Use builder for Flight server in Java

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

@@ -20,7 +20,6 @@ package org.apache.arrow.flight;
 import static io.grpc.stub.ClientCalls.asyncClientStreamingCall;
 import static io.grpc.stub.ClientCalls.asyncServerStreamingCall;
 
-import java.net.URI;
 import java.net.URISyntaxException;
 import java.util.Iterator;
 import java.util.concurrent.TimeUnit;
@@ -50,8 +49,8 @@ import com.google.common.util.concurrent.SettableFuture;
 
 import io.grpc.ClientCall;
 import io.grpc.ManagedChannel;
-import io.grpc.ManagedChannelBuilder;
 import io.grpc.MethodDescriptor;
+import io.grpc.netty.NettyChannelBuilder;
 import io.grpc.stub.ClientCallStreamObserver;
 import io.grpc.stub.ClientResponseObserver;
 import io.grpc.stub.StreamObserver;
@@ -71,17 +70,9 @@ public class FlightClient implements AutoCloseable {
   private final MethodDescriptor<Flight.Ticket, ArrowMessage> doGetDescriptor;
   private final MethodDescriptor<ArrowMessage, Flight.PutResult> doPutDescriptor;
 
-  /**
-   * Construct client for accessing RouteGuide server using the existing channel.
-   */
-  public FlightClient(BufferAllocator incomingAllocator, Location location) {
-    final ManagedChannelBuilder<?> channelBuilder =
-        ManagedChannelBuilder.forAddress(location.getUri().getHost(), location.getUri().getPort())
-            .maxTraceEvents(MAX_CHANNEL_TRACE_EVENTS)
-            .maxInboundMessageSize(FlightServer.MAX_GRPC_MESSAGE_SIZE)
-            .usePlaintext();
+  private FlightClient(BufferAllocator incomingAllocator, ManagedChannel channel) {
     this.allocator = incomingAllocator.newChildAllocator("flight-client", 0, Long.MAX_VALUE);
-    channel = channelBuilder.build();
+    this.channel = channel;
     blockingStub = FlightServiceGrpc.newBlockingStub(channel).withInterceptors(authInterceptor);
     asyncStub = FlightServiceGrpc.newStub(channel).withInterceptors(authInterceptor);
     doGetDescriptor = FlightBindingService.getDoGetDescriptor(allocator);
@@ -324,4 +315,64 @@ public class FlightClient implements AutoCloseable {
     allocator.close();
   }
 
+  /**
+   * Create a builder for a Flight client.
+   * @param allocator The allocator to use for the client.
+   * @param location The location to connect to.
+   */
+  public static Builder builder(BufferAllocator allocator, Location location) {
+    return new Builder(allocator, location);
+  }
+
+  /**
+   * A builder for Flight clients.
+   */
+  public static final class Builder {
+
+    private final BufferAllocator allocator;
+    private final Location location;
+    private boolean forceTls = false;
+
+    private Builder(BufferAllocator allocator, Location location) {
+      this.allocator = allocator;
+      this.location = location;
+    }
+
+    /**
+     * Force the client to connect over TLS.
+     */
+    public Builder useTls() {
+      this.forceTls = true;
+      return this;
+    }
+
+    /**
+     * Create the client from this builder.
+     */
+    public FlightClient build() {
+      final NettyChannelBuilder builder;
+
+      switch (location.getUri().getScheme()) {
+        case LocationSchemes.GRPC:
+        case LocationSchemes.GRPC_INSECURE:
+        case LocationSchemes.GRPC_TLS: {
+          builder = NettyChannelBuilder.forAddress(location.getUri().getHost(), location.getUri().getPort());
+          break;
+        }
+        default:
+          throw new IllegalArgumentException("Scheme is not supported: " + location.getUri().getScheme());
+      }
+
+      if (this.forceTls || LocationSchemes.GRPC_TLS.equals(location.getUri().getScheme())) {
+        builder.useTransportSecurity();
+      } else {
+        builder.usePlaintext();
+      }
+
+      builder
+          .maxTraceEvents(MAX_CHANNEL_TRACE_EVENTS)
+          .maxInboundMessageSize(FlightServer.MAX_GRPC_MESSAGE_SIZE);
+      return new FlightClient(allocator, builder.build());
+    }
+  }
 }

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

@@ -17,8 +17,13 @@
 
 package org.apache.arrow.flight;
 
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileNotFoundException;
 import java.io.IOException;
+import java.io.InputStream;
 import java.net.InetSocketAddress;
+import java.util.concurrent.ForkJoinPool;
 import java.util.concurrent.TimeUnit;
 
 import org.apache.arrow.flight.auth.ServerAuthHandler;
@@ -39,41 +44,9 @@ public class FlightServer implements AutoCloseable {
   /** The maximum size of an individual gRPC message. This effectively disables the limit. */
   static final int MAX_GRPC_MESSAGE_SIZE = Integer.MAX_VALUE;
 
-  /**
-   * Constructs a new instance.
-   *
-   * @param allocator The allocator to use for storing/copying arrow data.
-   * @param location The location to serve on.
-   * @param producer The underlying business logic for the server.
-   * @param authHandler The authorization handler for the server.
-   */
-  public FlightServer(
-      BufferAllocator allocator,
-      Location location,
-      FlightProducer producer,
-      ServerAuthHandler authHandler) {
-    final NettyServerBuilder builder;
-    switch (location.getUri().getScheme()) {
-      case LocationSchemes.GRPC_DOMAIN_SOCKET: {
-        // TODO: need reflection to check if domain sockets are available
-        throw new UnsupportedOperationException("Domain sockets are not available.");
-      }
-      case LocationSchemes.GRPC:
-      case LocationSchemes.GRPC_INSECURE: {
-        builder = NettyServerBuilder
-            .forAddress(new InetSocketAddress(location.getUri().getHost(), location.getUri().getPort()));
-        break;
-      }
-      default:
-        throw new IllegalArgumentException("Scheme is not supported: " + location.getUri().getScheme());
-    }
-    this.server = builder
-        .maxInboundMessageSize(MAX_GRPC_MESSAGE_SIZE)
-        .addService(
-            ServerInterceptors.intercept(
-                new FlightBindingService(allocator, producer, authHandler),
-                new ServerAuthInterceptor(authHandler)))
-        .build();
+  /** Create a new instance from a gRPC server. For internal use only. */
+  private FlightServer(Server server) {
+    this.server = server;
   }
 
   /** Start the server. */
@@ -115,20 +88,86 @@ public class FlightServer implements AutoCloseable {
     }
   }
 
-  public interface OutputFlight {
-    void sendData(int count);
-
-    void done();
-
-    void fail(Throwable t);
+  public static Builder builder(BufferAllocator allocator, Location location, FlightProducer producer) {
+    return new Builder(allocator, location, producer);
   }
 
-  public interface FlightServerHandler {
+  public static final class Builder {
 
-    public FlightInfo getFlightInfo(String descriptor) throws Exception;
+    private final BufferAllocator allocator;
+    private final Location location;
+    private final FlightProducer producer;
+    private ServerAuthHandler authHandler = ServerAuthHandler.NO_OP;
 
-    public OutputFlight setupFlight(VectorSchemaRoot root);
+    private InputStream certChain;
+    private InputStream key;
 
+    Builder(BufferAllocator allocator, Location location, FlightProducer producer) {
+      this.allocator = allocator;
+      this.location = location;
+      this.producer = producer;
+    }
+
+    public FlightServer build() {
+      final NettyServerBuilder builder;
+      switch (location.getUri().getScheme()) {
+        case LocationSchemes.GRPC_DOMAIN_SOCKET: {
+          // TODO: need reflection to check if domain sockets are available
+          throw new UnsupportedOperationException("Domain sockets are not available.");
+        }
+        case LocationSchemes.GRPC:
+        case LocationSchemes.GRPC_INSECURE: {
+          builder = NettyServerBuilder
+              .forAddress(new InetSocketAddress(location.getUri().getHost(), location.getUri().getPort()));
+          break;
+        }
+        case LocationSchemes.GRPC_TLS: {
+          if (certChain == null) {
+            throw new IllegalArgumentException("Must provide a certificate and key to serve gRPC over TLS");
+          }
+          builder = NettyServerBuilder
+              .forAddress(new InetSocketAddress(location.getUri().getHost(), location.getUri().getPort()));
+          break;
+        }
+        default:
+          throw new IllegalArgumentException("Scheme is not supported: " + location.getUri().getScheme());
+      }
+
+      if (certChain != null) {
+        builder.useTransportSecurity(certChain, key);
+      }
+
+      final Server server = builder
+          .executor(new ForkJoinPool())
+          .maxInboundMessageSize(MAX_GRPC_MESSAGE_SIZE)
+          .addService(
+              ServerInterceptors.intercept(
+                  new FlightBindingService(allocator, producer, authHandler),
+                  new ServerAuthInterceptor(authHandler)))
+          .build();
+      return new FlightServer(server);
+    }
+
+    /**
+     * Enable TLS on the server.
+     * @param certChain The certificate chain to use.
+     * @param key The private key to use.
+     */
+    public Builder useTls(final File certChain, final File key) throws IOException {
+      this.certChain = new FileInputStream(certChain);
+      this.key = new FileInputStream(key);
+      return this;
+    }
+
+    public Builder useTls(final InputStream certChain, final InputStream key) {
+      this.certChain = certChain;
+      this.key = key;
+      return this;
+    }
+
+    public Builder setAuthHandler(ServerAuthHandler authHandler) {
+      this.authHandler = authHandler;
+      return this;
+    }
   }
-
 }

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

@@ -46,7 +46,7 @@ public class ExampleFlightServer implements AutoCloseable {
     this.allocator = allocator.newChildAllocator("flight-server", 0, Long.MAX_VALUE);
     this.location = location;
     this.mem = new InMemoryStore(this.allocator, location);
-    this.flightServer = new FlightServer(allocator, location, mem, ServerAuthHandler.NO_OP);
+    this.flightServer = FlightServer.builder(allocator, location, mem).build();
   }
 
   public Location getLocation() {

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

@@ -82,7 +82,7 @@ class IntegrationTestClient {
 
     final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
     final Location defaultLocation = Location.forGrpcInsecure(host, port);
-    final FlightClient client = new FlightClient(allocator, defaultLocation);
+    final FlightClient client = FlightClient.builder(allocator, defaultLocation).build();
 
     final String inputPath = cmd.getOptionValue("j");
 
@@ -120,7 +120,7 @@ class IntegrationTestClient {
       }
       for (Location location : locations) {
         System.out.println("Verifying location " + location.getUri());
-        FlightClient readClient = new FlightClient(allocator, location);
+        FlightClient readClient = FlightClient.builder(allocator, location).build();
         FlightStream stream = readClient.getStream(endpoint.getTicket());
         VectorSchemaRoot downloadedRoot;
         try (VectorSchemaRoot root = stream.getRoot()) {

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

@@ -48,7 +48,7 @@ public class TestBackPressure {
         final BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
         final PerformanceTestServer server = FlightTestUtil.getStartedServer(
             (port) -> (new PerformanceTestServer(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, port))));
-        final FlightClient client = new FlightClient(a, server.getLocation())
+        final FlightClient client = FlightClient.builder(a, server.getLocation()).build()
     ) {
       FlightStream fs1 = client.getStream(client.getInfo(
           TestPerf.getPerfFlightDescriptor(110L * BATCH_SIZE, BATCH_SIZE, 1))
@@ -123,11 +123,13 @@ public class TestBackPressure {
           BufferAllocator serverAllocator = allocator.newChildAllocator("server", 0, Long.MAX_VALUE);
           FlightServer server =
               FlightTestUtil.getStartedServer(
-                  (port) -> new FlightServer(serverAllocator, Location.forGrpcInsecure("localhost", port), producer,
-                      ServerAuthHandler.NO_OP));
+                  (port) -> FlightServer.builder(serverAllocator, Location.forGrpcInsecure("localhost", port), producer)
+                      .build());
           BufferAllocator clientAllocator = allocator.newChildAllocator("client", 0, Long.MAX_VALUE);
           FlightClient client =
-              new FlightClient(clientAllocator, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, server.getPort()))
+              FlightClient
+                  .builder(clientAllocator, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, server.getPort()))
+                  .build()
       ) {
         FlightStream stream = client.getStream(new Ticket(new byte[1]));
         VectorSchemaRoot root = stream.getRoot();

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

@@ -140,11 +140,12 @@ public class TestBasicOperation {
         Producer producer = new Producer(a);
         FlightServer s =
             FlightTestUtil.getStartedServer(
-                (port) -> new FlightServer(a, Location.forGrpcInsecure("localhost", port), producer,
-                    ServerAuthHandler.NO_OP))) {
+                (port) -> FlightServer.builder(a, Location.forGrpcInsecure("localhost", port), producer).build()
+            )) {
 
       try (
-          FlightClient c = new FlightClient(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, s.getPort()));
+          FlightClient c = FlightClient.builder(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, s.getPort()))
+              .build()
       ) {
         try (BufferAllocator testAllocator = a.newChildAllocator("testcase", 0, Long.MAX_VALUE)) {
           consumer.accept(c, testAllocator);

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

@@ -71,9 +71,10 @@ public class TestCallOptions {
         Producer producer = new Producer(a);
         FlightServer s =
             FlightTestUtil.getStartedServer(
-                (port) -> new FlightServer(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, port), producer,
-                    ServerAuthHandler.NO_OP));
-        FlightClient client = new FlightClient(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, s.getPort()))) {
+                (port) -> FlightServer.builder(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, port), producer)
+                    .build());
+        FlightClient client = FlightClient.builder(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, s.getPort()))
+            .build()) {
       testFn.accept(client);
     } catch (InterruptedException | IOException e) {
       throw new RuntimeException(e);

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

@@ -45,10 +45,10 @@ public class TestLargeMessage {
          final Producer producer = new Producer(a);
          final FlightServer s =
              FlightTestUtil.getStartedServer(
-                 (port) -> new FlightServer(a, Location.forGrpcInsecure("localhost", port), producer,
-                     ServerAuthHandler.NO_OP))) {
+                 (port) -> FlightServer.builder(a, Location.forGrpcInsecure("localhost", port), producer).build())) {
 
-      try (FlightClient client = new FlightClient(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, s.getPort()))) {
+      try (FlightClient client = FlightClient
+          .builder(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, s.getPort())).build()) {
         FlightStream stream = client.getStream(new Ticket(new byte[]{}));
         try (VectorSchemaRoot root = stream.getRoot()) {
           while (stream.next()) {
@@ -76,10 +76,11 @@ public class TestLargeMessage {
          final Producer producer = new Producer(a);
          final FlightServer s =
              FlightTestUtil.getStartedServer(
-                 (port) -> new FlightServer(a, Location.forGrpcInsecure("localhost", port), producer,
-                     ServerAuthHandler.NO_OP))) {
+                 (port) -> FlightServer.builder(a, Location.forGrpcInsecure("localhost", port), producer).build()
+             )) {
 
-      try (FlightClient client = new FlightClient(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, s.getPort()));
+      try (FlightClient client = FlightClient
+          .builder(a, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, s.getPort())).build();
            BufferAllocator testAllocator = a.newChildAllocator("testcase", 0, Long.MAX_VALUE);
            VectorSchemaRoot root = generateData(testAllocator)) {
         final FlightClient.ClientStreamListener listener = client.startPut(FlightDescriptor.path("hello"), root);

java/flight/src/test/java/org/apache/arrow/flight/auth/TestAuth.java

@@ -119,7 +119,7 @@ public class TestAuth {
       }
     };
 
-    server = FlightTestUtil.getStartedServer((port) -> new FlightServer(
+    server = FlightTestUtil.getStartedServer((port) -> FlightServer.builder(
         allocator,
         Location.forGrpcInsecure("localhost", port),
         new NoOpFlightProducer() {
@@ -149,9 +149,9 @@ public class TestAuth {
             root.clear();
             listener.completed();
           }
-        },
-        new BasicServerAuthHandler(validator)));
-    client = new FlightClient(allocator, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, server.getPort()));
+        }).setAuthHandler(new BasicServerAuthHandler(validator)).build());
+    client = FlightClient.builder(allocator, Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, server.getPort()))
+        .build();
   }
 
   @After

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

