gRPC closes the stream it passes to a marshaller as soon as parsing returns, so a buffer + * merely borrowed through {@link HasByteBuffer} cannot outlive {@code parse()}. {@link + * Detachable#detach()} transfers ownership of the buffer to us; we wrap that memory as a foreign + * allocation whose release closes the detached stream, and hand out retained slices of it. The + * memory goes back to gRPC once the message and every slice are closed. + */ + private static final class OwnedFrame extends InputStream { + private final ArrowBuf frame; + private final int length; + private int position; + + private OwnedFrame(ArrowBuf frame, int length) { + this.frame = frame; + this.length = length; + } + + /** + * Take ownership of the stream's buffer if gRPC exposes it and it holds the entire message. + * + * @return the owned frame, or null to fall back to copying. + */ + static OwnedFrame tryTakeOwnership(BufferAllocator allocator, InputStream stream) { + if (!(stream instanceof Detachable) + || !(stream instanceof HasByteBuffer) + || !(stream instanceof KnownLength) + || !((HasByteBuffer) stream).byteBufferSupported()) { + return null; + } + final int size; + try { + size = stream.available(); + } catch (IOException e) { + throw new RuntimeException("Failed to query gRPC stream length", e); + } + final ByteBuffer peek = ((HasByteBuffer) stream).getByteBuffer(); + if (peek == null || !peek.isDirect() || size == 0 || peek.remaining() != size) { + // Empty, on-heap, or split across several buffers: only a single direct buffer holding + // the whole message can be wrapped. + return null; + } + + final InputStream owner = ((Detachable) stream).detach(); + final long address; + try { + if (!(owner instanceof HasByteBuffer) || !((HasByteBuffer) owner).byteBufferSupported()) { + throw new IllegalStateException("Detached gRPC stream does not expose its buffer"); + } + final ByteBuffer owned = ((HasByteBuffer) owner).getByteBuffer(); + if (owned == null || !owned.isDirect() || owned.remaining() != size) { + throw new IllegalStateException("Detached gRPC buffer differs from the one inspected"); + } + address = MemoryUtil.getByteBufferAddress(owned) + owned.position(); + } catch (RuntimeException | Error e) { + AutoCloseables.closeNoChecked(owner); + throw e; + } + + final ArrowBuf frame; + try { + frame = + allocator.wrapForeignAllocation( + new ForeignAllocation(size, address) { + @Override + protected void release0() { + AutoCloseables.closeNoChecked(owner); + } + }); + } catch (RuntimeException | Error e) { + // wrapForeignAllocation failed before adopting the allocation, so release0 will not run. + AutoCloseables.closeNoChecked(owner); + throw e; + } + return new OwnedFrame(frame, size); + } + + /** A zero-copy view of the next {@code size} bytes; the caller owns the returned buffer. */ + ArrowBuf slice(int size) throws IOException { + if (size < 0 || size > available()) { + throw new IOException("Unexpected end of gRPC message"); + } + final int offset = position; + position += size; + // A slice shares the frame's reference count, so give it its own reference first. + frame.getReferenceManager().retain(); + try { + return frame.slice(offset, size); + } catch (RuntimeException | Error e) { + frame.getReferenceManager().release(); + throw e; + } + } + + @Override + public int read() { + if (position >= length) { + return -1; + } + return frame.getByte(position++) & 0xFF; + } + + @Override + public int read(byte[] b, int off, int len) { + if (position >= length) { + return -1; + } + final int n = Math.min(len, available()); + frame.getBytes(position, b, off, n); + position += n; + return n; + } + + @Override + public long skip(long n) { + final int skipped = (int) Math.min(n, available()); + position += skipped; + return skipped; + } + + @Override + public int available() { + return length - position; + } + + @Override + public void close() { + frame.close(); + } + } + private static int readRawVarint32(int firstByte, InputStream is) throws IOException { return CodedInputStream.readRawVarint32(firstByte, is); } diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java index 15cfd6ba85..9195eb6780 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightStream.java @@ -480,6 +480,7 @@ public void onNext(ArrowMessage msg) { ex = new UnsupportedOperationException( "Unable to handle message of type: " + msg.getMessageType()); + AutoCloseables.closeNoChecked(msg); enqueue(DONE_EX); } } diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java index fcba88d212..341f10b16e 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java @@ -16,84 +16,81 @@ */ package org.apache.arrow.flight.grpc; -import com.google.common.base.Throwables; import com.google.common.io.ByteStreams; -import io.grpc.internal.ReadableBuffer; +import io.grpc.HasByteBuffer; import java.io.IOException; import java.io.InputStream; -import java.lang.reflect.Field; +import java.nio.ByteBuffer; import org.apache.arrow.memory.ArrowBuf; /** - * Enable access to ReadableBuffer directly to copy data from a BufferInputStream into a target - * ByteBuffer/ByteBuf. + * Reads a gRPC message body into an {@link ArrowBuf}. * - *
This could be solved by BufferInputStream exposing Drainable. + *
When the stream gRPC hands us exposes its backing buffers through {@link HasByteBuffer}, the + * bytes are copied straight from those buffers into the target, skipping the intermediate heap + * array. Otherwise the stream is read into a heap array first. */ public class GetReadableBuffer { - private static final Field READABLE_BUFFER; - private static final Class BUFFER_INPUT_STREAM; - - static { - Field tmpField = null; - Class tmpClazz = null; - try { - Class clazz = Class.forName("io.grpc.internal.ReadableBuffers$BufferInputStream"); - - Field f = clazz.getDeclaredField("buffer"); - f.setAccessible(true); - // don't set until we've gotten past all exception cases. - tmpField = f; - tmpClazz = clazz; - } catch (Exception e) { - new RuntimeException("Failed to initialize GetReadableBuffer, falling back to slow path", e) - .printStackTrace(); - } - READABLE_BUFFER = tmpField; - BUFFER_INPUT_STREAM = tmpClazz; - } - - /** - * Extracts the ReadableBuffer for the given input stream. - * - * @param is Must be an instance of io.grpc.internal.ReadableBuffers$BufferInputStream or null - * will be returned. - */ - public static ReadableBuffer getReadableBuffer(InputStream is) { - - if (BUFFER_INPUT_STREAM == null || !is.getClass().equals(BUFFER_INPUT_STREAM)) { - return null; - } - - try { - return (ReadableBuffer) READABLE_BUFFER.get(is); - } catch (Exception ex) { - throw Throwables.propagate(ex); - } - } + private GetReadableBuffer() {} /** - * Helper method to read a gRPC-provided InputStream into an ArrowBuf. + * Read exactly {@code size} bytes from {@code stream} into {@code buf}. * - * @param stream The stream to read from. Should be an instance of {@link #BUFFER_INPUT_STREAM}. - * @param buf The buffer to read into. + * @param stream The stream to read from. + * @param buf The buffer to read into. Its writer index is set to {@code size} on success. * @param size The number of bytes to read. - * @param fastPath Whether to enable the fast path (i.e. detect whether the stream is a {@link - * #BUFFER_INPUT_STREAM}). - * @throws IOException if there is an error reading form the stream + * @param fastPath Whether to copy directly from the stream's backing buffers when it exposes + * them. + * @throws IOException if the stream ends early or cannot be read. */ public static void readIntoBuffer( final InputStream stream, final ArrowBuf buf, final int size, final boolean fastPath) throws IOException { - ReadableBuffer readableBuffer = fastPath ? getReadableBuffer(stream) : null; - byte[] heapBytes = new byte[size]; - if (readableBuffer != null) { - readableBuffer.readBytes(heapBytes, 0, size); + if (fastPath + && stream instanceof HasByteBuffer + && ((HasByteBuffer) stream).byteBufferSupported()) { + readFromByteBuffers((HasByteBuffer) stream, stream, buf, size); } else { + final byte[] heapBytes = new byte[size]; ByteStreams.readFully(stream, heapBytes); + buf.setBytes(0, heapBytes); } - buf.writeBytes(heapBytes); buf.writerIndex(size); } + + private static void readFromByteBuffers( + final HasByteBuffer source, final InputStream stream, final ArrowBuf buf, final int size) + throws IOException { + int copied = 0; + while (copied < size) { + final ByteBuffer chunk = source.getByteBuffer(); + if (chunk == null || !chunk.hasRemaining()) { + throw new IOException( + "Unexpected end of gRPC stream: expected " + size + " bytes, got " + copied); + } + final int toRead = Math.min(size - copied, chunk.remaining()); + buf.setBytes(copied, chunk, chunk.position(), toRead); + // getByteBuffer() does not advance the stream; skip() consumes the bytes we just copied. + consume(stream, toRead); + copied += toRead; + } + } + + private static void consume(final InputStream stream, final int count) throws IOException { + int skipped = 0; + while (skipped < count) { + final long n = stream.skip(count - skipped); + if (n > 0) { + skipped += (int) n; + continue; + } + // InputStream.skip is allowed to return 0 before the end of the stream. Force progress with + // a single-byte read, which does distinguish end of stream. + if (stream.read() == -1) { + throw new IOException("Unexpected end of gRPC stream while skipping copied bytes"); + } + skipped++; + } + } } diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessageDetachable.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessageDetachable.java new file mode 100644 index 0000000000..0d4661723f --- /dev/null +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessageDetachable.java @@ -0,0 +1,282 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.flight; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import com.google.protobuf.ByteString; +import com.google.protobuf.CodedOutputStream; +import com.google.protobuf.WireFormat; +import io.grpc.Detachable; +import io.grpc.HasByteBuffer; +import io.grpc.KnownLength; +import java.io.ByteArrayOutputStream; +import java.io.InputStream; +import java.nio.ByteBuffer; +import org.apache.arrow.flight.impl.Flight.FlightData; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Tests for parsing FlightData without copying, by taking ownership of gRPC's buffer. */ +public class TestArrowMessageDetachable { + + private BufferAllocator allocator; + + @BeforeEach + public void setUp() { + allocator = new RootAllocator(Long.MAX_VALUE); + } + + @AfterEach + public void tearDown() { + assertEquals(0, allocator.getAllocatedMemory()); + allocator.close(); + } + + @Test + public void contiguousDirectBufferIsDetachedAndWrappedWithoutCopy() throws Exception { + final byte[] metadata = payload(16); + final byte[] body = payload(64); + final byte[] serialized = flightData(metadata, body); + final MockGrpcInputStream stream = MockGrpcInputStream.direct(serialized); + + try (ArrowMessage message = ArrowMessage.createMarshaller(allocator).parse(stream)) { + // gRPC closes the stream it handed us as soon as parse() returns. + stream.close(); + + assertEquals(1, stream.state.detachCount, "must take ownership through detach()"); + assertEquals( + serialized.length, + allocator.getAllocatedMemory(), + "the whole gRPC frame is wrapped, not a copy of the body"); + assertArrayEquals(metadata, contents(message.getApplicationMetadata())); + assertArrayEquals(body, contents(message.getBufs().iterator().next())); + assertEquals( + 0, stream.state.detachedCloseCount, "gRPC memory must stay alive with the message"); + } + assertEquals( + 1, stream.state.detachedCloseCount, "closing the message releases gRPC's buffer once"); + } + + @Test + public void heapBufferFallsBackToCopyingWithoutDetaching() throws Exception { + final byte[] metadata = payload(16); + final byte[] body = payload(64); + final MockGrpcInputStream stream = MockGrpcInputStream.heap(flightData(metadata, body)); + + try (ArrowMessage message = ArrowMessage.createMarshaller(allocator).parse(stream)) { + stream.close(); + + assertEquals(0, stream.state.detachCount, "on-heap buffers cannot be wrapped"); + assertEquals( + metadata.length + body.length, + allocator.getAllocatedMemory(), + "copying path allocates exactly the two fields"); + assertArrayEquals(metadata, contents(message.getApplicationMetadata())); + assertArrayEquals(body, contents(message.getBufs().iterator().next())); + } + assertEquals(0, stream.state.detachedCloseCount); + } + + @Test + public void messageSplitAcrossBuffersFallsBackToCopying() throws Exception { + final byte[] metadata = payload(16); + final byte[] body = payload(64); + final byte[] serialized = flightData(metadata, body); + final MockGrpcInputStream stream = MockGrpcInputStream.directFragmented(serialized, 20); + + try (ArrowMessage message = ArrowMessage.createMarshaller(allocator).parse(stream)) { + stream.close(); + + assertEquals(0, stream.state.detachCount, "only a single contiguous buffer is wrapped"); + assertEquals(metadata.length + body.length, allocator.getAllocatedMemory()); + assertArrayEquals(metadata, contents(message.getApplicationMetadata())); + assertArrayEquals(body, contents(message.getBufs().iterator().next())); + } + } + + @Test + public void parseFailureReleasesTheDetachedBuffer() throws Exception { + // A complete body, then a descriptor whose declared length runs past the end of the frame. + final ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + final CodedOutputStream out = CodedOutputStream.newInstance(bytes); + out.writeBytes(FlightData.DATA_BODY_FIELD_NUMBER, ByteString.copyFrom(payload(8))); + out.writeTag(FlightData.FLIGHT_DESCRIPTOR_FIELD_NUMBER, WireFormat.WIRETYPE_LENGTH_DELIMITED); + out.writeUInt32NoTag(10); + out.writeRawByte(1); + out.flush(); + final MockGrpcInputStream stream = MockGrpcInputStream.direct(bytes.toByteArray()); + + assertThrows( + RuntimeException.class, () -> ArrowMessage.createMarshaller(allocator).parse(stream)); + stream.close(); + + assertEquals(1, stream.state.detachCount); + assertEquals(1, stream.state.detachedCloseCount, "the detached buffer must not leak"); + assertEquals(0, allocator.getAllocatedMemory()); + } + + @Test + public void lastOfDuplicateBodyFieldsWinsWithoutLeaking() throws Exception { + final byte[] first = payload(8); + final byte[] second = payload(24); + final ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + final CodedOutputStream out = CodedOutputStream.newInstance(bytes); + out.writeBytes(FlightData.DATA_BODY_FIELD_NUMBER, ByteString.copyFrom(first)); + out.writeBytes(FlightData.DATA_BODY_FIELD_NUMBER, ByteString.copyFrom(second)); + out.flush(); + final byte[] serialized = bytes.toByteArray(); + final MockGrpcInputStream stream = MockGrpcInputStream.direct(serialized); + + try (ArrowMessage message = ArrowMessage.createMarshaller(allocator).parse(stream)) { + stream.close(); + assertEquals(1, stream.state.detachCount); + assertArrayEquals(second, contents(message.getBufs().iterator().next())); + assertEquals(serialized.length, allocator.getAllocatedMemory()); + } + assertEquals(1, stream.state.detachedCloseCount); + } + + private static byte[] flightData(byte[] metadata, byte[] body) { + return FlightData.newBuilder() + .setAppMetadata(ByteString.copyFrom(metadata)) + .setDataBody(ByteString.copyFrom(body)) + .build() + .toByteArray(); + } + + private static byte[] payload(int size) { + final byte[] bytes = new byte[size]; + for (int i = 0; i < size; i++) { + bytes[i] = (byte) (i * 13 + 7); + } + return bytes; + } + + private static byte[] contents(ArrowBuf buf) { + final byte[] out = new byte[(int) buf.writerIndex()]; + buf.getBytes(0, out); + return out; + } + + /** + * A stand-in for the stream gRPC's Netty transport hands to a marshaller: one buffer, exposed + * through the three public capabilities, with counters for detach and close. + */ + static final class MockGrpcInputStream extends InputStream + implements Detachable, HasByteBuffer, KnownLength { + + static final class State { + int detachCount; + int detachedCloseCount; + } + + final State state; + private ByteBuffer buffer; + private final boolean detached; + + /** How many bytes getByteBuffer() exposes at once; gRPC may hold a message in pieces. */ + private final int exposeLimit; + + static MockGrpcInputStream direct(byte[] bytes) { + return directFragmented(bytes, Integer.MAX_VALUE); + } + + static MockGrpcInputStream directFragmented(byte[] bytes, int exposeLimit) { + final ByteBuffer buffer = ByteBuffer.allocateDirect(bytes.length); + buffer.put(bytes).flip(); + return new MockGrpcInputStream(buffer, new State(), false, exposeLimit); + } + + static MockGrpcInputStream heap(byte[] bytes) { + return new MockGrpcInputStream(ByteBuffer.wrap(bytes), new State(), false, Integer.MAX_VALUE); + } + + private MockGrpcInputStream(ByteBuffer buffer, State state, boolean detached, int exposeLimit) { + this.buffer = buffer; + this.state = state; + this.detached = detached; + this.exposeLimit = exposeLimit; + } + + @Override + public int read() { + return buffer.hasRemaining() ? buffer.get() & 0xFF : -1; + } + + @Override + public int read(byte[] b, int off, int len) { + if (!buffer.hasRemaining()) { + return -1; + } + final int n = Math.min(len, buffer.remaining()); + buffer.get(b, off, n); + return n; + } + + @Override + public long skip(long n) { + final int skipped = (int) Math.min(n, buffer.remaining()); + buffer.position(buffer.position() + skipped); + return skipped; + } + + @Override + public int available() { + return buffer.remaining(); + } + + @Override + public boolean byteBufferSupported() { + return true; + } + + @Override + public ByteBuffer getByteBuffer() { + if (!buffer.hasRemaining()) { + return null; + } + final ByteBuffer view = buffer.duplicate(); + if (exposeLimit < view.remaining()) { + view.limit(view.position() + exposeLimit); + } + return view; + } + + @Override + public InputStream detach() { + state.detachCount++; + final MockGrpcInputStream owner = new MockGrpcInputStream(buffer, state, true, exposeLimit); + buffer = ByteBuffer.allocate(0); + return owner; + } + + @Override + public void close() { + if (detached) { + state.detachedCloseCount++; + } + buffer = ByteBuffer.allocate(0); + } + } +} diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestFlightZeroCopyRead.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestFlightZeroCopyRead.java new file mode 100644 index 0000000000..4431f301df --- /dev/null +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestFlightZeroCopyRead.java @@ -0,0 +1,129 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.flight; + +import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; +import static org.apache.arrow.flight.Location.forGrpcInsecure; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.google.common.io.ByteStreams; +import io.grpc.Drainable; +import java.io.ByteArrayOutputStream; +import java.io.InputStream; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.VectorUnloader; +import org.apache.arrow.vector.ipc.message.IpcOption; +import org.junit.jupiter.api.Test; + +/** + * Proves the zero-copy read path is taken over a real gRPC transport, not only against a mock. + * + *
The oracle is allocator accounting. When the client takes ownership of gRPC's buffer, it is
+ * charged for the whole wire frame (header, tags and body). When it copies, it is charged for a
+ * buffer of the body's size only. The two are never equal, so the number distinguishes the paths.
+ */
+public class TestFlightZeroCopyRead {
+
+ private static final int ROWS = 1000;
+
+ @Test
+ public void recordBatchIsReadWithoutCopyingOverRealTransport() throws Exception {
+ assertTrue(ArrowMessage.ENABLE_ZERO_COPY_READ, "zero-copy reads must be enabled for this test");
+ try (BufferAllocator serverAllocator = new RootAllocator(Long.MAX_VALUE);
+ BufferAllocator clientAllocator = new RootAllocator(Long.MAX_VALUE);
+ FlightServer server =
+ FlightServer.builder(
+ serverAllocator,
+ forGrpcInsecure(LOCALHOST, 0),
+ new OneBatchProducer(serverAllocator))
+ .build()
+ .start();
+ FlightClient client = FlightClient.builder(clientAllocator, server.getLocation()).build()) {
+ final long wireSize = wireSizeOfTheBatch(clientAllocator);
+ assertEquals(0, clientAllocator.getAllocatedMemory());
+
+ try (FlightStream stream = client.getStream(new Ticket(new byte[0]))) {
+ assertTrue(stream.next());
+ final VectorSchemaRoot root = stream.getRoot();
+ final IntVector values = (IntVector) root.getVector("c1");
+ assertEquals(ROWS, root.getRowCount());
+ for (int i = 0; i < ROWS; i++) {
+ assertEquals(i * 3, values.get(i));
+ }
+ assertEquals(
+ wireSize,
+ clientAllocator.getAllocatedMemory(),
+ "the client should be holding gRPC's whole frame, not a copy of the body");
+ assertFalse(stream.next());
+ }
+ assertEquals(0, clientAllocator.getAllocatedMemory(), "gRPC's buffer must be released");
+ }
+ }
+
+ /** The exact number of bytes the server puts on the wire for the batch the producer sends. */
+ private static long wireSizeOfTheBatch(BufferAllocator allocator) throws Exception {
+ final ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+ // Same ownership convention as the server: the batch's buffer references belong to the
+ // message, which is closed once the wire stream is done with them.
+ try (VectorSchemaRoot root = theBatch(allocator);
+ ArrowMessage message =
+ new ArrowMessage(
+ new VectorUnloader(root).getRecordBatch(), null, false, IpcOption.DEFAULT);
+ InputStream wire = ArrowMessage.createMarshaller(allocator).stream(message)) {
+ if (wire instanceof Drainable) {
+ ((Drainable) wire).drainTo(bytes);
+ } else {
+ ByteStreams.copy(wire, bytes);
+ }
+ }
+ return bytes.size();
+ }
+
+ private static VectorSchemaRoot theBatch(BufferAllocator allocator) {
+ final IntVector values = new IntVector("c1", allocator);
+ final VectorSchemaRoot root = VectorSchemaRoot.of(values);
+ root.allocateNew();
+ for (int i = 0; i < ROWS; i++) {
+ values.set(i, i * 3);
+ }
+ values.setValueCount(ROWS);
+ root.setRowCount(ROWS);
+ return root;
+ }
+
+ private static final class OneBatchProducer extends NoOpFlightProducer {
+ private final BufferAllocator allocator;
+
+ OneBatchProducer(BufferAllocator allocator) {
+ this.allocator = allocator;
+ }
+
+ @Override
+ public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) {
+ try (VectorSchemaRoot root = theBatch(allocator)) {
+ listener.start(root);
+ listener.putNext();
+ listener.completed();
+ }
+ }
+ }
+}
diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java
new file mode 100644
index 0000000000..589ce7df17
--- /dev/null
+++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/grpc/TestGetReadableBuffer.java
@@ -0,0 +1,271 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.flight.grpc;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import io.grpc.HasByteBuffer;
+import io.grpc.KnownLength;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.util.ArrayDeque;
+import java.util.Arrays;
+import java.util.Deque;
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/** Tests for reading gRPC message bodies into an {@link ArrowBuf} through gRPC's public API. */
+public class TestGetReadableBuffer {
+
+ private BufferAllocator allocator;
+
+ @BeforeEach
+ public void setUp() {
+ allocator = new RootAllocator(Long.MAX_VALUE);
+ }
+
+ @AfterEach
+ public void tearDown() {
+ assertEquals(0, allocator.getAllocatedMemory());
+ allocator.close();
+ }
+
+ @Test
+ public void fastPathReadsThroughByteBuffersNotThroughHeapReads() throws IOException {
+ final byte[] payload = payload(64);
+ try (ChunkedStream stream = new ChunkedStream(payload);
+ ArrowBuf buf = allocator.buffer(payload.length)) {
+ GetReadableBuffer.readIntoBuffer(stream, buf, payload.length, true);
+
+ assertArrayEquals(payload, contents(buf));
+ assertEquals(0, stream.available());
+ assertEquals(0, stream.heapReads, "fast path must not go through InputStream.read");
+ assertTrue(stream.byteBufferPeeks > 0, "fast path must use HasByteBuffer.getByteBuffer");
+ }
+ }
+
+ @Test
+ public void fastPathToleratesSkipReturningZero() throws IOException {
+ final byte[] payload = payload(48);
+ try (ChunkedStream stream = new ChunkedStream(true, payload);
+ ArrowBuf buf = allocator.buffer(payload.length)) {
+ GetReadableBuffer.readIntoBuffer(stream, buf, payload.length, true);
+
+ assertArrayEquals(payload, contents(buf));
+ assertEquals(0, stream.available());
+ }
+ }
+
+ @Test
+ public void fastPathCopiesAcrossChunkBoundaries() throws IOException {
+ final byte[] payload = payload(70);
+ try (ChunkedStream stream = new ChunkedStream(split(payload, 10, 1, 32, 27));
+ ArrowBuf buf = allocator.buffer(payload.length)) {
+ GetReadableBuffer.readIntoBuffer(stream, buf, payload.length, true);
+
+ assertArrayEquals(payload, contents(buf));
+ assertEquals(0, stream.available());
+ assertEquals(0, stream.heapReads);
+ }
+ }
+
+ @Test
+ public void fastPathReadsOnlyTheRequestedBytes() throws IOException {
+ final byte[] payload = payload(40);
+ try (ChunkedStream stream = new ChunkedStream(split(payload, 16, 24));
+ ArrowBuf buf = allocator.buffer(20)) {
+ GetReadableBuffer.readIntoBuffer(stream, buf, 20, true);
+
+ assertArrayEquals(Arrays.copyOf(payload, 20), contents(buf));
+ assertEquals(20, stream.available());
+ }
+ }
+
+ @Test
+ public void fallsBackToHeapReadWhenByteBuffersAreNotSupported() throws IOException {
+ final byte[] payload = payload(24);
+ try (ChunkedStream stream = ChunkedStream.withoutByteBufferSupport(split(payload, 9, 15));
+ ArrowBuf buf = allocator.buffer(payload.length)) {
+ GetReadableBuffer.readIntoBuffer(stream, buf, payload.length, true);
+
+ assertArrayEquals(payload, contents(buf));
+ assertEquals(0, stream.available());
+ assertEquals(0, stream.byteBufferPeeks, "must not call getByteBuffer when unsupported");
+ assertTrue(stream.heapReads > 0);
+ }
+ }
+
+ @Test
+ public void fastPathDisabledUsesHeapRead() throws IOException {
+ final byte[] payload = payload(24);
+ try (ChunkedStream stream = new ChunkedStream(payload);
+ ArrowBuf buf = allocator.buffer(payload.length)) {
+ GetReadableBuffer.readIntoBuffer(stream, buf, payload.length, false);
+
+ assertArrayEquals(payload, contents(buf));
+ assertEquals(0, stream.byteBufferPeeks, "fastPath=false must not touch HasByteBuffer");
+ assertTrue(stream.heapReads > 0);
+ }
+ }
+
+ @Test
+ public void fastPathThrowsWhenStreamEndsEarly() throws IOException {
+ final byte[] payload = payload(10);
+ try (ChunkedStream stream = new ChunkedStream(payload);
+ ArrowBuf buf = allocator.buffer(16)) {
+ assertThrows(
+ IOException.class, () -> GetReadableBuffer.readIntoBuffer(stream, buf, 16, true));
+ }
+ }
+
+ private static byte[][] split(byte[] payload, int... sizes) {
+ final byte[][] pieces = new byte[sizes.length][];
+ int offset = 0;
+ for (int i = 0; i < sizes.length; i++) {
+ pieces[i] = Arrays.copyOfRange(payload, offset, offset + sizes[i]);
+ offset += sizes[i];
+ }
+ if (offset != payload.length) {
+ throw new IllegalArgumentException("sizes must sum to payload length");
+ }
+ return pieces;
+ }
+
+ private static byte[] payload(int size) {
+ final byte[] bytes = new byte[size];
+ for (int i = 0; i < size; i++) {
+ bytes[i] = (byte) (i * 7 + 3);
+ }
+ return bytes;
+ }
+
+ private static byte[] contents(ArrowBuf buf) {
+ final byte[] out = new byte[(int) buf.writerIndex()];
+ buf.getBytes(0, out);
+ return out;
+ }
+
+ /**
+ * A stand-in for gRPC's message stream: a chain of direct buffers exposed through {@link
+ * HasByteBuffer} and {@link KnownLength}, with counters for how it was read.
+ */
+ static final class ChunkedStream extends InputStream implements HasByteBuffer, KnownLength {
+ private final Deque