GH-939: Remove reflection for gRPC buffers

flight/flight-core/pom.xml

@@ -61,6 +61,7 @@ under the License.
     <dependency>
       <groupId>io.grpc</groupId>
       <artifactId>grpc-core</artifactId>
+      <scope>runtime</scope>
     </dependency>
     <dependency>
       <groupId>io.grpc</groupId>
@@ -144,6 +145,12 @@ under the License.
       <type>test-jar</type>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.apache.commons</groupId>
+      <artifactId>commons-lang3</artifactId>
+      <version>3.20.0</version>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>

flight/flight-core/src/main/java/module-info.java

@@ -30,7 +30,6 @@
   requires com.google.protobuf;
   requires com.google.protobuf.util;
   requires io.grpc;
-  requires io.grpc.internal;
   requires io.grpc.netty;
   requires io.grpc.protobuf;
   requires io.grpc.stub;

flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java

@@ -19,11 +19,9 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.protobuf.ByteString;
-import com.google.protobuf.CodedInputStream;
 import com.google.protobuf.CodedOutputStream;
 import com.google.protobuf.WireFormat;
 import io.grpc.Drainable;
-import io.grpc.KnownLength;
 import io.grpc.MethodDescriptor.Marshaller;
 import io.grpc.protobuf.ProtoUtils;
 import io.netty.buffer.ByteBuf;
@@ -38,10 +36,11 @@
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.List;
+import org.apache.arrow.flight.FlightDataParser.ArrowBufReader;
+import org.apache.arrow.flight.FlightDataParser.FlightDataReader;
+import org.apache.arrow.flight.FlightDataParser.InputStreamReader;
 import org.apache.arrow.flight.grpc.AddWritableBuffer;
-import org.apache.arrow.flight.grpc.GetReadableBuffer;
 import org.apache.arrow.flight.impl.Flight.FlightData;
 import org.apache.arrow.flight.impl.Flight.FlightDescriptor;
 import org.apache.arrow.memory.ArrowBuf;
