Change getStreams to use multiple clients

java/flight/flight-jdbc-driver/pom.xml

@@ -143,6 +143,11 @@
             <artifactId>bcpkix-jdk15on</artifactId>
             <version>1.61</version>
         </dependency>
+        <dependency>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-pool2</artifactId>
+            <version>2.10.0</version>
+        </dependency>
     </dependencies>
 
     <build>

java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java

@@ -20,14 +20,17 @@
 import java.io.IOException;
 import java.security.GeneralSecurityException;
 import java.sql.SQLException;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
+import java.util.NoSuchElementException;
 import java.util.Set;
-import java.util.stream.Collectors;
 
 import org.apache.arrow.driver.jdbc.client.utils.ClientAuthenticationUtils;
+import org.apache.arrow.driver.jdbc.client.utils.KeyedFlightSqlClientObjectPool;
+import org.apache.arrow.driver.jdbc.client.utils.KeyedFlightSqlClientObjectPoolFactory;
 import org.apache.arrow.flight.CallOption;
 import org.apache.arrow.flight.FlightClient;
 import org.apache.arrow.flight.FlightClientMiddleware;
@@ -49,31 +52,39 @@
 import org.apache.arrow.util.Preconditions;
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.apache.calcite.avatica.Meta.StatementType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * A {@link FlightSqlClient} handler.
  */
 public final class ArrowFlightSqlClientHandler implements AutoCloseable {
 
+  private static final Logger logger = LoggerFactory.getLogger(ArrowFlightSqlClientHandler.class);
   private final FlightSqlClient sqlClient;
   private final Set<CallOption> options = new HashSet<>();
+  private final KeyedFlightSqlClientObjectPool flightSqlClientPool;
 
   ArrowFlightSqlClientHandler(final FlightSqlClient sqlClient,
-                              final Collection<CallOption> options) {
+                              BufferAllocator allocator, final Collection<CallOption> options) {
     this.options.addAll(options);
+    this.flightSqlClientPool = new KeyedFlightSqlClientObjectPool(
+        new KeyedFlightSqlClientObjectPoolFactory(allocator));
     this.sqlClient = Preconditions.checkNotNull(sqlClient);
   }
 
   /**
    * Creates a new {@link ArrowFlightSqlClientHandler} from the provided {@code client} and {@code options}.
    *
    * @param client  the {@link FlightClient} to manage under a {@link FlightSqlClient} wrapper.
+   * @param allocator the {@link BufferAllocator} used to create the client.
    * @param options the {@link CallOption}s to persist in between subsequent client calls.
    * @return a new {@link ArrowFlightSqlClientHandler}.
    */
   public static ArrowFlightSqlClientHandler createNewHandler(final FlightClient client,
+                                                             final BufferAllocator allocator,
                                                              final Collection<CallOption> options) {
-    return new ArrowFlightSqlClientHandler(new FlightSqlClient(client), options);
+    return new ArrowFlightSqlClientHandler(new FlightSqlClient(client), allocator, options);
   }
 
   /**
@@ -92,11 +103,36 @@ private CallOption[] getOptions() {
    * @param flightInfo The {@link FlightInfo} instance from which to fetch results.
    * @return a {@code FlightStream} of results.
    */
-  public List<FlightStream> getStreams(final FlightInfo flightInfo) {
-    return flightInfo.getEndpoints().stream()
-        .map(FlightEndpoint::getTicket)
-        .map(ticket -> sqlClient.getStream(ticket, getOptions()))
-        .collect(Collectors.toList());
+  public List<FlightStream> getStreams(final FlightInfo flightInfo) throws SQLException {
+    final List<FlightStream> allStreams = new ArrayList<>();
+
+    for (final FlightEndpoint endpoint : flightInfo.getEndpoints()) {
+      List<Location> locations = endpoint.getLocations();
+      if (locations.isEmpty()) {
+        allStreams.add(sqlClient.getStream(endpoint.getTicket(), getOptions()));
+        continue;
+      }
+      final Location location = locations.get(0); // purposefully discard other locations
+
+      logger.info(String.format("Getting a client for location %s.", location));
+      FlightSqlClient flightSqlClient;
+      try {
+        flightSqlClient = flightSqlClientPool.borrowObject(location);
+      } catch (final NoSuchElementException e) {
+        try {
+          flightSqlClientPool.addObject(location);
+          flightSqlClient = flightSqlClientPool.borrowObject(location);
+        } catch (final Exception addObjectEx) {
+          throw new SQLException("Failed to create and get a new FlightSqlClient in the ObjectPool.", addObjectEx);
+        }
+      } catch (final Exception borrowObjectEx) {
+        throw new SQLException("Failed to borrow a FlightSqlClient from the ObjectPool", borrowObjectEx);
+      }
+
+      allStreams.add(flightSqlClient.getStream(endpoint.getTicket(), getOptions()));
+      flightSqlClientPool.returnObject(location, flightSqlClient);
+    }
+    return allStreams;
   }
 
   /**
@@ -114,6 +150,7 @@ public FlightInfo getInfo(final String query) {
   public void close() throws SQLException {
     try {
       AutoCloseables.close(sqlClient);
+      flightSqlClientPool.close();
     } catch (final Exception e) {
       throw new SQLException("Failed to clean up client resources.", e);
     }
@@ -529,7 +566,7 @@ public ArrowFlightSqlClientHandler build() throws SQLException {
               ClientAuthenticationUtils.getAuthenticate(
                   client, new CredentialCallOption(new BearerCredentialWriter(token))));
         }
-        return ArrowFlightSqlClientHandler.createNewHandler(client, options);
+        return ArrowFlightSqlClientHandler.createNewHandler(client, allocator, options);
 
       } catch (final IllegalArgumentException | GeneralSecurityException | IOException | FlightRuntimeException e) {
         final SQLException originalException = new SQLException(e);

java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/client/utils/KeyedFlightSqlClientObjectPool.java

@@ -0,0 +1,39 @@
+/*
+ * 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.driver.jdbc.client.utils;
+
+import org.apache.arrow.flight.Location;
+import org.apache.arrow.flight.sql.FlightSqlClient;
+import org.apache.commons.pool2.KeyedPooledObjectFactory;
+import org.apache.commons.pool2.impl.GenericKeyedObjectPool;
+
+/**
+ * Draft.
+ */
+public class KeyedFlightSqlClientObjectPool extends GenericKeyedObjectPool<Location, FlightSqlClient> {
+
+  public KeyedFlightSqlClientObjectPool(KeyedPooledObjectFactory<Location, FlightSqlClient> factory) {
+    super(factory);
+  }
+
+  @Override
+  public synchronized void close() {
+    ((KeyedFlightSqlClientObjectPoolFactory) getFactory()).closeAllocator();
+    super.close();
+  }
+}

java/flight/flight-jdbc-driver/src/main/java/org/apache/arrow/driver/jdbc/client/utils/KeyedFlightSqlClientObjectPoolFactory.java

@@ -0,0 +1,78 @@
+/*
+ * 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.driver.jdbc.client.utils;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+import org.apache.arrow.flight.FlightClient;
+import org.apache.arrow.flight.Location;
+import org.apache.arrow.flight.sql.FlightSqlClient;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.commons.pool2.BaseKeyedPooledObjectFactory;
+import org.apache.commons.pool2.PooledObject;
+import org.apache.commons.pool2.impl.DefaultPooledObject;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Draft.
+ */
+public class KeyedFlightSqlClientObjectPoolFactory extends BaseKeyedPooledObjectFactory<Location, FlightSqlClient> {
+  private static final Logger logger = LoggerFactory.getLogger(KeyedFlightSqlClientObjectPoolFactory.class);
+  private final BufferAllocator parentAllocator;
+  private final AtomicInteger clientCounter = new AtomicInteger();
+
+  /**
+   * Draft.
+   * @param parentAllocator allocator.
+   */
+  public KeyedFlightSqlClientObjectPoolFactory(final BufferAllocator parentAllocator) {
+    super();
+    this.parentAllocator = parentAllocator
+        .newChildAllocator("KeyedFlightSqlClientObjectPoolFactory", 0, parentAllocator.getLimit());
+  }
+
+  @Override
+  public FlightSqlClient create(Location key) throws Exception {
+    logger.info("Trying to create a new FlightSqlClient.");
+    return new FlightSqlClient(
+        FlightClient.builder(
+            parentAllocator.newChildAllocator(
+                "flight-sql-client-pool_id-" + clientCounter.getAndIncrement(),
+                0,
+                parentAllocator.getLimit()),
+            key).build());
+  }
+
+  public void closeAllocator() {
+    parentAllocator.getChildAllocators().forEach(BufferAllocator::close);
+    parentAllocator.close();
+  }
+
+  @Override
+  public PooledObject<FlightSqlClient> wrap(FlightSqlClient value) {
+    logger.info("Wrapping an existing FlightSqlClient.");
+    return new DefaultPooledObject<>(value);
+  }
+
+  @Override
+  public void destroyObject(Location key, PooledObject<FlightSqlClient> p) throws Exception {
+    logger.info("Closing a client.");
+    p.getObject().close();
+  }
+}