From bb7e1267c5236f122fef0d7d50ad8eb4232eb504 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Tue, 24 Feb 2026 14:46:12 +0200 Subject: [PATCH 01/39] initial commit --- diff | 0 .../elasticsearch/search/SearchService.java | 23 ++- .../search/fetch/FetchSearchResult.java | 13 +- .../transport/OutboundHandler.java | 14 +- .../transport/TransportResponse.java | 43 ++++++ .../transport/TransportService.java | 7 +- .../SearchServiceCircuitBreakerTests.java | 114 ++++++++++++--- .../transport/OutboundHandlerTests.java | 137 +++++++++++++++--- 8 files changed, 297 insertions(+), 54 deletions(-) create mode 100644 diff diff --git a/diff b/diff new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/server/src/main/java/org/elasticsearch/search/SearchService.java b/server/src/main/java/org/elasticsearch/search/SearchService.java index 035ce4ff8b7fc..a5f38f0cfa2be 100644 --- a/server/src/main/java/org/elasticsearch/search/SearchService.java +++ b/server/src/main/java/org/elasticsearch/search/SearchService.java @@ -139,6 +139,7 @@ import org.elasticsearch.threadpool.ThreadPool; import org.elasticsearch.threadpool.ThreadPool.Names; import org.elasticsearch.transport.TransportRequest; +import org.elasticsearch.transport.TransportResponse; import org.elasticsearch.transport.Transports; import java.io.IOException; @@ -1637,21 +1638,29 @@ private static void checkKeepAliveLimit(long keepAlive, long maxKeepAlive) { } /** - * Wraps a listener to release circuit breaker bytes from a FetchSearchResult after the response is sent. - * The fetchResultExtractor function extracts the FetchSearchResult from the response type. + * Wraps a listener to defer circuit breaker release until after the response has been fully written to the network. + * Sets an {@code afterSendRelease} on the {@link TransportResponse} so that {@code OutboundHandler} releases + * the circuit breaker bytes only when the Netty write promise completes. */ - private ActionListener releaseCircuitBreakerOnResponse( + private ActionListener releaseCircuitBreakerOnResponse( ActionListener listener, Function fetchResultExtractor ) { return ActionListener.wrap(response -> { + FetchSearchResult fetchResult = fetchResultExtractor.apply(response); + if (fetchResult != null && fetchResult.getSearchHitsSizeBytes() > 0) { + response.setAfterSendRelease(() -> fetchResult.releaseCircuitBreakerBytes(circuitBreaker)); + } try { listener.onResponse(response); } finally { - // Release bytes after the response handler completes - FetchSearchResult fetchResult = fetchResultExtractor.apply(response); - if (fetchResult != null) { - fetchResult.releaseCircuitBreakerBytes(circuitBreaker); + // If the transport layer consumed afterSendRelease, it owns the release and will + // fire it on write completion. Otherwise, release now as a safety net (e.g. the + // channel doesn't support the mechanism, or listener.onResponse() threw before + // the transport could consume it). + Releasable unconsumed = response.consumeAfterSendRelease(); + if (unconsumed != null) { + Releasables.closeExpectNoException(unconsumed); } } }, listener::onFailure); diff --git a/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java b/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java index d9477a06f9bfe..dd23472482c27 100644 --- a/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java +++ b/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java @@ -23,12 +23,13 @@ import org.elasticsearch.transport.LeakTracker; import java.io.IOException; +import java.util.concurrent.atomic.AtomicLong; public final class FetchSearchResult extends SearchPhaseResult { private SearchHits hits; - private transient long searchHitsSizeBytes = 0L; + private final transient AtomicLong searchHitsSizeBytes = new AtomicLong(0L); // client side counter private transient int counter; @@ -88,17 +89,17 @@ public SearchHits hits() { } public void setSearchHitsSizeBytes(long bytes) { - this.searchHitsSizeBytes = bytes; + this.searchHitsSizeBytes.set(bytes); } public long getSearchHitsSizeBytes() { - return searchHitsSizeBytes; + return searchHitsSizeBytes.get(); } public void releaseCircuitBreakerBytes(CircuitBreaker circuitBreaker) { - if (searchHitsSizeBytes > 0L) { - circuitBreaker.addWithoutBreaking(-searchHitsSizeBytes); - searchHitsSizeBytes = 0L; + long bytes = searchHitsSizeBytes.getAndSet(0L); + if (bytes > 0L) { + circuitBreaker.addWithoutBreaking(-bytes); } } diff --git a/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java b/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java index 5427209b30616..4cdcbae8ac41c 100644 --- a/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java +++ b/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java @@ -146,6 +146,10 @@ void sendResponse( ) { assert assertValidTransportVersion(transportVersion); assert response.hasReferences(); + final Releasable afterSendRelease = response.consumeAfterSendRelease(); + final Releasable onAfter = afterSendRelease != null + ? Releasables.wrap(() -> messageListener.onResponseSent(requestId, action), afterSendRelease) + : () -> messageListener.onResponseSent(requestId, action); try { sendMessage( channel, @@ -157,9 +161,10 @@ void sendResponse( compressionScheme, transportVersion, responseStatsConsumer, - () -> messageListener.onResponseSent(requestId, action) + onAfter ); } catch (Exception ex) { + Releasables.closeExpectNoException(afterSendRelease); if (isHandshake) { logger.error( () -> format( @@ -217,6 +222,11 @@ public enum MessageDirection { RESPONSE_ERROR } + /** + * @param onAfter released when the channel write completes (success or failure). On serialization failure, {@code onAfter} + * is not released by this method; the caller is responsible for any cleanup (e.g. releasing resources + * that were composed into {@code onAfter}). + */ private void sendMessage( TcpChannel channel, MessageDirection messageDirection, @@ -253,7 +263,7 @@ private void sendMessage( throw e; } finally { if (serializeSuccess == false) { - Releasables.close(byteStreamOutput, onAfter); + Releasables.close(byteStreamOutput); } } responseStatsConsumer.addResponseStats(message.length()); diff --git a/server/src/main/java/org/elasticsearch/transport/TransportResponse.java b/server/src/main/java/org/elasticsearch/transport/TransportResponse.java index 13083521a710d..54c88759f5610 100644 --- a/server/src/main/java/org/elasticsearch/transport/TransportResponse.java +++ b/server/src/main/java/org/elasticsearch/transport/TransportResponse.java @@ -9,9 +9,52 @@ package org.elasticsearch.transport; +import org.elasticsearch.core.Releasable; +import org.elasticsearch.core.Releasables; + +import java.util.concurrent.atomic.AtomicReferenceFieldUpdater; + public abstract class TransportResponse extends TransportMessage { + + @SuppressWarnings("rawtypes") + private static final AtomicReferenceFieldUpdater AFTER_SEND_RELEASE_UPDATER = + AtomicReferenceFieldUpdater.newUpdater(TransportResponse.class, Releasable.class, "afterSendRelease"); + + private transient volatile Releasable afterSendRelease; + /** * Constructs a new empty transport response */ protected TransportResponse() {} + + /** + * Sets a {@link Releasable} that will be released after the response has been fully written to the network. + * This allows callers to defer resource cleanup (e.g. circuit breaker release) until the serialized bytes + * have actually been sent, rather than releasing as soon as the send is queued. + * + *

If a releasable has already been set (and not yet consumed), the new releasable is composed with + * the existing one so that both are released together when the transport layer consumes the after-send + * releasable. This supports multiple wrapper layers each attaching their own cleanup. + * + *

Thread-safe: uses an {@link AtomicReferenceFieldUpdater} so that a set on one thread + * is visible to a consume on another (e.g. search-service thread sets, transport thread consumes), + * without allocating an {@code AtomicReference} per instance. + */ + public void setAfterSendRelease(Releasable afterSendRelease) { + while (true) { + Releasable existing = AFTER_SEND_RELEASE_UPDATER.get(this); + Releasable combined = existing == null ? afterSendRelease : Releasables.wrap(existing, afterSendRelease); + if (AFTER_SEND_RELEASE_UPDATER.compareAndSet(this, existing, combined)) { + return; + } + } + } + + /** + * Atomically returns and clears the after-send {@link Releasable}. Returns {@code null} if none was set. + * Called by the transport layer to extract the releasable and compose it into the send-completion callback. + */ + public Releasable consumeAfterSendRelease() { + return AFTER_SEND_RELEASE_UPDATER.getAndSet(this, null); + } } diff --git a/server/src/main/java/org/elasticsearch/transport/TransportService.java b/server/src/main/java/org/elasticsearch/transport/TransportService.java index 60170125d65c0..4b27f0d4c848b 100644 --- a/server/src/main/java/org/elasticsearch/transport/TransportService.java +++ b/server/src/main/java/org/elasticsearch/transport/TransportService.java @@ -1570,7 +1570,7 @@ public String getProfileName() { @Override public void sendResponse(TransportResponse response) { - service.onResponseSent(requestId, action); + final Releasable afterSendRelease = response.consumeAfterSendRelease(); try (var shutdownBlock = service.pendingDirectHandlers.withRef()) { if (shutdownBlock == null) { // already shutting down, the handler will be completed by sendRequestInternal or doStop @@ -1603,6 +1603,11 @@ public String toString() { } }); } + } finally { + service.onResponseSent(requestId, action); + if (afterSendRelease != null) { + afterSendRelease.close(); + } } } diff --git a/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java b/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java index 323a53bb6de2c..48ce11af61658 100644 --- a/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java @@ -12,11 +12,14 @@ import org.elasticsearch.common.breaker.CircuitBreaker; import org.elasticsearch.common.breaker.CircuitBreakingException; import org.elasticsearch.common.breaker.NoopCircuitBreaker; +import org.elasticsearch.core.Releasable; +import org.elasticsearch.core.Releasables; import org.elasticsearch.search.fetch.FetchSearchResult; import org.elasticsearch.search.fetch.QueryFetchSearchResult; import org.elasticsearch.search.fetch.ScrollQueryFetchSearchResult; import org.elasticsearch.search.query.QuerySearchResult; import org.elasticsearch.test.ESTestCase; +import org.elasticsearch.transport.TransportResponse; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicLong; @@ -24,6 +27,8 @@ import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.notNullValue; +import static org.hamcrest.Matchers.nullValue; /** * Unit tests for circuit breaker release logic in SearchService. @@ -47,8 +52,14 @@ public void testReleaseCircuitBreakerForFetchResult() { assertThat(successCalled.get(), is(true)); assertThat(failureCalled.get(), is(false)); + // Safety-net finally block consumed and released afterSendRelease because the + // tracking listener doesn't simulate a transport layer that consumes it assertThat(breakerUsed.get(), equalTo(0L)); assertThat(result.getSearchHitsSizeBytes(), equalTo(0L)); + + // afterSendRelease was already consumed by the safety-net finally block + Releasable afterSend = result.consumeAfterSendRelease(); + assertThat(afterSend, nullValue()); } finally { result.decRef(); } @@ -74,6 +85,9 @@ public void testReleaseCircuitBreakerForQueryFetchResult() { assertThat(failureCalled.get(), is(false)); assertThat(breakerUsed.get(), equalTo(0L)); assertThat(fetchResult.getSearchHitsSizeBytes(), equalTo(0L)); + + Releasable afterSend = queryFetchResult.consumeAfterSendRelease(); + assertThat(afterSend, nullValue()); } finally { if (queryFetchResult != null) { queryFetchResult.decRef(); @@ -104,6 +118,9 @@ public void testReleaseCircuitBreakerForScrollResult() { assertThat(failureCalled.get(), is(false)); assertThat(breakerUsed.get(), equalTo(0L)); assertThat(fetchResult.getSearchHitsSizeBytes(), equalTo(0L)); + + Releasable afterSend = scrollResult.consumeAfterSendRelease(); + assertThat(afterSend, nullValue()); } finally { if (scrollResult != null) { scrollResult.decRef(); @@ -134,11 +151,61 @@ public void testExtractorReturnsNull() { AtomicBoolean successCalled = new AtomicBoolean(false); AtomicBoolean failureCalled = new AtomicBoolean(false); - querySearchResultListener(successCalled, failureCalled, breaker).onResponse(new QuerySearchResult()); + QuerySearchResult result = new QuerySearchResult(); + querySearchResultListener(successCalled, failureCalled, breaker).onResponse(result); assertThat(successCalled.get(), is(true)); assertThat(failureCalled.get(), is(false)); - // No breaker to release, should complete normally + assertThat(result.consumeAfterSendRelease(), nullValue()); + } + + /** + * Simulates the network path where the transport layer consumes afterSendRelease + * during listener.onResponse(). The safety-net finally block should see that + * afterSendRelease was already consumed and not release the bytes itself. + * The transport layer (simulated here) owns the release and fires it later + * when the network write completes. + */ + public void testTransportConsumesAfterSendRelease() { + AtomicLong breakerUsed = new AtomicLong(5000); + CircuitBreaker breaker = new TestCircuitBreaker(breakerUsed); + + FetchSearchResult result = new FetchSearchResult(); + try { + result.setSearchHitsSizeBytes(5000L); + AtomicBoolean successCalled = new AtomicBoolean(false); + + // Simulate a transport layer listener that consumes afterSendRelease + // (like OutboundHandler.sendResponse does for the network path) + Releasable[] captured = new Releasable[1]; + ActionListener transportSimulator = new ActionListener<>() { + @Override + public void onResponse(FetchSearchResult response) { + captured[0] = response.consumeAfterSendRelease(); + successCalled.set(true); + } + + @Override + public void onFailure(Exception e) { + fail("unexpected failure"); + } + }; + + fetchSearchResultListenerWithCustomInner(transportSimulator, breaker).onResponse(result); + + assertThat(successCalled.get(), is(true)); + // Bytes must still be reserved: the transport owns the release + assertThat(breakerUsed.get(), equalTo(5000L)); + assertThat(result.getSearchHitsSizeBytes(), equalTo(5000L)); + assertThat(captured[0], notNullValue()); + + // Simulate Netty write completion + captured[0].close(); + assertThat(breakerUsed.get(), equalTo(0L)); + assertThat(result.getSearchHitsSizeBytes(), equalTo(0L)); + } finally { + result.decRef(); + } } public void testMultipleReleasesAreIdempotent() { @@ -180,6 +247,9 @@ public void testLargeAllocation() { assertThat(successCalled.get(), is(true)); assertThat(breakerUsed.get(), equalTo(0L)); assertThat(result.getSearchHitsSizeBytes(), equalTo(0L)); + + Releasable afterSend = result.consumeAfterSendRelease(); + assertThat(afterSend, nullValue()); } finally { result.decRef(); } @@ -231,31 +301,28 @@ public void onFailure(Exception e) { } /** - * Wrap a listener with circuit breaker release. + * Wrap a listener with deferred circuit breaker release via afterSendRelease on the response. + * Mirrors the behavior of {@code SearchService.releaseCircuitBreakerOnResponse()}. */ - private ActionListener withCircuitBreakerRelease( + private ActionListener withCircuitBreakerRelease( ActionListener listener, CircuitBreaker breaker, Function fetchResultExtractor ) { - return new ActionListener<>() { - @Override - public void onResponse(T response) { - try { - listener.onResponse(response); - } finally { - FetchSearchResult fetchResult = fetchResultExtractor.apply(response); - if (fetchResult != null) { - fetchResult.releaseCircuitBreakerBytes(breaker); - } - } + return ActionListener.wrap(response -> { + FetchSearchResult fetchResult = fetchResultExtractor.apply(response); + if (fetchResult != null && fetchResult.getSearchHitsSizeBytes() > 0) { + response.setAfterSendRelease(() -> fetchResult.releaseCircuitBreakerBytes(breaker)); } - - @Override - public void onFailure(Exception e) { - listener.onFailure(e); + try { + listener.onResponse(response); + } finally { + Releasable unconsumed = response.consumeAfterSendRelease(); + if (unconsumed != null) { + Releasables.closeExpectNoException(unconsumed); + } } - }; + }, listener::onFailure); } private ActionListener querySearchResultListener( @@ -274,6 +341,13 @@ private ActionListener fetchSearchResultListener( return withCircuitBreakerRelease(trackingListener(successCalled, failureCalled), breaker, Function.identity()); } + private ActionListener fetchSearchResultListenerWithCustomInner( + ActionListener inner, + CircuitBreaker breaker + ) { + return withCircuitBreakerRelease(inner, breaker, Function.identity()); + } + private ActionListener queryFetchSearchResultListener( AtomicBoolean successCalled, AtomicBoolean failureCalled, diff --git a/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java b/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java index d76d4491a2b5d..e3e22920fa53d 100644 --- a/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java +++ b/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java @@ -272,6 +272,119 @@ public void onResponseSent(long requestId, String action) { assertEquals("header_value", header.getHeaders().v1().get("header")); } + /** + * Verifies that {@code afterSendRelease} set on a {@link TransportResponse} is NOT released when + * {@code sendResponse} returns (the send is only queued), but IS released when the write-completion + * listener fires with success. + */ + public void testAfterSendReleaseOnWriteSuccess() { + TransportVersion version = TransportVersionUtils.randomCompatibleVersion(); + String action = randomAlphaOfLength(10); + long requestId = randomLongBetween(0, 300); + TestResponse response = new TestResponse(randomAlphaOfLength(10)); + + AtomicBoolean released = new AtomicBoolean(false); + response.setAfterSendRelease(() -> assertTrue(released.compareAndSet(false, true))); + + AtomicBoolean onResponseSentCalled = new AtomicBoolean(false); + handler.setMessageListener(new TransportMessageListener() { + @Override + public void onResponseSent(long requestId, String action) { + onResponseSentCalled.set(true); + } + }); + + Compression.Scheme compress = randomFrom(compressionScheme, null); + handler.sendResponse(version, channel, requestId, action, response, compress, false, ResponseStatsConsumer.NONE); + + assertFalse("afterSendRelease must not be released before write completes", released.get()); + assertFalse("onResponseSent must not fire before write completes", onResponseSentCalled.get()); + assertNull("afterSendRelease should have been consumed from the response", response.consumeAfterSendRelease()); + + ActionListener sendListener = channel.getListenerCaptor().get(); + sendListener.onResponse(null); + + assertTrue("afterSendRelease must be released on write success", released.get()); + assertTrue("onResponseSent must fire on write success", onResponseSentCalled.get()); + } + + /** + * Verifies that {@code afterSendRelease} set on a {@link TransportResponse} is released when the + * write-completion listener fires with a failure (e.g. network error during write). + */ + public void testAfterSendReleaseOnWriteFailure() { + TransportVersion version = TransportVersionUtils.randomCompatibleVersion(); + String action = randomAlphaOfLength(10); + long requestId = randomLongBetween(0, 300); + TestResponse response = new TestResponse(randomAlphaOfLength(10)); + + AtomicBoolean released = new AtomicBoolean(false); + response.setAfterSendRelease(() -> assertTrue(released.compareAndSet(false, true))); + + AtomicBoolean onResponseSentCalled = new AtomicBoolean(false); + handler.setMessageListener(new TransportMessageListener() { + @Override + public void onResponseSent(long requestId, String action) { + onResponseSentCalled.set(true); + } + }); + + Compression.Scheme compress = randomFrom(compressionScheme, null); + handler.sendResponse(version, channel, requestId, action, response, compress, false, ResponseStatsConsumer.NONE); + + assertFalse("afterSendRelease must not be released before write completes", released.get()); + assertFalse("onResponseSent must not fire before write completes", onResponseSentCalled.get()); + + ActionListener sendListener = channel.getListenerCaptor().get(); + sendListener.onFailure(new IOException("write failed")); + + assertTrue("afterSendRelease must be released on write failure", released.get()); + assertTrue("onResponseSent must fire on write failure", onResponseSentCalled.get()); + } + + /** + * Verifies that {@code afterSendRelease} set on a {@link TransportResponse} is released immediately + * when serialization fails (before the write is ever attempted), via the catch block in + * {@link OutboundHandler#sendResponse}. + */ + public void testAfterSendReleaseOnSerializationFailure() { + TransportVersion version = TransportVersionUtils.randomCompatibleVersion(); + String action = randomAlphaOfLength(10); + long requestId = randomLongBetween(0, 300); + + var response = new ReleasbleTestResponse(randomAlphaOfLength(10)) { + @Override + public void writeTo(StreamOutput out) { + throw new CircuitBreakingException("simulated cbe", CircuitBreaker.Durability.TRANSIENT); + } + }; + + AtomicBoolean released = new AtomicBoolean(false); + response.setAfterSendRelease(() -> assertTrue(released.compareAndSet(false, true))); + + handler.setMessageListener(new TransportMessageListener() { + @Override + public void onResponseSent(long requestId, String action) { + throw new AssertionError("onResponseSent(success) must not be called on serialization failure"); + } + + @Override + public void onResponseSent(long requestId, String action, Exception error) { + // expected — error response is sent instead + } + }); + + Compression.Scheme compress = randomFrom(compressionScheme, null); + try { + handler.sendResponse(version, channel, requestId, action, response, compress, false, ResponseStatsConsumer.NONE); + } finally { + response.decRef(); + } + + assertTrue("afterSendRelease must be released on serialization failure", released.get()); + assertTrue("response resources must be released on serialization failure", response.released.get()); + } + public void testErrorResponse() throws IOException { ThreadContext threadContext = threadPool.getThreadContext(); TransportVersion version = TransportVersionUtils.randomCompatibleVersion(); @@ -353,8 +466,8 @@ public void onResponseSent(long requestId, String action) { @Override public void onResponseSent(long requestId, String action, Exception error) { assertNotNull(channel.getMessageCaptor().get()); - assertThat(requestIdRef.get(), equalTo(requestId)); - assertThat(actionRef.get(), equalTo(action)); + requestIdRef.set(requestId); + actionRef.set(action); exceptionRef.set(error); } }); @@ -404,18 +517,14 @@ public void sendMessage(BytesReference reference, ActionListener listener) handler.setMessageListener(new TransportMessageListener() { @Override public void onResponseSent(long requestId, String action) { - assertNull(channel.getMessageCaptor().get()); - assertThat(requestIdRef.get(), equalTo(0L)); - requestIdRef.set(requestId); - assertNull(actionRef.get()); - actionRef.set(action); + throw new AssertionError("onResponseSent(success) must not be called on serialization failure"); } @Override public void onResponseSent(long requestId, String action, Exception error) { assertNull(channel.getMessageCaptor().get()); - assertThat(requestIdRef.get(), equalTo(requestId)); - assertThat(actionRef.get(), equalTo(action)); + requestIdRef.set(requestId); + actionRef.set(action); exceptionRef.set(error); } }); @@ -453,16 +562,10 @@ public void writeTo(StreamOutput out) { } }; - AtomicLong requestIdRef = new AtomicLong(); - AtomicReference actionRef = new AtomicReference<>(); handler.setMessageListener(new TransportMessageListener() { @Override public void onResponseSent(long requestId, String action) { - assertNull(channel.getMessageCaptor().get()); - assertThat(requestIdRef.get(), equalTo(0L)); - requestIdRef.set(requestId); - assertNull(actionRef.get()); - actionRef.set(action); + throw new AssertionError("onResponseSent(success) must not be called on serialization failure"); } @Override @@ -478,8 +581,6 @@ public void onResponseSent(long requestId, String action, Exception error) { } assertNull(channel.getMessageCaptor().get()); assertNull(channel.getListenerCaptor().get()); - assertEquals(requestId, requestIdRef.get()); - assertEquals(action, actionRef.get()); assertTrue(response.released.get()); assertFalse(channel.isOpen()); assertTrue(closeListener.isDone()); From 80991f254fc45ee7182171bede75cfdbbd474345 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Tue, 24 Feb 2026 15:44:57 +0200 Subject: [PATCH 02/39] update after review --- diff | 0 .../transport/OutboundHandler.java | 5 ----- .../transport/TransportResponse.java | 18 ++++++------------ 3 files changed, 6 insertions(+), 17 deletions(-) delete mode 100644 diff diff --git a/diff b/diff deleted file mode 100644 index e69de29bb2d1d..0000000000000 diff --git a/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java b/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java index 4cdcbae8ac41c..53d2735b056c5 100644 --- a/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java +++ b/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java @@ -222,11 +222,6 @@ public enum MessageDirection { RESPONSE_ERROR } - /** - * @param onAfter released when the channel write completes (success or failure). On serialization failure, {@code onAfter} - * is not released by this method; the caller is responsible for any cleanup (e.g. releasing resources - * that were composed into {@code onAfter}). - */ private void sendMessage( TcpChannel channel, MessageDirection messageDirection, diff --git a/server/src/main/java/org/elasticsearch/transport/TransportResponse.java b/server/src/main/java/org/elasticsearch/transport/TransportResponse.java index 54c88759f5610..76e0544909f88 100644 --- a/server/src/main/java/org/elasticsearch/transport/TransportResponse.java +++ b/server/src/main/java/org/elasticsearch/transport/TransportResponse.java @@ -10,8 +10,6 @@ package org.elasticsearch.transport; import org.elasticsearch.core.Releasable; -import org.elasticsearch.core.Releasables; - import java.util.concurrent.atomic.AtomicReferenceFieldUpdater; public abstract class TransportResponse extends TransportMessage { @@ -22,6 +20,8 @@ public abstract class TransportResponse extends TransportMessage { private transient volatile Releasable afterSendRelease; + + /** * Constructs a new empty transport response */ @@ -32,21 +32,15 @@ protected TransportResponse() {} * This allows callers to defer resource cleanup (e.g. circuit breaker release) until the serialized bytes * have actually been sent, rather than releasing as soon as the send is queued. * - *

If a releasable has already been set (and not yet consumed), the new releasable is composed with - * the existing one so that both are released together when the transport layer consumes the after-send - * releasable. This supports multiple wrapper layers each attaching their own cleanup. - * *

Thread-safe: uses an {@link AtomicReferenceFieldUpdater} so that a set on one thread * is visible to a consume on another (e.g. search-service thread sets, transport thread consumes), * without allocating an {@code AtomicReference} per instance. + * + * @throws IllegalStateException if a releasable has already been set and not yet consumed */ public void setAfterSendRelease(Releasable afterSendRelease) { - while (true) { - Releasable existing = AFTER_SEND_RELEASE_UPDATER.get(this); - Releasable combined = existing == null ? afterSendRelease : Releasables.wrap(existing, afterSendRelease); - if (AFTER_SEND_RELEASE_UPDATER.compareAndSet(this, existing, combined)) { - return; - } + if (AFTER_SEND_RELEASE_UPDATER.compareAndSet(this, null, afterSendRelease) == false) { + throw new IllegalStateException("afterSendRelease already set"); } } From 4348a8d193a0d748db8f4ec26333d2bc6c030ed2 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Thu, 26 Feb 2026 11:32:12 +0200 Subject: [PATCH 03/39] update code --- .../elasticsearch/search/SearchService.java | 15 ++++++------ .../transport/OutboundHandler.java | 1 + .../transport/TransportResponse.java | 16 ++++--------- .../SearchServiceCircuitBreakerTests.java | 24 +------------------ .../transport/OutboundHandlerTests.java | 14 ----------- 5 files changed, 15 insertions(+), 55 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/search/SearchService.java b/server/src/main/java/org/elasticsearch/search/SearchService.java index a5f38f0cfa2be..df5e4e70b2972 100644 --- a/server/src/main/java/org/elasticsearch/search/SearchService.java +++ b/server/src/main/java/org/elasticsearch/search/SearchService.java @@ -1638,9 +1638,10 @@ private static void checkKeepAliveLimit(long keepAlive, long maxKeepAlive) { } /** - * Wraps a listener to defer circuit breaker release until after the response has been fully written to the network. - * Sets an {@code afterSendRelease} on the {@link TransportResponse} so that {@code OutboundHandler} releases - * the circuit breaker bytes only when the Netty write promise completes. + * Wraps a listener so that circuit breaker bytes are released after the response bytes have been + * written to the network, rather than immediately after the send is queued. The release is attached + * to the {@link TransportResponse} via {@code afterSendRelease} and executed by the transport layer + * on write completion. */ private ActionListener releaseCircuitBreakerOnResponse( ActionListener listener, @@ -1654,10 +1655,10 @@ private ActionListener releaseCircuitBreakerOnR try { listener.onResponse(response); } finally { - // If the transport layer consumed afterSendRelease, it owns the release and will - // fire it on write completion. Otherwise, release now as a safety net (e.g. the - // channel doesn't support the mechanism, or listener.onResponse() threw before - // the transport could consume it). + // The transport layer normally consumes afterSendRelease and releases it + // when the network write completes. If it was not consumed (e.g. an + // exception prevented the response from reaching the transport), release + // it here to avoid leaking circuit breaker bytes. Releasable unconsumed = response.consumeAfterSendRelease(); if (unconsumed != null) { Releasables.closeExpectNoException(unconsumed); diff --git a/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java b/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java index 53d2735b056c5..b413f155957df 100644 --- a/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java +++ b/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java @@ -222,6 +222,7 @@ public enum MessageDirection { RESPONSE_ERROR } + private void sendMessage( TcpChannel channel, MessageDirection messageDirection, diff --git a/server/src/main/java/org/elasticsearch/transport/TransportResponse.java b/server/src/main/java/org/elasticsearch/transport/TransportResponse.java index 76e0544909f88..e3b471de1d0d7 100644 --- a/server/src/main/java/org/elasticsearch/transport/TransportResponse.java +++ b/server/src/main/java/org/elasticsearch/transport/TransportResponse.java @@ -20,21 +20,15 @@ public abstract class TransportResponse extends TransportMessage { private transient volatile Releasable afterSendRelease; - - /** * Constructs a new empty transport response */ protected TransportResponse() {} /** - * Sets a {@link Releasable} that will be released after the response has been fully written to the network. - * This allows callers to defer resource cleanup (e.g. circuit breaker release) until the serialized bytes - * have actually been sent, rather than releasing as soon as the send is queued. - * - *

Thread-safe: uses an {@link AtomicReferenceFieldUpdater} so that a set on one thread - * is visible to a consume on another (e.g. search-service thread sets, transport thread consumes), - * without allocating an {@code AtomicReference} per instance. + * Attaches a {@link Releasable} to this response that the transport layer will release once the + * response bytes have been fully written to the network. This is used to defer resource cleanup + * (e.g. circuit breaker release) until after the write completes. * * @throws IllegalStateException if a releasable has already been set and not yet consumed */ @@ -45,8 +39,8 @@ public void setAfterSendRelease(Releasable afterSendRelease) { } /** - * Atomically returns and clears the after-send {@link Releasable}. Returns {@code null} if none was set. - * Called by the transport layer to extract the releasable and compose it into the send-completion callback. + * Atomically returns and clears the after-send {@link Releasable}, or {@code null} if none was set. + * The transport layer calls this to take ownership of the releasable and release it on write completion. */ public Releasable consumeAfterSendRelease() { return AFTER_SEND_RELEASE_UPDATER.getAndSet(this, null); diff --git a/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java b/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java index 48ce11af61658..c9390538cd539 100644 --- a/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java @@ -52,12 +52,9 @@ public void testReleaseCircuitBreakerForFetchResult() { assertThat(successCalled.get(), is(true)); assertThat(failureCalled.get(), is(false)); - // Safety-net finally block consumed and released afterSendRelease because the - // tracking listener doesn't simulate a transport layer that consumes it assertThat(breakerUsed.get(), equalTo(0L)); assertThat(result.getSearchHitsSizeBytes(), equalTo(0L)); - // afterSendRelease was already consumed by the safety-net finally block Releasable afterSend = result.consumeAfterSendRelease(); assertThat(afterSend, nullValue()); } finally { @@ -159,13 +156,6 @@ public void testExtractorReturnsNull() { assertThat(result.consumeAfterSendRelease(), nullValue()); } - /** - * Simulates the network path where the transport layer consumes afterSendRelease - * during listener.onResponse(). The safety-net finally block should see that - * afterSendRelease was already consumed and not release the bytes itself. - * The transport layer (simulated here) owns the release and fires it later - * when the network write completes. - */ public void testTransportConsumesAfterSendRelease() { AtomicLong breakerUsed = new AtomicLong(5000); CircuitBreaker breaker = new TestCircuitBreaker(breakerUsed); @@ -175,8 +165,6 @@ public void testTransportConsumesAfterSendRelease() { result.setSearchHitsSizeBytes(5000L); AtomicBoolean successCalled = new AtomicBoolean(false); - // Simulate a transport layer listener that consumes afterSendRelease - // (like OutboundHandler.sendResponse does for the network path) Releasable[] captured = new Releasable[1]; ActionListener transportSimulator = new ActionListener<>() { @Override @@ -191,15 +179,13 @@ public void onFailure(Exception e) { } }; - fetchSearchResultListenerWithCustomInner(transportSimulator, breaker).onResponse(result); + withCircuitBreakerRelease(transportSimulator, breaker, Function.identity()).onResponse(result); assertThat(successCalled.get(), is(true)); - // Bytes must still be reserved: the transport owns the release assertThat(breakerUsed.get(), equalTo(5000L)); assertThat(result.getSearchHitsSizeBytes(), equalTo(5000L)); assertThat(captured[0], notNullValue()); - // Simulate Netty write completion captured[0].close(); assertThat(breakerUsed.get(), equalTo(0L)); assertThat(result.getSearchHitsSizeBytes(), equalTo(0L)); @@ -302,7 +288,6 @@ public void onFailure(Exception e) { /** * Wrap a listener with deferred circuit breaker release via afterSendRelease on the response. - * Mirrors the behavior of {@code SearchService.releaseCircuitBreakerOnResponse()}. */ private ActionListener withCircuitBreakerRelease( ActionListener listener, @@ -341,13 +326,6 @@ private ActionListener fetchSearchResultListener( return withCircuitBreakerRelease(trackingListener(successCalled, failureCalled), breaker, Function.identity()); } - private ActionListener fetchSearchResultListenerWithCustomInner( - ActionListener inner, - CircuitBreaker breaker - ) { - return withCircuitBreakerRelease(inner, breaker, Function.identity()); - } - private ActionListener queryFetchSearchResultListener( AtomicBoolean successCalled, AtomicBoolean failureCalled, diff --git a/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java b/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java index e3e22920fa53d..ff84648f38494 100644 --- a/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java +++ b/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java @@ -272,11 +272,6 @@ public void onResponseSent(long requestId, String action) { assertEquals("header_value", header.getHeaders().v1().get("header")); } - /** - * Verifies that {@code afterSendRelease} set on a {@link TransportResponse} is NOT released when - * {@code sendResponse} returns (the send is only queued), but IS released when the write-completion - * listener fires with success. - */ public void testAfterSendReleaseOnWriteSuccess() { TransportVersion version = TransportVersionUtils.randomCompatibleVersion(); String action = randomAlphaOfLength(10); @@ -308,10 +303,6 @@ public void onResponseSent(long requestId, String action) { assertTrue("onResponseSent must fire on write success", onResponseSentCalled.get()); } - /** - * Verifies that {@code afterSendRelease} set on a {@link TransportResponse} is released when the - * write-completion listener fires with a failure (e.g. network error during write). - */ public void testAfterSendReleaseOnWriteFailure() { TransportVersion version = TransportVersionUtils.randomCompatibleVersion(); String action = randomAlphaOfLength(10); @@ -342,11 +333,6 @@ public void onResponseSent(long requestId, String action) { assertTrue("onResponseSent must fire on write failure", onResponseSentCalled.get()); } - /** - * Verifies that {@code afterSendRelease} set on a {@link TransportResponse} is released immediately - * when serialization fails (before the write is ever attempted), via the catch block in - * {@link OutboundHandler#sendResponse}. - */ public void testAfterSendReleaseOnSerializationFailure() { TransportVersion version = TransportVersionUtils.randomCompatibleVersion(); String action = randomAlphaOfLength(10); From 3f6f58108cf60ae54bce368129a4689ddb0ead34 Mon Sep 17 00:00:00 2001 From: elasticsearchmachine Date: Thu, 26 Feb 2026 10:17:57 +0000 Subject: [PATCH 04/39] [CI] Auto commit changes from spotless --- .../java/org/elasticsearch/transport/OutboundHandler.java | 1 - .../java/org/elasticsearch/transport/TransportResponse.java | 5 +++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java b/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java index b413f155957df..53d2735b056c5 100644 --- a/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java +++ b/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java @@ -222,7 +222,6 @@ public enum MessageDirection { RESPONSE_ERROR } - private void sendMessage( TcpChannel channel, MessageDirection messageDirection, diff --git a/server/src/main/java/org/elasticsearch/transport/TransportResponse.java b/server/src/main/java/org/elasticsearch/transport/TransportResponse.java index e3b471de1d0d7..fc1b70c5a6c58 100644 --- a/server/src/main/java/org/elasticsearch/transport/TransportResponse.java +++ b/server/src/main/java/org/elasticsearch/transport/TransportResponse.java @@ -10,13 +10,14 @@ package org.elasticsearch.transport; import org.elasticsearch.core.Releasable; + import java.util.concurrent.atomic.AtomicReferenceFieldUpdater; public abstract class TransportResponse extends TransportMessage { @SuppressWarnings("rawtypes") - private static final AtomicReferenceFieldUpdater AFTER_SEND_RELEASE_UPDATER = - AtomicReferenceFieldUpdater.newUpdater(TransportResponse.class, Releasable.class, "afterSendRelease"); + private static final AtomicReferenceFieldUpdater AFTER_SEND_RELEASE_UPDATER = AtomicReferenceFieldUpdater + .newUpdater(TransportResponse.class, Releasable.class, "afterSendRelease"); private transient volatile Releasable afterSendRelease; From 5a37cde091ae3958fe674e0a60007fdd8d6dc63b Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Thu, 5 Mar 2026 17:00:46 +0200 Subject: [PATCH 05/39] update codee --- .../action/search/SearchTransportService.java | 95 +++++- .../elasticsearch/search/SearchService.java | 39 +-- .../search/fetch/FetchSearchResult.java | 20 ++ .../transport/OutboundHandler.java | 7 +- .../transport/TransportResponse.java | 32 -- .../transport/TransportService.java | 7 +- .../action/search/AsBytesResponseTests.java | 275 ++++++++++++++++++ .../SearchServiceCircuitBreakerTests.java | 160 ++++++---- .../transport/OutboundHandlerTests.java | 99 ------- 9 files changed, 501 insertions(+), 233 deletions(-) create mode 100644 server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java index 74ca0c9080ee6..62a4fae651ef8 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java +++ b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java @@ -22,6 +22,7 @@ import org.elasticsearch.client.internal.node.NodeClient; import org.elasticsearch.cluster.node.DiscoveryNode; import org.elasticsearch.common.io.stream.NamedWriteableRegistry; +import org.elasticsearch.common.io.stream.RecyclerBytesStreamOutput; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.common.io.stream.Writeable; @@ -31,6 +32,7 @@ import org.elasticsearch.common.util.concurrent.ThrottledTaskRunner; import org.elasticsearch.core.Nullable; import org.elasticsearch.core.Releasable; +import org.elasticsearch.core.Releasables; import org.elasticsearch.search.SearchPhaseResult; import org.elasticsearch.search.SearchService; import org.elasticsearch.search.dfs.DfsSearchResult; @@ -50,9 +52,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; @@ -66,6 +71,7 @@ import java.util.Objects; import java.util.concurrent.Executor; import java.util.function.BiFunction; +import java.util.function.Consumer; /** * An encapsulation of {@link SearchService} operations exposed through @@ -457,7 +463,11 @@ public static void registerRequestHandler( (request, channel, task) -> searchService.executeQueryPhase( request, (SearchShardTask) task, - new ChannelActionListener<>(channel) + asBytesResponse(transportService, channel, response -> { + if (response instanceof QueryFetchSearchResult qfr) { + qfr.fetchResult().releaseCircuitBreakerBytes(searchService.getCircuitBreaker()); + } + }) ) ); TransportActionProxy.registerProxyActionWithDynamicResponseType( @@ -513,7 +523,11 @@ public static void registerRequestHandler( (request, channel, task) -> searchService.executeFetchPhase( request, (SearchShardTask) task, - new ChannelActionListener<>(channel) + asBytesResponse( + transportService, + channel, + response -> response.result().fetchResult().releaseCircuitBreakerBytes(searchService.getCircuitBreaker()) + ) ) ); TransportActionProxy.registerProxyAction( @@ -541,7 +555,15 @@ public static void registerRequestHandler( ); final TransportRequestHandler shardFetchRequestHandler = (request, channel, task) -> searchService - .executeFetchPhase(request, (SearchShardTask) task, new ChannelActionListener<>(channel)); + .executeFetchPhase( + request, + (SearchShardTask) task, + asBytesResponse( + transportService, + channel, + response -> response.releaseCircuitBreakerBytes(searchService.getCircuitBreaker()) + ) + ); transportService.registerRequestHandler( FETCH_ID_SCROLL_ACTION_NAME, EsExecutors.DIRECT_EXECUTOR_SERVICE, @@ -672,6 +694,73 @@ private boolean assertNodePresent() { } } + /** + * Creates a listener that releases circuit-breaker reservations (via {@code afterSerialize}) and sends the + * response over the channel. For network channels the response is pre-serialized into a + * {@link BytesTransportResponse} so that the original in-memory data can be freed by the caller's + * {@code respondAndRelease}. For local (direct) channels the original response is forwarded as-is, + * because TransportService.DirectResponseChannel} passes the object directly to the response + * handler without serialization/deserialization. + *

+ * The caller is responsible for calling {@code response.decRef()} (typically via + * {@link ActionListener#respondAndRelease}); this method never calls it. + * + * @param afterSerialize called after serialization (network) or after the handler processes the response (local) + * to release resources such as circuit-breaker reservations + */ + static ActionListener asBytesResponse( + TransportService transportService, + TransportChannel channel, + Consumer afterSerialize + ) { + TransportChannel innerChannel = channel instanceof TaskTransportChannel ttc ? ttc.getChannel() : channel; + if (TransportService.isDirectResponseChannel(innerChannel)) { + var channelListener = new ChannelActionListener(channel); + return new ActionListener<>() { + @Override + public void onResponse(T response) { + try { + channelListener.onResponse(response); + } finally { + afterSerialize.accept(response); + } + } + + @Override + public void onFailure(Exception e) { + channelListener.onFailure(e); + } + }; + } + + var channelListener = new ChannelActionListener(channel); + return new ActionListener<>() { + @Override + public void onResponse(T response) { + RecyclerBytesStreamOutput out = transportService.newNetworkBytesStream(); + try { + out.setTransportVersion(channel.getVersion()); + response.writeTo(out); + } catch (Exception e) { + Releasables.close(out); + afterSerialize.accept(response); + channelListener.onFailure(e); + return; + } + afterSerialize.accept(response); + ActionListener.respondAndRelease( + channelListener, + new BytesTransportResponse(out.moveToBytesReference(), out.getTransportVersion()) + ); + } + + @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); diff --git a/server/src/main/java/org/elasticsearch/search/SearchService.java b/server/src/main/java/org/elasticsearch/search/SearchService.java index efd1568a23613..3cc1f9f5fbf35 100644 --- a/server/src/main/java/org/elasticsearch/search/SearchService.java +++ b/server/src/main/java/org/elasticsearch/search/SearchService.java @@ -139,7 +139,6 @@ import org.elasticsearch.threadpool.ThreadPool; import org.elasticsearch.threadpool.ThreadPool.Names; import org.elasticsearch.transport.TransportRequest; -import org.elasticsearch.transport.TransportResponse; import org.elasticsearch.transport.Transports; import java.io.IOException; @@ -1210,13 +1209,7 @@ public void executeFetchPhase( // we handle the failure in the failure listener below throw e; } - }, - wrapFailureListener( - releaseCircuitBreakerOnResponse(listener, result -> result.result().fetchResult()), - readerContext, - markAsUsed - ) - ); + }, wrapFailureListener(listener, readerContext, markAsUsed)); } public void executeFetchPhase(ShardFetchRequest request, CancellableTask task, ActionListener listener) { @@ -1255,7 +1248,7 @@ public void executeFetchPhase(ShardFetchRequest request, CancellableTask task, A // we handle the failure in the failure listener below throw e; } - }, wrapFailureListener(releaseCircuitBreakerOnResponse(listener, result -> result), readerContext, markAsUsed)); + }, wrapFailureListener(listener, readerContext, markAsUsed)); })); } @@ -1645,30 +1638,26 @@ private static void checkKeepAliveLimit(long keepAlive, long maxKeepAlive) { } /** - * Wraps a listener so that circuit breaker bytes are released after the response bytes have been - * written to the network, rather than immediately after the send is queued. The release is attached - * to the {@link TransportResponse} via {@code afterSendRelease} and executed by the transport layer - * on write completion. + * Wraps a listener to release circuit breaker bytes from a FetchSearchResult after the response is consumed. + * Used for the query phase when a single-shard optimisation produces a {@code QueryFetchSearchResult} that is + * consumed locally (e.g. by {@code SearchQueryThenFetchAsyncAction}). For transport paths the circuit-breaker + * reservation is extracted (atomically zeroed) by {@code SearchTransportService.asBytesResponse} right after + * serialization and deferred until Netty completes the write. Because + * {@link FetchSearchResult#extractCircuitBreakerRelease} uses {@code AtomicLong.getAndSet(0)}, the release + * here is a safe no-op for transport paths — the bytes have already been taken. */ - private ActionListener releaseCircuitBreakerOnResponse( + private ActionListener releaseCircuitBreakerOnResponse( ActionListener listener, Function fetchResultExtractor ) { return ActionListener.wrap(response -> { - FetchSearchResult fetchResult = fetchResultExtractor.apply(response); - if (fetchResult != null && fetchResult.getSearchHitsSizeBytes() > 0) { - response.setAfterSendRelease(() -> fetchResult.releaseCircuitBreakerBytes(circuitBreaker)); - } try { listener.onResponse(response); } finally { - // The transport layer normally consumes afterSendRelease and releases it - // when the network write completes. If it was not consumed (e.g. an - // exception prevented the response from reaching the transport), release - // it here to avoid leaking circuit breaker bytes. - Releasable unconsumed = response.consumeAfterSendRelease(); - if (unconsumed != null) { - Releasables.closeExpectNoException(unconsumed); + // Release bytes after the response handler completes + FetchSearchResult fetchResult = fetchResultExtractor.apply(response); + if (fetchResult != null) { + fetchResult.releaseCircuitBreakerBytes(circuitBreaker); } } }, listener::onFailure); diff --git a/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java b/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java index dd23472482c27..d7d2266d31380 100644 --- a/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java +++ b/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java @@ -13,6 +13,7 @@ import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.core.RefCounted; +import org.elasticsearch.core.Releasable; import org.elasticsearch.core.SimpleRefCounted; import org.elasticsearch.search.SearchHit; import org.elasticsearch.search.SearchHits; @@ -27,6 +28,8 @@ public final class FetchSearchResult extends SearchPhaseResult { + private static final Releasable NOOP_RELEASE = () -> {}; + private SearchHits hits; private final transient AtomicLong searchHitsSizeBytes = new AtomicLong(0L); @@ -103,6 +106,23 @@ public void releaseCircuitBreakerBytes(CircuitBreaker circuitBreaker) { } } + /** + * Atomically takes the circuit-breaker reservation tracked by this result and returns a {@link Releasable} + * that will release those bytes when closed. Subsequent calls return a no-op releasable. + *

+ * This is used by the transport layer to defer the circuit-breaker release until after Netty + * completes the write, by attaching the returned releasable to the ref-counted serialized bytes. + * The returned releasable captures only the byte count and the breaker reference, not this result, + * so the result (and its search hits) can be freed independently. + */ + public Releasable extractCircuitBreakerRelease(CircuitBreaker circuitBreaker) { + long bytes = searchHitsSizeBytes.getAndSet(0L); + if (bytes > 0L) { + return () -> circuitBreaker.addWithoutBreaking(-bytes); + } + return NOOP_RELEASE; + } + public FetchSearchResult initCounter() { counter = 0; return this; diff --git a/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java b/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java index 53d2735b056c5..7b47367a50ce9 100644 --- a/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java +++ b/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java @@ -146,10 +146,6 @@ void sendResponse( ) { assert assertValidTransportVersion(transportVersion); assert response.hasReferences(); - final Releasable afterSendRelease = response.consumeAfterSendRelease(); - final Releasable onAfter = afterSendRelease != null - ? Releasables.wrap(() -> messageListener.onResponseSent(requestId, action), afterSendRelease) - : () -> messageListener.onResponseSent(requestId, action); try { sendMessage( channel, @@ -161,10 +157,9 @@ void sendResponse( compressionScheme, transportVersion, responseStatsConsumer, - onAfter + () -> messageListener.onResponseSent(requestId, action) ); } catch (Exception ex) { - Releasables.closeExpectNoException(afterSendRelease); if (isHandshake) { logger.error( () -> format( diff --git a/server/src/main/java/org/elasticsearch/transport/TransportResponse.java b/server/src/main/java/org/elasticsearch/transport/TransportResponse.java index fc1b70c5a6c58..13083521a710d 100644 --- a/server/src/main/java/org/elasticsearch/transport/TransportResponse.java +++ b/server/src/main/java/org/elasticsearch/transport/TransportResponse.java @@ -9,41 +9,9 @@ package org.elasticsearch.transport; -import org.elasticsearch.core.Releasable; - -import java.util.concurrent.atomic.AtomicReferenceFieldUpdater; - public abstract class TransportResponse extends TransportMessage { - - @SuppressWarnings("rawtypes") - private static final AtomicReferenceFieldUpdater AFTER_SEND_RELEASE_UPDATER = AtomicReferenceFieldUpdater - .newUpdater(TransportResponse.class, Releasable.class, "afterSendRelease"); - - private transient volatile Releasable afterSendRelease; - /** * Constructs a new empty transport response */ protected TransportResponse() {} - - /** - * Attaches a {@link Releasable} to this response that the transport layer will release once the - * response bytes have been fully written to the network. This is used to defer resource cleanup - * (e.g. circuit breaker release) until after the write completes. - * - * @throws IllegalStateException if a releasable has already been set and not yet consumed - */ - public void setAfterSendRelease(Releasable afterSendRelease) { - if (AFTER_SEND_RELEASE_UPDATER.compareAndSet(this, null, afterSendRelease) == false) { - throw new IllegalStateException("afterSendRelease already set"); - } - } - - /** - * Atomically returns and clears the after-send {@link Releasable}, or {@code null} if none was set. - * The transport layer calls this to take ownership of the releasable and release it on write completion. - */ - public Releasable consumeAfterSendRelease() { - return AFTER_SEND_RELEASE_UPDATER.getAndSet(this, null); - } } diff --git a/server/src/main/java/org/elasticsearch/transport/TransportService.java b/server/src/main/java/org/elasticsearch/transport/TransportService.java index 4b27f0d4c848b..60170125d65c0 100644 --- a/server/src/main/java/org/elasticsearch/transport/TransportService.java +++ b/server/src/main/java/org/elasticsearch/transport/TransportService.java @@ -1570,7 +1570,7 @@ public String getProfileName() { @Override public void sendResponse(TransportResponse response) { - final Releasable afterSendRelease = response.consumeAfterSendRelease(); + service.onResponseSent(requestId, action); try (var shutdownBlock = service.pendingDirectHandlers.withRef()) { if (shutdownBlock == null) { // already shutting down, the handler will be completed by sendRequestInternal or doStop @@ -1603,11 +1603,6 @@ public String toString() { } }); } - } finally { - service.onResponseSent(requestId, action); - if (afterSendRelease != null) { - afterSendRelease.close(); - } } } diff --git a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java new file mode 100644 index 0000000000000..ba866dc099606 --- /dev/null +++ b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java @@ -0,0 +1,275 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ +package org.elasticsearch.action.search; + +import org.apache.lucene.search.TotalHits; +import org.elasticsearch.action.ActionListener; +import org.elasticsearch.common.breaker.CircuitBreaker; +import org.elasticsearch.common.breaker.CircuitBreakingException; +import org.elasticsearch.common.breaker.NoopCircuitBreaker; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.search.SearchHits; +import org.elasticsearch.search.fetch.FetchSearchResult; +import org.elasticsearch.search.internal.ShardSearchContextId; +import org.elasticsearch.test.ESTestCase; +import org.elasticsearch.test.transport.MockTransport; +import org.elasticsearch.threadpool.TestThreadPool; +import org.elasticsearch.threadpool.ThreadPool; +import org.elasticsearch.transport.BytesTransportResponse; +import org.elasticsearch.transport.TestTransportChannel; +import org.elasticsearch.transport.TransportResponse; +import org.elasticsearch.transport.TransportService; +import org.junit.After; +import org.junit.Before; + +import java.io.IOException; +import java.util.Collections; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; + +import static org.elasticsearch.cluster.node.DiscoveryNodeUtils.builder; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.notNullValue; + +/** + * Tests for {@link SearchTransportService#asBytesResponse}, covering the network-channel path + * that pre-serializes responses into {@link BytesTransportResponse} and the associated + * {@code afterSerialize} callback. These tests exercise the actual method (not a mirror helper) + * and verify {@link BytesTransportResponse} creation, round-trip deserialization, and cleanup + * on serialization failure. + *

+ * {@link TestTransportChannel} is not a {@code DirectResponseChannel}, so the method takes the + * network (pre-serialization) path for all tests here. + */ +public class AsBytesResponseTests extends ESTestCase { + + private ThreadPool threadPool; + private TransportService transportService; + + @Override + @Before + public void setUp() throws Exception { + super.setUp(); + threadPool = new TestThreadPool(getTestName()); + var mockTransport = new MockTransport(); + transportService = mockTransport.createTransportService( + Settings.EMPTY, + threadPool, + TransportService.NOOP_TRANSPORT_INTERCEPTOR, + address -> builder("test-node").build(), + null, + Collections.emptySet() + ); + } + + @Override + @After + public void tearDown() throws Exception { + transportService.close(); + ThreadPool.terminate(threadPool, 30, TimeUnit.SECONDS); + super.tearDown(); + } + + public void testNetworkPathCallsAfterSerializeOnSuccess() { + var afterSerializeCalled = new AtomicBoolean(false); + var sentResponse = new AtomicReference(); + + var channel = new TestTransportChannel(ActionListener.wrap(resp -> { + resp.mustIncRef(); + sentResponse.set(resp); + }, e -> fail("unexpected failure: " + e))); + + ActionListener listener = SearchTransportService.asBytesResponse( + transportService, + channel, + response -> afterSerializeCalled.set(true) + ); + + listener.onResponse(new SimpleTestResponse("hello")); + + assertTrue("afterSerialize must be called after successful serialization", afterSerializeCalled.get()); + assertThat(sentResponse.get(), notNullValue()); + assertThat(sentResponse.get(), instanceOf(BytesTransportResponse.class)); + sentResponse.get().decRef(); + } + + public void testNetworkPathCallsAfterSerializeOnSerializationFailure() { + var afterSerializeCalled = new AtomicBoolean(false); + var sentException = new AtomicReference(); + + var channel = new TestTransportChannel( + ActionListener.wrap(resp -> fail("should not succeed when serialization fails"), sentException::set) + ); + + ActionListener listener = SearchTransportService.asBytesResponse( + transportService, + channel, + response -> afterSerializeCalled.set(true) + ); + + listener.onResponse(new FailingTestResponse()); + + assertTrue("afterSerialize must be called even on serialization failure", afterSerializeCalled.get()); + assertThat(sentException.get(), notNullValue()); + } + + public void testNetworkPathBytesResponseRoundTrip() throws Exception { + var sentResponse = new AtomicReference(); + + var channel = new TestTransportChannel(ActionListener.wrap(resp -> { + resp.mustIncRef(); + sentResponse.set(resp); + }, e -> fail("unexpected failure: " + e))); + + var original = new SimpleTestResponse("round-trip-test"); + ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, response -> {}); + + listener.onResponse(original); + + assertThat(sentResponse.get(), instanceOf(BytesTransportResponse.class)); + var bytesResp = (BytesTransportResponse) sentResponse.get(); + try (StreamInput in = bytesResp.streamInput()) { + var deserialized = new SimpleTestResponse(in); + assertThat(deserialized.value, equalTo("round-trip-test")); + } finally { + bytesResp.decRef(); + } + } + + public void testNetworkPathReleasesCircuitBreakerBytesAfterSerialization() { + var breakerUsed = new AtomicLong(5000); + CircuitBreaker breaker = new TestCircuitBreaker(breakerUsed); + + var sentResponse = new AtomicReference(); + + var channel = new TestTransportChannel(ActionListener.wrap(resp -> { + resp.mustIncRef(); + sentResponse.set(resp); + }, e -> fail("unexpected failure: " + e))); + + var contextId = new ShardSearchContextId("test-session", 1); + var fetchResult = new FetchSearchResult(contextId, null); + try { + var hits = SearchHits.empty(new TotalHits(0, TotalHits.Relation.EQUAL_TO), 0.0f); + fetchResult.shardResult(hits, null); + fetchResult.setSearchHitsSizeBytes(5000L); + + ActionListener listener = SearchTransportService.asBytesResponse( + transportService, + channel, + response -> response.releaseCircuitBreakerBytes(breaker) + ); + + listener.onResponse(fetchResult); + + assertThat("breaker bytes should be released after serialization", breakerUsed.get(), equalTo(0L)); + assertThat(fetchResult.getSearchHitsSizeBytes(), equalTo(0L)); + assertThat(sentResponse.get(), instanceOf(BytesTransportResponse.class)); + } finally { + fetchResult.decRef(); + if (sentResponse.get() != null) { + sentResponse.get().decRef(); + } + } + } + + public void testNetworkPathReleasesCircuitBreakerBytesOnSerializationFailure() { + var breakerUsed = new AtomicLong(3000); + CircuitBreaker breaker = new TestCircuitBreaker(breakerUsed); + + var sentException = new AtomicReference(); + + var channel = new TestTransportChannel( + ActionListener.wrap(resp -> fail("should not succeed when serialization fails"), sentException::set) + ); + + ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, response -> { + long bytes = breakerUsed.get(); + if (bytes > 0L) { + breakerUsed.addAndGet(-bytes); + } + }); + + listener.onResponse(new FailingTestResponse()); + + assertThat("breaker bytes should be released even on serialization failure", breakerUsed.get(), equalTo(0L)); + assertThat(sentException.get(), notNullValue()); + } + + public void testNetworkPathOnFailureDoesNotCallAfterSerialize() { + var afterSerializeCalled = new AtomicBoolean(false); + var sentException = new AtomicReference(); + + var channel = new TestTransportChannel(ActionListener.wrap(resp -> fail("should not succeed"), sentException::set)); + + ActionListener listener = SearchTransportService.asBytesResponse( + transportService, + channel, + response -> afterSerializeCalled.set(true) + ); + + listener.onFailure(new RuntimeException("upstream failure")); + + assertFalse("afterSerialize must not be called on upstream failure", afterSerializeCalled.get()); + assertThat(sentException.get(), notNullValue()); + } + + static class SimpleTestResponse extends TransportResponse { + final String value; + + SimpleTestResponse(String value) { + this.value = value; + } + + SimpleTestResponse(StreamInput in) throws IOException { + this.value = in.readString(); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeString(value); + } + } + + static class FailingTestResponse extends TransportResponse { + @Override + public void writeTo(StreamOutput out) throws IOException { + throw new IOException("simulated serialization failure"); + } + } + + private static class TestCircuitBreaker extends NoopCircuitBreaker { + private final AtomicLong used; + + TestCircuitBreaker(AtomicLong used) { + super("test"); + this.used = used; + } + + @Override + public void addEstimateBytesAndMaybeBreak(long bytes, String label) throws CircuitBreakingException { + used.addAndGet(bytes); + } + + @Override + public void addWithoutBreaking(long bytes) { + used.addAndGet(bytes); + } + + @Override + public long getUsed() { + return used.get(); + } + } +} diff --git a/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java b/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java index c9390538cd539..5094d568c15c3 100644 --- a/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java @@ -19,7 +19,6 @@ import org.elasticsearch.search.fetch.ScrollQueryFetchSearchResult; import org.elasticsearch.search.query.QuerySearchResult; import org.elasticsearch.test.ESTestCase; -import org.elasticsearch.transport.TransportResponse; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicLong; @@ -27,13 +26,18 @@ import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.is; -import static org.hamcrest.Matchers.notNullValue; -import static org.hamcrest.Matchers.nullValue; /** - * Unit tests for circuit breaker release logic in SearchService. - * Tests the generic helper method that releases circuit breaker bytes - * after search results are sent to the coordinator. + * Unit tests for circuit breaker release logic used by SearchService and SearchTransportService. + *

+ * For transport paths (fetch, scroll-fetch), the circuit-breaker reservation is atomically + * extracted right after serialization inside {@code SearchTransportService.asBytesResponse} and + * attached to the ref-counted serialized bytes. The actual release happens when the refcount + * drops to zero — i.e. after Netty completes the write. + *

+ * For the query phase (which may produce a {@code QueryFetchSearchResult} consumed locally by + * {@code SearchQueryThenFetchAsyncAction}), the circuit breaker is released after the listener + * consumes the response via {@code SearchService.releaseCircuitBreakerOnResponse}. */ public class SearchServiceCircuitBreakerTests extends ESTestCase { @@ -54,9 +58,6 @@ public void testReleaseCircuitBreakerForFetchResult() { assertThat(failureCalled.get(), is(false)); assertThat(breakerUsed.get(), equalTo(0L)); assertThat(result.getSearchHitsSizeBytes(), equalTo(0L)); - - Releasable afterSend = result.consumeAfterSendRelease(); - assertThat(afterSend, nullValue()); } finally { result.decRef(); } @@ -82,9 +83,6 @@ public void testReleaseCircuitBreakerForQueryFetchResult() { assertThat(failureCalled.get(), is(false)); assertThat(breakerUsed.get(), equalTo(0L)); assertThat(fetchResult.getSearchHitsSizeBytes(), equalTo(0L)); - - Releasable afterSend = queryFetchResult.consumeAfterSendRelease(); - assertThat(afterSend, nullValue()); } finally { if (queryFetchResult != null) { queryFetchResult.decRef(); @@ -115,9 +113,6 @@ public void testReleaseCircuitBreakerForScrollResult() { assertThat(failureCalled.get(), is(false)); assertThat(breakerUsed.get(), equalTo(0L)); assertThat(fetchResult.getSearchHitsSizeBytes(), equalTo(0L)); - - Releasable afterSend = scrollResult.consumeAfterSendRelease(); - assertThat(afterSend, nullValue()); } finally { if (scrollResult != null) { scrollResult.decRef(); @@ -153,48 +148,48 @@ public void testExtractorReturnsNull() { assertThat(successCalled.get(), is(true)); assertThat(failureCalled.get(), is(false)); - assertThat(result.consumeAfterSendRelease(), nullValue()); } - public void testTransportConsumesAfterSendRelease() { - AtomicLong breakerUsed = new AtomicLong(5000); + public void testMultipleReleasesAreIdempotent() { + AtomicLong breakerUsed = new AtomicLong(2000); CircuitBreaker breaker = new TestCircuitBreaker(breakerUsed); FetchSearchResult result = new FetchSearchResult(); try { - result.setSearchHitsSizeBytes(5000L); - AtomicBoolean successCalled = new AtomicBoolean(false); - - Releasable[] captured = new Releasable[1]; - ActionListener transportSimulator = new ActionListener<>() { - @Override - public void onResponse(FetchSearchResult response) { - captured[0] = response.consumeAfterSendRelease(); - successCalled.set(true); - } + result.setSearchHitsSizeBytes(2000L); - @Override - public void onFailure(Exception e) { - fail("unexpected failure"); - } - }; + result.releaseCircuitBreakerBytes(breaker); + assertThat(breakerUsed.get(), equalTo(0L)); + assertThat(result.getSearchHitsSizeBytes(), equalTo(0L)); - withCircuitBreakerRelease(transportSimulator, breaker, Function.identity()).onResponse(result); + result.releaseCircuitBreakerBytes(breaker); + assertThat(breakerUsed.get(), equalTo(0L)); + assertThat(result.getSearchHitsSizeBytes(), equalTo(0L)); + } finally { + result.decRef(); + } + } - assertThat(successCalled.get(), is(true)); - assertThat(breakerUsed.get(), equalTo(5000L)); - assertThat(result.getSearchHitsSizeBytes(), equalTo(5000L)); - assertThat(captured[0], notNullValue()); + public void testExtractCircuitBreakerRelease() { + AtomicLong breakerUsed = new AtomicLong(3000); + CircuitBreaker breaker = new TestCircuitBreaker(breakerUsed); - captured[0].close(); - assertThat(breakerUsed.get(), equalTo(0L)); + FetchSearchResult result = new FetchSearchResult(); + try { + result.setSearchHitsSizeBytes(3000L); + + Releasable release = result.extractCircuitBreakerRelease(breaker); assertThat(result.getSearchHitsSizeBytes(), equalTo(0L)); + assertThat(breakerUsed.get(), equalTo(3000L)); + + Releasables.close(release); + assertThat(breakerUsed.get(), equalTo(0L)); } finally { result.decRef(); } } - public void testMultipleReleasesAreIdempotent() { + public void testExtractCircuitBreakerReleaseIsIdempotent() { AtomicLong breakerUsed = new AtomicLong(2000); CircuitBreaker breaker = new TestCircuitBreaker(breakerUsed); @@ -202,15 +197,36 @@ public void testMultipleReleasesAreIdempotent() { try { result.setSearchHitsSizeBytes(2000L); - // First release - result.releaseCircuitBreakerBytes(breaker); - assertThat(breakerUsed.get(), equalTo(0L)); + Releasable first = result.extractCircuitBreakerRelease(breaker); + Releasable second = result.extractCircuitBreakerRelease(breaker); + assertThat(result.getSearchHitsSizeBytes(), equalTo(0L)); + assertThat(breakerUsed.get(), equalTo(2000L)); - // Next release - should be no-op + Releasables.close(first); + assertThat(breakerUsed.get(), equalTo(0L)); + + Releasables.close(second); + assertThat(breakerUsed.get(), equalTo(0L)); + } finally { + result.decRef(); + } + } + + public void testExtractAndDirectReleaseAreIdempotent() { + AtomicLong breakerUsed = new AtomicLong(4000); + CircuitBreaker breaker = new TestCircuitBreaker(breakerUsed); + + FetchSearchResult result = new FetchSearchResult(); + try { + result.setSearchHitsSizeBytes(4000L); + + Releasable release = result.extractCircuitBreakerRelease(breaker); result.releaseCircuitBreakerBytes(breaker); + assertThat(breakerUsed.get(), equalTo(4000L)); + + Releasables.close(release); assertThat(breakerUsed.get(), equalTo(0L)); - assertThat(result.getSearchHitsSizeBytes(), equalTo(0L)); } finally { result.decRef(); } @@ -233,9 +249,6 @@ public void testLargeAllocation() { assertThat(successCalled.get(), is(true)); assertThat(breakerUsed.get(), equalTo(0L)); assertThat(result.getSearchHitsSizeBytes(), equalTo(0L)); - - Releasable afterSend = result.consumeAfterSendRelease(); - assertThat(afterSend, nullValue()); } finally { result.decRef(); } @@ -287,35 +300,54 @@ public void onFailure(Exception e) { } /** - * Wrap a listener with deferred circuit breaker release via afterSendRelease on the response. + * Wrap a listener with circuit breaker release after {@code onResponse}, mirroring + * {@code SearchService.releaseCircuitBreakerOnResponse} (used for the query phase local path). */ - private ActionListener withCircuitBreakerRelease( + private ActionListener withCircuitBreakerReleaseAfterResponse( ActionListener listener, CircuitBreaker breaker, Function fetchResultExtractor ) { return ActionListener.wrap(response -> { - FetchSearchResult fetchResult = fetchResultExtractor.apply(response); - if (fetchResult != null && fetchResult.getSearchHitsSizeBytes() > 0) { - response.setAfterSendRelease(() -> fetchResult.releaseCircuitBreakerBytes(breaker)); - } try { listener.onResponse(response); } finally { - Releasable unconsumed = response.consumeAfterSendRelease(); - if (unconsumed != null) { - Releasables.closeExpectNoException(unconsumed); + FetchSearchResult fetchResult = fetchResultExtractor.apply(response); + if (fetchResult != null) { + fetchResult.releaseCircuitBreakerBytes(breaker); } } }, listener::onFailure); } + /** + * Wrap a listener with deferred circuit breaker release, mirroring the {@code asBytesResponse} + * callback in {@code SearchTransportService} (used for transport fetch paths). The reservation + * is atomically extracted before {@code onResponse} and released afterwards, simulating the + * Netty write-completion callback that would close the deferred releasable in production. + */ + private ActionListener withDeferredCircuitBreakerRelease( + ActionListener listener, + CircuitBreaker breaker, + Function fetchResultExtractor + ) { + return ActionListener.wrap(response -> { + FetchSearchResult fetchResult = fetchResultExtractor.apply(response); + Releasable cbRelease = (fetchResult != null) ? fetchResult.extractCircuitBreakerRelease(breaker) : () -> {}; + try { + listener.onResponse(response); + } finally { + Releasables.close(cbRelease); + } + }, listener::onFailure); + } + private ActionListener querySearchResultListener( AtomicBoolean successCalled, AtomicBoolean failureCalled, CircuitBreaker breaker ) { - return withCircuitBreakerRelease(trackingListener(successCalled, failureCalled), breaker, qr -> null); + return withCircuitBreakerReleaseAfterResponse(trackingListener(successCalled, failureCalled), breaker, qr -> null); } private ActionListener fetchSearchResultListener( @@ -323,7 +355,7 @@ private ActionListener fetchSearchResultListener( AtomicBoolean failureCalled, CircuitBreaker breaker ) { - return withCircuitBreakerRelease(trackingListener(successCalled, failureCalled), breaker, Function.identity()); + return withDeferredCircuitBreakerRelease(trackingListener(successCalled, failureCalled), breaker, Function.identity()); } private ActionListener queryFetchSearchResultListener( @@ -331,7 +363,11 @@ private ActionListener queryFetchSearchResultListener( AtomicBoolean failureCalled, CircuitBreaker breaker ) { - return withCircuitBreakerRelease(trackingListener(successCalled, failureCalled), breaker, QueryFetchSearchResult::fetchResult); + return withCircuitBreakerReleaseAfterResponse( + trackingListener(successCalled, failureCalled), + breaker, + QueryFetchSearchResult::fetchResult + ); } private ActionListener scrollQueryFetchSearchResultListener( @@ -339,7 +375,7 @@ private ActionListener scrollQueryFetchSearchResul AtomicBoolean failureCalled, CircuitBreaker breaker ) { - return withCircuitBreakerRelease(trackingListener(successCalled, failureCalled), breaker, sr -> sr.result().fetchResult()); + return withDeferredCircuitBreakerRelease(trackingListener(successCalled, failureCalled), breaker, sr -> sr.result().fetchResult()); } /** diff --git a/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java b/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java index ff84648f38494..32cf9237e030a 100644 --- a/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java +++ b/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java @@ -272,105 +272,6 @@ public void onResponseSent(long requestId, String action) { assertEquals("header_value", header.getHeaders().v1().get("header")); } - public void testAfterSendReleaseOnWriteSuccess() { - TransportVersion version = TransportVersionUtils.randomCompatibleVersion(); - String action = randomAlphaOfLength(10); - long requestId = randomLongBetween(0, 300); - TestResponse response = new TestResponse(randomAlphaOfLength(10)); - - AtomicBoolean released = new AtomicBoolean(false); - response.setAfterSendRelease(() -> assertTrue(released.compareAndSet(false, true))); - - AtomicBoolean onResponseSentCalled = new AtomicBoolean(false); - handler.setMessageListener(new TransportMessageListener() { - @Override - public void onResponseSent(long requestId, String action) { - onResponseSentCalled.set(true); - } - }); - - Compression.Scheme compress = randomFrom(compressionScheme, null); - handler.sendResponse(version, channel, requestId, action, response, compress, false, ResponseStatsConsumer.NONE); - - assertFalse("afterSendRelease must not be released before write completes", released.get()); - assertFalse("onResponseSent must not fire before write completes", onResponseSentCalled.get()); - assertNull("afterSendRelease should have been consumed from the response", response.consumeAfterSendRelease()); - - ActionListener sendListener = channel.getListenerCaptor().get(); - sendListener.onResponse(null); - - assertTrue("afterSendRelease must be released on write success", released.get()); - assertTrue("onResponseSent must fire on write success", onResponseSentCalled.get()); - } - - public void testAfterSendReleaseOnWriteFailure() { - TransportVersion version = TransportVersionUtils.randomCompatibleVersion(); - String action = randomAlphaOfLength(10); - long requestId = randomLongBetween(0, 300); - TestResponse response = new TestResponse(randomAlphaOfLength(10)); - - AtomicBoolean released = new AtomicBoolean(false); - response.setAfterSendRelease(() -> assertTrue(released.compareAndSet(false, true))); - - AtomicBoolean onResponseSentCalled = new AtomicBoolean(false); - handler.setMessageListener(new TransportMessageListener() { - @Override - public void onResponseSent(long requestId, String action) { - onResponseSentCalled.set(true); - } - }); - - Compression.Scheme compress = randomFrom(compressionScheme, null); - handler.sendResponse(version, channel, requestId, action, response, compress, false, ResponseStatsConsumer.NONE); - - assertFalse("afterSendRelease must not be released before write completes", released.get()); - assertFalse("onResponseSent must not fire before write completes", onResponseSentCalled.get()); - - ActionListener sendListener = channel.getListenerCaptor().get(); - sendListener.onFailure(new IOException("write failed")); - - assertTrue("afterSendRelease must be released on write failure", released.get()); - assertTrue("onResponseSent must fire on write failure", onResponseSentCalled.get()); - } - - public void testAfterSendReleaseOnSerializationFailure() { - TransportVersion version = TransportVersionUtils.randomCompatibleVersion(); - String action = randomAlphaOfLength(10); - long requestId = randomLongBetween(0, 300); - - var response = new ReleasbleTestResponse(randomAlphaOfLength(10)) { - @Override - public void writeTo(StreamOutput out) { - throw new CircuitBreakingException("simulated cbe", CircuitBreaker.Durability.TRANSIENT); - } - }; - - AtomicBoolean released = new AtomicBoolean(false); - response.setAfterSendRelease(() -> assertTrue(released.compareAndSet(false, true))); - - handler.setMessageListener(new TransportMessageListener() { - @Override - public void onResponseSent(long requestId, String action) { - throw new AssertionError("onResponseSent(success) must not be called on serialization failure"); - } - - @Override - public void onResponseSent(long requestId, String action, Exception error) { - // expected — error response is sent instead - } - }); - - Compression.Scheme compress = randomFrom(compressionScheme, null); - try { - handler.sendResponse(version, channel, requestId, action, response, compress, false, ResponseStatsConsumer.NONE); - } finally { - response.decRef(); - } - - assertTrue("afterSendRelease must be released on serialization failure", released.get()); - assertTrue("response resources must be released on serialization failure", response.released.get()); - } - public void testErrorResponse() throws IOException { ThreadContext threadContext = threadPool.getThreadContext(); TransportVersion version = TransportVersionUtils.randomCompatibleVersion(); From a26a1ec5c51f6f1766982c1208469331c0265db0 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Thu, 5 Mar 2026 17:08:31 +0200 Subject: [PATCH 06/39] update after review --- .../main/java/org/elasticsearch/transport/OutboundHandler.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java b/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java index 7b47367a50ce9..5427209b30616 100644 --- a/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java +++ b/server/src/main/java/org/elasticsearch/transport/OutboundHandler.java @@ -253,7 +253,7 @@ private void sendMessage( throw e; } finally { if (serializeSuccess == false) { - Releasables.close(byteStreamOutput); + Releasables.close(byteStreamOutput, onAfter); } } responseStatsConsumer.addResponseStats(message.length()); From 9487e49f988770023afcd48c772aa679ccae6161 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Thu, 5 Mar 2026 17:21:03 +0200 Subject: [PATCH 07/39] update after review --- .../elasticsearch/search/SearchService.java | 7 +- .../search/fetch/FetchSearchResult.java | 20 ---- .../SearchServiceCircuitBreakerTests.java | 94 ++----------------- 3 files changed, 9 insertions(+), 112 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/search/SearchService.java b/server/src/main/java/org/elasticsearch/search/SearchService.java index 3cc1f9f5fbf35..5652b0acb6b9f 100644 --- a/server/src/main/java/org/elasticsearch/search/SearchService.java +++ b/server/src/main/java/org/elasticsearch/search/SearchService.java @@ -1641,10 +1641,9 @@ private static void checkKeepAliveLimit(long keepAlive, long maxKeepAlive) { * Wraps a listener to release circuit breaker bytes from a FetchSearchResult after the response is consumed. * Used for the query phase when a single-shard optimisation produces a {@code QueryFetchSearchResult} that is * consumed locally (e.g. by {@code SearchQueryThenFetchAsyncAction}). For transport paths the circuit-breaker - * reservation is extracted (atomically zeroed) by {@code SearchTransportService.asBytesResponse} right after - * serialization and deferred until Netty completes the write. Because - * {@link FetchSearchResult#extractCircuitBreakerRelease} uses {@code AtomicLong.getAndSet(0)}, the release - * here is a safe no-op for transport paths — the bytes have already been taken. + * reservation is released by {@code SearchTransportService.asBytesResponse} right after serialization. + * Because {@link FetchSearchResult#releaseCircuitBreakerBytes} uses {@code AtomicLong.getAndSet(0)}, the + * release here is a safe no-op for transport paths — the bytes have already been released. */ private ActionListener releaseCircuitBreakerOnResponse( ActionListener listener, diff --git a/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java b/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java index d7d2266d31380..dd23472482c27 100644 --- a/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java +++ b/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java @@ -13,7 +13,6 @@ import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.core.RefCounted; -import org.elasticsearch.core.Releasable; import org.elasticsearch.core.SimpleRefCounted; import org.elasticsearch.search.SearchHit; import org.elasticsearch.search.SearchHits; @@ -28,8 +27,6 @@ public final class FetchSearchResult extends SearchPhaseResult { - private static final Releasable NOOP_RELEASE = () -> {}; - private SearchHits hits; private final transient AtomicLong searchHitsSizeBytes = new AtomicLong(0L); @@ -106,23 +103,6 @@ public void releaseCircuitBreakerBytes(CircuitBreaker circuitBreaker) { } } - /** - * Atomically takes the circuit-breaker reservation tracked by this result and returns a {@link Releasable} - * that will release those bytes when closed. Subsequent calls return a no-op releasable. - *

- * This is used by the transport layer to defer the circuit-breaker release until after Netty - * completes the write, by attaching the returned releasable to the ref-counted serialized bytes. - * The returned releasable captures only the byte count and the breaker reference, not this result, - * so the result (and its search hits) can be freed independently. - */ - public Releasable extractCircuitBreakerRelease(CircuitBreaker circuitBreaker) { - long bytes = searchHitsSizeBytes.getAndSet(0L); - if (bytes > 0L) { - return () -> circuitBreaker.addWithoutBreaking(-bytes); - } - return NOOP_RELEASE; - } - public FetchSearchResult initCounter() { counter = 0; return this; diff --git a/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java b/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java index 5094d568c15c3..a303230bf6c25 100644 --- a/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java @@ -12,8 +12,6 @@ import org.elasticsearch.common.breaker.CircuitBreaker; import org.elasticsearch.common.breaker.CircuitBreakingException; import org.elasticsearch.common.breaker.NoopCircuitBreaker; -import org.elasticsearch.core.Releasable; -import org.elasticsearch.core.Releasables; import org.elasticsearch.search.fetch.FetchSearchResult; import org.elasticsearch.search.fetch.QueryFetchSearchResult; import org.elasticsearch.search.fetch.ScrollQueryFetchSearchResult; @@ -170,68 +168,6 @@ public void testMultipleReleasesAreIdempotent() { } } - public void testExtractCircuitBreakerRelease() { - AtomicLong breakerUsed = new AtomicLong(3000); - CircuitBreaker breaker = new TestCircuitBreaker(breakerUsed); - - FetchSearchResult result = new FetchSearchResult(); - try { - result.setSearchHitsSizeBytes(3000L); - - Releasable release = result.extractCircuitBreakerRelease(breaker); - assertThat(result.getSearchHitsSizeBytes(), equalTo(0L)); - assertThat(breakerUsed.get(), equalTo(3000L)); - - Releasables.close(release); - assertThat(breakerUsed.get(), equalTo(0L)); - } finally { - result.decRef(); - } - } - - public void testExtractCircuitBreakerReleaseIsIdempotent() { - AtomicLong breakerUsed = new AtomicLong(2000); - CircuitBreaker breaker = new TestCircuitBreaker(breakerUsed); - - FetchSearchResult result = new FetchSearchResult(); - try { - result.setSearchHitsSizeBytes(2000L); - - Releasable first = result.extractCircuitBreakerRelease(breaker); - Releasable second = result.extractCircuitBreakerRelease(breaker); - - assertThat(result.getSearchHitsSizeBytes(), equalTo(0L)); - assertThat(breakerUsed.get(), equalTo(2000L)); - - Releasables.close(first); - assertThat(breakerUsed.get(), equalTo(0L)); - - Releasables.close(second); - assertThat(breakerUsed.get(), equalTo(0L)); - } finally { - result.decRef(); - } - } - - public void testExtractAndDirectReleaseAreIdempotent() { - AtomicLong breakerUsed = new AtomicLong(4000); - CircuitBreaker breaker = new TestCircuitBreaker(breakerUsed); - - FetchSearchResult result = new FetchSearchResult(); - try { - result.setSearchHitsSizeBytes(4000L); - - Releasable release = result.extractCircuitBreakerRelease(breaker); - result.releaseCircuitBreakerBytes(breaker); - assertThat(breakerUsed.get(), equalTo(4000L)); - - Releasables.close(release); - assertThat(breakerUsed.get(), equalTo(0L)); - } finally { - result.decRef(); - } - } - public void testLargeAllocation() { long largeBytes = randomLongBetween(1_000_000, 10_000_000); AtomicLong breakerUsed = new AtomicLong(largeBytes); @@ -320,28 +256,6 @@ private ActionListener withCircuitBreakerReleaseAfterResponse( }, listener::onFailure); } - /** - * Wrap a listener with deferred circuit breaker release, mirroring the {@code asBytesResponse} - * callback in {@code SearchTransportService} (used for transport fetch paths). The reservation - * is atomically extracted before {@code onResponse} and released afterwards, simulating the - * Netty write-completion callback that would close the deferred releasable in production. - */ - private ActionListener withDeferredCircuitBreakerRelease( - ActionListener listener, - CircuitBreaker breaker, - Function fetchResultExtractor - ) { - return ActionListener.wrap(response -> { - FetchSearchResult fetchResult = fetchResultExtractor.apply(response); - Releasable cbRelease = (fetchResult != null) ? fetchResult.extractCircuitBreakerRelease(breaker) : () -> {}; - try { - listener.onResponse(response); - } finally { - Releasables.close(cbRelease); - } - }, listener::onFailure); - } - private ActionListener querySearchResultListener( AtomicBoolean successCalled, AtomicBoolean failureCalled, @@ -355,7 +269,7 @@ private ActionListener fetchSearchResultListener( AtomicBoolean failureCalled, CircuitBreaker breaker ) { - return withDeferredCircuitBreakerRelease(trackingListener(successCalled, failureCalled), breaker, Function.identity()); + return withCircuitBreakerReleaseAfterResponse(trackingListener(successCalled, failureCalled), breaker, Function.identity()); } private ActionListener queryFetchSearchResultListener( @@ -375,7 +289,11 @@ private ActionListener scrollQueryFetchSearchResul AtomicBoolean failureCalled, CircuitBreaker breaker ) { - return withDeferredCircuitBreakerRelease(trackingListener(successCalled, failureCalled), breaker, sr -> sr.result().fetchResult()); + return withCircuitBreakerReleaseAfterResponse( + trackingListener(successCalled, failureCalled), + breaker, + sr -> sr.result().fetchResult() + ); } /** From 69d1e54d573082aeade272bf579f5bb000b5b831 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Thu, 5 Mar 2026 17:28:56 +0200 Subject: [PATCH 08/39] update after review --- .../action/search/SearchTransportService.java | 6 +- .../action/search/AsBytesResponseTests.java | 98 +++++++++++++++++-- .../SearchServiceCircuitBreakerTests.java | 12 +-- .../transport/TestDirectResponseChannel.java | 37 +++++++ 4 files changed, 135 insertions(+), 18 deletions(-) create mode 100644 test/framework/src/main/java/org/elasticsearch/transport/TestDirectResponseChannel.java diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java index 62a4fae651ef8..d9c611a90b7e8 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java +++ b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java @@ -463,11 +463,7 @@ public static void registerRequestHandler( (request, channel, task) -> searchService.executeQueryPhase( request, (SearchShardTask) task, - asBytesResponse(transportService, channel, response -> { - if (response instanceof QueryFetchSearchResult qfr) { - qfr.fetchResult().releaseCircuitBreakerBytes(searchService.getCircuitBreaker()); - } - }) + new ChannelActionListener<>(channel) ) ); TransportActionProxy.registerProxyActionWithDynamicResponseType( diff --git a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java index ba866dc099606..b35c15826db8e 100644 --- a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java @@ -24,6 +24,7 @@ import org.elasticsearch.threadpool.TestThreadPool; import org.elasticsearch.threadpool.ThreadPool; import org.elasticsearch.transport.BytesTransportResponse; +import org.elasticsearch.transport.TestDirectResponseChannel; import org.elasticsearch.transport.TestTransportChannel; import org.elasticsearch.transport.TransportResponse; import org.elasticsearch.transport.TransportService; @@ -43,14 +44,11 @@ import static org.hamcrest.Matchers.notNullValue; /** - * Tests for {@link SearchTransportService#asBytesResponse}, covering the network-channel path - * that pre-serializes responses into {@link BytesTransportResponse} and the associated - * {@code afterSerialize} callback. These tests exercise the actual method (not a mirror helper) - * and verify {@link BytesTransportResponse} creation, round-trip deserialization, and cleanup + * Tests for {@link SearchTransportService#asBytesResponse}, covering both the network-channel path + * (pre-serializes into {@link BytesTransportResponse}) and the direct-channel path (forwards the + * original response as-is for local/same-node requests). Tests verify {@link BytesTransportResponse} + * creation, round-trip deserialization, {@code afterSerialize} callback semantics, and cleanup * on serialization failure. - *

- * {@link TestTransportChannel} is not a {@code DirectResponseChannel}, so the method takes the - * network (pre-serialization) path for all tests here. */ public class AsBytesResponseTests extends ESTestCase { @@ -225,6 +223,92 @@ public void testNetworkPathOnFailureDoesNotCallAfterSerialize() { assertThat(sentException.get(), notNullValue()); } + public void testDirectPathForwardsOriginalResponse() { + var afterSerializeCalled = new AtomicBoolean(false); + var sentResponse = new AtomicReference(); + + var channel = new TestDirectResponseChannel(ActionListener.wrap(sentResponse::set, e -> fail("unexpected failure: " + e))); + + ActionListener listener = SearchTransportService.asBytesResponse( + transportService, + channel, + response -> afterSerializeCalled.set(true) + ); + + var original = new SimpleTestResponse("direct-test"); + listener.onResponse(original); + + assertTrue("afterSerialize must be called on direct path", afterSerializeCalled.get()); + assertSame("direct path must forward the original response, not BytesTransportResponse", original, sentResponse.get()); + } + + public void testDirectPathCallsAfterSerializeAfterOnResponse() { + var callOrder = new java.util.ArrayList(); + + var channel = new TestDirectResponseChannel( + ActionListener.wrap(resp -> callOrder.add("onResponse"), e -> fail("unexpected failure: " + e)) + ); + + ActionListener listener = SearchTransportService.asBytesResponse( + transportService, + channel, + response -> callOrder.add("afterSerialize") + ); + + listener.onResponse(new SimpleTestResponse("order-test")); + + assertThat(callOrder, equalTo(java.util.List.of("onResponse", "afterSerialize"))); + } + + public void testDirectPathOnFailureDoesNotCallAfterSerialize() { + var afterSerializeCalled = new AtomicBoolean(false); + var sentException = new AtomicReference(); + + var channel = new TestDirectResponseChannel(ActionListener.wrap(resp -> fail("should not succeed"), sentException::set)); + + ActionListener listener = SearchTransportService.asBytesResponse( + transportService, + channel, + response -> afterSerializeCalled.set(true) + ); + + listener.onFailure(new RuntimeException("upstream failure")); + + assertFalse("afterSerialize must not be called on upstream failure", afterSerializeCalled.get()); + assertThat(sentException.get(), notNullValue()); + } + + public void testDirectPathReleasesCircuitBreakerBytesAfterResponse() { + var breakerUsed = new AtomicLong(5000); + CircuitBreaker breaker = new TestCircuitBreaker(breakerUsed); + + var sentResponse = new AtomicReference(); + + var channel = new TestDirectResponseChannel(ActionListener.wrap(sentResponse::set, e -> fail("unexpected failure: " + e))); + + var contextId = new ShardSearchContextId("test-session", 1); + var fetchResult = new FetchSearchResult(contextId, null); + try { + var hits = SearchHits.empty(new TotalHits(0, TotalHits.Relation.EQUAL_TO), 0.0f); + fetchResult.shardResult(hits, null); + fetchResult.setSearchHitsSizeBytes(5000L); + + ActionListener listener = SearchTransportService.asBytesResponse( + transportService, + channel, + response -> response.releaseCircuitBreakerBytes(breaker) + ); + + listener.onResponse(fetchResult); + + assertThat("breaker bytes should be released after response on direct path", breakerUsed.get(), equalTo(0L)); + assertThat(fetchResult.getSearchHitsSizeBytes(), equalTo(0L)); + assertSame("direct path must forward original response", fetchResult, sentResponse.get()); + } finally { + fetchResult.decRef(); + } + } + static class SimpleTestResponse extends TransportResponse { final String value; diff --git a/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java b/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java index a303230bf6c25..1b06628472243 100644 --- a/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java @@ -28,14 +28,14 @@ /** * Unit tests for circuit breaker release logic used by SearchService and SearchTransportService. *

- * For transport paths (fetch, scroll-fetch), the circuit-breaker reservation is atomically - * extracted right after serialization inside {@code SearchTransportService.asBytesResponse} and - * attached to the ref-counted serialized bytes. The actual release happens when the refcount - * drops to zero — i.e. after Netty completes the write. - *

- * For the query phase (which may produce a {@code QueryFetchSearchResult} consumed locally by + * For transport paths (fetch, scroll-fetch), the circuit-breaker reservation is released right + * after serialization inside {@code SearchTransportService.asBytesResponse}. For the query phase + * (which may produce a {@code QueryFetchSearchResult} consumed locally by * {@code SearchQueryThenFetchAsyncAction}), the circuit breaker is released after the listener * consumes the response via {@code SearchService.releaseCircuitBreakerOnResponse}. + *

+ * Because {@code FetchSearchResult.releaseCircuitBreakerBytes} uses {@code AtomicLong.getAndSet(0)}, + * multiple releases from different paths are idempotent. */ public class SearchServiceCircuitBreakerTests extends ESTestCase { diff --git a/test/framework/src/main/java/org/elasticsearch/transport/TestDirectResponseChannel.java b/test/framework/src/main/java/org/elasticsearch/transport/TestDirectResponseChannel.java new file mode 100644 index 0000000000000..076f6023dc0b7 --- /dev/null +++ b/test/framework/src/main/java/org/elasticsearch/transport/TestDirectResponseChannel.java @@ -0,0 +1,37 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.transport; + +import org.elasticsearch.action.ActionListener; + +/** + * A test-only {@link TransportService.DirectResponseChannel} that delegates to an {@link ActionListener} + * instead of interacting with a real {@link TransportService}. This allows testing code paths that branch + * on {@link TransportService#isDirectResponseChannel}. + */ +public class TestDirectResponseChannel extends TransportService.DirectResponseChannel { + + private final ActionListener listener; + + public TestDirectResponseChannel(ActionListener listener) { + super(null, "test", 0, null); + this.listener = listener; + } + + @Override + public void sendResponse(TransportResponse response) { + listener.onResponse(response); + } + + @Override + public void sendResponse(Exception exception) { + listener.onFailure(exception); + } +} From 1ae96d2db9d1d2a8bdfb22b80c4ac6fe2c8d45b2 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Thu, 5 Mar 2026 17:53:58 +0200 Subject: [PATCH 09/39] update after review --- .../action/search/SearchTransportService.java | 31 +++++++++---- .../action/search/AsBytesResponseTests.java | 44 +++++++++++++++++++ .../transport/OutboundHandlerTests.java | 34 +++++++++++--- .../transport/TestTransportChannels.java | 5 +++ 4 files changed, 100 insertions(+), 14 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java index d9c611a90b7e8..90486ee2db1f8 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java +++ b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java @@ -695,14 +695,21 @@ private boolean assertNodePresent() { * response over the channel. For network channels the response is pre-serialized into a * {@link BytesTransportResponse} so that the original in-memory data can be freed by the caller's * {@code respondAndRelease}. For local (direct) channels the original response is forwarded as-is, - * because TransportService.DirectResponseChannel} passes the object directly to the response + * because {@code TransportService.DirectResponseChannel} passes the object directly to the response * handler without serialization/deserialization. *

* The caller is responsible for calling {@code response.decRef()} (typically via * {@link ActionListener#respondAndRelease}); this method never calls it. + *

+ * Direct-channel timing note: on the direct (local) path, {@code afterSerialize} runs after + * {@code channel.sendResponse()} returns. If the response handler uses a non-direct executor, + * {@code sendResponse()} may return before the handler fully processes the response, so + * {@code afterSerialize} can execute concurrently with the handler. This is safe because + * {@code afterSerialize} only releases circuit-breaker accounting (not the response data + * itself, which is ref-counted). * - * @param afterSerialize called after serialization (network) or after the handler processes the response (local) - * to release resources such as circuit-breaker reservations + * @param afterSerialize called after serialization (network) or after {@code channel.sendResponse()} returns + * (local) to release resources such as circuit-breaker reservations; must not throw */ static ActionListener asBytesResponse( TransportService transportService, @@ -739,15 +746,23 @@ public void onResponse(T response) { response.writeTo(out); } catch (Exception e) { Releasables.close(out); + try { + afterSerialize.accept(response); + } catch (Exception suppressed) { + e.addSuppressed(suppressed); + } + channelListener.onFailure(e); + return; + } + var bytesRef = out.moveToBytesReference(); + try { afterSerialize.accept(response); + } catch (Exception e) { + Releasables.close(bytesRef); channelListener.onFailure(e); return; } - afterSerialize.accept(response); - ActionListener.respondAndRelease( - channelListener, - new BytesTransportResponse(out.moveToBytesReference(), out.getTransportVersion()) - ); + ActionListener.respondAndRelease(channelListener, new BytesTransportResponse(bytesRef, out.getTransportVersion())); } @Override diff --git a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java index b35c15826db8e..914f699f5a78a 100644 --- a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java @@ -26,6 +26,7 @@ import org.elasticsearch.transport.BytesTransportResponse; import org.elasticsearch.transport.TestDirectResponseChannel; import org.elasticsearch.transport.TestTransportChannel; +import org.elasticsearch.transport.TestTransportChannels; import org.elasticsearch.transport.TransportResponse; import org.elasticsearch.transport.TransportService; import org.junit.After; @@ -309,6 +310,49 @@ public void testDirectPathReleasesCircuitBreakerBytesAfterResponse() { } } + public void testTaskTransportChannelUnwrapsToDirectPath() { + var afterSerializeCalled = new AtomicBoolean(false); + var sentResponse = new AtomicReference(); + + var directChannel = new TestDirectResponseChannel(ActionListener.wrap(sentResponse::set, e -> fail("unexpected failure: " + e))); + var taskChannel = TestTransportChannels.newTaskTransportChannel(directChannel, () -> {}); + + ActionListener listener = SearchTransportService.asBytesResponse( + transportService, + taskChannel, + response -> afterSerializeCalled.set(true) + ); + + var original = new SimpleTestResponse("task-wrapped-test"); + listener.onResponse(original); + + assertTrue("afterSerialize must be called on task-wrapped direct path", afterSerializeCalled.get()); + assertSame("task-wrapped direct channel must forward original response, not BytesTransportResponse", original, sentResponse.get()); + } + + public void testTaskTransportChannelUnwrapsToNetworkPath() { + var afterSerializeCalled = new AtomicBoolean(false); + var sentResponse = new AtomicReference(); + + var networkChannel = new TestTransportChannel(ActionListener.wrap(resp -> { + resp.mustIncRef(); + sentResponse.set(resp); + }, e -> fail("unexpected failure: " + e))); + var taskChannel = TestTransportChannels.newTaskTransportChannel(networkChannel, () -> {}); + + ActionListener listener = SearchTransportService.asBytesResponse( + transportService, + taskChannel, + response -> afterSerializeCalled.set(true) + ); + + listener.onResponse(new SimpleTestResponse("task-network-test")); + + assertTrue("afterSerialize must be called on task-wrapped network path", afterSerializeCalled.get()); + assertThat(sentResponse.get(), instanceOf(BytesTransportResponse.class)); + sentResponse.get().decRef(); + } + static class SimpleTestResponse extends TransportResponse { final String value; diff --git a/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java b/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java index 32cf9237e030a..06346ab95e290 100644 --- a/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java +++ b/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java @@ -340,6 +340,9 @@ public void writeTo(StreamOutput out) { AtomicLong requestIdRef = new AtomicLong(); AtomicReference actionRef = new AtomicReference<>(); AtomicReference exceptionRef = new AtomicReference<>(); + // When response serialization fails, OutboundHandler.sendMessage's finally block calls + // Releasables.close(byteStreamOutput, onAfter) which invokes onResponseSent(success) as a + // cleanup side-effect. Then sendErrorResponse sends the error and invokes onResponseSent(error). handler.setMessageListener(new TransportMessageListener() { @Override public void onResponseSent(long requestId, String action) { @@ -353,8 +356,8 @@ public void onResponseSent(long requestId, String action) { @Override public void onResponseSent(long requestId, String action, Exception error) { assertNotNull(channel.getMessageCaptor().get()); - requestIdRef.set(requestId); - actionRef.set(action); + assertThat(requestIdRef.get(), equalTo(requestId)); + assertThat(actionRef.get(), equalTo(action)); exceptionRef.set(error); } }); @@ -401,17 +404,25 @@ public void sendMessage(BytesReference reference, ActionListener listener) AtomicLong requestIdRef = new AtomicLong(); AtomicReference actionRef = new AtomicReference<>(); AtomicReference exceptionRef = new AtomicReference<>(); + // When response serialization fails, OutboundHandler.sendMessage's finally block calls + // Releasables.close(byteStreamOutput, onAfter) which invokes onResponseSent(success) as a + // cleanup side-effect. Then sendErrorResponse fires onResponseSent(error) -- but here the + // error response also fails to send (pipe broken), so it fires from the catch block instead. handler.setMessageListener(new TransportMessageListener() { @Override public void onResponseSent(long requestId, String action) { - throw new AssertionError("onResponseSent(success) must not be called on serialization failure"); + assertNull(channel.getMessageCaptor().get()); + assertThat(requestIdRef.get(), equalTo(0L)); + requestIdRef.set(requestId); + assertNull(actionRef.get()); + actionRef.set(action); } @Override public void onResponseSent(long requestId, String action, Exception error) { assertNull(channel.getMessageCaptor().get()); - requestIdRef.set(requestId); - actionRef.set(action); + assertThat(requestIdRef.get(), equalTo(requestId)); + assertThat(actionRef.get(), equalTo(action)); exceptionRef.set(error); } }); @@ -449,10 +460,19 @@ public void writeTo(StreamOutput out) { } }; + AtomicLong requestIdRef = new AtomicLong(); + AtomicReference actionRef = new AtomicReference<>(); + // When serialization fails, sendMessage's finally block invokes onResponseSent(success) as a + // cleanup side-effect via Releasables.close(onAfter). For handshakes no error response is sent, + // so onResponseSent(error) must not be called. handler.setMessageListener(new TransportMessageListener() { @Override public void onResponseSent(long requestId, String action) { - throw new AssertionError("onResponseSent(success) must not be called on serialization failure"); + assertNull(channel.getMessageCaptor().get()); + assertThat(requestIdRef.get(), equalTo(0L)); + requestIdRef.set(requestId); + assertNull(actionRef.get()); + actionRef.set(action); } @Override @@ -466,6 +486,8 @@ public void onResponseSent(long requestId, String action, Exception error) { } finally { response.decRef(); } + assertEquals(requestId, requestIdRef.get()); + assertEquals(action, actionRef.get()); assertNull(channel.getMessageCaptor().get()); assertNull(channel.getListenerCaptor().get()); assertTrue(response.released.get()); diff --git a/test/framework/src/main/java/org/elasticsearch/transport/TestTransportChannels.java b/test/framework/src/main/java/org/elasticsearch/transport/TestTransportChannels.java index effabd85591f9..ca627f96ee2dd 100644 --- a/test/framework/src/main/java/org/elasticsearch/transport/TestTransportChannels.java +++ b/test/framework/src/main/java/org/elasticsearch/transport/TestTransportChannels.java @@ -12,10 +12,15 @@ import org.elasticsearch.TransportVersion; import org.elasticsearch.common.network.HandlingTimeTracker; import org.elasticsearch.common.util.PageCacheRecycler; +import org.elasticsearch.core.Releasable; import org.elasticsearch.threadpool.ThreadPool; public class TestTransportChannels { + public static TaskTransportChannel newTaskTransportChannel(TransportChannel channel, Releasable onTaskFinished) { + return new TaskTransportChannel(1, channel, onTaskFinished); + } + public static TcpTransportChannel newFakeTcpTransportChannel( String nodeName, TcpChannel channel, From 5444e46929071272e9eaedd28f7c30076e97a92e Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Thu, 5 Mar 2026 19:58:05 +0200 Subject: [PATCH 10/39] update after review --- .../elasticsearch/action/search/SearchTransportService.java | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java index 90486ee2db1f8..cb05d65fdd715 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java +++ b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java @@ -463,7 +463,11 @@ public static void registerRequestHandler( (request, channel, task) -> searchService.executeQueryPhase( request, (SearchShardTask) task, - new ChannelActionListener<>(channel) + asBytesResponse(transportService, channel, response -> { + if (response instanceof QueryFetchSearchResult qfr) { + qfr.fetchResult().releaseCircuitBreakerBytes(searchService.getCircuitBreaker()); + } + }) ) ); TransportActionProxy.registerProxyActionWithDynamicResponseType( From ea8a10da20c92d4774a673567c7fdb2818fb4167 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Thu, 5 Mar 2026 20:11:30 +0200 Subject: [PATCH 11/39] udpate --- .../action/search/SearchTransportService.java | 41 +++++++++++-------- .../elasticsearch/search/SearchService.java | 16 +++++--- .../SearchServiceCircuitBreakerTests.java | 30 +++++++++----- 3 files changed, 54 insertions(+), 33 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java index cb05d65fdd715..a81c23603a094 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java +++ b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java @@ -695,25 +695,32 @@ private boolean assertNodePresent() { } /** - * Creates a listener that releases circuit-breaker reservations (via {@code afterSerialize}) and sends the - * response over the channel. For network channels the response is pre-serialized into a - * {@link BytesTransportResponse} so that the original in-memory data can be freed by the caller's - * {@code respondAndRelease}. For local (direct) channels the original response is forwarded as-is, - * because {@code TransportService.DirectResponseChannel} passes the object directly to the response - * handler without serialization/deserialization. - *

- * The caller is responsible for calling {@code response.decRef()} (typically via + * Wraps a transport channel so that the circuit-breaker bytes reserved for the response are released as + * early as possible -- right after the response has been converted to bytes, rather than after those bytes + * have been fully written to the network. + * + *

Why this exists: search responses (fetch results, query+fetch results) reserve circuit-breaker + * memory while their {@link org.elasticsearch.search.SearchHits} are held in Java objects. Without this + * wrapper the reservation is only freed after the full network send completes, keeping the breaker + * artificially inflated under high concurrency and causing spurious {@code CircuitBreakingException}s. + * + *

How it works -- two paths: + *

    + *
  1. Network (remote node): the response is serialized into a {@link BytesTransportResponse}. + * Once the bytes are produced, {@code afterSerialize} is called to release the breaker, and only + * then are the bytes sent over the wire. The original in-memory response can be freed immediately.
    2. + *
  2. Direct (same node): no serialization is needed -- the Java object is handed directly to + * the response handler. {@code afterSerialize} is called right after {@code channel.sendResponse()} + * returns. This is safe because the callback only updates breaker accounting; the actual + * response data is protected by ref-counting and stays alive until the handler is done with it.
    3. + *
+ * + *

The caller is responsible for calling {@code response.decRef()} (typically via * {@link ActionListener#respondAndRelease}); this method never calls it. - *

- * Direct-channel timing note: on the direct (local) path, {@code afterSerialize} runs after - * {@code channel.sendResponse()} returns. If the response handler uses a non-direct executor, - * {@code sendResponse()} may return before the handler fully processes the response, so - * {@code afterSerialize} can execute concurrently with the handler. This is safe because - * {@code afterSerialize} only releases circuit-breaker accounting (not the response data - * itself, which is ref-counted). * - * @param afterSerialize called after serialization (network) or after {@code channel.sendResponse()} returns - * (local) to release resources such as circuit-breaker reservations; must not throw + * @param afterSerialize callback invoked once the response bytes have been produced (network) or once + * {@code channel.sendResponse()} returns (direct); used to release circuit-breaker + * reservations. Must not throw. */ static ActionListener asBytesResponse( TransportService transportService, diff --git a/server/src/main/java/org/elasticsearch/search/SearchService.java b/server/src/main/java/org/elasticsearch/search/SearchService.java index 5652b0acb6b9f..3ca56f16867e5 100644 --- a/server/src/main/java/org/elasticsearch/search/SearchService.java +++ b/server/src/main/java/org/elasticsearch/search/SearchService.java @@ -1638,12 +1638,16 @@ private static void checkKeepAliveLimit(long keepAlive, long maxKeepAlive) { } /** - * Wraps a listener to release circuit breaker bytes from a FetchSearchResult after the response is consumed. - * Used for the query phase when a single-shard optimisation produces a {@code QueryFetchSearchResult} that is - * consumed locally (e.g. by {@code SearchQueryThenFetchAsyncAction}). For transport paths the circuit-breaker - * reservation is released by {@code SearchTransportService.asBytesResponse} right after serialization. - * Because {@link FetchSearchResult#releaseCircuitBreakerBytes} uses {@code AtomicLong.getAndSet(0)}, the - * release here is a safe no-op for transport paths — the bytes have already been released. + * Wraps a listener so that circuit-breaker bytes held by a {@link FetchSearchResult} are released after + * the response is handed off to the next listener in the chain. + * + *

When this matters: the single-shard query optimisation can produce a + * {@code QueryFetchSearchResult} that contains fetch data with a circuit-breaker reservation. On the + * transport path that reservation is already freed earlier by + * {@code SearchTransportService.asBytesResponse} (right after serialization). This wrapper acts as a + * second line of defence: because {@link FetchSearchResult#releaseCircuitBreakerBytes} uses + * {@code AtomicLong.getAndSet(0)}, calling it again here is a harmless no-op if the bytes were already + * released. */ private ActionListener releaseCircuitBreakerOnResponse( ActionListener listener, diff --git a/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java b/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java index 1b06628472243..1f7e1344d0ed3 100644 --- a/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java @@ -26,16 +26,26 @@ import static org.hamcrest.Matchers.is; /** - * Unit tests for circuit breaker release logic used by SearchService and SearchTransportService. - *

- * For transport paths (fetch, scroll-fetch), the circuit-breaker reservation is released right - * after serialization inside {@code SearchTransportService.asBytesResponse}. For the query phase - * (which may produce a {@code QueryFetchSearchResult} consumed locally by - * {@code SearchQueryThenFetchAsyncAction}), the circuit breaker is released after the listener - * consumes the response via {@code SearchService.releaseCircuitBreakerOnResponse}. - *

- * Because {@code FetchSearchResult.releaseCircuitBreakerBytes} uses {@code AtomicLong.getAndSet(0)}, - * multiple releases from different paths are idempotent. + * Tests that circuit-breaker bytes reserved for search fetch results are correctly released. + * + *

Background: when a fetch phase produces {@link org.elasticsearch.search.SearchHits}, the + * serialized size is reserved on the request circuit breaker. That reservation must be released once + * the response has been serialized (transport path) or consumed (local path), otherwise the breaker + * stays artificially inflated and blocks future requests. + * + *

Release happens in two places: + *

    + *
  1. {@code SearchTransportService.asBytesResponse} -- releases right after the response is + * serialized to bytes, before sending over the network. This covers fetch, scroll-fetch, + * and the single-shard query+fetch transport paths.
    2. + *
  2. {@code SearchService.releaseCircuitBreakerOnResponse} -- releases after the listener + * consumes the response. This is a safety net for the query phase; on transport paths the + * bytes are already gone so this call is a no-op.
    3. + *
+ * + *

Both paths are safe to call on the same result because + * {@code FetchSearchResult.releaseCircuitBreakerBytes} uses {@code AtomicLong.getAndSet(0)}, + * making repeated calls idempotent. */ public class SearchServiceCircuitBreakerTests extends ESTestCase { From 2fc3ee074cdf1398120703f37271c57d76da9886 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Thu, 5 Mar 2026 20:22:46 +0200 Subject: [PATCH 12/39] update code --- .../SearchServiceCircuitBreakerTests.java | 24 ++++++------------- .../transport/OutboundHandlerTests.java | 10 -------- 2 files changed, 7 insertions(+), 27 deletions(-) diff --git a/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java b/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java index 1f7e1344d0ed3..8100e3ca0f193 100644 --- a/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java @@ -151,8 +151,7 @@ public void testExtractorReturnsNull() { AtomicBoolean successCalled = new AtomicBoolean(false); AtomicBoolean failureCalled = new AtomicBoolean(false); - QuerySearchResult result = new QuerySearchResult(); - querySearchResultListener(successCalled, failureCalled, breaker).onResponse(result); + querySearchResultListener(successCalled, failureCalled, breaker).onResponse(new QuerySearchResult()); assertThat(successCalled.get(), is(true)); assertThat(failureCalled.get(), is(false)); @@ -246,10 +245,9 @@ public void onFailure(Exception e) { } /** - * Wrap a listener with circuit breaker release after {@code onResponse}, mirroring - * {@code SearchService.releaseCircuitBreakerOnResponse} (used for the query phase local path). + * Wrap a listener with circuit breaker release. */ - private ActionListener withCircuitBreakerReleaseAfterResponse( + private ActionListener withCircuitBreakerRelease( ActionListener listener, CircuitBreaker breaker, Function fetchResultExtractor @@ -271,7 +269,7 @@ private ActionListener querySearchResultListener( AtomicBoolean failureCalled, CircuitBreaker breaker ) { - return withCircuitBreakerReleaseAfterResponse(trackingListener(successCalled, failureCalled), breaker, qr -> null); + return withCircuitBreakerRelease(trackingListener(successCalled, failureCalled), breaker, qr -> null); } private ActionListener fetchSearchResultListener( @@ -279,7 +277,7 @@ private ActionListener fetchSearchResultListener( AtomicBoolean failureCalled, CircuitBreaker breaker ) { - return withCircuitBreakerReleaseAfterResponse(trackingListener(successCalled, failureCalled), breaker, Function.identity()); + return withCircuitBreakerRelease(trackingListener(successCalled, failureCalled), breaker, Function.identity()); } private ActionListener queryFetchSearchResultListener( @@ -287,11 +285,7 @@ private ActionListener queryFetchSearchResultListener( AtomicBoolean failureCalled, CircuitBreaker breaker ) { - return withCircuitBreakerReleaseAfterResponse( - trackingListener(successCalled, failureCalled), - breaker, - QueryFetchSearchResult::fetchResult - ); + return withCircuitBreakerRelease(trackingListener(successCalled, failureCalled), breaker, QueryFetchSearchResult::fetchResult); } private ActionListener scrollQueryFetchSearchResultListener( @@ -299,11 +293,7 @@ private ActionListener scrollQueryFetchSearchResul AtomicBoolean failureCalled, CircuitBreaker breaker ) { - return withCircuitBreakerReleaseAfterResponse( - trackingListener(successCalled, failureCalled), - breaker, - sr -> sr.result().fetchResult() - ); + return withCircuitBreakerRelease(trackingListener(successCalled, failureCalled), breaker, sr -> sr.result().fetchResult()); } /** diff --git a/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java b/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java index 06346ab95e290..0c7e75aa72e52 100644 --- a/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java +++ b/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java @@ -340,9 +340,6 @@ public void writeTo(StreamOutput out) { AtomicLong requestIdRef = new AtomicLong(); AtomicReference actionRef = new AtomicReference<>(); AtomicReference exceptionRef = new AtomicReference<>(); - // When response serialization fails, OutboundHandler.sendMessage's finally block calls - // Releasables.close(byteStreamOutput, onAfter) which invokes onResponseSent(success) as a - // cleanup side-effect. Then sendErrorResponse sends the error and invokes onResponseSent(error). handler.setMessageListener(new TransportMessageListener() { @Override public void onResponseSent(long requestId, String action) { @@ -404,10 +401,6 @@ public void sendMessage(BytesReference reference, ActionListener listener) AtomicLong requestIdRef = new AtomicLong(); AtomicReference actionRef = new AtomicReference<>(); AtomicReference exceptionRef = new AtomicReference<>(); - // When response serialization fails, OutboundHandler.sendMessage's finally block calls - // Releasables.close(byteStreamOutput, onAfter) which invokes onResponseSent(success) as a - // cleanup side-effect. Then sendErrorResponse fires onResponseSent(error) -- but here the - // error response also fails to send (pipe broken), so it fires from the catch block instead. handler.setMessageListener(new TransportMessageListener() { @Override public void onResponseSent(long requestId, String action) { @@ -462,9 +455,6 @@ public void writeTo(StreamOutput out) { AtomicLong requestIdRef = new AtomicLong(); AtomicReference actionRef = new AtomicReference<>(); - // When serialization fails, sendMessage's finally block invokes onResponseSent(success) as a - // cleanup side-effect via Releasables.close(onAfter). For handshakes no error response is sent, - // so onResponseSent(error) must not be called. handler.setMessageListener(new TransportMessageListener() { @Override public void onResponseSent(long requestId, String action) { From dc07e8229f65885ed9e4009c50d32d1920a93aa8 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Thu, 5 Mar 2026 20:26:40 +0200 Subject: [PATCH 13/39] update --- .../org/elasticsearch/transport/OutboundHandlerTests.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java b/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java index 0c7e75aa72e52..d76d4491a2b5d 100644 --- a/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java +++ b/server/src/test/java/org/elasticsearch/transport/OutboundHandlerTests.java @@ -476,10 +476,10 @@ public void onResponseSent(long requestId, String action, Exception error) { } finally { response.decRef(); } - assertEquals(requestId, requestIdRef.get()); - assertEquals(action, actionRef.get()); assertNull(channel.getMessageCaptor().get()); assertNull(channel.getListenerCaptor().get()); + assertEquals(requestId, requestIdRef.get()); + assertEquals(action, actionRef.get()); assertTrue(response.released.get()); assertFalse(channel.isOpen()); assertTrue(closeListener.isDone()); From 9b700a80edabffc2ccc11e96495154966145de27 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Thu, 5 Mar 2026 20:31:40 +0200 Subject: [PATCH 14/39] update --- .../action/search/AsBytesResponseTests.java | 1 + .../SearchServiceCircuitBreakerTests.java | 50 ++++++++----------- .../transport/TestDirectResponseChannel.java | 5 -- 3 files changed, 23 insertions(+), 33 deletions(-) diff --git a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java index 914f699f5a78a..96dec50012844 100644 --- a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java @@ -400,4 +400,5 @@ public long getUsed() { return used.get(); } } + } diff --git a/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java b/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java index 8100e3ca0f193..323a53bb6de2c 100644 --- a/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/SearchServiceCircuitBreakerTests.java @@ -26,26 +26,9 @@ import static org.hamcrest.Matchers.is; /** - * Tests that circuit-breaker bytes reserved for search fetch results are correctly released. - * - *

Background: when a fetch phase produces {@link org.elasticsearch.search.SearchHits}, the - * serialized size is reserved on the request circuit breaker. That reservation must be released once - * the response has been serialized (transport path) or consumed (local path), otherwise the breaker - * stays artificially inflated and blocks future requests. - * - *

Release happens in two places: - *

    - *
  1. {@code SearchTransportService.asBytesResponse} -- releases right after the response is - * serialized to bytes, before sending over the network. This covers fetch, scroll-fetch, - * and the single-shard query+fetch transport paths.
    2. - *
  2. {@code SearchService.releaseCircuitBreakerOnResponse} -- releases after the listener - * consumes the response. This is a safety net for the query phase; on transport paths the - * bytes are already gone so this call is a no-op.
    3. - *
- * - *

Both paths are safe to call on the same result because - * {@code FetchSearchResult.releaseCircuitBreakerBytes} uses {@code AtomicLong.getAndSet(0)}, - * making repeated calls idempotent. + * Unit tests for circuit breaker release logic in SearchService. + * Tests the generic helper method that releases circuit breaker bytes + * after search results are sent to the coordinator. */ public class SearchServiceCircuitBreakerTests extends ESTestCase { @@ -155,6 +138,7 @@ public void testExtractorReturnsNull() { assertThat(successCalled.get(), is(true)); assertThat(failureCalled.get(), is(false)); + // No breaker to release, should complete normally } public void testMultipleReleasesAreIdempotent() { @@ -165,10 +149,12 @@ public void testMultipleReleasesAreIdempotent() { try { result.setSearchHitsSizeBytes(2000L); + // First release result.releaseCircuitBreakerBytes(breaker); assertThat(breakerUsed.get(), equalTo(0L)); assertThat(result.getSearchHitsSizeBytes(), equalTo(0L)); + // Next release - should be no-op result.releaseCircuitBreakerBytes(breaker); assertThat(breakerUsed.get(), equalTo(0L)); assertThat(result.getSearchHitsSizeBytes(), equalTo(0L)); @@ -252,16 +238,24 @@ private ActionListener withCircuitBreakerRelease( CircuitBreaker breaker, Function fetchResultExtractor ) { - return ActionListener.wrap(response -> { - try { - listener.onResponse(response); - } finally { - FetchSearchResult fetchResult = fetchResultExtractor.apply(response); - if (fetchResult != null) { - fetchResult.releaseCircuitBreakerBytes(breaker); + return new ActionListener<>() { + @Override + public void onResponse(T response) { + try { + listener.onResponse(response); + } finally { + FetchSearchResult fetchResult = fetchResultExtractor.apply(response); + if (fetchResult != null) { + fetchResult.releaseCircuitBreakerBytes(breaker); + } } } - }, listener::onFailure); + + @Override + public void onFailure(Exception e) { + listener.onFailure(e); + } + }; } private ActionListener querySearchResultListener( diff --git a/test/framework/src/main/java/org/elasticsearch/transport/TestDirectResponseChannel.java b/test/framework/src/main/java/org/elasticsearch/transport/TestDirectResponseChannel.java index 076f6023dc0b7..505feafd41c3b 100644 --- a/test/framework/src/main/java/org/elasticsearch/transport/TestDirectResponseChannel.java +++ b/test/framework/src/main/java/org/elasticsearch/transport/TestDirectResponseChannel.java @@ -11,11 +11,6 @@ import org.elasticsearch.action.ActionListener; -/** - * A test-only {@link TransportService.DirectResponseChannel} that delegates to an {@link ActionListener} - * instead of interacting with a real {@link TransportService}. This allows testing code paths that branch - * on {@link TransportService#isDirectResponseChannel}. - */ public class TestDirectResponseChannel extends TransportService.DirectResponseChannel { private final ActionListener listener; From 1b6361ff30b22b41e0fe63284a3f99d4bbabf0f9 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Fri, 6 Mar 2026 10:50:04 +0200 Subject: [PATCH 15/39] update after review --- .../action/search/SearchTransportService.java | 163 ++++++++++-------- .../elasticsearch/search/SearchService.java | 8 - .../action/search/AsBytesResponseTests.java | 49 +++--- 3 files changed, 115 insertions(+), 105 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java index a81c23603a094..b2e06dd483cec 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java +++ b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java @@ -695,92 +695,111 @@ private boolean assertNodePresent() { } /** - * Wraps a transport channel so that the circuit-breaker bytes reserved for the response are released as - * early as possible -- right after the response has been converted to bytes, rather than after those bytes - * have been fully written to the network. + * Returns a listener that invokes {@code afterSerialize} as soon as the response has been serialized + * (network path) or handed off (direct/same-node path), so circuit-breaker bytes can be released before + * the network send completes. * - *

Why this exists: search responses (fetch results, query+fetch results) reserve circuit-breaker - * memory while their {@link org.elasticsearch.search.SearchHits} are held in Java objects. Without this - * wrapper the reservation is only freed after the full network send completes, keeping the breaker - * artificially inflated under high concurrency and causing spurious {@code CircuitBreakingException}s. - * - *

How it works -- two paths: - *

    - *
  1. Network (remote node): the response is serialized into a {@link BytesTransportResponse}. - * Once the bytes are produced, {@code afterSerialize} is called to release the breaker, and only - * then are the bytes sent over the wire. The original in-memory response can be freed immediately.
    2. - *
  2. Direct (same node): no serialization is needed -- the Java object is handed directly to - * the response handler. {@code afterSerialize} is called right after {@code channel.sendResponse()} - * returns. This is safe because the callback only updates breaker accounting; the actual - * response data is protected by ref-counting and stays alive until the handler is done with it.
    3. - *
- * - *

The caller is responsible for calling {@code response.decRef()} (typically via - * {@link ActionListener#respondAndRelease}); this method never calls it. - * - * @param afterSerialize callback invoked once the response bytes have been produced (network) or once - * {@code channel.sendResponse()} returns (direct); used to release circuit-breaker - * reservations. Must not throw. + *

On the network path the response is pre-serialized into a {@link BytesTransportResponse}; on the + * direct path the original object is forwarded as-is. In both cases {@code afterSerialize} runs before + * the bytes hit the wire. */ static ActionListener asBytesResponse( TransportService transportService, TransportChannel channel, Consumer afterSerialize ) { - TransportChannel innerChannel = channel instanceof TaskTransportChannel ttc ? ttc.getChannel() : channel; - if (TransportService.isDirectResponseChannel(innerChannel)) { - var channelListener = new ChannelActionListener(channel); - return new ActionListener<>() { - @Override - public void onResponse(T response) { - try { - channelListener.onResponse(response); - } finally { - afterSerialize.accept(response); - } - } + if (isDirectResponseChannel(channel)) { + return new DirectPathListener<>(channel, afterSerialize); + } + return new NetworkPathListener<>(transportService, channel, afterSerialize); + } - @Override - public void onFailure(Exception e) { - channelListener.onFailure(e); - } - }; + private static boolean isDirectResponseChannel(TransportChannel channel) { + if (channel instanceof TaskTransportChannel ttc) { + channel = ttc.getChannel(); } + return TransportService.isDirectResponseChannel(channel); + } - var channelListener = new ChannelActionListener(channel); - return new ActionListener<>() { - @Override - public void onResponse(T response) { - RecyclerBytesStreamOutput out = transportService.newNetworkBytesStream(); - try { - out.setTransportVersion(channel.getVersion()); - response.writeTo(out); - } catch (Exception e) { - Releasables.close(out); - try { - afterSerialize.accept(response); - } catch (Exception suppressed) { - e.addSuppressed(suppressed); - } - channelListener.onFailure(e); - return; - } - var bytesRef = out.moveToBytesReference(); - try { - afterSerialize.accept(response); - } catch (Exception e) { - Releasables.close(bytesRef); - channelListener.onFailure(e); - return; - } - ActionListener.respondAndRelease(channelListener, new BytesTransportResponse(bytesRef, out.getTransportVersion())); + /** + * Forwards the original Java response as-is and invokes {@code afterSerialize} immediately afterwards. + * */ + private static class DirectPathListener implements ActionListener { + private final ChannelActionListener channelListener; + private final Consumer afterSerialize; + + DirectPathListener(TransportChannel channel, Consumer afterSerialize) { + this.channelListener = new ChannelActionListener<>(channel); + this.afterSerialize = afterSerialize; + } + + @Override + public void onResponse(T response) { + try { + channelListener.onResponse(response); + } finally { + afterSerialize.accept(response); } + } - @Override - public void onFailure(Exception e) { + @Override + public void onFailure(Exception e) { + channelListener.onFailure(e); + } + } + + /** + * Serializes the response into a {@link BytesTransportResponse}, invokes {@code afterSerialize} to release circuit-breaker bytes, + * and then sends the pre-serialized bytes over the wire. + */ + private static class NetworkPathListener implements ActionListener { + private final TransportService transportService; + private final TransportChannel channel; + private final Consumer afterSerialize; + private final ChannelActionListener channelListener; + + NetworkPathListener(TransportService transportService, TransportChannel channel, Consumer afterSerialize) { + this.transportService = transportService; + this.channel = channel; + this.afterSerialize = afterSerialize; + this.channelListener = new ChannelActionListener<>(channel); + } + + @Override + public void onResponse(T response) { + RecyclerBytesStreamOutput out = transportService.newNetworkBytesStream(); + try { + out.setTransportVersion(channel.getVersion()); + response.writeTo(out); + } catch (Exception e) { + Releasables.close(out); + invokeAfterSerializeSafely(response, e); channelListener.onFailure(e); + return; } - }; + var bytesRef = out.moveToBytesReference(); + try { + afterSerialize.accept(response); + } catch (Exception e) { + Releasables.close(bytesRef); + channelListener.onFailure(e); + return; + } + ActionListener.respondAndRelease(channelListener, new BytesTransportResponse(bytesRef, out.getTransportVersion())); + } + + @Override + public void onFailure(Exception e) { + channelListener.onFailure(e); + } + + private void invokeAfterSerializeSafely(T response, Exception primary) { + try { + afterSerialize.accept(response); + } catch (Exception suppressed) { + primary.addSuppressed(suppressed); + } + } } public void cancelSearchTask(SearchTask task, String reason) { diff --git a/server/src/main/java/org/elasticsearch/search/SearchService.java b/server/src/main/java/org/elasticsearch/search/SearchService.java index 3ca56f16867e5..7f75b93bce41b 100644 --- a/server/src/main/java/org/elasticsearch/search/SearchService.java +++ b/server/src/main/java/org/elasticsearch/search/SearchService.java @@ -1640,14 +1640,6 @@ private static void checkKeepAliveLimit(long keepAlive, long maxKeepAlive) { /** * Wraps a listener so that circuit-breaker bytes held by a {@link FetchSearchResult} are released after * the response is handed off to the next listener in the chain. - * - *

When this matters: the single-shard query optimisation can produce a - * {@code QueryFetchSearchResult} that contains fetch data with a circuit-breaker reservation. On the - * transport path that reservation is already freed earlier by - * {@code SearchTransportService.asBytesResponse} (right after serialization). This wrapper acts as a - * second line of defence: because {@link FetchSearchResult#releaseCircuitBreakerBytes} uses - * {@code AtomicLong.getAndSet(0)}, calling it again here is a harmless no-op if the bytes were already - * released. */ private ActionListener releaseCircuitBreakerOnResponse( ActionListener listener, diff --git a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java index 96dec50012844..bfa04d4dc00c0 100644 --- a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java @@ -80,6 +80,29 @@ public void tearDown() throws Exception { super.tearDown(); } + public void testNetworkPathBytesResponseRoundTrip() throws Exception { + var sentResponse = new AtomicReference(); + + var channel = new TestTransportChannel(ActionListener.wrap(resp -> { + resp.mustIncRef(); + sentResponse.set(resp); + }, e -> fail("unexpected failure: " + e))); + + var original = new SimpleTestResponse("round-trip-test"); + ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, response -> {}); + + listener.onResponse(original); + + assertThat(sentResponse.get(), instanceOf(BytesTransportResponse.class)); + var bytesResp = (BytesTransportResponse) sentResponse.get(); + try (StreamInput in = bytesResp.streamInput()) { + var deserialized = new SimpleTestResponse(in); + assertThat(deserialized.value, equalTo("round-trip-test")); + } finally { + bytesResp.decRef(); + } + } + public void testNetworkPathCallsAfterSerializeOnSuccess() { var afterSerializeCalled = new AtomicBoolean(false); var sentResponse = new AtomicReference(); @@ -120,30 +143,7 @@ public void testNetworkPathCallsAfterSerializeOnSerializationFailure() { listener.onResponse(new FailingTestResponse()); assertTrue("afterSerialize must be called even on serialization failure", afterSerializeCalled.get()); - assertThat(sentException.get(), notNullValue()); - } - - public void testNetworkPathBytesResponseRoundTrip() throws Exception { - var sentResponse = new AtomicReference(); - - var channel = new TestTransportChannel(ActionListener.wrap(resp -> { - resp.mustIncRef(); - sentResponse.set(resp); - }, e -> fail("unexpected failure: " + e))); - - var original = new SimpleTestResponse("round-trip-test"); - ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, response -> {}); - - listener.onResponse(original); - - assertThat(sentResponse.get(), instanceOf(BytesTransportResponse.class)); - var bytesResp = (BytesTransportResponse) sentResponse.get(); - try (StreamInput in = bytesResp.streamInput()) { - var deserialized = new SimpleTestResponse(in); - assertThat(deserialized.value, equalTo("round-trip-test")); - } finally { - bytesResp.decRef(); - } + assertThat(sentException.get(), instanceOf(IOException.class)); } public void testNetworkPathReleasesCircuitBreakerBytesAfterSerialization() { @@ -400,5 +400,4 @@ public long getUsed() { return used.get(); } } - } From 2a1ef9a55b43683256b586c4756f96103acbf83e Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Fri, 6 Mar 2026 12:51:11 +0200 Subject: [PATCH 16/39] update after review --- .../action/search/SearchTransportService.java | 69 ++++---- .../search/fetch/FetchSearchResult.java | 13 ++ .../action/search/AsBytesResponseTests.java | 155 ++++++++---------- 3 files changed, 123 insertions(+), 114 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java index b2e06dd483cec..944d3d15f4701 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java +++ b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java @@ -21,6 +21,7 @@ import org.elasticsearch.client.internal.OriginSettingClient; import org.elasticsearch.client.internal.node.NodeClient; import org.elasticsearch.cluster.node.DiscoveryNode; +import org.elasticsearch.common.bytes.ReleasableBytesReference; import org.elasticsearch.common.io.stream.NamedWriteableRegistry; import org.elasticsearch.common.io.stream.RecyclerBytesStreamOutput; import org.elasticsearch.common.io.stream.StreamInput; @@ -71,7 +72,7 @@ import java.util.Objects; import java.util.concurrent.Executor; import java.util.function.BiFunction; -import java.util.function.Consumer; +import java.util.function.Function; /** * An encapsulation of {@link SearchService} operations exposed through @@ -465,8 +466,9 @@ public static void registerRequestHandler( (SearchShardTask) task, asBytesResponse(transportService, channel, response -> { if (response instanceof QueryFetchSearchResult qfr) { - qfr.fetchResult().releaseCircuitBreakerBytes(searchService.getCircuitBreaker()); + return qfr.fetchResult().detachCircuitBreakerReservation(searchService.getCircuitBreaker()); } + return () -> {}; }) ) ); @@ -526,7 +528,7 @@ public static void registerRequestHandler( asBytesResponse( transportService, channel, - response -> response.result().fetchResult().releaseCircuitBreakerBytes(searchService.getCircuitBreaker()) + response -> response.result().fetchResult().detachCircuitBreakerReservation(searchService.getCircuitBreaker()) ) ) ); @@ -561,7 +563,7 @@ public static void registerRequestHandler( asBytesResponse( transportService, channel, - response -> response.releaseCircuitBreakerBytes(searchService.getCircuitBreaker()) + response -> response.detachCircuitBreakerReservation(searchService.getCircuitBreaker()) ) ); transportService.registerRequestHandler( @@ -695,23 +697,26 @@ private boolean assertNodePresent() { } /** - * Returns a listener that invokes {@code afterSerialize} as soon as the response has been serialized - * (network path) or handed off (direct/same-node path), so circuit-breaker bytes can be released before - * the network send completes. + * Returns a listener that ensures circuit-breaker bytes are not released until the response + * has actually been written to the network. * - *

On the network path the response is pre-serialized into a {@link BytesTransportResponse}; on the - * direct path the original object is forwarded as-is. In both cases {@code afterSerialize} runs before - * the bytes hit the wire. + *

On the network path, the response is serialized into bytes up front. + * Then {@code detachRelease} is called to extract the breaker reservation from the response + * as a {@link Releasable}, and then is attached to the serialized bytes' ref-count, so it + * is only closed when Netty finishes writing those bytes to the wire. + * + *

On the direct (same-node) path the response is forwarded as-is and the {@code Releasable} + * is closed immediately after hand-off. */ static ActionListener asBytesResponse( TransportService transportService, TransportChannel channel, - Consumer afterSerialize + Function detachRelease ) { if (isDirectResponseChannel(channel)) { - return new DirectPathListener<>(channel, afterSerialize); + return new DirectPathListener<>(channel, detachRelease); } - return new NetworkPathListener<>(transportService, channel, afterSerialize); + return new NetworkPathListener<>(transportService, channel, detachRelease); } private static boolean isDirectResponseChannel(TransportChannel channel) { @@ -722,23 +727,26 @@ private static boolean isDirectResponseChannel(TransportChannel channel) { } /** - * Forwards the original Java response as-is and invokes {@code afterSerialize} immediately afterwards. - * */ + * Forwards the original Java response as-is and releases the circuit-breaker reservation immediately + * after the response has been handed off. + */ private static class DirectPathListener implements ActionListener { private final ChannelActionListener channelListener; - private final Consumer afterSerialize; + private final Function detachRelease; - DirectPathListener(TransportChannel channel, Consumer afterSerialize) { + DirectPathListener(TransportChannel channel, Function detachRelease) { this.channelListener = new ChannelActionListener<>(channel); - this.afterSerialize = afterSerialize; + this.detachRelease = detachRelease; } @Override public void onResponse(T response) { + Releasable breakerRelease = null; try { channelListener.onResponse(response); + breakerRelease = detachRelease.apply(response); } finally { - afterSerialize.accept(response); + Releasables.close(breakerRelease); } } @@ -749,19 +757,20 @@ public void onFailure(Exception e) { } /** - * Serializes the response into a {@link BytesTransportResponse}, invokes {@code afterSerialize} to release circuit-breaker bytes, - * and then sends the pre-serialized bytes over the wire. + * Serializes the response into a {@link BytesTransportResponse} and bundles the circuit-breaker + * release with the serialized bytes' ref-count. The breaker reservation stays alive as long as + * the bytes are on the heap, and is released when the ref-count drops to zero. */ private static class NetworkPathListener implements ActionListener { private final TransportService transportService; private final TransportChannel channel; - private final Consumer afterSerialize; + private final Function detachRelease; private final ChannelActionListener channelListener; - NetworkPathListener(TransportService transportService, TransportChannel channel, Consumer afterSerialize) { + NetworkPathListener(TransportService transportService, TransportChannel channel, Function detachRelease) { this.transportService = transportService; this.channel = channel; - this.afterSerialize = afterSerialize; + this.detachRelease = detachRelease; this.channelListener = new ChannelActionListener<>(channel); } @@ -773,19 +782,21 @@ public void onResponse(T response) { response.writeTo(out); } catch (Exception e) { Releasables.close(out); - invokeAfterSerializeSafely(response, e); + detachAndReleaseSafely(response, e); channelListener.onFailure(e); return; } var bytesRef = out.moveToBytesReference(); + Releasable breakerRelease; try { - afterSerialize.accept(response); + breakerRelease = detachRelease.apply(response); } catch (Exception e) { Releasables.close(bytesRef); channelListener.onFailure(e); return; } - ActionListener.respondAndRelease(channelListener, new BytesTransportResponse(bytesRef, out.getTransportVersion())); + var trackingBytes = new ReleasableBytesReference(bytesRef, Releasables.wrap(bytesRef, breakerRelease)); + ActionListener.respondAndRelease(channelListener, new BytesTransportResponse(trackingBytes, out.getTransportVersion())); } @Override @@ -793,9 +804,9 @@ public void onFailure(Exception e) { channelListener.onFailure(e); } - private void invokeAfterSerializeSafely(T response, Exception primary) { + private void detachAndReleaseSafely(T response, Exception primary) { try { - afterSerialize.accept(response); + Releasables.close(detachRelease.apply(response)); } catch (Exception suppressed) { primary.addSuppressed(suppressed); } diff --git a/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java b/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java index dd23472482c27..8ddc8764bf05b 100644 --- a/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java +++ b/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java @@ -13,6 +13,7 @@ import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.core.RefCounted; +import org.elasticsearch.core.Releasable; import org.elasticsearch.core.SimpleRefCounted; import org.elasticsearch.search.SearchHit; import org.elasticsearch.search.SearchHits; @@ -103,6 +104,18 @@ public void releaseCircuitBreakerBytes(CircuitBreaker circuitBreaker) { } } + /** + * Extracts the circuit-breaker reservation from this result and returns a + * {@link Releasable} that releases those bytes from the breaker when closed. + */ + public Releasable detachCircuitBreakerReservation(CircuitBreaker circuitBreaker) { + long bytes = searchHitsSizeBytes.getAndSet(0L); + if (bytes > 0L) { + return () -> circuitBreaker.addWithoutBreaking(-bytes); + } + return () -> {}; + } + public FetchSearchResult initCounter() { counter = 0; return this; diff --git a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java index bfa04d4dc00c0..a72cab8895808 100644 --- a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java @@ -44,13 +44,6 @@ import static org.hamcrest.Matchers.instanceOf; import static org.hamcrest.Matchers.notNullValue; -/** - * Tests for {@link SearchTransportService#asBytesResponse}, covering both the network-channel path - * (pre-serializes into {@link BytesTransportResponse}) and the direct-channel path (forwards the - * original response as-is for local/same-node requests). Tests verify {@link BytesTransportResponse} - * creation, round-trip deserialization, {@code afterSerialize} callback semantics, and cleanup - * on serialization failure. - */ public class AsBytesResponseTests extends ESTestCase { private ThreadPool threadPool; @@ -88,8 +81,12 @@ public void testNetworkPathBytesResponseRoundTrip() throws Exception { sentResponse.set(resp); }, e -> fail("unexpected failure: " + e))); - var original = new SimpleTestResponse("round-trip-test"); - ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, response -> {}); + var original = new SimpleTestResponse("test"); + ActionListener listener = SearchTransportService.asBytesResponse( + transportService, + channel, + response -> () -> {} + ); listener.onResponse(original); @@ -97,14 +94,15 @@ public void testNetworkPathBytesResponseRoundTrip() throws Exception { var bytesResp = (BytesTransportResponse) sentResponse.get(); try (StreamInput in = bytesResp.streamInput()) { var deserialized = new SimpleTestResponse(in); - assertThat(deserialized.value, equalTo("round-trip-test")); + assertThat(deserialized.value, equalTo("test")); } finally { bytesResp.decRef(); } } - public void testNetworkPathCallsAfterSerializeOnSuccess() { - var afterSerializeCalled = new AtomicBoolean(false); + public void testNetworkPathCallsDetachReleaseOnSuccess() { + var detachCalled = new AtomicBoolean(false); + var releasableClosed = new AtomicBoolean(false); var sentResponse = new AtomicReference(); var channel = new TestTransportChannel(ActionListener.wrap(resp -> { @@ -112,22 +110,24 @@ public void testNetworkPathCallsAfterSerializeOnSuccess() { sentResponse.set(resp); }, e -> fail("unexpected failure: " + e))); - ActionListener listener = SearchTransportService.asBytesResponse( - transportService, - channel, - response -> afterSerializeCalled.set(true) - ); + ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, response -> { + detachCalled.set(true); + return () -> releasableClosed.set(true); + }); listener.onResponse(new SimpleTestResponse("hello")); - assertTrue("afterSerialize must be called after successful serialization", afterSerializeCalled.get()); + assertTrue("detachRelease must be called after successful serialization", detachCalled.get()); + assertFalse("releasable must NOT be closed yet (deferred to bytes release)", releasableClosed.get()); assertThat(sentResponse.get(), notNullValue()); assertThat(sentResponse.get(), instanceOf(BytesTransportResponse.class)); + sentResponse.get().decRef(); + assertTrue("releasable must be closed after bytes are released", releasableClosed.get()); } - public void testNetworkPathCallsAfterSerializeOnSerializationFailure() { - var afterSerializeCalled = new AtomicBoolean(false); + public void testNetworkPathReleasesImmediatelyOnSerializationFailure() { + var releasableClosed = new AtomicBoolean(false); var sentException = new AtomicReference(); var channel = new TestTransportChannel( @@ -137,16 +137,17 @@ public void testNetworkPathCallsAfterSerializeOnSerializationFailure() { ActionListener listener = SearchTransportService.asBytesResponse( transportService, channel, - response -> afterSerializeCalled.set(true) + response -> () -> releasableClosed.set(true) ); listener.onResponse(new FailingTestResponse()); - assertTrue("afterSerialize must be called even on serialization failure", afterSerializeCalled.get()); + assertTrue("releasable must be closed immediately on serialization failure", releasableClosed.get()); + assertThat(sentException.get(), notNullValue()); assertThat(sentException.get(), instanceOf(IOException.class)); } - public void testNetworkPathReleasesCircuitBreakerBytesAfterSerialization() { + public void testNetworkPathDefersCircuitBreakerReleaseUntilBytesReleased() { var breakerUsed = new AtomicLong(5000); CircuitBreaker breaker = new TestCircuitBreaker(breakerUsed); @@ -167,26 +168,24 @@ public void testNetworkPathReleasesCircuitBreakerBytesAfterSerialization() { ActionListener listener = SearchTransportService.asBytesResponse( transportService, channel, - response -> response.releaseCircuitBreakerBytes(breaker) + response -> response.detachCircuitBreakerReservation(breaker) ); listener.onResponse(fetchResult); - assertThat("breaker bytes should be released after serialization", breakerUsed.get(), equalTo(0L)); - assertThat(fetchResult.getSearchHitsSizeBytes(), equalTo(0L)); + assertThat("detach must zero out the field on the result", fetchResult.getSearchHitsSizeBytes(), equalTo(0L)); + assertThat("breaker must still be reserved while bytes are in flight", breakerUsed.get(), equalTo(5000L)); assertThat(sentResponse.get(), instanceOf(BytesTransportResponse.class)); + + sentResponse.get().decRef(); + assertThat("breaker must be released after bytes are released", breakerUsed.get(), equalTo(0L)); } finally { fetchResult.decRef(); - if (sentResponse.get() != null) { - sentResponse.get().decRef(); - } } } public void testNetworkPathReleasesCircuitBreakerBytesOnSerializationFailure() { var breakerUsed = new AtomicLong(3000); - CircuitBreaker breaker = new TestCircuitBreaker(breakerUsed); - var sentException = new AtomicReference(); var channel = new TestTransportChannel( @@ -195,37 +194,39 @@ public void testNetworkPathReleasesCircuitBreakerBytesOnSerializationFailure() { ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, response -> { long bytes = breakerUsed.get(); - if (bytes > 0L) { - breakerUsed.addAndGet(-bytes); - } + return () -> { + if (bytes > 0L) { + breakerUsed.addAndGet(-bytes); + } + }; }); listener.onResponse(new FailingTestResponse()); - assertThat("breaker bytes should be released even on serialization failure", breakerUsed.get(), equalTo(0L)); + assertThat("breaker bytes should be released immediately on serialization failure", breakerUsed.get(), equalTo(0L)); assertThat(sentException.get(), notNullValue()); + assertThat(sentException.get(), instanceOf(IOException.class)); } - public void testNetworkPathOnFailureDoesNotCallAfterSerialize() { - var afterSerializeCalled = new AtomicBoolean(false); + public void testNetworkPathOnFailureDoesNotCallDetachRelease() { + var detachCalled = new AtomicBoolean(false); var sentException = new AtomicReference(); var channel = new TestTransportChannel(ActionListener.wrap(resp -> fail("should not succeed"), sentException::set)); - ActionListener listener = SearchTransportService.asBytesResponse( - transportService, - channel, - response -> afterSerializeCalled.set(true) - ); + ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, response -> { + detachCalled.set(true); + return () -> {}; + }); listener.onFailure(new RuntimeException("upstream failure")); - assertFalse("afterSerialize must not be called on upstream failure", afterSerializeCalled.get()); + assertFalse("detachRelease must not be called on upstream failure", detachCalled.get()); assertThat(sentException.get(), notNullValue()); } public void testDirectPathForwardsOriginalResponse() { - var afterSerializeCalled = new AtomicBoolean(false); + var releasableClosed = new AtomicBoolean(false); var sentResponse = new AtomicReference(); var channel = new TestDirectResponseChannel(ActionListener.wrap(sentResponse::set, e -> fail("unexpected failure: " + e))); @@ -233,53 +234,34 @@ public void testDirectPathForwardsOriginalResponse() { ActionListener listener = SearchTransportService.asBytesResponse( transportService, channel, - response -> afterSerializeCalled.set(true) + response -> () -> releasableClosed.set(true) ); var original = new SimpleTestResponse("direct-test"); listener.onResponse(original); - assertTrue("afterSerialize must be called on direct path", afterSerializeCalled.get()); + assertTrue("releasable must be closed immediately on direct path", releasableClosed.get()); assertSame("direct path must forward the original response, not BytesTransportResponse", original, sentResponse.get()); } - public void testDirectPathCallsAfterSerializeAfterOnResponse() { - var callOrder = new java.util.ArrayList(); - - var channel = new TestDirectResponseChannel( - ActionListener.wrap(resp -> callOrder.add("onResponse"), e -> fail("unexpected failure: " + e)) - ); - - ActionListener listener = SearchTransportService.asBytesResponse( - transportService, - channel, - response -> callOrder.add("afterSerialize") - ); - - listener.onResponse(new SimpleTestResponse("order-test")); - - assertThat(callOrder, equalTo(java.util.List.of("onResponse", "afterSerialize"))); - } - - public void testDirectPathOnFailureDoesNotCallAfterSerialize() { - var afterSerializeCalled = new AtomicBoolean(false); + public void testDirectPathOnFailureDoesNotCallDetachRelease() { + var detachCalled = new AtomicBoolean(false); var sentException = new AtomicReference(); var channel = new TestDirectResponseChannel(ActionListener.wrap(resp -> fail("should not succeed"), sentException::set)); - ActionListener listener = SearchTransportService.asBytesResponse( - transportService, - channel, - response -> afterSerializeCalled.set(true) - ); + ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, response -> { + detachCalled.set(true); + return () -> {}; + }); listener.onFailure(new RuntimeException("upstream failure")); - assertFalse("afterSerialize must not be called on upstream failure", afterSerializeCalled.get()); + assertFalse("detachRelease must not be called on upstream failure", detachCalled.get()); assertThat(sentException.get(), notNullValue()); } - public void testDirectPathReleasesCircuitBreakerBytesAfterResponse() { + public void testDirectPathReleasesCircuitBreakerBytesImmediately() { var breakerUsed = new AtomicLong(5000); CircuitBreaker breaker = new TestCircuitBreaker(breakerUsed); @@ -297,12 +279,12 @@ public void testDirectPathReleasesCircuitBreakerBytesAfterResponse() { ActionListener listener = SearchTransportService.asBytesResponse( transportService, channel, - response -> response.releaseCircuitBreakerBytes(breaker) + response -> response.detachCircuitBreakerReservation(breaker) ); listener.onResponse(fetchResult); - assertThat("breaker bytes should be released after response on direct path", breakerUsed.get(), equalTo(0L)); + assertThat("breaker bytes should be released immediately on direct path", breakerUsed.get(), equalTo(0L)); assertThat(fetchResult.getSearchHitsSizeBytes(), equalTo(0L)); assertSame("direct path must forward original response", fetchResult, sentResponse.get()); } finally { @@ -311,7 +293,7 @@ public void testDirectPathReleasesCircuitBreakerBytesAfterResponse() { } public void testTaskTransportChannelUnwrapsToDirectPath() { - var afterSerializeCalled = new AtomicBoolean(false); + var releasableClosed = new AtomicBoolean(false); var sentResponse = new AtomicReference(); var directChannel = new TestDirectResponseChannel(ActionListener.wrap(sentResponse::set, e -> fail("unexpected failure: " + e))); @@ -320,18 +302,20 @@ public void testTaskTransportChannelUnwrapsToDirectPath() { ActionListener listener = SearchTransportService.asBytesResponse( transportService, taskChannel, - response -> afterSerializeCalled.set(true) + response -> () -> releasableClosed.set(true) ); var original = new SimpleTestResponse("task-wrapped-test"); listener.onResponse(original); - assertTrue("afterSerialize must be called on task-wrapped direct path", afterSerializeCalled.get()); + assertTrue("releasable must be closed immediately on task-wrapped direct path", releasableClosed.get()); assertSame("task-wrapped direct channel must forward original response, not BytesTransportResponse", original, sentResponse.get()); + sentResponse.get().decRef(); } public void testTaskTransportChannelUnwrapsToNetworkPath() { - var afterSerializeCalled = new AtomicBoolean(false); + var detachCalled = new AtomicBoolean(false); + var releasableClosed = new AtomicBoolean(false); var sentResponse = new AtomicReference(); var networkChannel = new TestTransportChannel(ActionListener.wrap(resp -> { @@ -340,17 +324,18 @@ public void testTaskTransportChannelUnwrapsToNetworkPath() { }, e -> fail("unexpected failure: " + e))); var taskChannel = TestTransportChannels.newTaskTransportChannel(networkChannel, () -> {}); - ActionListener listener = SearchTransportService.asBytesResponse( - transportService, - taskChannel, - response -> afterSerializeCalled.set(true) - ); + ActionListener listener = SearchTransportService.asBytesResponse(transportService, taskChannel, response -> { + detachCalled.set(true); + return () -> releasableClosed.set(true); + }); listener.onResponse(new SimpleTestResponse("task-network-test")); - assertTrue("afterSerialize must be called on task-wrapped network path", afterSerializeCalled.get()); + assertTrue("detachRelease must be called on task-wrapped network path", detachCalled.get()); + assertFalse("releasable must NOT be closed yet on task-wrapped network path", releasableClosed.get()); assertThat(sentResponse.get(), instanceOf(BytesTransportResponse.class)); sentResponse.get().decRef(); + assertTrue("releasable must be closed after bytes are released", releasableClosed.get()); } static class SimpleTestResponse extends TransportResponse { From d89663edcc42bcfeb3290c3d418337ef9675df9b Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Fri, 6 Mar 2026 12:52:56 +0200 Subject: [PATCH 17/39] update after review --- .../main/java/org/elasticsearch/search/SearchService.java | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/search/SearchService.java b/server/src/main/java/org/elasticsearch/search/SearchService.java index 7f75b93bce41b..14b4e0dc3e4ba 100644 --- a/server/src/main/java/org/elasticsearch/search/SearchService.java +++ b/server/src/main/java/org/elasticsearch/search/SearchService.java @@ -1636,11 +1636,7 @@ private static void checkKeepAliveLimit(long keepAlive, long maxKeepAlive) { ); } } - - /** - * Wraps a listener so that circuit-breaker bytes held by a {@link FetchSearchResult} are released after - * the response is handed off to the next listener in the chain. - */ + private ActionListener releaseCircuitBreakerOnResponse( ActionListener listener, Function fetchResultExtractor From 4f5b108f3ffba3da6d58e24c3aed3f05f6bb2e3d Mon Sep 17 00:00:00 2001 From: elasticsearchmachine Date: Fri, 6 Mar 2026 11:04:02 +0000 Subject: [PATCH 18/39] [CI] Auto commit changes from spotless --- .../src/main/java/org/elasticsearch/search/SearchService.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/main/java/org/elasticsearch/search/SearchService.java b/server/src/main/java/org/elasticsearch/search/SearchService.java index 14b4e0dc3e4ba..c44371e68cb68 100644 --- a/server/src/main/java/org/elasticsearch/search/SearchService.java +++ b/server/src/main/java/org/elasticsearch/search/SearchService.java @@ -1636,7 +1636,7 @@ private static void checkKeepAliveLimit(long keepAlive, long maxKeepAlive) { ); } } - + private ActionListener releaseCircuitBreakerOnResponse( ActionListener listener, Function fetchResultExtractor From 3ce47fad2b127640a0a3aa5ba6112a0829698bd6 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Mon, 9 Mar 2026 10:53:01 +0200 Subject: [PATCH 19/39] update after review --- .../action/search/SearchTransportService.java | 28 +++++++++---- .../elasticsearch/transport/TcpTransport.java | 5 +++ .../elasticsearch/transport/Transport.java | 7 +++- .../transport/TransportService.java | 5 +++ .../action/search/AsBytesResponseTests.java | 41 ++++++++++++------- 5 files changed, 63 insertions(+), 23 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java index 944d3d15f4701..865e940f88057 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java +++ b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java @@ -21,6 +21,7 @@ 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.RecyclerBytesStreamOutput; @@ -469,7 +470,7 @@ public static void registerRequestHandler( return qfr.fetchResult().detachCircuitBreakerReservation(searchService.getCircuitBreaker()); } return () -> {}; - }) + }, searchService.getCircuitBreaker()) ) ); TransportActionProxy.registerProxyActionWithDynamicResponseType( @@ -528,7 +529,8 @@ public static void registerRequestHandler( asBytesResponse( transportService, channel, - response -> response.result().fetchResult().detachCircuitBreakerReservation(searchService.getCircuitBreaker()) + response -> response.result().fetchResult().detachCircuitBreakerReservation(searchService.getCircuitBreaker()), + searchService.getCircuitBreaker() ) ) ); @@ -563,7 +565,8 @@ public static void registerRequestHandler( asBytesResponse( transportService, channel, - response -> response.detachCircuitBreakerReservation(searchService.getCircuitBreaker()) + response -> response.detachCircuitBreakerReservation(searchService.getCircuitBreaker()), + searchService.getCircuitBreaker() ) ); transportService.registerRequestHandler( @@ -711,12 +714,13 @@ private boolean assertNodePresent() { static ActionListener asBytesResponse( TransportService transportService, TransportChannel channel, - Function detachRelease + Function detachRelease, + @Nullable CircuitBreaker circuitBreaker ) { if (isDirectResponseChannel(channel)) { return new DirectPathListener<>(channel, detachRelease); } - return new NetworkPathListener<>(transportService, channel, detachRelease); + return new NetworkPathListener<>(transportService, channel, detachRelease, circuitBreaker); } private static boolean isDirectResponseChannel(TransportChannel channel) { @@ -766,17 +770,25 @@ private static class NetworkPathListener implements private final TransportChannel channel; private final Function detachRelease; private final ChannelActionListener channelListener; - - NetworkPathListener(TransportService transportService, TransportChannel channel, Function detachRelease) { + @Nullable + private final CircuitBreaker circuitBreaker; + + NetworkPathListener( + TransportService transportService, + TransportChannel channel, + Function detachRelease, + @Nullable CircuitBreaker circuitBreaker + ) { this.transportService = transportService; this.channel = channel; this.detachRelease = detachRelease; this.channelListener = new ChannelActionListener<>(channel); + this.circuitBreaker = circuitBreaker; } @Override public void onResponse(T response) { - RecyclerBytesStreamOutput out = transportService.newNetworkBytesStream(); + RecyclerBytesStreamOutput out = transportService.newNetworkBytesStream(circuitBreaker); try { out.setTransportVersion(channel.getVersion()); response.writeTo(out); diff --git a/server/src/main/java/org/elasticsearch/transport/TcpTransport.java b/server/src/main/java/org/elasticsearch/transport/TcpTransport.java index 24570c544bdea..4e36d4b2e4d46 100644 --- a/server/src/main/java/org/elasticsearch/transport/TcpTransport.java +++ b/server/src/main/java/org/elasticsearch/transport/TcpTransport.java @@ -991,6 +991,11 @@ public RecyclerBytesStreamOutput newNetworkBytesStream() { return new RecyclerBytesStreamOutput(recycler); } + @Override + public RecyclerBytesStreamOutput newNetworkBytesStream(CircuitBreaker circuitBreaker) { + return new RecyclerBytesStreamOutput(recycler, circuitBreaker); + } + /** * Ensures this transport is still started / open * diff --git a/server/src/main/java/org/elasticsearch/transport/Transport.java b/server/src/main/java/org/elasticsearch/transport/Transport.java index 4bca3e917082c..c78e75db38f33 100644 --- a/server/src/main/java/org/elasticsearch/transport/Transport.java +++ b/server/src/main/java/org/elasticsearch/transport/Transport.java @@ -13,6 +13,7 @@ import org.elasticsearch.action.ActionListener; import org.elasticsearch.cluster.node.DiscoveryNode; import org.elasticsearch.common.component.LifecycleComponent; +import org.elasticsearch.common.breaker.CircuitBreaker; import org.elasticsearch.common.io.stream.RecyclerBytesStreamOutput; import org.elasticsearch.common.transport.BoundTransportAddress; import org.elasticsearch.common.transport.TransportAddress; @@ -89,7 +90,11 @@ default boolean isSecure() { RequestHandlers getRequestHandlers(); default RecyclerBytesStreamOutput newNetworkBytesStream() { - return new RecyclerBytesStreamOutput(NON_RECYCLING_INSTANCE); + return newNetworkBytesStream(null); + } + + default RecyclerBytesStreamOutput newNetworkBytesStream(CircuitBreaker circuitBreaker) { + return new RecyclerBytesStreamOutput(NON_RECYCLING_INSTANCE, circuitBreaker); } /** diff --git a/server/src/main/java/org/elasticsearch/transport/TransportService.java b/server/src/main/java/org/elasticsearch/transport/TransportService.java index 60170125d65c0..e5d70b799746a 100644 --- a/server/src/main/java/org/elasticsearch/transport/TransportService.java +++ b/server/src/main/java/org/elasticsearch/transport/TransportService.java @@ -22,6 +22,7 @@ import org.elasticsearch.cluster.project.ProjectResolver; import org.elasticsearch.common.ReferenceDocs; import org.elasticsearch.common.Strings; +import org.elasticsearch.common.breaker.CircuitBreaker; import org.elasticsearch.common.component.AbstractLifecycleComponent; import org.elasticsearch.common.io.stream.RecyclerBytesStreamOutput; import org.elasticsearch.common.io.stream.StreamInput; @@ -664,6 +665,10 @@ public RecyclerBytesStreamOutput newNetworkBytesStream() { return transport.newNetworkBytesStream(); } + public RecyclerBytesStreamOutput newNetworkBytesStream(CircuitBreaker circuitBreaker) { + return transport.newNetworkBytesStream(circuitBreaker); + } + static class HandshakeRequest extends AbstractTransportRequest { public static final HandshakeRequest INSTANCE = new HandshakeRequest(); diff --git a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java index a72cab8895808..4b4269248bdba 100644 --- a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java @@ -16,6 +16,7 @@ import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.common.util.PageCacheRecycler; import org.elasticsearch.search.SearchHits; import org.elasticsearch.search.fetch.FetchSearchResult; import org.elasticsearch.search.internal.ShardSearchContextId; @@ -85,7 +86,8 @@ public void testNetworkPathBytesResponseRoundTrip() throws Exception { ActionListener listener = SearchTransportService.asBytesResponse( transportService, channel, - response -> () -> {} + response -> () -> {}, + null ); listener.onResponse(original); @@ -113,7 +115,7 @@ public void testNetworkPathCallsDetachReleaseOnSuccess() { ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, response -> { detachCalled.set(true); return () -> releasableClosed.set(true); - }); + }, null); listener.onResponse(new SimpleTestResponse("hello")); @@ -137,7 +139,8 @@ public void testNetworkPathReleasesImmediatelyOnSerializationFailure() { ActionListener listener = SearchTransportService.asBytesResponse( transportService, channel, - response -> () -> releasableClosed.set(true) + response -> () -> releasableClosed.set(true), + null ); listener.onResponse(new FailingTestResponse()); @@ -148,7 +151,9 @@ public void testNetworkPathReleasesImmediatelyOnSerializationFailure() { } public void testNetworkPathDefersCircuitBreakerReleaseUntilBytesReleased() { - var breakerUsed = new AtomicLong(5000); + long responseBytes = 5000L; + long pageBytes = PageCacheRecycler.BYTE_PAGE_SIZE; + var breakerUsed = new AtomicLong(responseBytes); CircuitBreaker breaker = new TestCircuitBreaker(breakerUsed); var sentResponse = new AtomicReference(); @@ -163,18 +168,23 @@ public void testNetworkPathDefersCircuitBreakerReleaseUntilBytesReleased() { try { var hits = SearchHits.empty(new TotalHits(0, TotalHits.Relation.EQUAL_TO), 0.0f); fetchResult.shardResult(hits, null); - fetchResult.setSearchHitsSizeBytes(5000L); + fetchResult.setSearchHitsSizeBytes(responseBytes); ActionListener listener = SearchTransportService.asBytesResponse( transportService, channel, - response -> response.detachCircuitBreakerReservation(breaker) + response -> response.detachCircuitBreakerReservation(breaker), + breaker ); listener.onResponse(fetchResult); assertThat("detach must zero out the field on the result", fetchResult.getSearchHitsSizeBytes(), equalTo(0L)); - assertThat("breaker must still be reserved while bytes are in flight", breakerUsed.get(), equalTo(5000L)); + assertThat( + "breaker must account for both the detached response reservation and the serialized page bytes", + breakerUsed.get(), + equalTo(responseBytes + pageBytes) + ); assertThat(sentResponse.get(), instanceOf(BytesTransportResponse.class)); sentResponse.get().decRef(); @@ -199,7 +209,7 @@ public void testNetworkPathReleasesCircuitBreakerBytesOnSerializationFailure() { breakerUsed.addAndGet(-bytes); } }; - }); + }, null); listener.onResponse(new FailingTestResponse()); @@ -217,7 +227,7 @@ public void testNetworkPathOnFailureDoesNotCallDetachRelease() { ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, response -> { detachCalled.set(true); return () -> {}; - }); + }, null); listener.onFailure(new RuntimeException("upstream failure")); @@ -234,7 +244,8 @@ public void testDirectPathForwardsOriginalResponse() { ActionListener listener = SearchTransportService.asBytesResponse( transportService, channel, - response -> () -> releasableClosed.set(true) + response -> () -> releasableClosed.set(true), + null ); var original = new SimpleTestResponse("direct-test"); @@ -253,7 +264,7 @@ public void testDirectPathOnFailureDoesNotCallDetachRelease() { ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, response -> { detachCalled.set(true); return () -> {}; - }); + }, null); listener.onFailure(new RuntimeException("upstream failure")); @@ -279,7 +290,8 @@ public void testDirectPathReleasesCircuitBreakerBytesImmediately() { ActionListener listener = SearchTransportService.asBytesResponse( transportService, channel, - response -> response.detachCircuitBreakerReservation(breaker) + response -> response.detachCircuitBreakerReservation(breaker), + breaker ); listener.onResponse(fetchResult); @@ -302,7 +314,8 @@ public void testTaskTransportChannelUnwrapsToDirectPath() { ActionListener listener = SearchTransportService.asBytesResponse( transportService, taskChannel, - response -> () -> releasableClosed.set(true) + response -> () -> releasableClosed.set(true), + null ); var original = new SimpleTestResponse("task-wrapped-test"); @@ -327,7 +340,7 @@ public void testTaskTransportChannelUnwrapsToNetworkPath() { ActionListener listener = SearchTransportService.asBytesResponse(transportService, taskChannel, response -> { detachCalled.set(true); return () -> releasableClosed.set(true); - }); + }, null); listener.onResponse(new SimpleTestResponse("task-network-test")); From 320cedf79aebcc787013867cd66fa594658d0148 Mon Sep 17 00:00:00 2001 From: elasticsearchmachine Date: Mon, 9 Mar 2026 09:00:52 +0000 Subject: [PATCH 20/39] [CI] Auto commit changes from spotless --- server/src/main/java/org/elasticsearch/transport/Transport.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/main/java/org/elasticsearch/transport/Transport.java b/server/src/main/java/org/elasticsearch/transport/Transport.java index c78e75db38f33..a8a53017681a4 100644 --- a/server/src/main/java/org/elasticsearch/transport/Transport.java +++ b/server/src/main/java/org/elasticsearch/transport/Transport.java @@ -12,8 +12,8 @@ import org.elasticsearch.TransportVersion; import org.elasticsearch.action.ActionListener; import org.elasticsearch.cluster.node.DiscoveryNode; -import org.elasticsearch.common.component.LifecycleComponent; import org.elasticsearch.common.breaker.CircuitBreaker; +import org.elasticsearch.common.component.LifecycleComponent; import org.elasticsearch.common.io.stream.RecyclerBytesStreamOutput; import org.elasticsearch.common.transport.BoundTransportAddress; import org.elasticsearch.common.transport.TransportAddress; From b75717ddc3163550b3193e12e841ecabb6b059a0 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Mon, 9 Mar 2026 12:36:45 +0200 Subject: [PATCH 21/39] update after review --- .../action/search/TransportSearchIT.java | 11 +++++++++-- .../hotthreads/TransportNodesHotThreadsAction.java | 2 +- .../search/SearchQueryThenFetchAsyncAction.java | 4 ++-- .../cluster/coordination/JoinValidationService.java | 2 +- .../coordination/PublicationTransportHandler.java | 4 ++-- .../org/elasticsearch/transport/TcpTransport.java | 8 ++------ .../java/org/elasticsearch/transport/Transport.java | 7 ++----- .../org/elasticsearch/transport/TransportService.java | 6 +----- .../PublicationTransportHandlerTests.java | 7 ++++--- .../snapshots/SnapshotResiliencyTestHelper.java | 5 +++-- .../transport/TransportActionProxyTests.java | 4 ++-- .../coordination/AbstractCoordinatorTestCase.java | 5 +++-- .../test/transport/StubbableTransport.java | 6 ++++-- 13 files changed, 36 insertions(+), 35 deletions(-) diff --git a/server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java b/server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java index 2127a19a9af99..bcac1fb772777 100644 --- a/server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java +++ b/server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java @@ -77,6 +77,7 @@ import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked; import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount; import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertResponse; +import static org.hamcrest.Matchers.anyOf; import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.equalTo; @@ -503,7 +504,10 @@ public void onFailure(Exception e) { Exception.class, client.prepareSearch("test").addAggregation(new TestAggregationBuilder("test")) ); - assertThat(exc.getCause().getMessage(), containsString("")); + assertThat( + exc.getCause().getMessage(), + anyOf(containsString(""), containsString("RecyclerBytesStreamOutput")) + ); }); final AtomicArray exceptions = new AtomicArray<>(10); @@ -530,7 +534,10 @@ public void onFailure(Exception exc) { latch.await(); assertThat(exceptions.asList().size(), equalTo(10)); for (Exception exc : exceptions.asList()) { - assertThat(exc.getCause().getMessage(), containsString("")); + assertThat( + exc.getCause().getMessage(), + anyOf(containsString(""), containsString("RecyclerBytesStreamOutput")) + ); } assertBusy(() -> assertThat(requestBreakerUsed(), equalTo(0L))); } finally { diff --git a/server/src/main/java/org/elasticsearch/action/admin/cluster/node/hotthreads/TransportNodesHotThreadsAction.java b/server/src/main/java/org/elasticsearch/action/admin/cluster/node/hotthreads/TransportNodesHotThreadsAction.java index 228dd76f1e2ac..06d4065bcce15 100644 --- a/server/src/main/java/org/elasticsearch/action/admin/cluster/node/hotthreads/TransportNodesHotThreadsAction.java +++ b/server/src/main/java/org/elasticsearch/action/admin/cluster/node/hotthreads/TransportNodesHotThreadsAction.java @@ -87,7 +87,7 @@ protected NodeHotThreads nodeOperation(NodeRequest request, Task task) { .interval(request.requestOptions.interval()) .threadElementsSnapshotCount(request.requestOptions.snapshots()) .ignoreIdleThreads(request.requestOptions.ignoreIdleThreads()); - final var out = transportService.newNetworkBytesStream(); + final var out = transportService.newNetworkBytesStream(null); final var trackedResource = LeakTracker.wrap(out); var success = false; try { diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchQueryThenFetchAsyncAction.java b/server/src/main/java/org/elasticsearch/action/search/SearchQueryThenFetchAsyncAction.java index a44f65da1b381..f4397eaed337f 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchQueryThenFetchAsyncAction.java +++ b/server/src/main/java/org/elasticsearch/action/search/SearchQueryThenFetchAsyncAction.java @@ -874,7 +874,7 @@ void onShardDone() { return; } var channelListener = new ChannelActionListener<>(channel); - RecyclerBytesStreamOutput out = dependencies.transportService.newNetworkBytesStream(); + RecyclerBytesStreamOutput out = dependencies.transportService.newNetworkBytesStream(null); out.setTransportVersion(channel.getVersion()); boolean success = false; @@ -997,7 +997,7 @@ void bwcRespond() { } } final int resultCount = queryPhaseResultConsumer.getNumShards(); - out = dependencies.transportService.newNetworkBytesStream(); + out = dependencies.transportService.newNetworkBytesStream(null); out.setTransportVersion(channel.getVersion()); try { out.writeVInt(resultCount); diff --git a/server/src/main/java/org/elasticsearch/cluster/coordination/JoinValidationService.java b/server/src/main/java/org/elasticsearch/cluster/coordination/JoinValidationService.java index 36e2209bd404d..f6f1b2d21f753 100644 --- a/server/src/main/java/org/elasticsearch/cluster/coordination/JoinValidationService.java +++ b/server/src/main/java/org/elasticsearch/cluster/coordination/JoinValidationService.java @@ -390,7 +390,7 @@ private ReleasableBytesReference maybeSerializeClusterState( } assert clusterState.nodes().isLocalNodeElectedMaster(); - try (var bytesStream = transportService.newNetworkBytesStream()) { + try (var bytesStream = transportService.newNetworkBytesStream(null)) { try (var stream = CompressorFactory.COMPRESSOR.threadLocalStreamOutput(Streams.flushOnCloseStream(bytesStream))) { stream.setTransportVersion(version); clusterState.writeTo(stream); diff --git a/server/src/main/java/org/elasticsearch/cluster/coordination/PublicationTransportHandler.java b/server/src/main/java/org/elasticsearch/cluster/coordination/PublicationTransportHandler.java index 20201a9fae9c1..6288cdbbe4b58 100644 --- a/server/src/main/java/org/elasticsearch/cluster/coordination/PublicationTransportHandler.java +++ b/server/src/main/java/org/elasticsearch/cluster/coordination/PublicationTransportHandler.java @@ -249,7 +249,7 @@ public PublicationContext newPublicationContext(ClusterStatePublicationEvent clu } private ReleasableBytesReference serializeFullClusterState(ClusterState clusterState, DiscoveryNode node, TransportVersion version) { - try (RecyclerBytesStreamOutput bytesStream = transportService.newNetworkBytesStream()) { + try (RecyclerBytesStreamOutput bytesStream = transportService.newNetworkBytesStream(null)) { final long uncompressedBytes; try (StreamOutput stream = CompressorFactory.COMPRESSOR.threadLocalStreamOutput(Streams.flushOnCloseStream(bytesStream))) { stream.setTransportVersion(version); @@ -278,7 +278,7 @@ private ReleasableBytesReference serializeDiffClusterState( TransportVersion version ) { final long clusterStateVersion = newState.version(); - try (RecyclerBytesStreamOutput bytesStream = transportService.newNetworkBytesStream()) { + try (RecyclerBytesStreamOutput bytesStream = transportService.newNetworkBytesStream(null)) { final long uncompressedBytes; try (StreamOutput stream = CompressorFactory.COMPRESSOR.threadLocalStreamOutput(Streams.flushOnCloseStream(bytesStream))) { stream.setTransportVersion(version); diff --git a/server/src/main/java/org/elasticsearch/transport/TcpTransport.java b/server/src/main/java/org/elasticsearch/transport/TcpTransport.java index 4e36d4b2e4d46..a8ef3c34847aa 100644 --- a/server/src/main/java/org/elasticsearch/transport/TcpTransport.java +++ b/server/src/main/java/org/elasticsearch/transport/TcpTransport.java @@ -43,6 +43,7 @@ import org.elasticsearch.common.util.PageCacheRecycler; import org.elasticsearch.common.util.concurrent.ConcurrentCollections; import org.elasticsearch.common.util.concurrent.CountDown; +import org.elasticsearch.core.Nullable; import org.elasticsearch.core.TimeValue; import org.elasticsearch.indices.breaker.CircuitBreakerService; import org.elasticsearch.monitor.jvm.JvmInfo; @@ -987,12 +988,7 @@ final Set getAcceptedChannels() { } @Override - public RecyclerBytesStreamOutput newNetworkBytesStream() { - return new RecyclerBytesStreamOutput(recycler); - } - - @Override - public RecyclerBytesStreamOutput newNetworkBytesStream(CircuitBreaker circuitBreaker) { + public RecyclerBytesStreamOutput newNetworkBytesStream(@Nullable CircuitBreaker circuitBreaker) { return new RecyclerBytesStreamOutput(recycler, circuitBreaker); } diff --git a/server/src/main/java/org/elasticsearch/transport/Transport.java b/server/src/main/java/org/elasticsearch/transport/Transport.java index a8a53017681a4..2838c5d2c1366 100644 --- a/server/src/main/java/org/elasticsearch/transport/Transport.java +++ b/server/src/main/java/org/elasticsearch/transport/Transport.java @@ -19,6 +19,7 @@ import org.elasticsearch.common.transport.TransportAddress; import org.elasticsearch.common.util.Maps; import org.elasticsearch.common.util.concurrent.ConcurrentCollections; +import org.elasticsearch.core.Nullable; import org.elasticsearch.core.RefCounted; import org.elasticsearch.core.TimeValue; @@ -89,11 +90,7 @@ default boolean isSecure() { RequestHandlers getRequestHandlers(); - default RecyclerBytesStreamOutput newNetworkBytesStream() { - return newNetworkBytesStream(null); - } - - default RecyclerBytesStreamOutput newNetworkBytesStream(CircuitBreaker circuitBreaker) { + default RecyclerBytesStreamOutput newNetworkBytesStream(@Nullable CircuitBreaker circuitBreaker) { return new RecyclerBytesStreamOutput(NON_RECYCLING_INSTANCE, circuitBreaker); } diff --git a/server/src/main/java/org/elasticsearch/transport/TransportService.java b/server/src/main/java/org/elasticsearch/transport/TransportService.java index e5d70b799746a..60f5ad3399da7 100644 --- a/server/src/main/java/org/elasticsearch/transport/TransportService.java +++ b/server/src/main/java/org/elasticsearch/transport/TransportService.java @@ -661,11 +661,7 @@ public ConnectionManager getConnectionManager() { return connectionManager; } - public RecyclerBytesStreamOutput newNetworkBytesStream() { - return transport.newNetworkBytesStream(); - } - - public RecyclerBytesStreamOutput newNetworkBytesStream(CircuitBreaker circuitBreaker) { + public RecyclerBytesStreamOutput newNetworkBytesStream(@Nullable CircuitBreaker circuitBreaker) { return transport.newNetworkBytesStream(circuitBreaker); } diff --git a/server/src/test/java/org/elasticsearch/cluster/coordination/PublicationTransportHandlerTests.java b/server/src/test/java/org/elasticsearch/cluster/coordination/PublicationTransportHandlerTests.java index bb56bd3ce4bf6..2460485cb808e 100644 --- a/server/src/test/java/org/elasticsearch/cluster/coordination/PublicationTransportHandlerTests.java +++ b/server/src/test/java/org/elasticsearch/cluster/coordination/PublicationTransportHandlerTests.java @@ -25,6 +25,7 @@ import org.elasticsearch.cluster.node.DiscoveryNodes; import org.elasticsearch.cluster.service.BatchSummary; import org.elasticsearch.common.UUIDs; +import org.elasticsearch.common.breaker.CircuitBreaker; import org.elasticsearch.common.compress.Compressor; import org.elasticsearch.common.compress.CompressorFactory; import org.elasticsearch.common.io.stream.RecyclerBytesStreamOutput; @@ -84,7 +85,7 @@ public void testDiffSerializationFailure() { final TransportService transportService = mock(TransportService.class); final BytesRefRecycler recycler = new BytesRefRecycler(new MockPageCacheRecycler(Settings.EMPTY)); - when(transportService.newNetworkBytesStream()).then(invocation -> new RecyclerBytesStreamOutput(recycler)); + when(transportService.newNetworkBytesStream(any())).then(invocation -> new RecyclerBytesStreamOutput(recycler)); Transport.Connection connection = mock(Transport.Connection.class); when(connection.getTransportVersion()).thenReturn(TransportVersion.current()); when(transportService.getConnection(any())).thenReturn(connection); @@ -206,8 +207,8 @@ protected void onSendRequest(long requestId, String action, TransportRequest req } @Override - public RecyclerBytesStreamOutput newNetworkBytesStream() { - return new RecyclerBytesStreamOutput(recycler); + public RecyclerBytesStreamOutput newNetworkBytesStream(@Nullable CircuitBreaker circuitBreaker) { + return new RecyclerBytesStreamOutput(recycler, circuitBreaker); } }; diff --git a/server/src/test/java/org/elasticsearch/snapshots/SnapshotResiliencyTestHelper.java b/server/src/test/java/org/elasticsearch/snapshots/SnapshotResiliencyTestHelper.java index 0fa47dc9f1712..d2b96d1a03331 100644 --- a/server/src/test/java/org/elasticsearch/snapshots/SnapshotResiliencyTestHelper.java +++ b/server/src/test/java/org/elasticsearch/snapshots/SnapshotResiliencyTestHelper.java @@ -80,6 +80,7 @@ import org.elasticsearch.cluster.service.FakeThreadPoolMasterService; import org.elasticsearch.cluster.service.MasterService; import org.elasticsearch.cluster.version.CompatibilityVersionsUtils; +import org.elasticsearch.common.breaker.CircuitBreaker; import org.elasticsearch.common.io.stream.NamedWriteableRegistry; import org.elasticsearch.common.io.stream.RecyclerBytesStreamOutput; import org.elasticsearch.common.logging.activity.ActivityLogWriterProvider; @@ -558,9 +559,9 @@ protected NamedWriteableRegistry writeableRegistry() { } @Override - public RecyclerBytesStreamOutput newNetworkBytesStream() { + public RecyclerBytesStreamOutput newNetworkBytesStream(CircuitBreaker circuitBreaker) { // skip leak checks in these tests since they do indeed leak - return new RecyclerBytesStreamOutput(BytesRefRecycler.NON_RECYCLING_INSTANCE); + return new RecyclerBytesStreamOutput(BytesRefRecycler.NON_RECYCLING_INSTANCE, circuitBreaker); // TODO fix these leaks and implement leak checking } }; diff --git a/server/src/test/java/org/elasticsearch/transport/TransportActionProxyTests.java b/server/src/test/java/org/elasticsearch/transport/TransportActionProxyTests.java index c4146634f10c6..48281ad2edb94 100644 --- a/server/src/test/java/org/elasticsearch/transport/TransportActionProxyTests.java +++ b/server/src/test/java/org/elasticsearch/transport/TransportActionProxyTests.java @@ -344,7 +344,7 @@ public void testSendLocalRequestBytesTransportResponseSameVersion() throws Excep assertEquals(request.sourceNode, "TS_A"); SimpleTestResponse tsB = new SimpleTestResponse("TS_B"); - try (RecyclerBytesStreamOutput out = serviceB.newNetworkBytesStream()) { + try (RecyclerBytesStreamOutput out = serviceB.newNetworkBytesStream(null)) { out.setTransportVersion(transportVersion1); tsB.writeTo(out); // simulate what happens in SearchQueryThenFetchAsyncAction with NodeQueryResponse @@ -417,7 +417,7 @@ public void testSendLocalRequestBytesTransportResponseDifferentVersions() throws assertEquals(request.sourceNode, "TS_A"); SimpleTestResponse tsB = new SimpleTestResponse("TS_B"); - try (RecyclerBytesStreamOutput out = serviceB.newNetworkBytesStream()) { + try (RecyclerBytesStreamOutput out = serviceB.newNetworkBytesStream(null)) { out.setTransportVersion(transportVersion1); tsB.writeTo(out); // simulate what happens in SearchQueryThenFetchAsyncAction with NodeQueryResponse diff --git a/test/framework/src/main/java/org/elasticsearch/cluster/coordination/AbstractCoordinatorTestCase.java b/test/framework/src/main/java/org/elasticsearch/cluster/coordination/AbstractCoordinatorTestCase.java index c207044b1ec72..7b1f5f8416f0a 100644 --- a/test/framework/src/main/java/org/elasticsearch/cluster/coordination/AbstractCoordinatorTestCase.java +++ b/test/framework/src/main/java/org/elasticsearch/cluster/coordination/AbstractCoordinatorTestCase.java @@ -48,6 +48,7 @@ import org.elasticsearch.common.Randomness; import org.elasticsearch.common.Strings; import org.elasticsearch.common.UUIDs; +import org.elasticsearch.common.breaker.CircuitBreaker; import org.elasticsearch.common.io.stream.BytesStreamOutput; import org.elasticsearch.common.io.stream.NamedWriteableAwareStreamInput; import org.elasticsearch.common.io.stream.NamedWriteableRegistry; @@ -1096,8 +1097,8 @@ public String toString() { } @Override - public RecyclerBytesStreamOutput newNetworkBytesStream() { - return new RecyclerBytesStreamOutput(clearableRecycler); + public RecyclerBytesStreamOutput newNetworkBytesStream(@Nullable CircuitBreaker circuitBreaker) { + return new RecyclerBytesStreamOutput(clearableRecycler, circuitBreaker); } }; diff --git a/test/framework/src/main/java/org/elasticsearch/test/transport/StubbableTransport.java b/test/framework/src/main/java/org/elasticsearch/test/transport/StubbableTransport.java index c5200c816fb6e..57cbbf7b78384 100644 --- a/test/framework/src/main/java/org/elasticsearch/test/transport/StubbableTransport.java +++ b/test/framework/src/main/java/org/elasticsearch/test/transport/StubbableTransport.java @@ -12,6 +12,7 @@ import org.elasticsearch.TransportVersion; import org.elasticsearch.action.ActionListener; import org.elasticsearch.cluster.node.DiscoveryNode; +import org.elasticsearch.common.breaker.CircuitBreaker; import org.elasticsearch.common.component.Lifecycle; import org.elasticsearch.common.component.LifecycleListener; import org.elasticsearch.common.io.stream.RecyclerBytesStreamOutput; @@ -20,6 +21,7 @@ import org.elasticsearch.common.transport.TransportAddress; import org.elasticsearch.common.util.MockPageCacheRecycler; import org.elasticsearch.common.util.PageCacheRecycler; +import org.elasticsearch.core.Nullable; import org.elasticsearch.tasks.Task; import org.elasticsearch.transport.BytesRefRecycler; import org.elasticsearch.transport.ConnectionProfile; @@ -343,7 +345,7 @@ default void clearCallback() {} } @Override - public RecyclerBytesStreamOutput newNetworkBytesStream() { - return new RecyclerBytesStreamOutput(new BytesRefRecycler(recycler)); + public RecyclerBytesStreamOutput newNetworkBytesStream(@Nullable CircuitBreaker circuitBreaker) { + return new RecyclerBytesStreamOutput(new BytesRefRecycler(recycler), circuitBreaker); } } From c15da98d8d97829227e9bdf273b6d441ca149085 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Mon, 9 Mar 2026 12:50:38 +0200 Subject: [PATCH 22/39] update after review --- .../action/search/SearchTransportService.java | 7 ++----- .../action/search/AsBytesResponseTests.java | 15 ++++++--------- 2 files changed, 8 insertions(+), 14 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java index 865e940f88057..a9dd40f8d931d 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java +++ b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java @@ -22,7 +22,6 @@ 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.RecyclerBytesStreamOutput; import org.elasticsearch.common.io.stream.StreamInput; @@ -799,16 +798,14 @@ public void onResponse(T response) { return; } var bytesRef = out.moveToBytesReference(); - Releasable breakerRelease; try { - breakerRelease = detachRelease.apply(response); + Releasables.close(detachRelease.apply(response)); } catch (Exception e) { Releasables.close(bytesRef); channelListener.onFailure(e); return; } - var trackingBytes = new ReleasableBytesReference(bytesRef, Releasables.wrap(bytesRef, breakerRelease)); - ActionListener.respondAndRelease(channelListener, new BytesTransportResponse(trackingBytes, out.getTransportVersion())); + ActionListener.respondAndRelease(channelListener, new BytesTransportResponse(bytesRef, out.getTransportVersion())); } @Override diff --git a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java index 4b4269248bdba..e099d6268bec7 100644 --- a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java @@ -120,12 +120,10 @@ public void testNetworkPathCallsDetachReleaseOnSuccess() { listener.onResponse(new SimpleTestResponse("hello")); assertTrue("detachRelease must be called after successful serialization", detachCalled.get()); - assertFalse("releasable must NOT be closed yet (deferred to bytes release)", releasableClosed.get()); + assertTrue("releasable must be closed eagerly after serialization", releasableClosed.get()); assertThat(sentResponse.get(), notNullValue()); assertThat(sentResponse.get(), instanceOf(BytesTransportResponse.class)); - sentResponse.get().decRef(); - assertTrue("releasable must be closed after bytes are released", releasableClosed.get()); } public void testNetworkPathReleasesImmediatelyOnSerializationFailure() { @@ -150,7 +148,7 @@ public void testNetworkPathReleasesImmediatelyOnSerializationFailure() { assertThat(sentException.get(), instanceOf(IOException.class)); } - public void testNetworkPathDefersCircuitBreakerReleaseUntilBytesReleased() { + public void testNetworkPathReleasesResponseBreakerEagerlyAndPageBreakerOnSend() { long responseBytes = 5000L; long pageBytes = PageCacheRecycler.BYTE_PAGE_SIZE; var breakerUsed = new AtomicLong(responseBytes); @@ -181,14 +179,14 @@ public void testNetworkPathDefersCircuitBreakerReleaseUntilBytesReleased() { assertThat("detach must zero out the field on the result", fetchResult.getSearchHitsSizeBytes(), equalTo(0L)); assertThat( - "breaker must account for both the detached response reservation and the serialized page bytes", + "response reservation must be released eagerly, only page bytes remain on the breaker", breakerUsed.get(), - equalTo(responseBytes + pageBytes) + equalTo(pageBytes) ); assertThat(sentResponse.get(), instanceOf(BytesTransportResponse.class)); sentResponse.get().decRef(); - assertThat("breaker must be released after bytes are released", breakerUsed.get(), equalTo(0L)); + assertThat("breaker must be fully released after bytes are sent", breakerUsed.get(), equalTo(0L)); } finally { fetchResult.decRef(); } @@ -345,10 +343,9 @@ public void testTaskTransportChannelUnwrapsToNetworkPath() { listener.onResponse(new SimpleTestResponse("task-network-test")); assertTrue("detachRelease must be called on task-wrapped network path", detachCalled.get()); - assertFalse("releasable must NOT be closed yet on task-wrapped network path", releasableClosed.get()); + assertTrue("releasable must be closed eagerly on task-wrapped network path", releasableClosed.get()); assertThat(sentResponse.get(), instanceOf(BytesTransportResponse.class)); sentResponse.get().decRef(); - assertTrue("releasable must be closed after bytes are released", releasableClosed.get()); } static class SimpleTestResponse extends TransportResponse { From b8da157bed5c9a68d997a97e67e2dca13d8c5e5d Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Mon, 9 Mar 2026 13:34:36 +0200 Subject: [PATCH 23/39] update after review --- .../action/search/TransportSearchIT.java | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java b/server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java index bcac1fb772777..331d46a045c1d 100644 --- a/server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java +++ b/server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java @@ -21,8 +21,10 @@ import org.elasticsearch.action.support.WriteRequest; import org.elasticsearch.client.internal.Client; import org.elasticsearch.cluster.metadata.IndexMetadata; +import org.elasticsearch.ExceptionsHelper; 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; @@ -77,7 +79,6 @@ import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked; import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount; import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertResponse; -import static org.hamcrest.Matchers.anyOf; import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.equalTo; @@ -504,9 +505,9 @@ public void onFailure(Exception e) { Exception.class, client.prepareSearch("test").addAggregation(new TestAggregationBuilder("test")) ); - assertThat( - exc.getCause().getMessage(), - anyOf(containsString(""), containsString("RecyclerBytesStreamOutput")) + assertNotNull( + "root cause must be a CircuitBreakingException", + ExceptionsHelper.unwrap(exc, CircuitBreakingException.class) ); }); @@ -534,9 +535,9 @@ public void onFailure(Exception exc) { latch.await(); assertThat(exceptions.asList().size(), equalTo(10)); for (Exception exc : exceptions.asList()) { - assertThat( - exc.getCause().getMessage(), - anyOf(containsString(""), containsString("RecyclerBytesStreamOutput")) + assertNotNull( + "root cause must be a CircuitBreakingException", + ExceptionsHelper.unwrap(exc, CircuitBreakingException.class) ); } assertBusy(() -> assertThat(requestBreakerUsed(), equalTo(0L))); From 391165c6a45c99d595be3c107682250e04f19644 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Mon, 9 Mar 2026 13:37:52 +0200 Subject: [PATCH 24/39] Update newNetworkBytesStream method signature --- .../elasticsearch/snapshots/SnapshotResiliencyTestHelper.java | 4 ---- 1 file changed, 4 deletions(-) diff --git a/server/src/test/java/org/elasticsearch/snapshots/SnapshotResiliencyTestHelper.java b/server/src/test/java/org/elasticsearch/snapshots/SnapshotResiliencyTestHelper.java index 481298ea8a0d1..439c43ad43228 100644 --- a/server/src/test/java/org/elasticsearch/snapshots/SnapshotResiliencyTestHelper.java +++ b/server/src/test/java/org/elasticsearch/snapshots/SnapshotResiliencyTestHelper.java @@ -560,11 +560,7 @@ protected NamedWriteableRegistry writeableRegistry() { } @Override -<<<<<<< fix/release-breaker-after-send-response - public RecyclerBytesStreamOutput newNetworkBytesStream(CircuitBreaker circuitBreaker) { -======= public RecyclerBytesStreamOutput newNetworkBytesStream(@Nullable CircuitBreaker circuitBreaker) { ->>>>>>> main // skip leak checks in these tests since they do indeed leak return new RecyclerBytesStreamOutput(BytesRefRecycler.NON_RECYCLING_INSTANCE, circuitBreaker); // TODO fix these leaks and implement leak checking From b825df8bedaed5c026838d51a50c424d246b96d9 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Mon, 9 Mar 2026 13:38:33 +0200 Subject: [PATCH 25/39] Fix mock behavior for newNetworkBytesStream method --- .../coordination/PublicationTransportHandlerTests.java | 4 ---- 1 file changed, 4 deletions(-) diff --git a/server/src/test/java/org/elasticsearch/cluster/coordination/PublicationTransportHandlerTests.java b/server/src/test/java/org/elasticsearch/cluster/coordination/PublicationTransportHandlerTests.java index acc9edb3f9903..e903d324da240 100644 --- a/server/src/test/java/org/elasticsearch/cluster/coordination/PublicationTransportHandlerTests.java +++ b/server/src/test/java/org/elasticsearch/cluster/coordination/PublicationTransportHandlerTests.java @@ -85,13 +85,9 @@ public void testDiffSerializationFailure() { final TransportService transportService = mock(TransportService.class); final BytesRefRecycler recycler = new BytesRefRecycler(new MockPageCacheRecycler(Settings.EMPTY)); -<<<<<<< fix/release-breaker-after-send-response - when(transportService.newNetworkBytesStream(any())).then(invocation -> new RecyclerBytesStreamOutput(recycler)); -======= when(transportService.newNetworkBytesStream(any())).then( invocation -> new RecyclerBytesStreamOutput(recycler, invocation.getArgument(0)) ); ->>>>>>> main Transport.Connection connection = mock(Transport.Connection.class); when(connection.getTransportVersion()).thenReturn(TransportVersion.current()); when(transportService.getConnection(any())).thenReturn(connection); From a74f69df739eb0cbc9fbdf5713cfad64bff8a8f0 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Mon, 9 Mar 2026 13:39:15 +0200 Subject: [PATCH 26/39] Fix merge conflict in TransportService.java Removed merge conflict markers and cleaned up code. --- .../java/org/elasticsearch/transport/TransportService.java | 3 --- 1 file changed, 3 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/transport/TransportService.java b/server/src/main/java/org/elasticsearch/transport/TransportService.java index bb94acd2bf442..2519720fd5fa6 100644 --- a/server/src/main/java/org/elasticsearch/transport/TransportService.java +++ b/server/src/main/java/org/elasticsearch/transport/TransportService.java @@ -661,15 +661,12 @@ public ConnectionManager getConnectionManager() { return connectionManager; } -<<<<<<< fix/release-breaker-after-send-response -======= /** * @return a {@link RecyclerBytesStreamOutput} which allocates its pages with {@code org.elasticsearch.transport.netty4.NettyAllocator}, * tracking these allocations using the provided {@link CircuitBreaker} if this is not {@code null}. *

* In tests in which Netty is not in use, each page is allocated as a {@code new byte[]}. */ ->>>>>>> main public RecyclerBytesStreamOutput newNetworkBytesStream(@Nullable CircuitBreaker circuitBreaker) { return transport.newNetworkBytesStream(circuitBreaker); } From 77974547bc36b37d1134a8560aa7c4de55d73e69 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Mon, 9 Mar 2026 13:40:12 +0200 Subject: [PATCH 27/39] Remove default implementation for newNetworkBytesStream Removed default implementation of newNetworkBytesStream method. --- .../main/java/org/elasticsearch/transport/Transport.java | 6 ------ 1 file changed, 6 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/transport/Transport.java b/server/src/main/java/org/elasticsearch/transport/Transport.java index 9fcebddf619e0..156f419f8eb84 100644 --- a/server/src/main/java/org/elasticsearch/transport/Transport.java +++ b/server/src/main/java/org/elasticsearch/transport/Transport.java @@ -88,11 +88,6 @@ default boolean isSecure() { RequestHandlers getRequestHandlers(); -<<<<<<< fix/release-breaker-after-send-response - default RecyclerBytesStreamOutput newNetworkBytesStream(@Nullable CircuitBreaker circuitBreaker) { - return new RecyclerBytesStreamOutput(NON_RECYCLING_INSTANCE, circuitBreaker); - } -======= /** * @return a {@link RecyclerBytesStreamOutput} which allocates its pages with {@code org.elasticsearch.transport.netty4.NettyAllocator}, * tracking these allocations using the provided {@link CircuitBreaker} if this is not {@code null}. @@ -100,7 +95,6 @@ default RecyclerBytesStreamOutput newNetworkBytesStream(@Nullable CircuitBreaker * In tests in which Netty is not in use, each page is allocated as a {@code new byte[]}. */ RecyclerBytesStreamOutput newNetworkBytesStream(@Nullable CircuitBreaker circuitBreaker); ->>>>>>> main /** * A unidirectional connection to a {@link DiscoveryNode} From 00767dd759c18ff3bcf818b48dea079b443bc5a3 Mon Sep 17 00:00:00 2001 From: elasticsearchmachine Date: Mon, 9 Mar 2026 12:10:35 +0000 Subject: [PATCH 28/39] [CI] Auto commit changes from spotless --- .../java/org/elasticsearch/action/search/TransportSearchIT.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java b/server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java index 331d46a045c1d..a45237acd01eb 100644 --- a/server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java +++ b/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; @@ -21,7 +22,6 @@ import org.elasticsearch.action.support.WriteRequest; import org.elasticsearch.client.internal.Client; import org.elasticsearch.cluster.metadata.IndexMetadata; -import org.elasticsearch.ExceptionsHelper; import org.elasticsearch.common.Strings; import org.elasticsearch.common.breaker.CircuitBreaker; import org.elasticsearch.common.breaker.CircuitBreakingException; From e8023dbd957631e07314df9550c345d4b4b5f0d6 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Tue, 17 Mar 2026 16:45:17 +0200 Subject: [PATCH 29/39] update after review --- .../action/search/SearchTransportService.java | 100 ++------ .../elasticsearch/search/SearchService.java | 10 +- .../search/fetch/FetchSearchResult.java | 24 +- .../action/search/AsBytesResponseTests.java | 220 +----------------- 4 files changed, 38 insertions(+), 316 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java index a9dd40f8d931d..1c46df81117fb 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java +++ b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java @@ -72,7 +72,6 @@ import java.util.Objects; import java.util.concurrent.Executor; import java.util.function.BiFunction; -import java.util.function.Function; /** * An encapsulation of {@link SearchService} operations exposed through @@ -464,12 +463,7 @@ public static void registerRequestHandler( (request, channel, task) -> searchService.executeQueryPhase( request, (SearchShardTask) task, - asBytesResponse(transportService, channel, response -> { - if (response instanceof QueryFetchSearchResult qfr) { - return qfr.fetchResult().detachCircuitBreakerReservation(searchService.getCircuitBreaker()); - } - return () -> {}; - }, searchService.getCircuitBreaker()) + asBytesResponse(transportService, channel, searchService.getCircuitBreaker()) ) ); TransportActionProxy.registerProxyActionWithDynamicResponseType( @@ -525,12 +519,7 @@ public static void registerRequestHandler( (request, channel, task) -> searchService.executeFetchPhase( request, (SearchShardTask) task, - asBytesResponse( - transportService, - channel, - response -> response.result().fetchResult().detachCircuitBreakerReservation(searchService.getCircuitBreaker()), - searchService.getCircuitBreaker() - ) + asBytesResponse(transportService, channel, searchService.getCircuitBreaker()) ) ); TransportActionProxy.registerProxyAction( @@ -561,12 +550,7 @@ public static void registerRequestHandler( .executeFetchPhase( request, (SearchShardTask) task, - asBytesResponse( - transportService, - channel, - response -> response.detachCircuitBreakerReservation(searchService.getCircuitBreaker()), - searchService.getCircuitBreaker() - ) + asBytesResponse(transportService, channel, searchService.getCircuitBreaker()) ); transportService.registerRequestHandler( FETCH_ID_SCROLL_ACTION_NAME, @@ -699,27 +683,24 @@ private boolean assertNodePresent() { } /** - * Returns a listener that ensures circuit-breaker bytes are not released until the response - * has actually been written to the network. + * Returns a listener that serializes responses to bytes on the network path. * - *

On the network path, the response is serialized into bytes up front. - * Then {@code detachRelease} is called to extract the breaker reservation from the response - * as a {@link Releasable}, and then is attached to the serialized bytes' ref-count, so it - * is only closed when Netty finishes writing those bytes to the wire. + *

On the network path, the response is serialized into bytes using a + * circuit-breaker-aware stream and sent as a {@link BytesTransportResponse}. * - *

On the direct (same-node) path the response is forwarded as-is and the {@code Releasable} - * is closed immediately after hand-off. + *

On the direct (same-node) path the response is forwarded as-is. + * + *

Circuit-breaker accounting for response objects is handled by the caller. */ static ActionListener asBytesResponse( TransportService transportService, TransportChannel channel, - Function detachRelease, @Nullable CircuitBreaker circuitBreaker ) { if (isDirectResponseChannel(channel)) { - return new DirectPathListener<>(channel, detachRelease); + return new ChannelActionListener<>(channel); } - return new NetworkPathListener<>(transportService, channel, detachRelease, circuitBreaker); + return new NetworkPathListener<>(transportService, channel, circuitBreaker); } private static boolean isDirectResponseChannel(TransportChannel channel) { @@ -729,58 +710,21 @@ private static boolean isDirectResponseChannel(TransportChannel channel) { return TransportService.isDirectResponseChannel(channel); } - /** - * Forwards the original Java response as-is and releases the circuit-breaker reservation immediately - * after the response has been handed off. - */ - private static class DirectPathListener implements ActionListener { - private final ChannelActionListener channelListener; - private final Function detachRelease; - - DirectPathListener(TransportChannel channel, Function detachRelease) { - this.channelListener = new ChannelActionListener<>(channel); - this.detachRelease = detachRelease; - } - - @Override - public void onResponse(T response) { - Releasable breakerRelease = null; - try { - channelListener.onResponse(response); - breakerRelease = detachRelease.apply(response); - } finally { - Releasables.close(breakerRelease); - } - } - - @Override - public void onFailure(Exception e) { - channelListener.onFailure(e); - } - } - /** * Serializes the response into a {@link BytesTransportResponse} and bundles the circuit-breaker - * release with the serialized bytes' ref-count. The breaker reservation stays alive as long as - * the bytes are on the heap, and is released when the ref-count drops to zero. + * accounting with the serialized bytes' lifecycle. The bytes stay on the heap until the + * {@link BytesTransportResponse} is released. */ private static class NetworkPathListener implements ActionListener { private final TransportService transportService; private final TransportChannel channel; - private final Function detachRelease; private final ChannelActionListener channelListener; @Nullable private final CircuitBreaker circuitBreaker; - NetworkPathListener( - TransportService transportService, - TransportChannel channel, - Function detachRelease, - @Nullable CircuitBreaker circuitBreaker - ) { + NetworkPathListener(TransportService transportService, TransportChannel channel, @Nullable CircuitBreaker circuitBreaker) { this.transportService = transportService; this.channel = channel; - this.detachRelease = detachRelease; this.channelListener = new ChannelActionListener<>(channel); this.circuitBreaker = circuitBreaker; } @@ -793,18 +737,10 @@ public void onResponse(T response) { response.writeTo(out); } catch (Exception e) { Releasables.close(out); - detachAndReleaseSafely(response, e); channelListener.onFailure(e); return; } var bytesRef = out.moveToBytesReference(); - try { - Releasables.close(detachRelease.apply(response)); - } catch (Exception e) { - Releasables.close(bytesRef); - channelListener.onFailure(e); - return; - } ActionListener.respondAndRelease(channelListener, new BytesTransportResponse(bytesRef, out.getTransportVersion())); } @@ -812,14 +748,6 @@ public void onResponse(T response) { public void onFailure(Exception e) { channelListener.onFailure(e); } - - private void detachAndReleaseSafely(T response, Exception primary) { - try { - Releasables.close(detachRelease.apply(response)); - } catch (Exception suppressed) { - primary.addSuppressed(suppressed); - } - } } public void cancelSearchTask(SearchTask task, String reason) { diff --git a/server/src/main/java/org/elasticsearch/search/SearchService.java b/server/src/main/java/org/elasticsearch/search/SearchService.java index b8914259931aa..45a8d054dddc4 100644 --- a/server/src/main/java/org/elasticsearch/search/SearchService.java +++ b/server/src/main/java/org/elasticsearch/search/SearchService.java @@ -1211,7 +1211,13 @@ public void executeFetchPhase( // we handle the failure in the failure listener below throw e; } - }, wrapFailureListener(listener, readerContext, markAsUsed)); + }, + wrapFailureListener( + releaseCircuitBreakerOnResponse(listener, result -> result.result().fetchResult()), + readerContext, + markAsUsed + ) + ); } public void executeFetchPhase(ShardFetchRequest request, CancellableTask task, ActionListener listener) { @@ -1250,7 +1256,7 @@ public void executeFetchPhase(ShardFetchRequest request, CancellableTask task, A // we handle the failure in the failure listener below throw e; } - }, wrapFailureListener(listener, readerContext, markAsUsed)); + }, wrapFailureListener(releaseCircuitBreakerOnResponse(listener, result -> result), readerContext, markAsUsed)); })); } diff --git a/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java b/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java index 8ddc8764bf05b..4ce6ba44a52f9 100644 --- a/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java +++ b/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java @@ -30,7 +30,7 @@ public final class FetchSearchResult extends SearchPhaseResult { private SearchHits hits; - private final transient AtomicLong searchHitsSizeBytes = new AtomicLong(0L); + private long searchHitsSizeBytes = 0L; // client side counter private transient int counter; @@ -90,32 +90,20 @@ public SearchHits hits() { } public void setSearchHitsSizeBytes(long bytes) { - this.searchHitsSizeBytes.set(bytes); + this.searchHitsSizeBytes = bytes; } public long getSearchHitsSizeBytes() { - return searchHitsSizeBytes.get(); + return searchHitsSizeBytes; } public void releaseCircuitBreakerBytes(CircuitBreaker circuitBreaker) { - long bytes = searchHitsSizeBytes.getAndSet(0L); - if (bytes > 0L) { - circuitBreaker.addWithoutBreaking(-bytes); + if (searchHitsSizeBytes > 0L) { + circuitBreaker.addWithoutBreaking(-searchHitsSizeBytes); + searchHitsSizeBytes = 0L; } } - /** - * Extracts the circuit-breaker reservation from this result and returns a - * {@link Releasable} that releases those bytes from the breaker when closed. - */ - public Releasable detachCircuitBreakerReservation(CircuitBreaker circuitBreaker) { - long bytes = searchHitsSizeBytes.getAndSet(0L); - if (bytes > 0L) { - return () -> circuitBreaker.addWithoutBreaking(-bytes); - } - return () -> {}; - } - public FetchSearchResult initCounter() { counter = 0; return this; diff --git a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java index e099d6268bec7..003620e77ebda 100644 --- a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java @@ -8,18 +8,10 @@ */ package org.elasticsearch.action.search; -import org.apache.lucene.search.TotalHits; import org.elasticsearch.action.ActionListener; -import org.elasticsearch.common.breaker.CircuitBreaker; -import org.elasticsearch.common.breaker.CircuitBreakingException; -import org.elasticsearch.common.breaker.NoopCircuitBreaker; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.common.settings.Settings; -import org.elasticsearch.common.util.PageCacheRecycler; -import org.elasticsearch.search.SearchHits; -import org.elasticsearch.search.fetch.FetchSearchResult; -import org.elasticsearch.search.internal.ShardSearchContextId; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.test.transport.MockTransport; import org.elasticsearch.threadpool.TestThreadPool; @@ -36,8 +28,6 @@ import java.io.IOException; import java.util.Collections; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; import static org.elasticsearch.cluster.node.DiscoveryNodeUtils.builder; @@ -83,12 +73,7 @@ public void testNetworkPathBytesResponseRoundTrip() throws Exception { }, e -> fail("unexpected failure: " + e))); var original = new SimpleTestResponse("test"); - ActionListener listener = SearchTransportService.asBytesResponse( - transportService, - channel, - response -> () -> {}, - null - ); + ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, null); listener.onResponse(original); @@ -102,231 +87,74 @@ public void testNetworkPathBytesResponseRoundTrip() throws Exception { } } - public void testNetworkPathCallsDetachReleaseOnSuccess() { - var detachCalled = new AtomicBoolean(false); - var releasableClosed = new AtomicBoolean(false); - var sentResponse = new AtomicReference(); - - var channel = new TestTransportChannel(ActionListener.wrap(resp -> { - resp.mustIncRef(); - sentResponse.set(resp); - }, e -> fail("unexpected failure: " + e))); - - ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, response -> { - detachCalled.set(true); - return () -> releasableClosed.set(true); - }, null); - - listener.onResponse(new SimpleTestResponse("hello")); - - assertTrue("detachRelease must be called after successful serialization", detachCalled.get()); - assertTrue("releasable must be closed eagerly after serialization", releasableClosed.get()); - assertThat(sentResponse.get(), notNullValue()); - assertThat(sentResponse.get(), instanceOf(BytesTransportResponse.class)); - sentResponse.get().decRef(); - } - - public void testNetworkPathReleasesImmediatelyOnSerializationFailure() { - var releasableClosed = new AtomicBoolean(false); + public void testNetworkPathSerializationFailureSendsFailure() { var sentException = new AtomicReference(); var channel = new TestTransportChannel( ActionListener.wrap(resp -> fail("should not succeed when serialization fails"), sentException::set) ); - ActionListener listener = SearchTransportService.asBytesResponse( - transportService, - channel, - response -> () -> releasableClosed.set(true), - null - ); + ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, null); listener.onResponse(new FailingTestResponse()); - assertTrue("releasable must be closed immediately on serialization failure", releasableClosed.get()); assertThat(sentException.get(), notNullValue()); assertThat(sentException.get(), instanceOf(IOException.class)); } - public void testNetworkPathReleasesResponseBreakerEagerlyAndPageBreakerOnSend() { - long responseBytes = 5000L; - long pageBytes = PageCacheRecycler.BYTE_PAGE_SIZE; - var breakerUsed = new AtomicLong(responseBytes); - CircuitBreaker breaker = new TestCircuitBreaker(breakerUsed); - - var sentResponse = new AtomicReference(); - - var channel = new TestTransportChannel(ActionListener.wrap(resp -> { - resp.mustIncRef(); - sentResponse.set(resp); - }, e -> fail("unexpected failure: " + e))); - - var contextId = new ShardSearchContextId("test-session", 1); - var fetchResult = new FetchSearchResult(contextId, null); - try { - var hits = SearchHits.empty(new TotalHits(0, TotalHits.Relation.EQUAL_TO), 0.0f); - fetchResult.shardResult(hits, null); - fetchResult.setSearchHitsSizeBytes(responseBytes); - - ActionListener listener = SearchTransportService.asBytesResponse( - transportService, - channel, - response -> response.detachCircuitBreakerReservation(breaker), - breaker - ); - - listener.onResponse(fetchResult); - - assertThat("detach must zero out the field on the result", fetchResult.getSearchHitsSizeBytes(), equalTo(0L)); - assertThat( - "response reservation must be released eagerly, only page bytes remain on the breaker", - breakerUsed.get(), - equalTo(pageBytes) - ); - assertThat(sentResponse.get(), instanceOf(BytesTransportResponse.class)); - - sentResponse.get().decRef(); - assertThat("breaker must be fully released after bytes are sent", breakerUsed.get(), equalTo(0L)); - } finally { - fetchResult.decRef(); - } - } - - public void testNetworkPathReleasesCircuitBreakerBytesOnSerializationFailure() { - var breakerUsed = new AtomicLong(3000); - var sentException = new AtomicReference(); - - var channel = new TestTransportChannel( - ActionListener.wrap(resp -> fail("should not succeed when serialization fails"), sentException::set) - ); - - ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, response -> { - long bytes = breakerUsed.get(); - return () -> { - if (bytes > 0L) { - breakerUsed.addAndGet(-bytes); - } - }; - }, null); - - listener.onResponse(new FailingTestResponse()); - - assertThat("breaker bytes should be released immediately on serialization failure", breakerUsed.get(), equalTo(0L)); - assertThat(sentException.get(), notNullValue()); - assertThat(sentException.get(), instanceOf(IOException.class)); - } - - public void testNetworkPathOnFailureDoesNotCallDetachRelease() { - var detachCalled = new AtomicBoolean(false); + public void testNetworkPathOnFailureForwardsFailure() { var sentException = new AtomicReference(); var channel = new TestTransportChannel(ActionListener.wrap(resp -> fail("should not succeed"), sentException::set)); - ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, response -> { - detachCalled.set(true); - return () -> {}; - }, null); + ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, null); listener.onFailure(new RuntimeException("upstream failure")); - assertFalse("detachRelease must not be called on upstream failure", detachCalled.get()); assertThat(sentException.get(), notNullValue()); } public void testDirectPathForwardsOriginalResponse() { - var releasableClosed = new AtomicBoolean(false); var sentResponse = new AtomicReference(); var channel = new TestDirectResponseChannel(ActionListener.wrap(sentResponse::set, e -> fail("unexpected failure: " + e))); - ActionListener listener = SearchTransportService.asBytesResponse( - transportService, - channel, - response -> () -> releasableClosed.set(true), - null - ); + ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, null); var original = new SimpleTestResponse("direct-test"); listener.onResponse(original); - assertTrue("releasable must be closed immediately on direct path", releasableClosed.get()); assertSame("direct path must forward the original response, not BytesTransportResponse", original, sentResponse.get()); } - public void testDirectPathOnFailureDoesNotCallDetachRelease() { - var detachCalled = new AtomicBoolean(false); + public void testDirectPathOnFailureForwardsFailure() { var sentException = new AtomicReference(); var channel = new TestDirectResponseChannel(ActionListener.wrap(resp -> fail("should not succeed"), sentException::set)); - ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, response -> { - detachCalled.set(true); - return () -> {}; - }, null); + ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, null); listener.onFailure(new RuntimeException("upstream failure")); - assertFalse("detachRelease must not be called on upstream failure", detachCalled.get()); assertThat(sentException.get(), notNullValue()); } - public void testDirectPathReleasesCircuitBreakerBytesImmediately() { - var breakerUsed = new AtomicLong(5000); - CircuitBreaker breaker = new TestCircuitBreaker(breakerUsed); - - var sentResponse = new AtomicReference(); - - var channel = new TestDirectResponseChannel(ActionListener.wrap(sentResponse::set, e -> fail("unexpected failure: " + e))); - - var contextId = new ShardSearchContextId("test-session", 1); - var fetchResult = new FetchSearchResult(contextId, null); - try { - var hits = SearchHits.empty(new TotalHits(0, TotalHits.Relation.EQUAL_TO), 0.0f); - fetchResult.shardResult(hits, null); - fetchResult.setSearchHitsSizeBytes(5000L); - - ActionListener listener = SearchTransportService.asBytesResponse( - transportService, - channel, - response -> response.detachCircuitBreakerReservation(breaker), - breaker - ); - - listener.onResponse(fetchResult); - - assertThat("breaker bytes should be released immediately on direct path", breakerUsed.get(), equalTo(0L)); - assertThat(fetchResult.getSearchHitsSizeBytes(), equalTo(0L)); - assertSame("direct path must forward original response", fetchResult, sentResponse.get()); - } finally { - fetchResult.decRef(); - } - } - public void testTaskTransportChannelUnwrapsToDirectPath() { - var releasableClosed = new AtomicBoolean(false); var sentResponse = new AtomicReference(); var directChannel = new TestDirectResponseChannel(ActionListener.wrap(sentResponse::set, e -> fail("unexpected failure: " + e))); var taskChannel = TestTransportChannels.newTaskTransportChannel(directChannel, () -> {}); - ActionListener listener = SearchTransportService.asBytesResponse( - transportService, - taskChannel, - response -> () -> releasableClosed.set(true), - null - ); + ActionListener listener = SearchTransportService.asBytesResponse(transportService, taskChannel, null); var original = new SimpleTestResponse("task-wrapped-test"); listener.onResponse(original); - assertTrue("releasable must be closed immediately on task-wrapped direct path", releasableClosed.get()); assertSame("task-wrapped direct channel must forward original response, not BytesTransportResponse", original, sentResponse.get()); sentResponse.get().decRef(); } public void testTaskTransportChannelUnwrapsToNetworkPath() { - var detachCalled = new AtomicBoolean(false); - var releasableClosed = new AtomicBoolean(false); var sentResponse = new AtomicReference(); var networkChannel = new TestTransportChannel(ActionListener.wrap(resp -> { @@ -335,15 +163,10 @@ public void testTaskTransportChannelUnwrapsToNetworkPath() { }, e -> fail("unexpected failure: " + e))); var taskChannel = TestTransportChannels.newTaskTransportChannel(networkChannel, () -> {}); - ActionListener listener = SearchTransportService.asBytesResponse(transportService, taskChannel, response -> { - detachCalled.set(true); - return () -> releasableClosed.set(true); - }, null); + ActionListener listener = SearchTransportService.asBytesResponse(transportService, taskChannel, null); listener.onResponse(new SimpleTestResponse("task-network-test")); - assertTrue("detachRelease must be called on task-wrapped network path", detachCalled.get()); - assertTrue("releasable must be closed eagerly on task-wrapped network path", releasableClosed.get()); assertThat(sentResponse.get(), instanceOf(BytesTransportResponse.class)); sentResponse.get().decRef(); } @@ -372,27 +195,4 @@ public void writeTo(StreamOutput out) throws IOException { } } - private static class TestCircuitBreaker extends NoopCircuitBreaker { - private final AtomicLong used; - - TestCircuitBreaker(AtomicLong used) { - super("test"); - this.used = used; - } - - @Override - public void addEstimateBytesAndMaybeBreak(long bytes, String label) throws CircuitBreakingException { - used.addAndGet(bytes); - } - - @Override - public void addWithoutBreaking(long bytes) { - used.addAndGet(bytes); - } - - @Override - public long getUsed() { - return used.get(); - } - } } From 506a922fec5a0f0c8e5a118d5bf7bc29d8f7e714 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Tue, 17 Mar 2026 16:45:56 +0200 Subject: [PATCH 30/39] update after review --- .../java/org/elasticsearch/search/fetch/FetchSearchResult.java | 2 -- 1 file changed, 2 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java b/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java index 4ce6ba44a52f9..6703bb5bd1920 100644 --- a/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java +++ b/server/src/main/java/org/elasticsearch/search/fetch/FetchSearchResult.java @@ -13,7 +13,6 @@ import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.core.RefCounted; -import org.elasticsearch.core.Releasable; import org.elasticsearch.core.SimpleRefCounted; import org.elasticsearch.search.SearchHit; import org.elasticsearch.search.SearchHits; @@ -24,7 +23,6 @@ import org.elasticsearch.transport.LeakTracker; import java.io.IOException; -import java.util.concurrent.atomic.AtomicLong; public final class FetchSearchResult extends SearchPhaseResult { From ff2a4c98383fa142356fc5f2b0b99d956385a980 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Tue, 17 Mar 2026 16:46:22 +0200 Subject: [PATCH 31/39] update after review --- .../src/main/java/org/elasticsearch/search/SearchService.java | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/server/src/main/java/org/elasticsearch/search/SearchService.java b/server/src/main/java/org/elasticsearch/search/SearchService.java index 45a8d054dddc4..2bf8ab3313719 100644 --- a/server/src/main/java/org/elasticsearch/search/SearchService.java +++ b/server/src/main/java/org/elasticsearch/search/SearchService.java @@ -1645,6 +1645,10 @@ private static void checkKeepAliveLimit(long keepAlive, long maxKeepAlive) { } } + /** + * Wraps a listener to release circuit breaker bytes from a FetchSearchResult after the response is sent. + * The fetchResultExtractor function extracts the FetchSearchResult from the response type. + */ private ActionListener releaseCircuitBreakerOnResponse( ActionListener listener, Function fetchResultExtractor From 8e78ca88598d7bac6d97bef710949df7dd923b6c Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Tue, 17 Mar 2026 17:03:31 +0200 Subject: [PATCH 32/39] update after review --- .../org/elasticsearch/action/search/SearchTransportService.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java index 1c46df81117fb..f2580fd05c9c3 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java +++ b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java @@ -731,6 +731,7 @@ private static class NetworkPathListener implements @Override public void onResponse(T response) { + // The stream output owns breaker-accounted bytes until moveToBytesReference transfers ownership to the response. RecyclerBytesStreamOutput out = transportService.newNetworkBytesStream(circuitBreaker); try { out.setTransportVersion(channel.getVersion()); @@ -741,6 +742,7 @@ public void onResponse(T response) { return; } var bytesRef = out.moveToBytesReference(); + // respondAndRelease releases the bytes once the transport layer completes. ActionListener.respondAndRelease(channelListener, new BytesTransportResponse(bytesRef, out.getTransportVersion())); } From cf6c2156fed8a28ec5bb2525e19d62560ba592bc Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Wed, 18 Mar 2026 09:39:23 +0200 Subject: [PATCH 33/39] update after review --- .../action/search/SearchTransportService.java | 27 +++++++++---------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java index f2580fd05c9c3..e5c2815b42b52 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java +++ b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java @@ -17,13 +17,14 @@ import org.elasticsearch.action.OriginalIndices; import org.elasticsearch.action.admin.cluster.node.tasks.cancel.CancelTasksRequest; import org.elasticsearch.action.admin.cluster.node.tasks.get.TransportGetTaskAction; +import org.elasticsearch.TransportVersion; import org.elasticsearch.action.support.ChannelActionListener; 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.RecyclerBytesStreamOutput; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.common.io.stream.Writeable; @@ -33,7 +34,6 @@ import org.elasticsearch.common.util.concurrent.ThrottledTaskRunner; import org.elasticsearch.core.Nullable; import org.elasticsearch.core.Releasable; -import org.elasticsearch.core.Releasables; import org.elasticsearch.search.SearchPhaseResult; import org.elasticsearch.search.SearchService; import org.elasticsearch.search.dfs.DfsSearchResult; @@ -711,39 +711,38 @@ private static boolean isDirectResponseChannel(TransportChannel channel) { } /** - * Serializes the response into a {@link BytesTransportResponse} and bundles the circuit-breaker - * accounting with the serialized bytes' lifecycle. The bytes stay on the heap until the - * {@link BytesTransportResponse} is released. + * 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 implements ActionListener { private final TransportService transportService; - private final TransportChannel channel; + private final TransportVersion transportVersion; private final ChannelActionListener channelListener; @Nullable private final CircuitBreaker circuitBreaker; NetworkPathListener(TransportService transportService, TransportChannel channel, @Nullable CircuitBreaker circuitBreaker) { this.transportService = transportService; - this.channel = channel; + this.transportVersion = channel.getVersion(); this.channelListener = new ChannelActionListener<>(channel); this.circuitBreaker = circuitBreaker; } @Override public void onResponse(T response) { - // The stream output owns breaker-accounted bytes until moveToBytesReference transfers ownership to the response. - RecyclerBytesStreamOutput out = transportService.newNetworkBytesStream(circuitBreaker); - try { - out.setTransportVersion(channel.getVersion()); + // 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) { - Releasables.close(out); channelListener.onFailure(e); return; } - var bytesRef = out.moveToBytesReference(); // respondAndRelease releases the bytes once the transport layer completes. - ActionListener.respondAndRelease(channelListener, new BytesTransportResponse(bytesRef, out.getTransportVersion())); + ActionListener.respondAndRelease(channelListener, new BytesTransportResponse(bytesRef, transportVersion)); } @Override From c7a08d1cc06c47e04e812ebb7ec6d0561c8e3884 Mon Sep 17 00:00:00 2001 From: elasticsearchmachine Date: Wed, 18 Mar 2026 07:48:34 +0000 Subject: [PATCH 34/39] [CI] Auto commit changes from spotless --- .../org/elasticsearch/action/search/SearchTransportService.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java index e5c2815b42b52..c2ec505187d6d 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java +++ b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java @@ -11,13 +11,13 @@ 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; import org.elasticsearch.action.OriginalIndices; import org.elasticsearch.action.admin.cluster.node.tasks.cancel.CancelTasksRequest; import org.elasticsearch.action.admin.cluster.node.tasks.get.TransportGetTaskAction; -import org.elasticsearch.TransportVersion; import org.elasticsearch.action.support.ChannelActionListener; import org.elasticsearch.client.internal.OriginSettingClient; import org.elasticsearch.client.internal.node.NodeClient; From 3fade55342c3f41ae0c6e37705f739c8f520c6c2 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Thu, 19 Mar 2026 14:43:27 +0200 Subject: [PATCH 35/39] update naming and small fixes after review --- .../action/search/SearchTransportService.java | 8 +++---- .../action/search/AsBytesResponseTests.java | 23 ++++++++----------- 2 files changed, 14 insertions(+), 17 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java index c2ec505187d6d..9774ba54d6b90 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java +++ b/server/src/main/java/org/elasticsearch/action/search/SearchTransportService.java @@ -463,7 +463,7 @@ public static void registerRequestHandler( (request, channel, task) -> searchService.executeQueryPhase( request, (SearchShardTask) task, - asBytesResponse(transportService, channel, searchService.getCircuitBreaker()) + channelListener(transportService, channel, searchService.getCircuitBreaker()) ) ); TransportActionProxy.registerProxyActionWithDynamicResponseType( @@ -519,7 +519,7 @@ public static void registerRequestHandler( (request, channel, task) -> searchService.executeFetchPhase( request, (SearchShardTask) task, - asBytesResponse(transportService, channel, searchService.getCircuitBreaker()) + channelListener(transportService, channel, searchService.getCircuitBreaker()) ) ); TransportActionProxy.registerProxyAction( @@ -550,7 +550,7 @@ public static void registerRequestHandler( .executeFetchPhase( request, (SearchShardTask) task, - asBytesResponse(transportService, channel, searchService.getCircuitBreaker()) + channelListener(transportService, channel, searchService.getCircuitBreaker()) ); transportService.registerRequestHandler( FETCH_ID_SCROLL_ACTION_NAME, @@ -692,7 +692,7 @@ private boolean assertNodePresent() { * *

Circuit-breaker accounting for response objects is handled by the caller. */ - static ActionListener asBytesResponse( + static ActionListener channelListener( TransportService transportService, TransportChannel channel, @Nullable CircuitBreaker circuitBreaker diff --git a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java index 003620e77ebda..8984bfe7a9ed7 100644 --- a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java @@ -12,6 +12,7 @@ import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.common.unit.ByteSizeValue; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.test.transport.MockTransport; import org.elasticsearch.threadpool.TestThreadPool; @@ -40,10 +41,8 @@ public class AsBytesResponseTests extends ESTestCase { private ThreadPool threadPool; private TransportService transportService; - @Override @Before - public void setUp() throws Exception { - super.setUp(); + public void setUpResources() throws Exception { threadPool = new TestThreadPool(getTestName()); var mockTransport = new MockTransport(); transportService = mockTransport.createTransportService( @@ -56,12 +55,10 @@ public void setUp() throws Exception { ); } - @Override @After - public void tearDown() throws Exception { + public void tearDownResources() throws Exception { transportService.close(); ThreadPool.terminate(threadPool, 30, TimeUnit.SECONDS); - super.tearDown(); } public void testNetworkPathBytesResponseRoundTrip() throws Exception { @@ -73,7 +70,7 @@ public void testNetworkPathBytesResponseRoundTrip() throws Exception { }, e -> fail("unexpected failure: " + e))); var original = new SimpleTestResponse("test"); - ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, null); + ActionListener listener = SearchTransportService.channelListener(transportService, channel, newLimitedBreaker(ByteSizeValue.ofMb(100))); listener.onResponse(original); @@ -94,7 +91,7 @@ public void testNetworkPathSerializationFailureSendsFailure() { ActionListener.wrap(resp -> fail("should not succeed when serialization fails"), sentException::set) ); - ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, null); + ActionListener listener = SearchTransportService.channelListener(transportService, channel, newLimitedBreaker(ByteSizeValue.ofMb(100))); listener.onResponse(new FailingTestResponse()); @@ -107,7 +104,7 @@ public void testNetworkPathOnFailureForwardsFailure() { var channel = new TestTransportChannel(ActionListener.wrap(resp -> fail("should not succeed"), sentException::set)); - ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, null); + ActionListener listener = SearchTransportService.channelListener(transportService, channel, newLimitedBreaker(ByteSizeValue.ofMb(100))); listener.onFailure(new RuntimeException("upstream failure")); @@ -119,7 +116,7 @@ public void testDirectPathForwardsOriginalResponse() { var channel = new TestDirectResponseChannel(ActionListener.wrap(sentResponse::set, e -> fail("unexpected failure: " + e))); - ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, null); + ActionListener listener = SearchTransportService.channelListener(transportService, channel, newLimitedBreaker(ByteSizeValue.ofMb(100))); var original = new SimpleTestResponse("direct-test"); listener.onResponse(original); @@ -132,7 +129,7 @@ public void testDirectPathOnFailureForwardsFailure() { var channel = new TestDirectResponseChannel(ActionListener.wrap(resp -> fail("should not succeed"), sentException::set)); - ActionListener listener = SearchTransportService.asBytesResponse(transportService, channel, null); + ActionListener listener = SearchTransportService.channelListener(transportService, channel, newLimitedBreaker(ByteSizeValue.ofMb(100))); listener.onFailure(new RuntimeException("upstream failure")); @@ -145,7 +142,7 @@ public void testTaskTransportChannelUnwrapsToDirectPath() { var directChannel = new TestDirectResponseChannel(ActionListener.wrap(sentResponse::set, e -> fail("unexpected failure: " + e))); var taskChannel = TestTransportChannels.newTaskTransportChannel(directChannel, () -> {}); - ActionListener listener = SearchTransportService.asBytesResponse(transportService, taskChannel, null); + ActionListener listener = SearchTransportService.channelListener(transportService, taskChannel, newLimitedBreaker(ByteSizeValue.ofMb(100))); var original = new SimpleTestResponse("task-wrapped-test"); listener.onResponse(original); @@ -163,7 +160,7 @@ public void testTaskTransportChannelUnwrapsToNetworkPath() { }, e -> fail("unexpected failure: " + e))); var taskChannel = TestTransportChannels.newTaskTransportChannel(networkChannel, () -> {}); - ActionListener listener = SearchTransportService.asBytesResponse(transportService, taskChannel, null); + ActionListener listener = SearchTransportService.channelListener(transportService, taskChannel, newLimitedBreaker(ByteSizeValue.ofMb(100))); listener.onResponse(new SimpleTestResponse("task-network-test")); From a89f2d2e47593a97c2932692a4c7b343ccecffe7 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Thu, 19 Mar 2026 14:50:01 +0200 Subject: [PATCH 36/39] update after review --- .../action/search/AsBytesResponseTests.java | 43 ++++++++++++++++--- 1 file changed, 36 insertions(+), 7 deletions(-) diff --git a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java index 8984bfe7a9ed7..ce38a0c5d11c0 100644 --- a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java @@ -70,7 +70,11 @@ public void testNetworkPathBytesResponseRoundTrip() throws Exception { }, e -> fail("unexpected failure: " + e))); var original = new SimpleTestResponse("test"); - ActionListener listener = SearchTransportService.channelListener(transportService, channel, newLimitedBreaker(ByteSizeValue.ofMb(100))); + ActionListener listener = SearchTransportService.channelListener( + transportService, + channel, + newLimitedBreaker(ByteSizeValue.ofMb(100)) + ); listener.onResponse(original); @@ -91,12 +95,17 @@ public void testNetworkPathSerializationFailureSendsFailure() { ActionListener.wrap(resp -> fail("should not succeed when serialization fails"), sentException::set) ); - ActionListener listener = SearchTransportService.channelListener(transportService, channel, newLimitedBreaker(ByteSizeValue.ofMb(100))); + ActionListener listener = SearchTransportService.channelListener( + transportService, + channel, + 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")); } public void testNetworkPathOnFailureForwardsFailure() { @@ -104,7 +113,11 @@ public void testNetworkPathOnFailureForwardsFailure() { var channel = new TestTransportChannel(ActionListener.wrap(resp -> fail("should not succeed"), sentException::set)); - ActionListener listener = SearchTransportService.channelListener(transportService, channel, newLimitedBreaker(ByteSizeValue.ofMb(100))); + ActionListener listener = SearchTransportService.channelListener( + transportService, + channel, + newLimitedBreaker(ByteSizeValue.ofMb(100)) + ); listener.onFailure(new RuntimeException("upstream failure")); @@ -116,7 +129,11 @@ public void testDirectPathForwardsOriginalResponse() { var channel = new TestDirectResponseChannel(ActionListener.wrap(sentResponse::set, e -> fail("unexpected failure: " + e))); - ActionListener listener = SearchTransportService.channelListener(transportService, channel, newLimitedBreaker(ByteSizeValue.ofMb(100))); + ActionListener listener = SearchTransportService.channelListener( + transportService, + channel, + newLimitedBreaker(ByteSizeValue.ofMb(100)) + ); var original = new SimpleTestResponse("direct-test"); listener.onResponse(original); @@ -129,7 +146,11 @@ public void testDirectPathOnFailureForwardsFailure() { var channel = new TestDirectResponseChannel(ActionListener.wrap(resp -> fail("should not succeed"), sentException::set)); - ActionListener listener = SearchTransportService.channelListener(transportService, channel, newLimitedBreaker(ByteSizeValue.ofMb(100))); + ActionListener listener = SearchTransportService.channelListener( + transportService, + channel, + newLimitedBreaker(ByteSizeValue.ofMb(100)) + ); listener.onFailure(new RuntimeException("upstream failure")); @@ -142,7 +163,11 @@ public void testTaskTransportChannelUnwrapsToDirectPath() { var directChannel = new TestDirectResponseChannel(ActionListener.wrap(sentResponse::set, e -> fail("unexpected failure: " + e))); var taskChannel = TestTransportChannels.newTaskTransportChannel(directChannel, () -> {}); - ActionListener listener = SearchTransportService.channelListener(transportService, taskChannel, newLimitedBreaker(ByteSizeValue.ofMb(100))); + ActionListener listener = SearchTransportService.channelListener( + transportService, + taskChannel, + newLimitedBreaker(ByteSizeValue.ofMb(100)) + ); var original = new SimpleTestResponse("task-wrapped-test"); listener.onResponse(original); @@ -160,7 +185,11 @@ public void testTaskTransportChannelUnwrapsToNetworkPath() { }, e -> fail("unexpected failure: " + e))); var taskChannel = TestTransportChannels.newTaskTransportChannel(networkChannel, () -> {}); - ActionListener listener = SearchTransportService.channelListener(transportService, taskChannel, newLimitedBreaker(ByteSizeValue.ofMb(100))); + ActionListener listener = SearchTransportService.channelListener( + transportService, + taskChannel, + newLimitedBreaker(ByteSizeValue.ofMb(100)) + ); listener.onResponse(new SimpleTestResponse("task-network-test")); From 27a750d366daa64818b7a4dc6d8e2c5b37d5e04b Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Thu, 19 Mar 2026 14:53:10 +0200 Subject: [PATCH 37/39] update after review --- .../action/search/AsBytesResponseTests.java | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java index ce38a0c5d11c0..e87db0c25a571 100644 --- a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java @@ -35,6 +35,7 @@ import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.instanceOf; import static org.hamcrest.Matchers.notNullValue; +import static org.hamcrest.Matchers.sameInstance; public class AsBytesResponseTests extends ESTestCase { @@ -119,9 +120,11 @@ public void testNetworkPathOnFailureForwardsFailure() { newLimitedBreaker(ByteSizeValue.ofMb(100)) ); - listener.onFailure(new RuntimeException("upstream failure")); + RuntimeException e = new RuntimeException("upstream failure"); + listener.onFailure(e); assertThat(sentException.get(), notNullValue()); + assertThat(sentException.get(), sameInstance(e)); } public void testDirectPathForwardsOriginalResponse() { @@ -152,9 +155,11 @@ public void testDirectPathOnFailureForwardsFailure() { newLimitedBreaker(ByteSizeValue.ofMb(100)) ); - listener.onFailure(new RuntimeException("upstream failure")); + RuntimeException e = new RuntimeException("upstream failure"); + listener.onFailure(e); assertThat(sentException.get(), notNullValue()); + assertThat(sentException.get(), sameInstance(e)); } public void testTaskTransportChannelUnwrapsToDirectPath() { From 3be8c42e858e902749687963206190f8d25b76d8 Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Thu, 19 Mar 2026 15:09:47 +0200 Subject: [PATCH 38/39] update after review --- .../elasticsearch/action/search/AsBytesResponseTests.java | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java index e87db0c25a571..f7611a24e0fc8 100644 --- a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java @@ -34,6 +34,7 @@ import static org.elasticsearch.cluster.node.DiscoveryNodeUtils.builder; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.not; import static org.hamcrest.Matchers.notNullValue; import static org.hamcrest.Matchers.sameInstance; @@ -141,7 +142,8 @@ public void testDirectPathForwardsOriginalResponse() { var original = new SimpleTestResponse("direct-test"); listener.onResponse(original); - assertSame("direct path must forward the original response, not BytesTransportResponse", original, sentResponse.get()); + assertSame(original, sentResponse.get()); + assertThat(sentResponse.get(), not(instanceOf(BytesTransportResponse.class))); } public void testDirectPathOnFailureForwardsFailure() { @@ -177,7 +179,8 @@ public void testTaskTransportChannelUnwrapsToDirectPath() { var original = new SimpleTestResponse("task-wrapped-test"); listener.onResponse(original); - assertSame("task-wrapped direct channel must forward original response, not BytesTransportResponse", original, sentResponse.get()); + assertSame(original, sentResponse.get()); + assertThat(sentResponse.get(), not(instanceOf(BytesTransportResponse.class))); sentResponse.get().decRef(); } From cc26220de8179073d9376e9f89a59c7bb7bc914d Mon Sep 17 00:00:00 2001 From: Dimitris Rempapis Date: Thu, 19 Mar 2026 15:15:36 +0200 Subject: [PATCH 39/39] update class name --- ...sts.java => SearchTransportServiceChannelListenerTests.java} | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) rename server/src/test/java/org/elasticsearch/action/search/{AsBytesResponseTests.java => SearchTransportServiceChannelListenerTests.java} (99%) diff --git a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java b/server/src/test/java/org/elasticsearch/action/search/SearchTransportServiceChannelListenerTests.java similarity index 99% rename from server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java rename to server/src/test/java/org/elasticsearch/action/search/SearchTransportServiceChannelListenerTests.java index f7611a24e0fc8..297a48c0ace29 100644 --- a/server/src/test/java/org/elasticsearch/action/search/AsBytesResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/SearchTransportServiceChannelListenerTests.java @@ -38,7 +38,7 @@ import static org.hamcrest.Matchers.notNullValue; import static org.hamcrest.Matchers.sameInstance; -public class AsBytesResponseTests extends ESTestCase { +public class SearchTransportServiceChannelListenerTests extends ESTestCase { private ThreadPool threadPool; private TransportService transportService;