Use ArrowBuf instead of byte[] for Flight metadata

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

@@ -53,7 +53,6 @@ import com.google.protobuf.WireFormat;
 import io.grpc.Drainable;
 import io.grpc.MethodDescriptor;
 import io.grpc.MethodDescriptor.Marshaller;
-import io.grpc.internal.ReadableBuffer;
 import io.grpc.protobuf.ProtoUtils;
 
 import io.netty.buffer.ArrowBuf;
@@ -80,8 +79,8 @@ class ArrowMessage implements AutoCloseable {
 
   private static Marshaller<FlightData> NO_BODY_MARSHALLER = ProtoUtils.marshaller(FlightData.getDefaultInstance());
 
-  /** Get the application-specific metadata in this message. */
-  public byte[] getApplicationMetadata() {
+  /** Get the application-specific metadata in this message. The ArrowMessage retains ownership of the buffer. */
+  public ArrowBuf getApplicationMetadata() {
     return appMetadata;
   }
 
@@ -122,7 +121,7 @@ class ArrowMessage implements AutoCloseable {
 
   private final FlightDescriptor descriptor;
   private final Message message;
-  private final byte[] appMetadata;
+  private final ArrowBuf appMetadata;
   private final List<ArrowBuf> bufs;
 
   public ArrowMessage(FlightDescriptor descriptor, Schema schema) {
@@ -136,7 +135,12 @@ class ArrowMessage implements AutoCloseable {
     this.appMetadata = null;
   }
 
-  public ArrowMessage(ArrowRecordBatch batch, byte[] appMetadata) {
+  /**
+   * Create an ArrowMessage from a record batch and app metadata.
+   * @param batch The record batch.
+   * @param appMetadata The app metadata. May be null. Takes ownership of the buffer otherwise.
+   */
+  public ArrowMessage(ArrowRecordBatch batch, ArrowBuf appMetadata) {
     FlatBufferBuilder builder = new FlatBufferBuilder();
     int batchOffset = batch.writeTo(builder);
     ByteBuffer serializedMessage = MessageSerializer.serializeMessage(builder, MessageHeader.RecordBatch, batchOffset,
@@ -163,7 +167,7 @@ class ArrowMessage implements AutoCloseable {
     this.appMetadata = null;
   }
 
-  private ArrowMessage(FlightDescriptor descriptor, Message message, byte[] appMetadata, ArrowBuf buf) {
+  private ArrowMessage(FlightDescriptor descriptor, Message message, ArrowBuf appMetadata, ArrowBuf buf) {
     this.message = message;
     this.descriptor = descriptor;
     this.appMetadata = appMetadata;
@@ -196,6 +200,7 @@ class ArrowMessage implements AutoCloseable {
     RecordBatch recordBatch = new RecordBatch();
     message.header(recordBatch);
     ArrowBuf underlying = bufs.get(0);
+    underlying.getReferenceManager().retain();
     ArrowRecordBatch batch = MessageSerializer.deserializeRecordBatch(recordBatch, underlying);
     return batch;
   }
@@ -217,7 +222,7 @@ class ArrowMessage implements AutoCloseable {
       FlightDescriptor descriptor = null;
       Message header = null;
       ArrowBuf body = null;
-      byte[] appMetadata = null;
+      ArrowBuf appMetadata = null;
       while (stream.available() > 0) {
         int tag = readRawVarint32(stream);
         switch (tag) {
@@ -238,8 +243,8 @@ class ArrowMessage implements AutoCloseable {
           }
           case APP_METADATA_TAG: {
             int size = readRawVarint32(stream);
-            appMetadata = new byte[size];
-            ByteStreams.readFully(stream, appMetadata);
+            appMetadata = allocator.buffer(size);
+            GetReadableBuffer.readIntoBuffer(stream, appMetadata, size, FAST_PATH);
             break;
           }
           case BODY_TAG:
@@ -250,15 +255,7 @@ class ArrowMessage implements AutoCloseable {
             }
             int size = readRawVarint32(stream);
             body = allocator.buffer(size);
-            ReadableBuffer readableBuffer = FAST_PATH ? GetReadableBuffer.getReadableBuffer(stream) : null;
-            if (readableBuffer != null) {
-              readableBuffer.readBytes(body.nioBuffer(0, size));
-            } else {
-              byte[] heapBytes = new byte[size];
-              ByteStreams.readFully(stream, heapBytes);
-              body.writeBytes(heapBytes);
-            }
-            body.writerIndex(size);
+            GetReadableBuffer.readIntoBuffer(stream, body, size, FAST_PATH);
             break;
 
           default:
@@ -309,8 +306,10 @@ class ArrowMessage implements AutoCloseable {
       CodedOutputStream cos = CodedOutputStream.newInstance(baos);
       cos.writeBytes(FlightData.DATA_HEADER_FIELD_NUMBER, bytes);
 
-      if (appMetadata != null && appMetadata.length > 0) {
-        cos.writeByteArray(FlightData.APP_METADATA_FIELD_NUMBER, appMetadata);
+      if (appMetadata != null && appMetadata.capacity() > 0) {
+        cos.writeByteBuffer(FlightData.APP_METADATA_FIELD_NUMBER, appMetadata.asNettyBuffer().nioBuffer());
+        // This is weird, but implicitly, writing an ArrowMessage frees any references it has
+        appMetadata.getReferenceManager().release();
       }
 
       cos.writeTag(FlightData.DATA_BODY_FIELD_NUMBER, WireFormat.WIRETYPE_LENGTH_DELIMITED);
@@ -335,6 +334,7 @@ class ArrowMessage implements AutoCloseable {
       initialBuf.writeBytes(baos.toByteArray());
       final CompositeByteBuf bb = new CompositeByteBuf(allocator.getAsByteBufAllocator(), true, bufs.size() + 1,
           ImmutableList.<ByteBuf>builder().add(initialBuf.asNettyBuffer()).addAll(allBufs).build());
+      // Implicitly, transfer ownership of our buffers to the input stream (which will decrement the refcount when done)
       final ByteBufInputStream is = new DrainableByteBufInputStream(bb);
       return is;
     } catch (Exception ex) {
@@ -364,7 +364,7 @@ class ArrowMessage implements AutoCloseable {
     }
 
     @Override
-    public void close() throws IOException {
+    public void close() {
       buf.release();
     }
 
@@ -399,5 +399,8 @@ class ArrowMessage implements AutoCloseable {
   @Override
   public void close() throws Exception {
     AutoCloseables.close(bufs);
+    if (appMetadata != null) {
+      appMetadata.close();
+    }
   }
 }

java/flight/src/main/java/org/apache/arrow/flight/FlightClient.java

@@ -19,12 +19,7 @@ package org.apache.arrow.flight;
 
 import java.io.InputStream;
 import java.net.URISyntaxException;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashSet;
 import java.util.Iterator;
-import java.util.List;
-import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
@@ -42,17 +37,11 @@ import org.apache.arrow.flight.impl.FlightServiceGrpc;
 import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceBlockingStub;
 import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceStub;
 import org.apache.arrow.memory.BufferAllocator;
-import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.VectorUnloader;
-import org.apache.arrow.vector.dictionary.Dictionary;
 import org.apache.arrow.vector.dictionary.DictionaryProvider;
 import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider;
-import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch;
 import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
-import org.apache.arrow.vector.types.pojo.Field;
-import org.apache.arrow.vector.types.pojo.Schema;
-import org.apache.arrow.vector.util.DictionaryUtility;
 
 import com.google.common.base.Preconditions;
 import com.google.common.base.Throwables;
@@ -69,6 +58,7 @@ import io.grpc.stub.ClientCalls;
 import io.grpc.stub.ClientResponseObserver;
 import io.grpc.stub.StreamObserver;
 
+import io.netty.buffer.ArrowBuf;
 import io.netty.channel.EventLoopGroup;
 import io.netty.channel.ServerChannel;
 import io.netty.handler.ssl.SslContextBuilder;
@@ -169,9 +159,11 @@ public class FlightClient implements AutoCloseable {
 
   /**
    * Create or append a descriptor with another stream.
+   *
    * @param descriptor FlightDescriptor the descriptor for the data
    * @param root VectorSchemaRoot the root containing data
-   * @param metadataListener A handler for metadata messages from the server.
+   * @param metadataListener A handler for metadata messages from the server. This will be passed buffers that will be
+   *     freed after {@link StreamListener#onNext(Object)} is called!
    * @param options RPC-layer hints for this call.
    * @return ClientStreamListener an interface to control uploading data
    */
@@ -193,7 +185,7 @@ public class FlightClient implements AutoCloseable {
     Preconditions.checkNotNull(descriptor);
     Preconditions.checkNotNull(root);
 
-    SetStreamObserver resultObserver = new SetStreamObserver(metadataListener);
+    SetStreamObserver resultObserver = new SetStreamObserver(allocator, metadataListener);
     final io.grpc.CallOptions callOptions = CallOptions.wrapStub(asyncStub, options).getCallOptions();
     ClientCallStreamObserver<ArrowMessage> observer = (ClientCallStreamObserver<ArrowMessage>)
         ClientCalls.asyncBidiStreamingCall(
@@ -268,33 +260,32 @@ public class FlightClient implements AutoCloseable {
   private static class SetStreamObserver implements StreamObserver<Flight.PutResult> {
 
     private final CompletableFuture<Void> result;
+    private final BufferAllocator allocator;
     private final StreamListener<PutResult> listener;
 
-    SetStreamObserver(FlightProducer.StreamListener<PutResult> listener) {
-      this.listener = listener;
+    SetStreamObserver(BufferAllocator allocator,
+        StreamListener<PutResult> listener) {
+      this.allocator = allocator;
+      this.listener = listener == null ? NoOpStreamListener.getInstance() : listener;
       result = new CompletableFuture<>();
     }
 
     @Override
     public void onNext(Flight.PutResult value) {
-      if (listener != null) {
-        listener.onNext(PutResult.fromProtocol(value));
+      try (final PutResult message = PutResult.fromProtocol(allocator, value)) {
+        listener.onNext(message);
       }
     }
 
     @Override
     public void onError(Throwable t) {
       result.completeExceptionally(t);
-      if (listener != null) {
-        listener.onError(t);
-      }
+      listener.onError(t);
     }
 
     @Override
     public void onCompleted() {
-      if (listener != null) {
-        listener.onCompleted();
-      }
+      listener.onCompleted();
       result.complete(null);
     }
 
@@ -322,12 +313,13 @@ public class FlightClient implements AutoCloseable {
     }
 
     @Override
-    public void putNext(byte[] appMetadata) {
+    public void putNext(ArrowBuf appMetadata) {
       ArrowRecordBatch batch = unloader.getRecordBatch();
       // Check the futureResult in case server sent an exception
       while (!observer.isReady() && !futureResult.isDone()) {
         /* busy wait */
       }
+      // Takes ownership of appMetadata
       observer.onNext(new ArrowMessage(batch, appMetadata));
     }
 
@@ -361,7 +353,11 @@ public class FlightClient implements AutoCloseable {
      */
     void putNext();
 
-    void putNext(byte[] appMetadata);
+    /**
+     * Send the current data in the corresponding {@link VectorSchemaRoot} to the server, along with
+     * application-specific metadata. This takes ownership of the buffer.
+     */
+    void putNext(ArrowBuf appMetadata);
 
     /**
      * Indicate an error to the server. Terminates the stream; do not call {@link #completed()}.
@@ -376,7 +372,6 @@ public class FlightClient implements AutoCloseable {
      * happened during the upload.
      */
     void getResult();
-
   }
 
   /**

java/flight/src/main/java/org/apache/arrow/flight/FlightProducer.java

@@ -20,6 +20,8 @@ package org.apache.arrow.flight;
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.dictionary.DictionaryProvider;
 
+import io.netty.buffer.ArrowBuf;
+
 /**
  * API to Implement an Arrow Flight producer.
  */
@@ -59,7 +61,7 @@ public interface FlightProducer {
    * @param context Per-call context.
    * @param flightStream The data stream being uploaded.
    */
-  public Runnable acceptPut(CallContext context,
+  Runnable acceptPut(CallContext context,
       FlightStream flightStream, StreamListener<PutResult> ackStream);
 
   /**
@@ -82,7 +84,7 @@ public interface FlightProducer {
   /**
    * An interface for sending Arrow data back to a client.
    */
-  public interface ServerStreamListener {
+  interface ServerStreamListener {
 
     /**
      * Check whether the call has been cancelled. If so, stop sending data.
@@ -115,8 +117,9 @@ public interface FlightProducer {
 
     /**
      * Send the current contents of the associated {@link VectorSchemaRoot} alongside application-defined metadata.
+     * @param metadata The metadata to send. Ownership of the buffer is transferred to the Flight implementation.
      */
-    void putNext(byte[] metadata);
+    void putNext(ArrowBuf metadata);
 
     /**
      * Indicate an error to the client. Terminates the stream; do not call {@link #completed()} afterwards.

java/flight/src/main/java/org/apache/arrow/flight/FlightService.java

@@ -17,11 +17,6 @@
 
 package org.apache.arrow.flight;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.function.BooleanSupplier;
@@ -37,17 +32,10 @@ import org.apache.arrow.flight.impl.Flight.HandshakeRequest;
 import org.apache.arrow.flight.impl.Flight.HandshakeResponse;
 import org.apache.arrow.flight.impl.FlightServiceGrpc.FlightServiceImplBase;
 import org.apache.arrow.memory.BufferAllocator;
-import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.VectorUnloader;
-import org.apache.arrow.vector.dictionary.Dictionary;
 import org.apache.arrow.vector.dictionary.DictionaryProvider;
 import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider;
-import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch;
-import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
-import org.apache.arrow.vector.types.pojo.Field;
-import org.apache.arrow.vector.types.pojo.Schema;
-import org.apache.arrow.vector.util.DictionaryUtility;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -55,6 +43,7 @@ import com.google.common.base.Preconditions;
 
 import io.grpc.stub.ServerCallStreamObserver;
 import io.grpc.stub.StreamObserver;
+import io.netty.buffer.ArrowBuf;
 
 /**
  * GRPC service implementation for a flight server.
@@ -167,7 +156,7 @@ class FlightService extends FlightServiceImplBase {
     }
 
     @Override
-    public void putNext(byte[] metadata) {
+    public void putNext(ArrowBuf metadata) {
       Preconditions.checkNotNull(unloader);
       responseObserver.onNext(new ArrowMessage(unloader.getRecordBatch(), metadata));
     }

java/flight/src/main/java/org/apache/arrow/flight/FlightStream.java

@@ -46,6 +46,7 @@ import com.google.common.collect.Iterables;
 import com.google.common.util.concurrent.SettableFuture;
 
 import io.grpc.stub.StreamObserver;
+import io.netty.buffer.ArrowBuf;
 
 /**
  * An adaptor between protobuf streams and flight data streams.
@@ -72,7 +73,7 @@ public class FlightStream implements AutoCloseable {
   private volatile Throwable ex;
   private volatile FlightDescriptor descriptor;
   private volatile Schema schema;
-  private volatile byte[] applicationMetadata = null;
+  private volatile ArrowBuf applicationMetadata = null;
 
   /**
    * Constructs a new instance.
@@ -117,6 +118,9 @@ public class FlightStream implements AutoCloseable {
         .collect(Collectors.toList());
 
     AutoCloseables.close(Iterables.concat(closeables, ImmutableList.of(root.get())));
+    if (applicationMetadata != null) {
+      applicationMetadata.close();
+    }
   }
 
   /**
@@ -149,31 +153,38 @@ public class FlightStream implements AutoCloseable {
           throw new Exception(ex);
         }
       } else {
-        ArrowMessage msg = ((ArrowMessage) data);
-        if (msg.getMessageType() == HeaderType.RECORD_BATCH) {
-          try (ArrowRecordBatch arb = msg.asRecordBatch()) {
-            loader.load(arb);
-          }
-          this.applicationMetadata = msg.getApplicationMetadata();
-        } else if (msg.getMessageType() == HeaderType.DICTIONARY_BATCH) {
-          try (ArrowDictionaryBatch arb = msg.asDictionaryBatch()) {
-            final long id = arb.getDictionaryId();
-            final Dictionary dictionary = dictionaries.lookup(id);
-            if (dictionary == null) {
-              throw new IllegalArgumentException("Dictionary not defined in schema: ID " + id);
+        try (ArrowMessage msg = ((ArrowMessage) data)) {
+          if (msg.getMessageType() == HeaderType.RECORD_BATCH) {
+            try (ArrowRecordBatch arb = msg.asRecordBatch()) {
+              loader.load(arb);
             }
+            if (this.applicationMetadata != null) {
+              this.applicationMetadata.close();
+            }
+            this.applicationMetadata = msg.getApplicationMetadata();
+            if (this.applicationMetadata != null) {
+              this.applicationMetadata.getReferenceManager().retain();
+            }
+          } else if (msg.getMessageType() == HeaderType.DICTIONARY_BATCH) {
+            try (ArrowDictionaryBatch arb = msg.asDictionaryBatch()) {
+              final long id = arb.getDictionaryId();
+              final Dictionary dictionary = dictionaries.lookup(id);
+              if (dictionary == null) {
+                throw new IllegalArgumentException("Dictionary not defined in schema: ID " + id);
+              }
 
-            final FieldVector vector = dictionary.getVector();
-            final VectorSchemaRoot dictionaryRoot = new VectorSchemaRoot(Collections.singletonList(vector.getField()),
-                Collections.singletonList(vector), 0);
-            final VectorLoader dictionaryLoader = new VectorLoader(dictionaryRoot);
-            dictionaryLoader.load(arb.getDictionary());
+              final FieldVector vector = dictionary.getVector();
+              final VectorSchemaRoot dictionaryRoot = new VectorSchemaRoot(Collections.singletonList(vector.getField()),
+                  Collections.singletonList(vector), 0);
+              final VectorLoader dictionaryLoader = new VectorLoader(dictionaryRoot);
+              dictionaryLoader.load(arb.getDictionary());
+            }
+            return next();
+          } else {
+            throw new UnsupportedOperationException("Message type is unsupported: " + msg.getMessageType());
           }
-          return next();
-        } else {
-          throw new UnsupportedOperationException("Message type is unsupported: " + msg.getMessageType());
+          return true;
         }
-        return true;
       }
     } catch (Exception e) {
       throw new RuntimeException(e);
@@ -191,11 +202,12 @@ public class FlightStream implements AutoCloseable {
 
   /**
    * Get the most recent metadata sent from the server. This may be cleared by calls to {@link #next()} if the server
-   * sends a message without metadata.
+   * sends a message without metadata. This does NOT take ownership of the buffer - call retain() to create a reference
+   * if you need the buffer after a call to {@link #next()}.
    *
    * @return the application metadata. May be null.
    */
-  public byte[] getLatestMetadata() {
+  public ArrowBuf getLatestMetadata() {
     return applicationMetadata;
   }
 

java/flight/src/main/java/org/apache/arrow/flight/NoOpStreamListener.java

@@ -0,0 +1,49 @@
+/*
+ * 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 org.apache.arrow.flight.FlightProducer.StreamListener;
+
+/**
+ * A {@link StreamListener} that does nothing for all callbacks.
+ * @param <T> The type of the callback object.
+ */
+public class NoOpStreamListener<T> implements StreamListener<T> {
+  private static NoOpStreamListener INSTANCE = new NoOpStreamListener();
+
+  /** Ignores the value received. */
+  @Override
+  public void onNext(T val) {
+  }
+
+  /** Ignores the error received. */
+  @Override
+  public void onError(Throwable t) {
+  }
+
+  /** Ignores the stream completion event. */
+  @Override
+  public void onCompleted() {
+  }
+
+  @SuppressWarnings("unchecked")
+  public static <T> StreamListener<T> getInstance() {
+    // Safe because we never use T
+    return (StreamListener<T>) INSTANCE;
+  }
+}

java/flight/src/main/java/org/apache/arrow/flight/PutResult.java

@@ -17,29 +17,32 @@
 
 package org.apache.arrow.flight;
 
-import java.nio.ByteBuffer;
-
 import org.apache.arrow.flight.impl.Flight;
+import org.apache.arrow.memory.BufferAllocator;
 
 import com.google.protobuf.ByteString;
 
+import io.netty.buffer.ArrowBuf;
+
 /**
  * A message from the server during a DoPut operation.
+ *
+ * <p>This object owns an {@link ArrowBuf} and should be closed when you are done with it.
  */
-public class PutResult {
+public class PutResult implements AutoCloseable {
 
-  private ByteBuffer applicationMetadata;
+  private ArrowBuf applicationMetadata;
 
-  private PutResult(ByteBuffer metadata) {
+  private PutResult(ArrowBuf metadata) {
     applicationMetadata = metadata;
   }
 
   /** Create a PutResult with application-specific metadata. */
-  public static PutResult metadata(byte[] metadata) {
+  public static PutResult metadata(ArrowBuf metadata) {
     if (metadata == null) {
       return empty();
     }
-    return new PutResult(ByteBuffer.wrap(metadata));
+    return new PutResult(metadata);
   }
 
   /** Create an empty PutResult. */
@@ -48,7 +51,7 @@ public class PutResult {
   }
 
   /** Get the metadata in this message. May be null. */
-  public ByteBuffer getApplicationMetadata() {
+  public ArrowBuf getApplicationMetadata() {
     return applicationMetadata;
   }
 
@@ -56,10 +59,29 @@ public class PutResult {
     if (applicationMetadata == null) {
       return Flight.PutResult.getDefaultInstance();
     }
-    return Flight.PutResult.newBuilder().setAppMetadata(ByteString.copyFrom(applicationMetadata)).build();
+    return Flight.PutResult.newBuilder().setAppMetadata(ByteString.copyFrom(applicationMetadata.nioBuffer())).build();
   }
 
-  static PutResult fromProtocol(Flight.PutResult message) {
-    return new PutResult(message.getAppMetadata().asReadOnlyByteBuffer());
+  /**
+   * Construct a PutResult from a Protobuf message.
+   *
+   * @param allocator The allocator to use for allocating application metadata memory. The result object owns the
+   *     allocated buffer, if any.
+   * @param message The gRPC/Protobuf message.
+   */
+  static PutResult fromProtocol(BufferAllocator allocator, Flight.PutResult message) {
+    final ArrowBuf buf = allocator.buffer(message.getAppMetadata().size());
+    message.getAppMetadata().asReadOnlyByteBufferList().forEach(bb -> {
+      buf.setBytes(buf.writerIndex(), bb);
+      buf.writerIndex(buf.writerIndex() + bb.limit());
+    });
+    return new PutResult(buf);
+  }
+
+  @Override
+  public void close() {
+    if (applicationMetadata != null) {
+      applicationMetadata.close();
+    }
   }
 }

java/flight/src/main/java/org/apache/arrow/flight/example/Stream.java

@@ -36,6 +36,8 @@ import org.apache.arrow.vector.types.pojo.Schema;
 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
 
+import io.netty.buffer.ArrowBuf;
+
 /**
  * A collection of Arrow record batches.
  */
@@ -91,8 +93,12 @@ public class Stream implements AutoCloseable, Iterable<ArrowRecordBatch> {
       final VectorLoader loader = new VectorLoader(root);
       int counter = 0;
       for (ArrowRecordBatch batch : batches) {
+        final byte[] rawMetadata = Integer.toString(counter).getBytes(StandardCharsets.UTF_8);
+        final ArrowBuf metadata = allocator.buffer(rawMetadata.length);
+        metadata.writeBytes(rawMetadata);
         loader.load(batch);
-        listener.putNext(Integer.toString(counter).getBytes(StandardCharsets.UTF_8));
+        // Transfers ownership of the buffer - do not free buffer ourselves
+        listener.putNext(metadata);
         counter++;
       }
       listener.completed();

java/flight/src/main/java/org/apache/arrow/flight/example/integration/IntegrationTestClient.java

@@ -22,6 +22,10 @@ import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 import org.apache.arrow.flight.FlightClient;
 import org.apache.arrow.flight.FlightDescriptor;
@@ -44,6 +48,8 @@ import org.apache.commons.cli.DefaultParser;
 import org.apache.commons.cli.Options;
 import org.apache.commons.cli.ParseException;
 
+import io.netty.buffer.ArrowBuf;
+
 /**
  * An Example Flight Server that provides access to the InMemoryStore.
  */
@@ -92,17 +98,19 @@ class IntegrationTestClient {
     FlightDescriptor descriptor = FlightDescriptor.path(inputPath);
     VectorSchemaRoot jsonRoot;
     try (JsonFileReader reader = new JsonFileReader(new File(inputPath), allocator);
-         VectorSchemaRoot root = VectorSchemaRoot.create(reader.start(), allocator)) {
+        VectorSchemaRoot root = VectorSchemaRoot.create(reader.start(), allocator)) {
       jsonRoot = VectorSchemaRoot.create(root.getSchema(), allocator);
       VectorUnloader unloader = new VectorUnloader(root);
       VectorLoader jsonLoader = new VectorLoader(jsonRoot);
+      final CompletableFuture<Void> completed = new CompletableFuture<>();
       FlightClient.ClientStreamListener stream = client.startPut(descriptor, root, reader,
           new StreamListener<PutResult>() {
             int counter = 0;
 
             @Override
             public void onNext(PutResult val) {
-              final String metadata = StandardCharsets.UTF_8.decode(val.getApplicationMetadata()).toString();
+              final String metadata = StandardCharsets.UTF_8.decode(val.getApplicationMetadata().nioBuffer())
+                  .toString();
               if (!Integer.toString(counter).equals(metadata)) {
                 throw new RuntimeException(
                     String.format("Invalid ACK from server. Expected '%d' but got '%s'.", counter, metadata));
@@ -112,15 +120,20 @@ class IntegrationTestClient {
 
             @Override
             public void onError(Throwable t) {
+              completed.completeExceptionally(t);
             }
 
             @Override
             public void onCompleted() {
+              completed.complete(null);
             }
           });
       int counter = 0;
       while (reader.read(root)) {
-        stream.putNext(Integer.toString(counter).getBytes(StandardCharsets.UTF_8));
+        final ArrowBuf metadata = allocator.buffer(0);
+        metadata.writeBytes(Integer.toString(counter).getBytes(StandardCharsets.UTF_8));
+        // Transfers ownership of the buffer, so do not release it ourselves
+        stream.putNext(metadata);
         jsonLoader.load(unloader.getRecordBatch());
         root.clear();
         counter++;
@@ -128,6 +141,9 @@ class IntegrationTestClient {
       stream.completed();
       // Need to call this, or exceptions from the server get swallowed
       stream.getResult();
+      completed.get(30, TimeUnit.SECONDS);
+    } catch (InterruptedException | ExecutionException | TimeoutException e) {
+      throw new RuntimeException(e);
     }
 
     // 2. Get the ticket for the data.

java/flight/src/main/java/org/apache/arrow/flight/grpc/GetReadableBuffer.java

@@ -17,12 +17,15 @@
 
 package org.apache.arrow.flight.grpc;
 
+import java.io.IOException;
 import java.io.InputStream;
 import java.lang.reflect.Field;
 
 import com.google.common.base.Throwables;
+import com.google.common.io.ByteStreams;
 
 import io.grpc.internal.ReadableBuffer;
+import io.netty.buffer.ArrowBuf;
 
 /**
  * Enable access to ReadableBuffer directly to copy data from an BufferInputStream into a target
@@ -72,4 +75,24 @@ public class GetReadableBuffer {
     }
   }
 
+  /**
+   * Helper method to read a gRPC-provided InputStream into an ArrowBuf.
+   * @param stream The stream to read from. Should be an instance of {@link #BUFFER_INPUT_STREAM}.
+   * @param buf The buffer to read into.
+   * @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
+   */
+  public static void readIntoBuffer(final InputStream stream, final ArrowBuf buf, final int size,
+      final boolean fastPath) throws IOException {
+    ReadableBuffer readableBuffer = fastPath ? getReadableBuffer(stream) : null;
+    if (readableBuffer != null) {
+      readableBuffer.readBytes(buf.nioBuffer(0, size));
+    } else {
+      byte[] heapBytes = new byte[size];
+      ByteStreams.readFully(stream, heapBytes);
+      buf.writeBytes(heapBytes);
+    }
+    buf.writerIndex(size);
+  }
 }

java/flight/src/test/java/org/apache/arrow/flight/TestApplicationMetadata.java

@@ -17,11 +17,9 @@
 
 package org.apache.arrow.flight;
 
-import java.util.Arrays;
 import java.util.Collections;
 
 import org.apache.arrow.flight.FlightProducer.StreamListener;
-import org.apache.arrow.flight.auth.ServerAuthHandler;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.IntVector;
@@ -34,6 +32,9 @@ import org.junit.Assert;
 import org.junit.Ignore;
 import org.junit.Test;
 
+import io.grpc.Status;
+import io.netty.buffer.ArrowBuf;
+
 /**
  * Tests for application-specific metadata support in Flight.
  */
@@ -57,7 +58,7 @@ public class TestApplicationMetadata {
         final IntVector vector = (IntVector) stream.getRoot().getVector("a");
         Assert.assertEquals(1, vector.getValueCount());
         Assert.assertEquals(10, vector.get(0));
-        Assert.assertArrayEquals(new byte[]{i}, stream.getLatestMetadata());
+        Assert.assertEquals(i, stream.getLatestMetadata().getByte(0));
         i++;
       }
     } catch (Exception e) {
@@ -85,7 +86,7 @@ public class TestApplicationMetadata {
         @Override
         public void onNext(PutResult val) {
           Assert.assertNotNull(val);
-          Assert.assertEquals(counter, val.getApplicationMetadata().get(0));
+          Assert.assertEquals(counter, val.getApplicationMetadata().getByte(0));
           counter++;
         }
 
@@ -104,10 +105,12 @@ public class TestApplicationMetadata {
       root.allocateNew();
       for (byte i = 0; i < 10; i++) {
         final IntVector vector = (IntVector) root.getVector("a");
+        final ArrowBuf metadata = a.buffer(1);
+        metadata.writeByte(i);
         vector.set(0, 10);
         vector.setValueCount(1);
         root.setRowCount(1);
-        writer.putNext(new byte[]{i});
+        writer.putNext(metadata);
       }
       writer.completed();
       // Must attempt to retrieve the result to get any server-side errors.
@@ -139,7 +142,9 @@ public class TestApplicationMetadata {
           vector.set(0, 10);
           vector.setValueCount(1);
           root.setRowCount(1);
-          listener.putNext(new byte[]{i});
+          final ArrowBuf metadata = allocator.buffer(1);
+          metadata.writeByte(i);
+          listener.putNext(metadata);
         }
         listener.completed();
       }
@@ -151,11 +156,14 @@ public class TestApplicationMetadata {
         try (FlightStream flightStream = stream) {
           byte current = 0;
           while (flightStream.next()) {
-            byte[] metadata = flightStream.getLatestMetadata();
-            if (!Arrays.equals(new byte[]{current}, metadata)) {
-              throw new IllegalArgumentException("Metadata does not match expected value.");
+            final ArrowBuf metadata = flightStream.getLatestMetadata();
+            if (current != metadata.getByte(0)) {
+              ackStream.onError(Status.INVALID_ARGUMENT.withDescription(String
+                  .format("Metadata does not match expected value; got %d but expected %d.", metadata.getByte(0),
+                      current)).asRuntimeException());
+              return;
             }
-            ackStream.onNext(PutResult.metadata(new byte[]{current}));
+            ackStream.onNext(PutResult.metadata(metadata));
             current++;
           }
           if (current != 10) {

java/flight/src/test/java/org/apache/arrow/flight/TestBackPressure.java

@@ -40,7 +40,6 @@ public class TestBackPressure {
   /**
    * Make sure that failing to consume one stream doesn't block other streams.
    */
-  @Ignore
   @Test
   public void ensureIndependentSteams() throws Exception {
     try (
@@ -49,26 +48,24 @@ public class TestBackPressure {
             (location) -> (new PerformanceTestServer(a, location)));
         final FlightClient client = FlightClient.builder(a, server.getLocation()).build()
     ) {
-      FlightStream fs1 = client.getStream(client.getInfo(
+      try (FlightStream fs1 = client.getStream(client.getInfo(
           TestPerf.getPerfFlightDescriptor(110L * BATCH_SIZE, BATCH_SIZE, 1))
-          .getEndpoints().get(0).getTicket());
-      consume(fs1, 10);
+          .getEndpoints().get(0).getTicket())) {
+        consume(fs1, 10);
 
-      // stop consuming fs1 but make sure we can consume a large amount of fs2.
-      FlightStream fs2 = client.getStream(client.getInfo(
-          TestPerf.getPerfFlightDescriptor(200L * BATCH_SIZE, BATCH_SIZE, 1))
-          .getEndpoints().get(0).getTicket());
-      consume(fs2, 100);
+        // stop consuming fs1 but make sure we can consume a large amount of fs2.
+        try (FlightStream fs2 = client.getStream(client.getInfo(
+            TestPerf.getPerfFlightDescriptor(200L * BATCH_SIZE, BATCH_SIZE, 1))
+            .getEndpoints().get(0).getTicket())) {
+          consume(fs2, 100);
 
-      consume(fs1, 100);
-      consume(fs2, 100);
-
-      consume(fs1);
-      consume(fs2);
-
-      fs1.close();
-      fs2.close();
+          consume(fs1, 100);
+          consume(fs2, 100);
 
+          consume(fs1);
+          consume(fs2);
+        }
+      }
     }
   }
 
@@ -92,27 +89,28 @@ public class TestBackPressure {
             ServerStreamListener listener) {
           int batches = 0;
           final Schema pojoSchema = new Schema(ImmutableList.of(Field.nullable("a", MinorType.BIGINT.getType())));
-          VectorSchemaRoot root = VectorSchemaRoot.create(pojoSchema, allocator);
-          listener.start(root);
-          while (true) {
-            while (!listener.isReady()) {
-              try {
-                Thread.sleep(1);
-                sleepTime.addAndGet(1L);
-              } catch (InterruptedException ignore) {
+          try (VectorSchemaRoot root = VectorSchemaRoot.create(pojoSchema, allocator)) {
+            listener.start(root);
+            while (true) {
+              while (!listener.isReady()) {
+                try {
+                  Thread.sleep(1);
+                  sleepTime.addAndGet(1L);
+                } catch (InterruptedException ignore) {
+                }
               }
-            }
 
-            if (batches > 100) {
-              root.clear();
-              listener.completed();
-              return;
-            }
+              if (batches > 100) {
+                root.clear();
+                listener.completed();
+                return;
+              }
 
-            root.allocateNew();
-            root.setRowCount(4095);
-            listener.putNext();
-            batches++;
+              root.allocateNew();
+              root.setRowCount(4095);
+              listener.putNext();
+              batches++;
+            }
           }
         }
       };
@@ -127,9 +125,9 @@ public class TestBackPressure {
           FlightClient client =
               FlightClient
                   .builder(clientAllocator, server.getLocation())
-                  .build()
+                  .build();
+          FlightStream stream = client.getStream(new Ticket(new byte[1]))
       ) {
-        FlightStream stream = client.getStream(new Ticket(new byte[1]));
         VectorSchemaRoot root = stream.getRoot();
         root.clear();
         Thread.sleep(wait);

java/flight/src/test/java/org/apache/arrow/flight/TestBasicOperation.java

@@ -97,19 +97,8 @@ public class TestBasicOperation {
       IntVector iv = new IntVector("c1", a);
 
       VectorSchemaRoot root = VectorSchemaRoot.of(iv);
-      ClientStreamListener listener = c.startPut(FlightDescriptor.path("hello"), root, new StreamListener<PutResult>() {
-        @Override
-        public void onNext(PutResult val) {
-        }
-
-        @Override
-        public void onError(Throwable t) {
-        }
-
-        @Override
-        public void onCompleted() {
-        }
-      });
+      ClientStreamListener listener = c
+          .startPut(FlightDescriptor.path("hello"), root, NoOpStreamListener.getInstance());
 
       //batch 1
       root.allocateNew();

java/flight/src/test/java/org/apache/arrow/flight/TestLargeMessage.java

@@ -21,7 +21,6 @@ import java.util.Arrays;
 import java.util.List;
 import java.util.stream.Stream;
 
-import org.apache.arrow.flight.FlightProducer.StreamListener;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.IntVector;
@@ -78,22 +77,7 @@ public class TestLargeMessage {
            BufferAllocator testAllocator = a.newChildAllocator("testcase", 0, Long.MAX_VALUE);
            VectorSchemaRoot root = generateData(testAllocator)) {
         final FlightClient.ClientStreamListener listener = client.startPut(FlightDescriptor.path("hello"), root,
-            new StreamListener<PutResult>() {
-              @Override
-              public void onNext(PutResult val) {
-
-              }
-
-              @Override
-              public void onError(Throwable t) {
-
-              }
-
-              @Override
-              public void onCompleted() {
-
-              }
-            });
+            NoOpStreamListener.getInstance());
         listener.putNext();
         listener.completed();
         listener.getResult();

java/flight/src/test/java/org/apache/arrow/flight/example/TestExampleServer.java

@@ -27,6 +27,7 @@ import org.apache.arrow.flight.FlightProducer.StreamListener;
 import org.apache.arrow.flight.FlightStream;
 import org.apache.arrow.flight.FlightTestUtil;
 import org.apache.arrow.flight.Location;
+import org.apache.arrow.flight.NoOpStreamListener;
 import org.apache.arrow.flight.PutResult;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
@@ -40,7 +41,6 @@ import org.junit.Test;
 /**
  * Ensure that example server supports get and put.
  */
-@org.junit.Ignore
 public class TestExampleServer {
 
   private BufferAllocator allocator;
@@ -78,19 +78,7 @@ public class TestExampleServer {
 
     VectorSchemaRoot root = VectorSchemaRoot.of(iv);
     ClientStreamListener listener = client.startPut(FlightDescriptor.path("hello"), root,
-        new StreamListener<PutResult>() {
-          @Override
-          public void onNext(PutResult val) {
-          }
-
-          @Override
-          public void onError(Throwable t) {
-          }
-
-          @Override
-          public void onCompleted() {
-          }
-        });
+        NoOpStreamListener.getInstance());
 
     //batch 1
     root.allocateNew();
@@ -117,10 +105,13 @@ public class TestExampleServer {
     listener.getResult();
 
     FlightInfo info = client.getInfo(FlightDescriptor.path("hello"));
-    FlightStream stream = client.getStream(info.getEndpoints().get(0).getTicket());
-    VectorSchemaRoot newRoot = stream.getRoot();
-    while (stream.next()) {
-      newRoot.clear();
+    try (final FlightStream stream = client.getStream(info.getEndpoints().get(0).getTicket())) {
+      VectorSchemaRoot newRoot = stream.getRoot();
+      while (stream.next()) {
+        newRoot.clear();
+      }
+    } catch (Exception e) {
+      throw new RuntimeException(e);
     }
   }
 }

java/flight/src/test/java/org/apache/arrow/flight/perf/TestPerf.java

@@ -28,7 +28,6 @@ import org.apache.arrow.flight.FlightDescriptor;
 import org.apache.arrow.flight.FlightInfo;
 import org.apache.arrow.flight.FlightStream;
 import org.apache.arrow.flight.FlightTestUtil;
-import org.apache.arrow.flight.Location;
 import org.apache.arrow.flight.Ticket;
 import org.apache.arrow.flight.perf.impl.PerfOuterClass.Perf;
 import org.apache.arrow.memory.BufferAllocator;
@@ -92,11 +91,13 @@ public class TestPerf {
             .map(t -> pool.submit(t))
             .collect(Collectors.toList());
 
-        Futures.whenAllSucceed(results);
-        Result r = new Result();
-        for (ListenableFuture<Result> f : results) {
-          r.add(f.get());
-        }
+        final Result r = Futures.whenAllSucceed(results).call(() -> {
+          Result res = new Result();
+          for (ListenableFuture<Result> f : results) {
+            res.add(f.get());
+          }
+          return res;
+        }).get();
 
         double seconds = r.nanos * 1.0d / 1000 / 1000 / 1000;
         System.out.println(String.format(
@@ -126,28 +127,29 @@ public class TestPerf {
     public Result call() throws Exception {
       final Result r = new Result();
       Stopwatch watch = Stopwatch.createStarted();
-      FlightStream stream = client.getStream(ticket);
-      final VectorSchemaRoot root = stream.getRoot();
-      try {
-        BigIntVector a = (BigIntVector) root.getVector("a");
-        while (stream.next()) {
-          int rows = root.getRowCount();
-          long aSum = r.aSum;
-          for (int i = 0; i < rows; i++) {
-            if (VALIDATE) {
-              aSum += a.get(i);
+      try (final FlightStream stream = client.getStream(ticket)) {
+        final VectorSchemaRoot root = stream.getRoot();
+        try {
+          BigIntVector a = (BigIntVector) root.getVector("a");
+          while (stream.next()) {
+            int rows = root.getRowCount();
+            long aSum = r.aSum;
+            for (int i = 0; i < rows; i++) {
+              if (VALIDATE) {
+                aSum += a.get(i);
+              }
             }
+            r.bytes += rows * 32;
+            r.rows += rows;
+            r.aSum = aSum;
+            r.batches++;
           }
-          r.bytes += rows * 32;
-          r.rows += rows;
-          r.aSum = aSum;
-          r.batches++;
-        }
 
-        r.nanos = watch.elapsed(TimeUnit.NANOSECONDS);
-        return r;
-      } finally {
-        root.clear();
+          r.nanos = watch.elapsed(TimeUnit.NANOSECONDS);
+          return r;
+        } finally {
+          root.clear();
+        }
       }
     }