Fix reader context leak when query response serialization fails

muted-tests.yml

@@ -318,9 +318,6 @@ tests:
 - class: org.elasticsearch.xpack.esql.qa.multi_node.GenerativeIT
   method: test
   issue: https://github.com/elastic/elasticsearch/issues/144587
-- class: org.elasticsearch.action.search.TransportSearchIT
-  method: testCircuitBreakerReduceFail
-  issue: https://github.com/elastic/elasticsearch/issues/144598
 - class: org.elasticsearch.xpack.esql.qa.single_node.GenerativeIT
   method: test
   issue: https://github.com/elastic/elasticsearch/issues/144587

server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java

@@ -40,6 +40,7 @@
 import org.elasticsearch.plugins.SearchPlugin;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.search.DocValueFormat;
+import org.elasticsearch.search.MockSearchService;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.SearchService;
 import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
@@ -541,13 +542,28 @@ public void onFailure(Exception exc) {
                 );
             }
             assertBusy(() -> assertThat(requestBreakerUsed(), equalTo(0L)));
+            assertBusy(MockSearchService::assertNoInFlightContext);
         } finally {
             updateClusterSettings(
                 Settings.builder().putNull("indices.breaker.request.limit").putNull(SearchService.BATCHED_QUERY_PHASE.getKey())
             );
         }
     }
 
+    public void testReaderContextFreedOnSerializationFailure() throws Exception {
+        String coordinatingNode = internalCluster().startCoordinatingOnlyNode(Settings.EMPTY);
+        indexSomeDocs("test", 1, 3);
+        ensureGreen("test");
+
+        updateClusterSettings(Settings.builder().put("indices.breaker.request.limit", "1b"));
+        try {
+            expectThrows(Exception.class, client(coordinatingNode).prepareSearch("test")::get);
+            assertBusy(MockSearchService::assertNoInFlightContext);
+        } finally {
+            updateClusterSettings(Settings.builder().putNull("indices.breaker.request.limit"));
+        }
+    }
+
     public void testCircuitBreakerFetchFail() throws Exception {
         int numShards = randomIntBetween(1, 10);
         int numDocs = numShards * 10;

server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java

@@ -11,6 +11,7 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionListenerResponseHandler;
@@ -738,8 +739,8 @@ public void onResponse(T response) {
                 response.writeTo(out);
                 bytesRef = out.moveToBytesReference();
             } catch (Exception e) {
-                channelListener.onFailure(e);
-                return;
+                // Propagate to caller so wrapFailureListener in SearchService can free the reader context.
+                throw ExceptionsHelper.convertToRuntime(e);
             }
             // respondAndRelease releases the bytes once the transport layer completes.
             ActionListener.respondAndRelease(channelListener, new BytesTransportResponse(bytesRef, transportVersion));

server/src/main/java/org/elasticsearch/search/SearchService.java

@@ -156,6 +156,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.LongSupplier;
@@ -754,11 +755,31 @@ public void executeQueryPhase(ShardSearchRequest request, CancellableTask task,
                 }
                 // TODO: i think it makes sense to always do a canMatch here and
                 // return an empty response (not null response) in case canMatch is false?
-                ensureAfterSeqNoRefreshed(shard, orig, () -> executeQueryPhase(orig, task), l);
+                executeQueryPhaseAsync(shard, orig, task, l);
             })
         );
     }
 