@@ -59,7 +59,7 @@ public class TestExampleServer {
     } else {
       System.out.println("Skipping server startup.");
     }
-    client = new FlightClient(allocator, l);
+    client = FlightClient.builder(allocator, l).build();
     caseAllocator = allocator.newChildAllocator("test-case", 0, Long.MAX_VALUE);
   }
 

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

@@ -64,7 +64,7 @@ public class PerformanceTestServer implements AutoCloseable {
     this.allocator = incomingAllocator.newChildAllocator("perf-server", 0, Long.MAX_VALUE);
     this.location = location;
     this.producer = new PerfProducer();
-    this.flightServer = new FlightServer(this.allocator, location, producer, ServerAuthHandler.NO_OP);
+    this.flightServer = FlightServer.builder(this.allocator, location, producer).build();
   }
 
   public Location getLocation() {

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

@@ -83,7 +83,7 @@ public class TestPerf {
           final PerformanceTestServer server =
               FlightTestUtil.getStartedServer((port) -> new PerformanceTestServer(a,
                   Location.forGrpcInsecure(FlightTestUtil.LOCALHOST, port)));
-          final FlightClient client = new FlightClient(a, server.getLocation());
+          final FlightClient client = FlightClient.builder(a, server.getLocation()).build();
       ) {
         final FlightInfo info = client.getInfo(getPerfFlightDescriptor(50_000_000L, 4095, 2));
         ListeningExecutorService pool = MoreExecutors.listeningDecorator(Executors.newFixedThreadPool(4));