@@ -55,10 +54,14 @@
 import org.apache.arrow.vector.ipc.message.MessageSerializer;
 import org.apache.arrow.vector.types.MetadataVersion;
 import org.apache.arrow.vector.types.pojo.Schema;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /** The in-memory representation of FlightData used to manage a stream of Arrow messages. */
 class ArrowMessage implements AutoCloseable {
 
+  private static final Logger LOG = LoggerFactory.getLogger(ArrowMessage.class);
+
   // If true, deserialize Arrow data by giving Arrow a reference to the underlying gRPC buffer
   // instead of copying the data. Defaults to true.
   public static final boolean ENABLE_ZERO_COPY_READ;
@@ -79,15 +82,6 @@ class ArrowMessage implements AutoCloseable {
     ENABLE_ZERO_COPY_WRITE = "true".equalsIgnoreCase(zeroCopyWriteFlag);
   }
 
-  private static final int DESCRIPTOR_TAG =
-      (FlightData.FLIGHT_DESCRIPTOR_FIELD_NUMBER << 3) | WireFormat.WIRETYPE_LENGTH_DELIMITED;
-  private static final int BODY_TAG =
-      (FlightData.DATA_BODY_FIELD_NUMBER << 3) | WireFormat.WIRETYPE_LENGTH_DELIMITED;
-  private static final int HEADER_TAG =
-      (FlightData.DATA_HEADER_FIELD_NUMBER << 3) | WireFormat.WIRETYPE_LENGTH_DELIMITED;
-  private static final int APP_METADATA_TAG =
-      (FlightData.APP_METADATA_FIELD_NUMBER << 3) | WireFormat.WIRETYPE_LENGTH_DELIMITED;
-
   private static final Marshaller<FlightData> NO_BODY_MARSHALLER =
       ProtoUtils.marshaller(FlightData.getDefaultInstance());
 
@@ -143,6 +137,8 @@ public static HeaderType getHeader(byte b) {
   private final ArrowBuf appMetadata;
   private final List<ArrowBuf> bufs;
   private final boolean tryZeroCopyWrite;
+  // For zero-copy reads, this releases the message-scoped allocator after local buffers close.
+  private final BufferAllocator messageAllocator;
 
   public ArrowMessage(FlightDescriptor descriptor, Schema schema, IpcOption option) {
     this.writeOption = option;
@@ -153,6 +149,7 @@ public ArrowMessage(FlightDescriptor descriptor, Schema schema, IpcOption option
     this.descriptor = descriptor;
     this.appMetadata = null;
     this.tryZeroCopyWrite = false;
+    this.messageAllocator = null;
   }
 
   /**
@@ -172,6 +169,7 @@ public ArrowMessage(
     this.descriptor = null;
     this.appMetadata = appMetadata;
     this.tryZeroCopyWrite = tryZeroCopy;
+    this.messageAllocator = null;
   }
 
   public ArrowMessage(ArrowDictionaryBatch batch, IpcOption option) {
@@ -185,6 +183,7 @@ public ArrowMessage(ArrowDictionaryBatch batch, IpcOption option) {
     this.descriptor = null;
     this.appMetadata = null;
     this.tryZeroCopyWrite = false;
+    this.messageAllocator = null;
   }
 
   /**
@@ -200,6 +199,7 @@ public ArrowMessage(ArrowBuf appMetadata) {
     this.descriptor = null;
     this.appMetadata = appMetadata;
     this.tryZeroCopyWrite = false;
+    this.messageAllocator = null;
   }
 
   public ArrowMessage(FlightDescriptor descriptor) {
@@ -210,13 +210,23 @@ public ArrowMessage(FlightDescriptor descriptor) {
     this.descriptor = descriptor;
     this.appMetadata = null;
     this.tryZeroCopyWrite = false;
+    this.messageAllocator = null;
   }
 
-  private ArrowMessage(
+  ArrowMessage(
       FlightDescriptor descriptor,
       MessageMetadataResult message,
       ArrowBuf appMetadata,
       ArrowBuf buf) {
+    this(descriptor, message, appMetadata, buf, null);
+  }
+
+  ArrowMessage(
+      FlightDescriptor descriptor,
+      MessageMetadataResult message,
+      ArrowBuf appMetadata,
+      ArrowBuf buf,
+      BufferAllocator messageAllocator) {
     // No need to take IpcOption as this is used for deserialized ArrowMessage coming from the wire.
     this.writeOption =
         message != null
@@ -229,6 +239,7 @@ private ArrowMessage(
     this.appMetadata = appMetadata;
     this.bufs = buf == null ? ImmutableList.of() : ImmutableList.of(buf);
     this.tryZeroCopyWrite = false;
+    this.messageAllocator = messageAllocator;
   }
 
   public MessageMetadataResult asSchemaMessage() {
@@ -269,6 +280,7 @@ public ArrowDictionaryBatch asDictionaryBatch() throws IOException {
         bufs.size() == 1, "A batch can only be consumed if it contains a single ArrowBuf.");
     Preconditions.checkArgument(getMessageType() == HeaderType.DICTIONARY_BATCH);
     ArrowBuf underlying = bufs.get(0);
+
     // Retain a reference to keep the batch alive when the message is closed
     underlying.getReferenceManager().retain();
     // Do not set drained - we still want to release our reference
@@ -280,165 +292,16 @@ public Iterable<ArrowBuf> getBufs() {
   }
 
   private static ArrowMessage frame(BufferAllocator allocator, final InputStream stream) {
-
-    ArrowBuf body = null;
-    ArrowBuf appMetadata = null;
-    try {
-      FlightDescriptor descriptor = null;
-      MessageMetadataResult header = null;
-      while (stream.available() > 0) {
-        final int tagFirstByte = stream.read();
-        if (tagFirstByte == -1) {
-          break;
-        }
-        int tag = readRawVarint32(tagFirstByte, stream);
-        switch (tag) {
-          case DESCRIPTOR_TAG:
-            {
-              byte[] bytes = readFieldBytes(stream);
-              descriptor = FlightDescriptor.parseFrom(bytes);
-              break;
-            }
-          case HEADER_TAG:
-            {
-              byte[] bytes = readFieldBytes(stream);
-              header = MessageMetadataResult.create(ByteBuffer.wrap(bytes), bytes.length);
-              break;
-            }
-          case APP_METADATA_TAG:
-            {
-              if (appMetadata != null) {
-                // only read last app metadata.
-                appMetadata.close();
-                appMetadata = null;
-              }
-              appMetadata = readFieldBuffer(allocator, stream);
-              break;
-            }
-          case BODY_TAG:
-            if (body != null) {
-              // only read last body.
-              body.getReferenceManager().release();
-              body = null;
-            }
-            body = readFieldBuffer(allocator, stream);
-            break;
-
-          default:
-            // ignore unknown fields.
-        }
+    FlightDataReader reader;
+    if (ENABLE_ZERO_COPY_READ) {
+      reader = ArrowBufReader.tryArrowBufReader(allocator, stream);
+      if (reader != null) {
+        return reader.toMessage();
       }
-      // Protobuf implementations can omit empty fields, such as body; for some message types, like
-      // RecordBatch,
-      // this will fail later as we still expect an empty buffer. In those cases only, fill in an
-      // empty buffer here -
-      // in other cases, like Schema, having an unexpected empty buffer will also cause failures.
-      // We don't fill in defaults for fields like header, for which there is no reasonable default,
-      // or for appMetadata
-      // or descriptor, which are intended to be empty in some cases.
-      if (header != null) {
-        switch (HeaderType.getHeader(header.headerType())) {
-          case SCHEMA:
-            // Ignore 0-length buffers in case a Protobuf implementation wrote it out
-            if (body != null && body.capacity() == 0) {
-              body.close();
-              body = null;
-            }
-            break;
-          case DICTIONARY_BATCH:
-          case RECORD_BATCH:
-            // A Protobuf implementation can skip 0-length bodies, so ensure we fill it in here
-            if (body == null) {
-              body = allocator.getEmpty();
-            }
-            break;
-          case NONE:
-          case TENSOR:
-          default:
-            // Do nothing
-            break;
-        }
-      }
-      return new ArrowMessage(descriptor, header, appMetadata, body);
-    } catch (Exception ioe) {
-      // No ArrowMessage takes ownership of the buffers read so far, so release them here.
-      AutoCloseables.close(ioe, appMetadata, body);
-      throw new RuntimeException(ioe);
     }
-  }
-
-  private static int readRawVarint32(InputStream is) throws IOException {
-    int firstByte = is.read();
-    return readRawVarint32(firstByte, is);
-  }
-
-  private static int readRawVarint32(int firstByte, InputStream is) throws IOException {
-    return CodedInputStream.readRawVarint32(firstByte, is);
-  }
-
-  /**
-   * Read and validate the length prefix of a length-delimited field.
-   *
-   * <p>The length is read straight off the wire, so it must not size an allocation unchecked. It
-   * can only be compared to the bytes left in the message when the stream is {@link KnownLength}:
-   * {@code available()} means something else on other streams (the decompressing stream gRPC uses
-   * for compressed messages reports 1 until EOF), so for those the read itself bounds the
-   * allocation, see {@link #readFieldBytes}.
-   */
-  private static int readFieldLength(InputStream stream) throws IOException {
-    final int size = readRawVarint32(stream);
-    if (size < 0) {
-      throw new IOException("Malformed FlightData frame: negative field length " + size);
-    }
-    if (stream instanceof KnownLength && size > stream.available()) {
-      throw fieldTooLong(size, stream.available());
-    }
-    return size;
-  }
-
-  /**
-   * Read a length-delimited field into a byte array.
-   *
-   * <p>{@link InputStream#readNBytes(int)} grows the array as bytes arrive, so a length prefix
-   * larger than the message is rejected without allocating the declared length first.
-   */
-  private static byte[] readFieldBytes(InputStream stream) throws IOException {
-    final int size = readFieldLength(stream);
-    final byte[] bytes = stream.readNBytes(size);
-    if (bytes.length != size) {
-      throw fieldTooLong(size, bytes.length);
-    }
-    return bytes;
-  }
-
-  /** Read a length-delimited field into a new buffer. The caller must release the buffer. */
-  private static ArrowBuf readFieldBuffer(BufferAllocator allocator, InputStream stream)
-      throws IOException {
-    if (!(stream instanceof KnownLength)) {
-      // The length can't be checked up front, so only allocate once the bytes have arrived.
-      final byte[] bytes = readFieldBytes(stream);
-      final ArrowBuf buf = allocator.buffer(bytes.length);
-      buf.writeBytes(bytes);
-      return buf;
-    }
-    final int size = readFieldLength(stream);
-    final ArrowBuf buf = allocator.buffer(size);
-    try {
-      GetReadableBuffer.readIntoBuffer(stream, buf, size, ENABLE_ZERO_COPY_READ);
-    } catch (IOException | RuntimeException e) {
-      buf.close();
-      throw e;
-    }
-    return buf;
-  }
 
-  private static IOException fieldTooLong(int size, int remaining) {
-    return new IOException(
-        "Malformed FlightData frame: field length "
-            + size
-            + " exceeds "
-            + remaining
-            + " bytes remaining in the message");
+    reader = new InputStreamReader(allocator, stream);
+    return reader.toMessage();
   }
 
   /**
@@ -650,6 +513,7 @@ public ArrowMessage parse(InputStream stream) {
 
   @Override
   public void close() throws Exception {
-    AutoCloseables.close(Iterables.concat(bufs, Collections.singletonList(appMetadata)));
+    AutoCloseables.close(
+        Iterables.concat(bufs, AutoCloseables.iter(appMetadata, messageAllocator)));
   }
 }