Inline CompletableFuture in Flight acceptPut

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

@@ -194,7 +194,7 @@ public class FlightClient implements AutoCloseable {
     DictionaryUtils.generateSchemaMessages(root.getSchema(), descriptor, provider, observer::onNext);
     return new PutObserver(new VectorUnloader(
         root, true /* include # of nulls in vectors */, true /* must align buffers to be C++-compatible */),
-        observer, resultObserver.getFuture());
+        observer, resultObserver);
   }
 
   /**
@@ -257,17 +257,14 @@ public class FlightClient implements AutoCloseable {
     return stream;
   }
 
-  private static class SetStreamObserver implements StreamObserver<Flight.PutResult> {
-
-    private final CompletableFuture<Void> result;
+  private static class SetStreamObserver extends CompletableFuture<Void> implements StreamObserver<Flight.PutResult> {
     private final BufferAllocator allocator;
     private final StreamListener<PutResult> listener;
 
-    SetStreamObserver(BufferAllocator allocator,
-        StreamListener<PutResult> listener) {
+    SetStreamObserver(BufferAllocator allocator, StreamListener<PutResult> listener) {
+      super();
       this.allocator = allocator;
       this.listener = listener == null ? NoOpStreamListener.getInstance() : listener;
-      result = new CompletableFuture<>();
     }
 
     @Override
@@ -279,18 +276,14 @@ public class FlightClient implements AutoCloseable {
 
     @Override
     public void onError(Throwable t) {
-      result.completeExceptionally(t);
       listener.onError(t);
+      completeExceptionally(t);
     }
 
     @Override
     public void onCompleted() {
       listener.onCompleted();
-      result.complete(null);
-    }
-
-    public CompletableFuture<Void> getFuture() {
-      return result;
+      complete(null);
     }
   }
 

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

@@ -117,10 +117,10 @@ public class FlightStream implements AutoCloseable {
         .map(t -> ((AutoCloseable) t))
         .collect(Collectors.toList());
 
-    AutoCloseables.close(Iterables.concat(closeables, ImmutableList.of(root.get())));
-    if (applicationMetadata != null) {
-      applicationMetadata.close();
-    }
+    // Must check for null since ImmutableList doesn't accept nulls
+    AutoCloseables.close(Iterables.concat(closeables,
+        applicationMetadata != null ? ImmutableList.of(root.get(), applicationMetadata)
+            : ImmutableList.of(root.get())));
   }
 
   /**

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

@@ -19,6 +19,7 @@ package org.apache.arrow.flight;
 
 import org.apache.arrow.flight.impl.Flight;
 import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.ReferenceManager;
 
 import com.google.protobuf.ByteString;
 
@@ -37,7 +38,11 @@ public class PutResult implements AutoCloseable {
     applicationMetadata = metadata;
   }
 
-  /** Create a PutResult with application-specific metadata. */
+  /**
+   * Create a PutResult with application-specific metadata.
+   *
+   * <p>This method assumes ownership of the {@link ArrowBuf}.
+   */
   public static PutResult metadata(ArrowBuf metadata) {
     if (metadata == null) {
       return empty();
@@ -50,7 +55,12 @@ public class PutResult implements AutoCloseable {
     return new PutResult(null);
   }
 
-  /** Get the metadata in this message. May be null. */
+  /**
+   * Get the metadata in this message. May be null.
+   *
+   * <p>Ownership of the {@link ArrowBuf} is retained by this object. Call {@link ReferenceManager#retain()} to preserve
+   * a reference.
+   */
   public ArrowBuf getApplicationMetadata() {
     return applicationMetadata;
   }

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

@@ -44,8 +44,8 @@ public class TestLargeMessage {
              FlightTestUtil.getStartedServer((location) -> FlightServer.builder(a, location, producer).build())) {
 
       try (FlightClient client = FlightClient.builder(a, s.getLocation()).build()) {
-        FlightStream stream = client.getStream(new Ticket(new byte[]{}));
-        try (VectorSchemaRoot root = stream.getRoot()) {
+        try (FlightStream stream = client.getStream(new Ticket(new byte[]{}));
+            VectorSchemaRoot root = stream.getRoot()) {
           while (stream.next()) {
             for (final Field field : root.getSchema().getFields()) {
               int value = 0;
@@ -57,7 +57,6 @@ public class TestLargeMessage {
             }
           }
         }
-        stream.close();
       }
     }
   }