From bce10bb2665273bab627c0bdf25fe53b77dc2702 Mon Sep 17 00:00:00 2001 From: Laura Trotta Date: Thu, 17 Sep 2026 15:59:33 +0200 Subject: [PATCH 1/7] keep feeding the encoder while it accepts bytes in multibuffer entity --- .../rest5_client/MultiBufferEntity.java | 29 +++++++++++++------ .../rest_client/MultiBufferEntity.java | 21 +++++++++----- 2 files changed, 33 insertions(+), 17 deletions(-) diff --git a/java-client/src/main/java/co/elastic/clients/transport/rest5_client/MultiBufferEntity.java b/java-client/src/main/java/co/elastic/clients/transport/rest5_client/MultiBufferEntity.java index 742753036b..f1f8955fcf 100644 --- a/java-client/src/main/java/co/elastic/clients/transport/rest5_client/MultiBufferEntity.java +++ b/java-client/src/main/java/co/elastic/clients/transport/rest5_client/MultiBufferEntity.java @@ -108,7 +108,8 @@ public void writeTo(OutputStream out) throws IOException { @Override public int available() { - return currentBuffer.remaining(); + ByteBuffer buffer = currentBuffer; + return buffer == null ? 0 : buffer.remaining(); } @Override @@ -118,14 +119,24 @@ public void produce(DataStreamChannel channel) throws IOException { return; } - channel.write(currentBuffer); - - if (!currentBuffer.hasRemaining()) { - if (iterator.hasNext()) { - currentBuffer = iterator.next().duplicate(); - } else { - currentBuffer = null; - channel.endStream(); + // Keep feeding buffers to the channel as long as it accepts them, so a single output event + // ships many small buffers (e.g. NDJSON bulk lines) instead of one buffer per event. Over + // HTTP/1.1 the chunk encoder accumulates them into its session buffer and flushes it when + // full. A zero return with bytes still pending means the channel cannot take more right now + // (socket buffer full, or HTTP/2 frame/flow-control window exhausted): stop and wait for + // the next output event. + while (currentBuffer != null) { + int written = channel.write(currentBuffer); + + if (!currentBuffer.hasRemaining()) { + if (iterator.hasNext()) { + currentBuffer = iterator.next().duplicate(); + } else { + currentBuffer = null; + channel.endStream(); + } + } else if (written == 0) { + break; } } } diff --git a/java-client/src/main/java/co/elastic/clients/transport/rest_client/MultiBufferEntity.java b/java-client/src/main/java/co/elastic/clients/transport/rest_client/MultiBufferEntity.java index cebbebd55e..f3ffea68b4 100644 --- a/java-client/src/main/java/co/elastic/clients/transport/rest_client/MultiBufferEntity.java +++ b/java-client/src/main/java/co/elastic/clients/transport/rest_client/MultiBufferEntity.java @@ -107,14 +107,19 @@ public void produceContent(ContentEncoder encoder, IOControl ioControl) throws I return; } - encoder.write(currentBuffer); - - if (!currentBuffer.hasRemaining()) { - if (iterator.hasNext()) { - currentBuffer = iterator.next().duplicate(); - } else { - currentBuffer = null; - encoder.complete(); + // Keep feeding buffers to the encoder as long as it accepts them. + while (currentBuffer != null) { + int written = encoder.write(currentBuffer); + + if (!currentBuffer.hasRemaining()) { + if (iterator.hasNext()) { + currentBuffer = iterator.next().duplicate(); + } else { + currentBuffer = null; + encoder.complete(); + } + } else if (written == 0) { + break; } } } From e8e46e87d5e689cf9bdb2981de347f5c5bce25dc Mon Sep 17 00:00:00 2001 From: Laura Trotta Date: Thu, 17 Sep 2026 16:55:25 +0200 Subject: [PATCH 2/7] rest5 MultiBufferEntity: coalesce small buffers into a 16 KiB staging buffer to fix bug for http2 --- .../rest5_client/MultiBufferEntity.java | 110 +++++++++++++----- 1 file changed, 83 insertions(+), 27 deletions(-) diff --git a/java-client/src/main/java/co/elastic/clients/transport/rest5_client/MultiBufferEntity.java b/java-client/src/main/java/co/elastic/clients/transport/rest5_client/MultiBufferEntity.java index f1f8955fcf..05c03ad125 100644 --- a/java-client/src/main/java/co/elastic/clients/transport/rest5_client/MultiBufferEntity.java +++ b/java-client/src/main/java/co/elastic/clients/transport/rest5_client/MultiBufferEntity.java @@ -39,10 +39,23 @@ */ class MultiBufferEntity extends AbstractHttpEntity implements AsyncEntityProducer { + /** + * Size of the buffer small input buffers are coalesced into before being handed to the channel. + *

+ * Each {@link DataStreamChannel#write} becomes at most one HTTP/2 DATA frame, and httpcore5's + * multiplexer emits a single frame per output event, so the frame payload size (16 KiB by default, + * RFC 9113 §4.2) is the largest write that can be fully used. Over HTTP/1.1 it also amortizes the + * per-write chunk header and session buffer maintenance. + */ + static final int STAGING_BUFFER_SIZE = 16 * 1024; + private final Iterable buffers; private Iterator iterator; + /** Input buffer currently being consumed, {@code null} once the input is exhausted. Never empty. */ private volatile ByteBuffer currentBuffer; + /** Coalesced bytes not yet accepted by the channel, always in read mode. Allocated on first use. */ + private volatile ByteBuffer staging; private final AtomicReference exceptionRef; @@ -61,11 +74,24 @@ public void close() throws IOException { private void init() { this.iterator = buffers.iterator(); - if (this.iterator.hasNext()) { - this.currentBuffer = this.iterator.next().duplicate(); - } else { - this.currentBuffer = null; + this.currentBuffer = null; + advance(); + ByteBuffer staging = this.staging; + if (staging != null) { + staging.position(0).limit(0); + } + } + + /** Moves {@link #currentBuffer} to the next non-empty input buffer, or {@code null} if there is none. */ + private void advance() { + while (iterator.hasNext()) { + ByteBuffer next = iterator.next(); + if (next.hasRemaining()) { + currentBuffer = next.duplicate(); + return; + } } + currentBuffer = null; } @Override @@ -108,37 +134,67 @@ public void writeTo(OutputStream out) throws IOException { @Override public int available() { - ByteBuffer buffer = currentBuffer; - return buffer == null ? 0 : buffer.remaining(); + ByteBuffer staging = this.staging; + ByteBuffer current = this.currentBuffer; + return (staging == null ? 0 : staging.remaining()) + (current == null ? 0 : current.remaining()); } @Override public void produce(DataStreamChannel channel) throws IOException { - if (currentBuffer == null) { - channel.endStream(); - return; + // The input is typically many small buffers (e.g. one per NDJSON line of a bulk request). Each + // write() below is at most one HTTP/2 DATA frame and, over HTTP/1.1, one chunk, so writing + // them one by one is very expensive (issue #1339). Coalesce them into a staging buffer and + // write that instead; input buffers that are already large are written directly. + // + // Keep going as long as the channel accepts everything: over HTTP/1.1 the encoder buffers up + // to the socket send buffer, over HTTP/2 the multiplexer takes exactly one frame per output + // event. Once a write is refused or partial (socket buffer full, frame or flow-control window + // exhausted), return and wait for the next output event. + while (true) { + ByteBuffer out; + if (staging != null && staging.hasRemaining()) { + out = staging; + } else if (currentBuffer == null) { + channel.endStream(); + return; + } else if (currentBuffer.remaining() >= STAGING_BUFFER_SIZE) { + out = currentBuffer; + } else { + out = fillStaging(); + } + + channel.write(out); + + if (out == currentBuffer && !out.hasRemaining()) { + advance(); + } + if (out.hasRemaining()) { + return; + } } + } - // Keep feeding buffers to the channel as long as it accepts them, so a single output event - // ships many small buffers (e.g. NDJSON bulk lines) instead of one buffer per event. Over - // HTTP/1.1 the chunk encoder accumulates them into its session buffer and flushes it when - // full. A zero return with bytes still pending means the channel cannot take more right now - // (socket buffer full, or HTTP/2 frame/flow-control window exhausted): stop and wait for - // the next output event. - while (currentBuffer != null) { - int written = channel.write(currentBuffer); - - if (!currentBuffer.hasRemaining()) { - if (iterator.hasNext()) { - currentBuffer = iterator.next().duplicate(); - } else { - currentBuffer = null; - channel.endStream(); - } - } else if (written == 0) { - break; + /** Copies as much input as fits into the (empty) staging buffer and returns it, ready to be read. */ + private ByteBuffer fillStaging() { + ByteBuffer staging = this.staging; + if (staging == null) { + staging = ByteBuffer.allocate(STAGING_BUFFER_SIZE); + this.staging = staging; + } + staging.clear(); + while (currentBuffer != null && staging.hasRemaining()) { + ByteBuffer src = currentBuffer; + int n = Math.min(src.remaining(), staging.remaining()); + int srcLimit = src.limit(); + src.limit(src.position() + n); + staging.put(src); + src.limit(srcLimit); + if (!src.hasRemaining()) { + advance(); } } + staging.flip(); + return staging; } @Override From 09095a93ceb36d2adc9112c9a73212927d9b5091 Mon Sep 17 00:00:00 2001 From: Laura Trotta Date: Thu, 5 Nov 2026 11:48:27 +0100 Subject: [PATCH 3/7] avoid buffering single body requests in staging buffer --- .../transport/rest5_client/MultiBufferEntity.java | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/java-client/src/main/java/co/elastic/clients/transport/rest5_client/MultiBufferEntity.java b/java-client/src/main/java/co/elastic/clients/transport/rest5_client/MultiBufferEntity.java index 05c03ad125..858947d027 100644 --- a/java-client/src/main/java/co/elastic/clients/transport/rest5_client/MultiBufferEntity.java +++ b/java-client/src/main/java/co/elastic/clients/transport/rest5_client/MultiBufferEntity.java @@ -40,12 +40,9 @@ class MultiBufferEntity extends AbstractHttpEntity implements AsyncEntityProducer { /** - * Size of the buffer small input buffers are coalesced into before being handed to the channel. - *

- * Each {@link DataStreamChannel#write} becomes at most one HTTP/2 DATA frame, and httpcore5's - * multiplexer emits a single frame per output event, so the frame payload size (16 KiB by default, - * RFC 9113 §4.2) is the largest write that can be fully used. Over HTTP/1.1 it also amortizes the - * per-write chunk header and session buffer maintenance. + * Small input buffers are coalesced up to this size before writing. Matches the default HTTP/2 + * frame payload size (RFC 9113 §4.2), since each {@link DataStreamChannel#write} becomes one frame. + * The staging buffer is never allocated for single-buffer bodies, such as most non-bulk requests. */ static final int STAGING_BUFFER_SIZE = 16 * 1024; @@ -144,7 +141,7 @@ public void produce(DataStreamChannel channel) throws IOException { // The input is typically many small buffers (e.g. one per NDJSON line of a bulk request). Each // write() below is at most one HTTP/2 DATA frame and, over HTTP/1.1, one chunk, so writing // them one by one is very expensive (issue #1339). Coalesce them into a staging buffer and - // write that instead; input buffers that are already large are written directly. + // write that instead; input buffers that are already large, or the last one, are written directly. // // Keep going as long as the channel accepts everything: over HTTP/1.1 the encoder buffers up // to the socket send buffer, over HTTP/2 the multiplexer takes exactly one frame per output @@ -157,7 +154,7 @@ public void produce(DataStreamChannel channel) throws IOException { } else if (currentBuffer == null) { channel.endStream(); return; - } else if (currentBuffer.remaining() >= STAGING_BUFFER_SIZE) { + } else if (currentBuffer.remaining() >= STAGING_BUFFER_SIZE || !iterator.hasNext()) { out = currentBuffer; } else { out = fillStaging(); From 5a7db3174f07ab158674f1f0964f139e0790b038 Mon Sep 17 00:00:00 2001 From: Laura Trotta Date: Thu, 5 Nov 2026 12:16:02 +0100 Subject: [PATCH 4/7] better explain rest5 mbe produce --- .../transport/rest5_client/MultiBufferEntity.java | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/java-client/src/main/java/co/elastic/clients/transport/rest5_client/MultiBufferEntity.java b/java-client/src/main/java/co/elastic/clients/transport/rest5_client/MultiBufferEntity.java index 858947d027..49cee1f7ae 100644 --- a/java-client/src/main/java/co/elastic/clients/transport/rest5_client/MultiBufferEntity.java +++ b/java-client/src/main/java/co/elastic/clients/transport/rest5_client/MultiBufferEntity.java @@ -138,15 +138,12 @@ public int available() { @Override public void produce(DataStreamChannel channel) throws IOException { - // The input is typically many small buffers (e.g. one per NDJSON line of a bulk request). Each - // write() below is at most one HTTP/2 DATA frame and, over HTTP/1.1, one chunk, so writing - // them one by one is very expensive (issue #1339). Coalesce them into a staging buffer and - // write that instead; input buffers that are already large, or the last one, are written directly. - // - // Keep going as long as the channel accepts everything: over HTTP/1.1 the encoder buffers up - // to the socket send buffer, over HTTP/2 the multiplexer takes exactly one frame per output - // event. Once a write is refused or partial (socket buffer full, frame or flow-control window - // exhausted), return and wait for the next output event. + // Each write() is one HTTP/2 frame or HTTP/1.1 chunk, so small buffers are coalesced. Each iteration: + // 1. Staging has unsent bytes: send those first, they come before anything in currentBuffer. + // 2. Staging is empty and there is no current buffer: the body is complete, end the stream. + // 3. The current buffer is large, or it's the last one: write it directly, copying is useless. + // 4. Otherwise small buffers follow: coalesce them into staging and write that. + // Loops while the channel accepts everything, returns on a refused or partial write to resume later. while (true) { ByteBuffer out; if (staging != null && staging.hasRemaining()) { From 6d88a1892a00267eb61a748ec6ba484119c4a200 Mon Sep 17 00:00:00 2001 From: Laura Trotta Date: Thu, 5 Nov 2026 12:41:47 +0100 Subject: [PATCH 5/7] unit tests --- .../rest5_client/MultiBufferEntityTest.java | 244 ++++++++++++++++++ .../rest_client/MultiBufferEntityTest.java | 169 ++++++++++++ 2 files changed, 413 insertions(+) diff --git a/java-client/src/test/java/co/elastic/clients/transport/rest5_client/MultiBufferEntityTest.java b/java-client/src/test/java/co/elastic/clients/transport/rest5_client/MultiBufferEntityTest.java index 62d4c2ce7e..18fa6658ed 100644 --- a/java-client/src/test/java/co/elastic/clients/transport/rest5_client/MultiBufferEntityTest.java +++ b/java-client/src/test/java/co/elastic/clients/transport/rest5_client/MultiBufferEntityTest.java @@ -25,14 +25,23 @@ import co.elastic.clients.transport.rest5_client.low_level.Rest5Client; import com.sun.net.httpserver.HttpServer; import org.apache.commons.io.IOUtils; +import org.apache.hc.core5.http.ContentType; +import org.apache.hc.core5.http.Header; import org.apache.hc.core5.http.HttpHost; +import org.apache.hc.core5.http.nio.DataStreamChannel; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import java.io.ByteArrayOutputStream; import java.io.IOException; import java.net.InetAddress; import java.net.InetSocketAddress; +import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; public class MultiBufferEntityTest { @@ -96,4 +105,239 @@ public void testBulkRequest() throws IOException { str ); } + + // Scenario 1: bytes left in staging by a partial write are sent before anything else. + @Test + public void testUnsentStagedBytesAreSentFirst() throws IOException { + List body = bulkBody(2000); + + // Like a socket that only takes 1000 bytes per output event + FakeChannel channel = new FakeChannel(1000, Integer.MAX_VALUE); + drain(entity(body), channel, body); + + // Like the HTTP/2 multiplexer: one frame per output event + channel = new FakeChannel(MultiBufferEntity.STAGING_BUFFER_SIZE, 1); + drain(entity(body), channel, body); + + assertUntouched(body); + } + + // Scenario 2: the stream is ended once, and only after the whole body has been written. + @Test + public void testEndOfStream() throws IOException { + for (List body : Arrays.asList( + Collections.emptyList(), + Arrays.asList(piece(0, 0), piece(0, 1)) + )) { + FakeChannel channel = new FakeChannel(Integer.MAX_VALUE, Integer.MAX_VALUE); + MultiBufferEntity entity = entity(body); + Assertions.assertEquals(0, entity.available()); + entity.produce(channel); + Assertions.assertEquals(1, channel.endStreamCount); + Assertions.assertTrue(channel.writeSizes.isEmpty()); + } + + List body = bulkBody(2000); + FakeChannel channel = new FakeChannel(MultiBufferEntity.STAGING_BUFFER_SIZE, 1); + drain(entity(body), channel, body); + } + + // Scenario 3: large buffers and the last buffer are written directly, without copying. + @Test + public void testLargeAndLastBuffersAreWrittenDirectly() throws IOException { + ByteBuffer large = piece(20_000, 0); + ByteBuffer last = piece(100, 1); + List body = Arrays.asList(large, last); + + FakeChannel channel = new FakeChannel(Integer.MAX_VALUE, Integer.MAX_VALUE); + drain(entity(body), channel, body); + Assertions.assertEquals(Arrays.asList(20_000, 100), channel.writeSizes); + Assertions.assertSame(large.array(), channel.writtenArrays.get(0)); + Assertions.assertSame(last.array(), channel.writtenArrays.get(1)); + + // Over HTTP/2 the large buffer is only partially accepted, and resumed on the next event + drain(entity(body), new FakeChannel(MultiBufferEntity.STAGING_BUFFER_SIZE, 1), body); + + // A single small buffer never goes through staging + ByteBuffer single = piece(200, 2); + body = Collections.singletonList(single); + channel = new FakeChannel(Integer.MAX_VALUE, Integer.MAX_VALUE); + drain(entity(body), channel, body); + Assertions.assertEquals(Collections.singletonList(200), channel.writeSizes); + Assertions.assertSame(single.array(), channel.writtenArrays.get(0)); + + assertUntouched(body); + } + + // Scenario 4: small buffers are coalesced into full staging-size writes. + @Test + public void testSmallBuffersAreCoalesced() throws IOException { + List body = bulkBody(2000); + int total = concat(body).length; + int expectedWrites = (total + MultiBufferEntity.STAGING_BUFFER_SIZE - 1) / MultiBufferEntity.STAGING_BUFFER_SIZE; + + FakeChannel channel = new FakeChannel(Integer.MAX_VALUE, Integer.MAX_VALUE); + drain(entity(body), channel, body); + Assertions.assertEquals(expectedWrites, channel.writeSizes.size()); + for (int i = 0; i < expectedWrites - 1; i++) { + Assertions.assertEquals(MultiBufferEntity.STAGING_BUFFER_SIZE, channel.writeSizes.get(i)); + } + + // Over HTTP/2, each output event carries a full frame + channel = new FakeChannel(MultiBufferEntity.STAGING_BUFFER_SIZE, 1); + Assertions.assertEquals(expectedWrites, drain(entity(body), channel, body)); + + assertUntouched(body); + } + + @Test + public void testRetryResendsFullBody() throws IOException { + List body = bulkBody(2000); + MultiBufferEntity entity = entity(body); + + // Attempts interrupted by close() and failed() are both reset + for (int attempt = 0; attempt < 2; attempt++) { + FakeChannel interrupted = new FakeChannel(1000, Integer.MAX_VALUE); + for (int i = 0; i < 3; i++) { + interrupted.newEvent(); + entity.produce(interrupted); + } + if (attempt == 0) { + entity.close(); + } else { + entity.failed(new IOException("test")); + } + drain(entity, new FakeChannel(1000, Integer.MAX_VALUE), body); + } + assertUntouched(body); + } + + //--------------------------------------------------------------------------------------------- + + private static final ByteBuffer SEPARATOR = ByteBuffer.wrap(new byte[]{'\n'}); + + /** + * Items of varying size, each followed by the same shared separator, like a bulk request body. Empty + * buffers are added at the start, in the middle and at the end, and must be skipped. + */ + private static List bulkBody(int items) { + List body = new ArrayList<>(); + for (int i = 0; i < items; i++) { + if (i % 500 == 0) { + body.add(piece(0, 0)); + } + body.add(piece(20 + i % 100, i)); + body.add(SEPARATOR); + } + body.add(piece(0, 0)); + return body; + } + + private static ByteBuffer piece(int size, int seed) { + byte[] bytes = new byte[size]; + for (int i = 0; i < size; i++) { + bytes[i] = (byte) (seed + i); + } + return ByteBuffer.wrap(bytes); + } + + private static byte[] concat(List body) { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + for (ByteBuffer buffer : body) { + out.write(buffer.array(), buffer.arrayOffset() + buffer.position(), buffer.remaining()); + } + return out.toByteArray(); + } + + private static MultiBufferEntity entity(List body) { + return new MultiBufferEntity(body, ContentType.APPLICATION_JSON); + } + + /** Input buffers must not be consumed, as they are shared (separator) and reused on retries. */ + private static void assertUntouched(List body) { + for (ByteBuffer buffer : body) { + Assertions.assertEquals(0, buffer.position()); + Assertions.assertEquals(buffer.capacity(), buffer.limit()); + } + } + + /** + * Calls {@code produce()} once per output event until the stream is ended, checks that the received bytes + * are the body, and returns the number of output events. + */ + private static int drain(MultiBufferEntity entity, FakeChannel channel, List body) throws IOException { + byte[] expected = concat(body); + int events = 0; + while (channel.endStreamCount == 0) { + Assertions.assertTrue(++events < 1_000_000, "no progress"); + // The HTTP/2 stream handler only asks for output while data is available + if (entity.available() == 0) { + Assertions.assertEquals(expected.length, channel.received.size()); + } + channel.newEvent(); + entity.produce(channel); + } + Assertions.assertEquals(1, channel.endStreamCount); + Assertions.assertArrayEquals(expected, channel.received.toByteArray()); + Assertions.assertEquals(0, entity.available()); + Assertions.assertArrayEquals(expected, IOUtils.toByteArray(entity.getContent())); + return events; + } + + /** + * A channel accepting at most {@code bytesPerEvent} bytes in at most {@code writesPerEvent} writes per output + * event, like a socket that fills up (HTTP/1.1) or the HTTP/2 multiplexer, which takes one frame per event. + */ + private static class FakeChannel implements DataStreamChannel { + private final int bytesPerEvent; + private final int writesPerEvent; + private int bytesLeft; + private int writesLeft; + + final ByteArrayOutputStream received = new ByteArrayOutputStream(); + final List writeSizes = new ArrayList<>(); + final List writtenArrays = new ArrayList<>(); + int endStreamCount; + + FakeChannel(int bytesPerEvent, int writesPerEvent) { + this.bytesPerEvent = bytesPerEvent; + this.writesPerEvent = writesPerEvent; + newEvent(); + } + + void newEvent() { + bytesLeft = bytesPerEvent; + writesLeft = writesPerEvent; + } + + @Override + public int write(ByteBuffer src) { + Assertions.assertEquals(0, endStreamCount, "write after end of stream"); + int n = Math.min(src.remaining(), bytesLeft); + if (writesLeft == 0 || n == 0) { + return 0; + } + writesLeft--; + bytesLeft -= n; + writeSizes.add(n); + writtenArrays.add(src.array()); + received.write(src.array(), src.arrayOffset() + src.position(), n); + src.position(src.position() + n); + return n; + } + + @Override + public void requestOutput() { + } + + @Override + public void endStream() { + endStreamCount++; + } + + @Override + public void endStream(List trailers) { + endStream(); + } + } } diff --git a/java-client/src/test/java/co/elastic/clients/transport/rest_client/MultiBufferEntityTest.java b/java-client/src/test/java/co/elastic/clients/transport/rest_client/MultiBufferEntityTest.java index b241c19052..af2660eae8 100644 --- a/java-client/src/test/java/co/elastic/clients/transport/rest_client/MultiBufferEntityTest.java +++ b/java-client/src/test/java/co/elastic/clients/transport/rest_client/MultiBufferEntityTest.java @@ -25,14 +25,22 @@ import com.sun.net.httpserver.HttpServer; import org.apache.commons.io.IOUtils; import org.apache.http.HttpHost; +import org.apache.http.entity.ContentType; +import org.apache.http.nio.ContentEncoder; import org.elasticsearch.client.RestClient; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import java.io.ByteArrayOutputStream; import java.io.IOException; import java.net.InetAddress; import java.net.InetSocketAddress; +import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; public class MultiBufferEntityTest { @@ -96,4 +104,165 @@ public void testBulkRequest() throws IOException { str ); } + + // Scenario 1: a buffer partially accepted by the encoder is resumed from where it stopped. + @Test + public void testPartiallyWrittenBufferIsResumed() throws IOException { + List body = bulkBody(2000); + drain(entity(body), new FakeEncoder(7), body); + assertUntouched(body); + } + + // Scenario 2: the encoder is completed once, and only after the whole body has been written. + @Test + public void testEndOfStream() throws IOException { + for (List body : Arrays.asList( + Collections.emptyList(), + Arrays.asList(piece(0, 0), piece(0, 1)) + )) { + FakeEncoder encoder = new FakeEncoder(Integer.MAX_VALUE); + entity(body).produceContent(encoder, null); + Assertions.assertEquals(1, encoder.completeCount); + Assertions.assertEquals(0, encoder.received.size()); + } + + List body = bulkBody(2000); + drain(entity(body), new FakeEncoder(1000), body); + } + + // Scenario 3: buffers are written in a single call as long as the encoder accepts them. + @Test + public void testWritesAllBuffersWhileEncoderAccepts() throws IOException { + List body = bulkBody(2000); + Assertions.assertEquals(1, drain(entity(body), new FakeEncoder(Integer.MAX_VALUE), body)); + assertUntouched(body); + } + + // Scenario 4: writing stops when the encoder is full, and every output event fills it. + @Test + public void testStopsWhenEncoderIsFull() throws IOException { + List body = bulkBody(2000); + int total = concat(body).length; + Assertions.assertEquals((total + 999) / 1000, drain(entity(body), new FakeEncoder(1000), body)); + } + + @Test + public void testRetryResendsFullBody() throws IOException { + List body = bulkBody(2000); + MultiBufferEntity entity = entity(body); + + FakeEncoder failed = new FakeEncoder(1000); + for (int i = 0; i < 3; i++) { + failed.newEvent(); + entity.produceContent(failed, null); + } + entity.close(); + + drain(entity, new FakeEncoder(1000), body); + assertUntouched(body); + } + + //--------------------------------------------------------------------------------------------- + + private static final ByteBuffer SEPARATOR = ByteBuffer.wrap(new byte[]{'\n'}); + + /** + * Items of varying size, each followed by the same shared separator, like a bulk request body. Empty + * buffers are added at the start, in the middle and at the end, and must be skipped. + */ + private static List bulkBody(int items) { + List body = new ArrayList<>(); + for (int i = 0; i < items; i++) { + if (i % 500 == 0) { + body.add(piece(0, 0)); + } + body.add(piece(20 + i % 100, i)); + body.add(SEPARATOR); + } + body.add(piece(0, 0)); + return body; + } + + private static ByteBuffer piece(int size, int seed) { + byte[] bytes = new byte[size]; + for (int i = 0; i < size; i++) { + bytes[i] = (byte) (seed + i); + } + return ByteBuffer.wrap(bytes); + } + + private static byte[] concat(List body) { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + for (ByteBuffer buffer : body) { + out.write(buffer.array(), buffer.arrayOffset() + buffer.position(), buffer.remaining()); + } + return out.toByteArray(); + } + + private static MultiBufferEntity entity(List body) { + return new MultiBufferEntity(body, ContentType.APPLICATION_JSON); + } + + /** Input buffers must not be consumed, as they are shared (separator) and reused on retries. */ + private static void assertUntouched(List body) { + for (ByteBuffer buffer : body) { + Assertions.assertEquals(0, buffer.position()); + Assertions.assertEquals(buffer.capacity(), buffer.limit()); + } + } + + /** + * Calls {@code produceContent()} once per output event until the encoder is completed, checks that the + * received bytes are the body, and returns the number of output events. + */ + private static int drain(MultiBufferEntity entity, FakeEncoder encoder, List body) throws IOException { + int events = 0; + while (!encoder.isCompleted()) { + Assertions.assertTrue(++events < 1_000_000, "no progress"); + encoder.newEvent(); + entity.produceContent(encoder, null); + } + Assertions.assertEquals(1, encoder.completeCount); + Assertions.assertArrayEquals(concat(body), encoder.received.toByteArray()); + Assertions.assertArrayEquals(concat(body), IOUtils.toByteArray(entity.getContent())); + return events; + } + + /** An encoder accepting at most {@code bytesPerEvent} bytes per output event, like a socket that fills up. */ + private static class FakeEncoder implements ContentEncoder { + private final int bytesPerEvent; + private int bytesLeft; + + final ByteArrayOutputStream received = new ByteArrayOutputStream(); + int completeCount; + + FakeEncoder(int bytesPerEvent) { + this.bytesPerEvent = bytesPerEvent; + newEvent(); + } + + void newEvent() { + bytesLeft = bytesPerEvent; + } + + @Override + public int write(ByteBuffer src) { + Assertions.assertEquals(0, completeCount, "write after completion"); + int n = Math.min(src.remaining(), bytesLeft); + bytesLeft -= n; + received.write(src.array(), src.arrayOffset() + src.position(), n); + src.position(src.position() + n); + return n; + } + + @Override + public void complete() { + completeCount++; + } + + @Override + public boolean isCompleted() { + return completeCount > 0; + } + } } From 31067b2963eab80ba3f26e1903af426591b260e0 Mon Sep 17 00:00:00 2001 From: Laura Trotta Date: Thu, 5 Nov 2026 12:52:33 +0100 Subject: [PATCH 6/7] apply same staging buffer logic to legacy rest client --- .../rest_client/MultiBufferEntity.java | 97 +++++++++++++++---- .../rest_client/MultiBufferEntityTest.java | 59 +++++++++-- 2 files changed, 126 insertions(+), 30 deletions(-) diff --git a/java-client/src/main/java/co/elastic/clients/transport/rest_client/MultiBufferEntity.java b/java-client/src/main/java/co/elastic/clients/transport/rest_client/MultiBufferEntity.java index f3ffea68b4..4c7eaf2a91 100644 --- a/java-client/src/main/java/co/elastic/clients/transport/rest_client/MultiBufferEntity.java +++ b/java-client/src/main/java/co/elastic/clients/transport/rest_client/MultiBufferEntity.java @@ -39,10 +39,19 @@ */ class MultiBufferEntity extends AbstractHttpEntity implements HttpAsyncContentProducer { + /** + * Small input buffers are coalesced up to this size before writing, since each write is one HTTP chunk. + * The staging buffer is never allocated for single-buffer bodies, such as most non-bulk requests. + */ + static final int STAGING_BUFFER_SIZE = 16 * 1024; + private final Iterable buffers; private Iterator iterator; + /** Input buffer currently being consumed, {@code null} once the input is exhausted. Never empty. */ private volatile ByteBuffer currentBuffer; + /** Coalesced bytes not yet accepted by the encoder, always in read mode. Allocated on first use. */ + private volatile ByteBuffer staging; MultiBufferEntity(Iterable buffers, ContentType contentType) { this.buffers = buffers; @@ -61,12 +70,27 @@ public void close() throws IOException { private void init() { Iterator localIterator = this.buffers.iterator(); - if (localIterator.hasNext()) { - this.currentBuffer = localIterator.next().duplicate(); - } else { - this.currentBuffer = null; - } + this.currentBuffer = nextNonEmpty(localIterator); this.iterator = localIterator; + ByteBuffer staging = this.staging; + if (staging != null) { + staging.position(0).limit(0); + } + } + + /** Moves {@link #currentBuffer} to the next non-empty input buffer, or {@code null} if there is none. */ + private void advance() { + currentBuffer = nextNonEmpty(iterator); + } + + private static ByteBuffer nextNonEmpty(Iterator iterator) { + while (iterator.hasNext()) { + ByteBuffer next = iterator.next(); + if (next.hasRemaining()) { + return next.duplicate(); + } + } + return null; } @Override @@ -102,25 +126,56 @@ public void writeTo(OutputStream out) throws IOException { @Override public void produceContent(ContentEncoder encoder, IOControl ioControl) throws IOException { - if (currentBuffer == null) { - encoder.complete(); - return; + // Each write() is one HTTP chunk, so small buffers are coalesced. Each iteration: + // 1. Staging has unsent bytes: send those first, they come before anything in currentBuffer. + // 2. Staging is empty and there is no current buffer: the body is complete, complete the encoder. + // 3. The current buffer is large, or it's the last one: write it directly, copying is useless. + // 4. Otherwise small buffers follow: coalesce them into staging and write that. + // Loops while the encoder accepts everything, returns on a refused or partial write to resume later. + while (true) { + ByteBuffer out; + if (staging != null && staging.hasRemaining()) { + out = staging; + } else if (currentBuffer == null) { + encoder.complete(); + return; + } else if (currentBuffer.remaining() >= STAGING_BUFFER_SIZE || !iterator.hasNext()) { + out = currentBuffer; + } else { + out = fillStaging(); + } + + encoder.write(out); + + if (out == currentBuffer && !out.hasRemaining()) { + advance(); + } + if (out.hasRemaining()) { + return; + } } + } - // Keep feeding buffers to the encoder as long as it accepts them. - while (currentBuffer != null) { - int written = encoder.write(currentBuffer); - - if (!currentBuffer.hasRemaining()) { - if (iterator.hasNext()) { - currentBuffer = iterator.next().duplicate(); - } else { - currentBuffer = null; - encoder.complete(); - } - } else if (written == 0) { - break; + /** Copies as much input as fits into the (empty) staging buffer and returns it, ready to be read. */ + private ByteBuffer fillStaging() { + ByteBuffer staging = this.staging; + if (staging == null) { + staging = ByteBuffer.allocate(STAGING_BUFFER_SIZE); + this.staging = staging; + } + staging.clear(); + while (currentBuffer != null && staging.hasRemaining()) { + ByteBuffer src = currentBuffer; + int n = Math.min(src.remaining(), staging.remaining()); + int srcLimit = src.limit(); + src.limit(src.position() + n); + staging.put(src); + src.limit(srcLimit); + if (!src.hasRemaining()) { + advance(); } } + staging.flip(); + return staging; } } diff --git a/java-client/src/test/java/co/elastic/clients/transport/rest_client/MultiBufferEntityTest.java b/java-client/src/test/java/co/elastic/clients/transport/rest_client/MultiBufferEntityTest.java index af2660eae8..5988b370ab 100644 --- a/java-client/src/test/java/co/elastic/clients/transport/rest_client/MultiBufferEntityTest.java +++ b/java-client/src/test/java/co/elastic/clients/transport/rest_client/MultiBufferEntityTest.java @@ -105,11 +105,16 @@ public void testBulkRequest() throws IOException { ); } - // Scenario 1: a buffer partially accepted by the encoder is resumed from where it stopped. + // Scenario 1: bytes left in staging by a partial write are sent before anything else. @Test - public void testPartiallyWrittenBufferIsResumed() throws IOException { + public void testUnsentStagedBytesAreSentFirst() throws IOException { List body = bulkBody(2000); + int total = concat(body).length; + drain(entity(body), new FakeEncoder(7), body); + // Every output event fills the encoder + Assertions.assertEquals((total + 999) / 1000, drain(entity(body), new FakeEncoder(1000), body)); + assertUntouched(body); } @@ -130,20 +135,49 @@ public void testEndOfStream() throws IOException { drain(entity(body), new FakeEncoder(1000), body); } - // Scenario 3: buffers are written in a single call as long as the encoder accepts them. + // Scenario 3: large buffers and the last buffer are written directly, without copying. @Test - public void testWritesAllBuffersWhileEncoderAccepts() throws IOException { - List body = bulkBody(2000); - Assertions.assertEquals(1, drain(entity(body), new FakeEncoder(Integer.MAX_VALUE), body)); + public void testLargeAndLastBuffersAreWrittenDirectly() throws IOException { + ByteBuffer large = piece(20_000, 0); + ByteBuffer last = piece(100, 1); + List body = Arrays.asList(large, last); + + FakeEncoder encoder = new FakeEncoder(Integer.MAX_VALUE); + drain(entity(body), encoder, body); + Assertions.assertEquals(Arrays.asList(20_000, 100), encoder.writeSizes); + Assertions.assertSame(large.array(), encoder.writtenArrays.get(0)); + Assertions.assertSame(last.array(), encoder.writtenArrays.get(1)); + + // The large buffer is only partially accepted, and resumed on the next event + drain(entity(body), new FakeEncoder(1000), body); + + // A single small buffer never goes through staging + ByteBuffer single = piece(200, 2); + body = Collections.singletonList(single); + encoder = new FakeEncoder(Integer.MAX_VALUE); + drain(entity(body), encoder, body); + Assertions.assertEquals(Collections.singletonList(200), encoder.writeSizes); + Assertions.assertSame(single.array(), encoder.writtenArrays.get(0)); + assertUntouched(body); } - // Scenario 4: writing stops when the encoder is full, and every output event fills it. + // Scenario 4: small buffers are coalesced into full staging-size writes, in a single call while the encoder + // accepts them. @Test - public void testStopsWhenEncoderIsFull() throws IOException { + public void testSmallBuffersAreCoalesced() throws IOException { List body = bulkBody(2000); int total = concat(body).length; - Assertions.assertEquals((total + 999) / 1000, drain(entity(body), new FakeEncoder(1000), body)); + int expectedWrites = (total + MultiBufferEntity.STAGING_BUFFER_SIZE - 1) / MultiBufferEntity.STAGING_BUFFER_SIZE; + + FakeEncoder encoder = new FakeEncoder(Integer.MAX_VALUE); + Assertions.assertEquals(1, drain(entity(body), encoder, body)); + Assertions.assertEquals(expectedWrites, encoder.writeSizes.size()); + for (int i = 0; i < expectedWrites - 1; i++) { + Assertions.assertEquals(MultiBufferEntity.STAGING_BUFFER_SIZE, encoder.writeSizes.get(i)); + } + + assertUntouched(body); } @Test @@ -234,6 +268,8 @@ private static class FakeEncoder implements ContentEncoder { private int bytesLeft; final ByteArrayOutputStream received = new ByteArrayOutputStream(); + final List writeSizes = new ArrayList<>(); + final List writtenArrays = new ArrayList<>(); int completeCount; FakeEncoder(int bytesPerEvent) { @@ -249,7 +285,12 @@ void newEvent() { public int write(ByteBuffer src) { Assertions.assertEquals(0, completeCount, "write after completion"); int n = Math.min(src.remaining(), bytesLeft); + if (n == 0) { + return 0; + } bytesLeft -= n; + writeSizes.add(n); + writtenArrays.add(src.array()); received.write(src.array(), src.arrayOffset() + src.position(), n); src.position(src.position() + n); return n; From 3f2e4b96228cf0e1ccab2b7e716feab59cc3fdfa Mon Sep 17 00:00:00 2001 From: Laura Trotta Date: Thu, 5 Nov 2026 13:23:26 +0100 Subject: [PATCH 7/7] rest5 MultiBufferEntityTest: use a fresh entity per retry attempt --- .../clients/transport/rest5_client/MultiBufferEntityTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/java-client/src/test/java/co/elastic/clients/transport/rest5_client/MultiBufferEntityTest.java b/java-client/src/test/java/co/elastic/clients/transport/rest5_client/MultiBufferEntityTest.java index 18fa6658ed..5da1f34a97 100644 --- a/java-client/src/test/java/co/elastic/clients/transport/rest5_client/MultiBufferEntityTest.java +++ b/java-client/src/test/java/co/elastic/clients/transport/rest5_client/MultiBufferEntityTest.java @@ -193,10 +193,10 @@ public void testSmallBuffersAreCoalesced() throws IOException { @Test public void testRetryResendsFullBody() throws IOException { List body = bulkBody(2000); - MultiBufferEntity entity = entity(body); // Attempts interrupted by close() and failed() are both reset for (int attempt = 0; attempt < 2; attempt++) { + MultiBufferEntity entity = entity(body); FakeChannel interrupted = new FakeChannel(1000, Integer.MAX_VALUE); for (int i = 0; i < 3; i++) { interrupted.newEvent();