Defer circuit breaker release until transport write completes

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

@@ -11,6 +11,7 @@
 
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.ScoreMode;
+import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.DocWriteResponse;
@@ -23,6 +24,7 @@
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.common.breaker.CircuitBreakingException;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.settings.Settings;
@@ -503,7 +505,10 @@ public void onFailure(Exception e) {
                     Exception.class,
                     client.prepareSearch("test").addAggregation(new TestAggregationBuilder("test"))
                 );
-                assertThat(exc.getCause().getMessage(), containsString("<reduce_aggs>"));
+                assertNotNull(
+                    "root cause must be a CircuitBreakingException",
+                    ExceptionsHelper.unwrap(exc, CircuitBreakingException.class)
+                );
             });
 
             final AtomicArray<Exception> exceptions = new AtomicArray<>(10);
@@ -530,7 +535,10 @@ public void onFailure(Exception exc) {
             latch.await();
             assertThat(exceptions.asList().size(), equalTo(10));
             for (Exception exc : exceptions.asList()) {
-                assertThat(exc.getCause().getMessage(), containsString("<reduce_aggs>"));
+                assertNotNull(
+                    "root cause must be a CircuitBreakingException",
+                    ExceptionsHelper.unwrap(exc, CircuitBreakingException.class)
+                );
             }
             assertBusy(() -> assertThat(requestBreakerUsed(), equalTo(0L)));
         } finally {

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.TransportVersion;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionListenerResponseHandler;
 import org.elasticsearch.action.ActionResponse;
@@ -21,6 +22,8 @@
 import org.elasticsearch.client.internal.OriginSettingClient;
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.common.bytes.ReleasableBytesReference;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
@@ -50,9 +53,12 @@
 import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.AbstractTransportRequest;
+import org.elasticsearch.transport.BytesTransportResponse;
 import org.elasticsearch.transport.RemoteClusterService;
+import org.elasticsearch.transport.TaskTransportChannel;
 import org.elasticsearch.transport.Transport;
 import org.elasticsearch.transport.TransportActionProxy;
+import org.elasticsearch.transport.TransportChannel;
 import org.elasticsearch.transport.TransportException;
 import org.elasticsearch.transport.TransportRequestHandler;
 import org.elasticsearch.transport.TransportRequestOptions;
@@ -457,7 +463,7 @@ public static void registerRequestHandler(
             (request, channel, task) -> searchService.executeQueryPhase(
                 request,
                 (SearchShardTask) task,
-                new ChannelActionListener<>(channel)
+                channelListener(transportService, channel, searchService.getCircuitBreaker())
             )
         );
         TransportActionProxy.registerProxyActionWithDynamicResponseType(
@@ -513,7 +519,7 @@ public static void registerRequestHandler(
             (request, channel, task) -> searchService.executeFetchPhase(
                 request,
                 (SearchShardTask) task,
-                new ChannelActionListener<>(channel)
+                channelListener(transportService, channel, searchService.getCircuitBreaker())
             )
         );
         TransportActionProxy.registerProxyAction(
@@ -541,7 +547,11 @@ public static void registerRequestHandler(
         );
 
         final TransportRequestHandler<ShardFetchRequest> shardFetchRequestHandler = (request, channel, task) -> searchService
-            .executeFetchPhase(request, (SearchShardTask) task, new ChannelActionListener<>(channel));
+            .executeFetchPhase(
+                request,
+                (SearchShardTask) task,
+                channelListener(transportService, channel, searchService.getCircuitBreaker())
+            );
         transportService.registerRequestHandler(
             FETCH_ID_SCROLL_ACTION_NAME,
             EsExecutors.DIRECT_EXECUTOR_SERVICE,
@@ -672,6 +682,75 @@ private boolean assertNodePresent() {
         }
     }
 
+    /**
+     * Returns a listener that serializes responses to bytes on the network path.
+     *
+     * <p>On the <b>network path</b>, the response is serialized into bytes using a
+     * circuit-breaker-aware stream and sent as a {@link BytesTransportResponse}.
+     *
+     * <p>On the <b>direct (same-node) path</b> the response is forwarded as-is.
+     *
+     * <p>Circuit-breaker accounting for response objects is handled by the caller.
+     */
+    static <T extends TransportResponse> ActionListener<T> channelListener(
+        TransportService transportService,
+        TransportChannel channel,
+        @Nullable CircuitBreaker circuitBreaker
+    ) {
+        if (isDirectResponseChannel(channel)) {
+            return new ChannelActionListener<>(channel);
+        }
+        return new NetworkPathListener<>(transportService, channel, circuitBreaker);
+    }
+
+    private static boolean isDirectResponseChannel(TransportChannel channel) {
+        if (channel instanceof TaskTransportChannel ttc) {
+            channel = ttc.getChannel();
+        }
+        return TransportService.isDirectResponseChannel(channel);
+    }
+
+    /**
+     * Serializes the response into a {@link BytesTransportResponse} while keeping the breaker-accounted
+     * bytes alive for the response lifecycle. Captures the transport version from the channel at
+     * construction time and reuses it for serialization and the response metadata.
+     */
+    private static class NetworkPathListener<T extends TransportResponse> implements ActionListener<T> {
+        private final TransportService transportService;
+        private final TransportVersion transportVersion;
+        private final ChannelActionListener<BytesTransportResponse> channelListener;
+        @Nullable
+        private final CircuitBreaker circuitBreaker;
+
+        NetworkPathListener(TransportService transportService, TransportChannel channel, @Nullable CircuitBreaker circuitBreaker) {
+            this.transportService = transportService;
+            this.transportVersion = channel.getVersion();
+            this.channelListener = new ChannelActionListener<>(channel);
+            this.circuitBreaker = circuitBreaker;
+        }
+
+        @Override
+        public void onResponse(T response) {
+            // The bytes reference keeps breaker-accounted bytes; the stream output closes after serialization.
+            final ReleasableBytesReference bytesRef;
+            try (var out = transportService.newNetworkBytesStream(circuitBreaker)) {
+                out.setTransportVersion(transportVersion);
+                response.writeTo(out);
+                bytesRef = out.moveToBytesReference();
+            } catch (Exception e) {
+                channelListener.onFailure(e);
+                return;
+            }
+            // respondAndRelease releases the bytes once the transport layer completes.
+            ActionListener.respondAndRelease(channelListener, new BytesTransportResponse(bytesRef, transportVersion));
+        }
+
+        @Override
+        public void onFailure(Exception e) {
+            channelListener.onFailure(e);
+        }
+    }
+
     public void cancelSearchTask(SearchTask task, String reason) {
         CancelTasksRequest req = new CancelTasksRequest().setTargetTaskId(new TaskId(client.getLocalNodeId(), task.getId()))
             .setReason("Fatal failure during search: " + reason);

server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java

@@ -28,7 +28,7 @@ public final class FetchSearchResult extends SearchPhaseResult {
 
     private SearchHits hits;
 
-    private transient long searchHitsSizeBytes = 0L;
+    private long searchHitsSizeBytes = 0L;
 
     // client side counter
     private transient int counter;