+    private void executeQueryPhaseAsync(
+        IndexShard shard,
+        ShardSearchRequest request,
+        CancellableTask task,
+        ActionListener<SearchPhaseResult> listener
+    ) {
+        // wrapFailureListener requires readerContext and markAsUsed, but those are created inside the supplier
+        // lambda below. The ActionListener.wrap callbacks are constructed (before the supplier runs) and must
+        // therefore read the listener indirectly. completionListenerRef starts as the plain listener so that any
+        // failure before the supplier runs is still forwarded. Once the supplier sets it to the wrapped version,
+        // the ActionListener.wrap callbacks will invoke the one that handles readerContext cleanup on failure.
+        final var completionListenerRef = new AtomicReference<>(listener);
+        ensureAfterSeqNoRefreshed(shard, request, () -> {
+            final ReaderContext readerContext = createOrGetReaderContext(request);
+            final Releasable markAsUsed = readerContext.markAsUsed(getKeepAlive(request));
+            completionListenerRef.set(wrapFailureListener(listener, readerContext, markAsUsed));
+            return executeQueryPhase(request, task, readerContext);
+        }, ActionListener.wrap(result -> completionListenerRef.get().onResponse(result), e -> completionListenerRef.get().onFailure(e)));
+    }
+
     private <T extends RefCounted> void ensureAfterSeqNoRefreshed(
         IndexShard shard,
         ShardSearchRequest request,
@@ -908,11 +929,10 @@ private static <T extends RefCounted> void runAsync(
      * It is the responsibility of the caller to ensure that the ref count is correctly decremented
      * when the object is no longer needed.
      */
-    private SearchPhaseResult executeQueryPhase(ShardSearchRequest request, CancellableTask task) throws Exception {
-        final ReaderContext readerContext = createOrGetReaderContext(request);
+    private SearchPhaseResult executeQueryPhase(ShardSearchRequest request, CancellableTask task, ReaderContext readerContext)
+        throws Exception {
         try (
             Releasable scope = tracer.withScope(task);
-            Releasable ignored = readerContext.markAsUsed(getKeepAlive(request));
             SearchContext context = createContext(readerContext, request, task, ResultsType.QUERY, true)
         ) {
             tracer.startTrace("executeQueryPhase", Map.of());
@@ -967,7 +987,6 @@ private SearchPhaseResult executeQueryPhase(ShardSearchRequest request, Cancella
                     : new ElasticsearchException(e.getCause());
             }
             logger.trace("Query phase failed", e);
-            processFailure(readerContext, e);
             throw e;
         }
     }
@@ -1651,36 +1670,36 @@ private <T> ActionListener<T> releaseCircuitBreakerOnResponse(
         ActionListener<T> listener,
         Function<T, FetchSearchResult> fetchResultExtractor
     ) {
-        return ActionListener.wrap(response -> {
-            try {
-                listener.onResponse(response);
-            } finally {
-                // Release bytes after the response handler completes
-                FetchSearchResult fetchResult = fetchResultExtractor.apply(response);
-                if (fetchResult != null) {
-                    fetchResult.releaseCircuitBreakerBytes(circuitBreaker);
-                }
-            }
-        }, listener::onFailure);
-    }
-
-    private <T> ActionListener<T> wrapFailureListener(ActionListener<T> listener, ReaderContext context, Releasable releasable) {
         return new ActionListener<>() {
             @Override
-            public void onResponse(T resp) {
-                Releasables.close(releasable);
-                listener.onResponse(resp);
+            public void onResponse(T response) {
+                try {
+                    listener.onResponse(response);
+                } finally {
+                    // Release bytes after the response handler completes, even if it throws.
+                    // Exceptions are intentionally allowed to propagate so that wrapFailureListener
+                    // can observe them and free the reader context via processFailure.
+                    FetchSearchResult fetchResult = fetchResultExtractor.apply(response);
+                    if (fetchResult != null) {
+                        fetchResult.releaseCircuitBreakerBytes(circuitBreaker);
+                    }
+                }
             }
 
             @Override
-            public void onFailure(Exception exc) {
-                processFailure(context, exc);
-                Releasables.close(releasable);
-                listener.onFailure(exc);
+            public void onFailure(Exception e) {
+                listener.onFailure(e);
             }
         };
     }
 
+    private <T> ActionListener<T> wrapFailureListener(ActionListener<T> listener, ReaderContext context, Releasable releasable) {
+        return ActionListener.releaseAfter(ActionListener.wrap(listener::onResponse, e -> {
+            processFailure(context, e);
+            listener.onFailure(e);
+        }), releasable);
+    }
+
     private static boolean isScrollContext(ReaderContext context) {
         return context instanceof LegacyReaderContext && context.singleSession() == false;
     }

server/src/test/java/org/elasticsearch/action/search/SearchTransportServiceChannelListenerTests.java

@@ -90,11 +90,12 @@ public void testNetworkPathBytesResponseRoundTrip() throws Exception {
         }
     }
 
-    public void testNetworkPathSerializationFailureSendsFailure() {
-        var sentException = new AtomicReference<Exception>();
-
+    public void testNetworkPathSerializationFailurePropagates() {
         var channel = new TestTransportChannel(
-            ActionListener.wrap(resp -> fail("should not succeed when serialization fails"), sentException::set)
+            ActionListener.wrap(
+                resp -> fail("should not succeed when serialization fails"),
+                e -> fail("should not send failure to channel; caller handles it")
+            )
         );
 
         ActionListener<FailingTestResponse> listener = SearchTransportService.channelListener(
@@ -103,11 +104,9 @@ public void testNetworkPathSerializationFailureSendsFailure() {
             newLimitedBreaker(ByteSizeValue.ofMb(100))
         );
 
-        listener.onResponse(new FailingTestResponse());
-
-        assertThat(sentException.get(), notNullValue());
-        assertThat(sentException.get(), instanceOf(IOException.class));
-        assertThat(sentException.get().getMessage(), equalTo("simulated serialization failure"));
+        var ex = expectThrows(RuntimeException.class, () -> listener.onResponse(new FailingTestResponse()));
+        assertThat(ex.getCause(), instanceOf(IOException.class));
+        assertThat(ex.getCause().getMessage(), equalTo("simulated serialization failure"));
     }
 
     public void testNetworkPathOnFailureForwardsFailure() {

x-pack/plugin/ql/test-fixtures/src/main/java/org/elasticsearch/xpack/ql/TestUtils.java

@@ -75,6 +75,7 @@
 
 import static java.util.Collections.emptyMap;
 import static org.elasticsearch.cluster.ClusterState.VERSION_INTRODUCING_TRANSPORT_VERSIONS;
+import static org.elasticsearch.test.ESTestCase.assertBusy;
 import static org.elasticsearch.test.ESTestCase.between;
 import static org.elasticsearch.test.ESTestCase.randomAlphaOfLength;
 import static org.elasticsearch.test.ESTestCase.randomBoolean;
@@ -174,15 +175,17 @@ public static EsRelation relation() {
     // Common methods / assertions
     //
 
-    public static void assertNoSearchContexts(RestClient client) throws IOException {
-        Map<String, Object> stats = searchStats(client);
-        @SuppressWarnings("unchecked")
-        Map<String, Object> indicesStats = (Map<String, Object>) stats.get("indices");
-        for (String index : indicesStats.keySet()) {
-            if (index.startsWith(".") == false) { // We are not interested in internal indices
-                assertEquals(index + " should have no search contexts", 0, getOpenContexts(stats, index));
+    public static void assertNoSearchContexts(RestClient client) throws Exception {
+        assertBusy(() -> {
+            Map<String, Object> stats = searchStats(client);
+            @SuppressWarnings("unchecked")
+            Map<String, Object> indicesStats = (Map<String, Object>) stats.get("indices");
+            for (String index : indicesStats.keySet()) {
+                if (index.startsWith(".") == false) { // We are not interested in internal indices
+                    assertEquals(index + " should have no search contexts", 0, getOpenContexts(stats, index));
+                }
             }
-        }
+        });
     }
 
     public static int getNumberOfSearchContexts(RestClient client, String index) throws IOException {