Skip to content

[Backport 8.19] MultiBufferEntity add loop to produce content. (#1366) #1370

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
l-trotta merged 1 commit into from
Nov 5, 2026
Merged

[Backport 8.19] MultiBufferEntity add loop to produce content. (#1366) #1370

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
92 changes: 76 additions & 16 deletions java-client/src/main/java/co/elastic/clients/transport/rest_client/MultiBufferEntity.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<ByteBuffer> buffers;

private Iterator<ByteBuffer> 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<ByteBuffer> buffers, ContentType contentType) {
this.buffers = buffers;
Expand All @@ -61,12 +70,27 @@ public void close() throws IOException {

private void init() {
Iterator<ByteBuffer> 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<ByteBuffer> iterator) {
while (iterator.hasNext()) {
ByteBuffer next = iterator.next();
if (next.hasRemaining()) {
return next.duplicate();
}
}
return null;
}

@Override
Expand Down Expand Up @@ -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;
}
}
210 changes: 210 additions & 0 deletions ...-client/src/test/java/co/elastic/clients/transport/rest_client/MultiBufferEntityTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {

Expand Down Expand Up @@ -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<ByteBuffer> 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<ByteBuffer> body : Arrays.asList(
Collections.<ByteBuffer>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<ByteBuffer> 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<ByteBuffer> 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<ByteBuffer> 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<ByteBuffer> 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<ByteBuffer> bulkBody(int items) {
List<ByteBuffer> 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<ByteBuffer> 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<ByteBuffer> 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<ByteBuffer> 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<ByteBuffer> 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<Integer> writeSizes = new ArrayList<>();
final List<byte[]> 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;
}
}
}
Loading