From c7c3ea1e7254f36c762625e1580d1554e5f1afd8 Mon Sep 17 00:00:00 2001 From: Laura Trotta <153528055+l-trotta@users.noreply.github.com> Date: Thu, 5 Nov 2026 14:41:12 +0100 Subject: [PATCH] MultiBufferEntity add loop to produce content. (#1366) * keep feeding the encoder while it accepts bytes in multibuffer entity * rest5 MultiBufferEntity: coalesce small buffers into a 16 KiB staging buffer to fix bug for http2 * avoid buffering single body requests in staging buffer * better explain rest5 mbe produce * unit tests * apply same staging buffer logic to legacy rest client * rest5 MultiBufferEntityTest: use a fresh entity per retry attempt --- .../rest_client/MultiBufferEntity.java | 92 ++++++-- .../rest_client/MultiBufferEntityTest.java | 210 ++++++++++++++++++ 2 files changed, 286 insertions(+), 16 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 cebbebd55e..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,20 +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(currentBuffer); + encoder.write(out); - if (!currentBuffer.hasRemaining()) { - if (iterator.hasNext()) { - currentBuffer = iterator.next().duplicate(); - } else { - currentBuffer = null; - encoder.complete(); + if (out == currentBuffer && !out.hasRemaining()) { + advance(); + } + if (out.hasRemaining()) { + return; + } + } + } + + /** 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 b241c19052..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 @@ -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,206 @@ 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); + 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); + } + + // 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: 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); + + 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: small buffers are coalesced into full staging-size writes, in a single call while the encoder + // accepts them. + @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; + + 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 + 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(); + final List writeSizes = new ArrayList<>(); + final List writtenArrays = new ArrayList<>(); + 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); + 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; + } + + @Override + public void complete() { + completeCount++; + } + + @Override + public boolean isCompleted() { + return completeCount > 0; + } + } }