GH-661: [Flight] JDBC: cache failed locations

Instead of repeatedly trying to access locations which have failed
before, cache them and move them to the end of the list of
locations to attempt on subsequent accesses.

Fixes #661.

Author: lidavidm

flight/flight-sql-jdbc-core/pom.xml

@@ -143,6 +143,12 @@ under the License.
       <groupId>org.checkerframework</groupId>
       <artifactId>checker-qual</artifactId>
     </dependency>
+
+    <dependency>
+      <groupId>com.github.ben-manes.caffeine</groupId>
+      <artifactId>caffeine</artifactId>
+      <version>3.1.8</version>
+    </dependency>
   </dependencies>
 
   <build>

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java

@@ -24,6 +24,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler;
+import org.apache.arrow.driver.jdbc.client.utils.FlightClientCache;
 import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl;
 import org.apache.arrow.flight.FlightClient;
 import org.apache.arrow.memory.BufferAllocator;
@@ -113,6 +114,7 @@ private static ArrowFlightSqlClientHandler createNewClientHandler(
           .withRetainCookies(config.retainCookies())
           .withRetainAuth(config.retainAuth())
           .withCatalog(config.getCatalog())
+          .withClientCache(config.useClientCache() ? new FlightClientCache() : null)
           .withConnectTimeout(config.getConnectTimeout())
           .build();
     } catch (final SQLException e) {

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

@@ -33,6 +33,8 @@
 import java.util.Optional;
 import java.util.Set;
 import org.apache.arrow.driver.jdbc.client.utils.ClientAuthenticationUtils;
+import org.apache.arrow.driver.jdbc.client.utils.FlightClientCache;
+import org.apache.arrow.driver.jdbc.client.utils.FlightLocationQueue;
 import org.apache.arrow.flight.CallOption;
 import org.apache.arrow.flight.CallStatus;
 import org.apache.arrow.flight.CloseSessionRequest;
@@ -75,21 +77,27 @@ public final class ArrowFlightSqlClientHandler implements AutoCloseable {
   // JDBC connection string query parameter
   private static final String CATALOG = "catalog";
 
+  private final String cacheKey;
   private final FlightSqlClient sqlClient;
   private final Set<CallOption> options = new HashSet<>();
   private final Builder builder;
   private final Optional<String> catalog;
+  private final @Nullable FlightClientCache flightClientCache;
 
   ArrowFlightSqlClientHandler(
+      final String cacheKey,
       final FlightSqlClient sqlClient,
       final Builder builder,
       final Collection<CallOption> credentialOptions,
-      final Optional<String> catalog) {
+      final Optional<String> catalog,
+      final @Nullable FlightClientCache flightClientCache) {
     this.options.addAll(builder.options);
     this.options.addAll(credentialOptions);
+    this.cacheKey = Preconditions.checkNotNull(cacheKey);
     this.sqlClient = Preconditions.checkNotNull(sqlClient);
     this.builder = builder;
     this.catalog = catalog;
+    this.flightClientCache = flightClientCache;
   }
 
   /**
@@ -101,12 +109,15 @@ public final class ArrowFlightSqlClientHandler implements AutoCloseable {
    * @return a new {@link ArrowFlightSqlClientHandler}.
    */
   static ArrowFlightSqlClientHandler createNewHandler(
+      final String cacheKey,
       final FlightClient client,
       final Builder builder,
       final Collection<CallOption> options,
-      final Optional<String> catalog) {
+      final Optional<String> catalog,
+      final @Nullable FlightClientCache flightClientCache) {
     final ArrowFlightSqlClientHandler handler =
-        new ArrowFlightSqlClientHandler(new FlightSqlClient(client), builder, options, catalog);
+        new ArrowFlightSqlClientHandler(
+            cacheKey, new FlightSqlClient(client), builder, options, catalog, flightClientCache);
     handler.setSetCatalogInSessionIfPresent();
     return handler;
   }
@@ -148,9 +159,14 @@ public List<CloseableEndpointStreamPair> getStreams(final FlightInfo flightInfo)
           // location
           // is the same as the original connection's Location and skip creating a FlightClient in
           // that scenario.
+          // Also copy the cache to the client so we can share a cache. Cache needs to cache
+          // negative attempts too.
           List<Exception> exceptions = new ArrayList<>();
           CloseableEndpointStreamPair stream = null;
-          for (Location location : endpoint.getLocations()) {
+          FlightLocationQueue locations =
+              new FlightLocationQueue(flightClientCache, endpoint.getLocations());
+          while (locations.hasNext()) {
+            Location location = locations.next();
             final URI endpointUri = location.getUri();
             if (endpointUri.getScheme().equals(LocationSchemes.REUSE_CONNECTION)) {
               stream =
@@ -163,6 +179,7 @@ public List<CloseableEndpointStreamPair> getStreams(final FlightInfo flightInfo)
                     .withHost(endpointUri.getHost())
                     .withPort(endpointUri.getPort())
                     .withEncryption(endpointUri.getScheme().equals(LocationSchemes.GRPC_TLS))
+                    .withClientCache(flightClientCache)
                     .withConnectTimeout(builder.connectTimeout);
 
             ArrowFlightSqlClientHandler endpointHandler = null;
@@ -177,12 +194,29 @@ public List<CloseableEndpointStreamPair> getStreams(final FlightInfo flightInfo)
               stream.getStream().getSchema();
             } catch (Exception ex) {
               if (endpointHandler != null) {
+                // If the exception is related to connectivity, mark the client as a dud.
+                if (flightClientCache != null) {
+                  if (ex instanceof FlightRuntimeException
+                      && ((FlightRuntimeException) ex).status().code()
+                          == FlightStatusCode.UNAVAILABLE
+                      &&
+                      // IOException covers SocketException and Netty's (private)
+                      // AnnotatedSocketException
+                      // We are looking for things like "Network is unreachable"
+                      ex.getCause() instanceof IOException) {
+                    flightClientCache.markLocationAsDud(location.toString());
+                  }
+                }
+
                 AutoCloseables.close(endpointHandler);
               }
               exceptions.add(ex);
               continue;
             }
 
+            if (flightClientCache != null) {
+              flightClientCache.markLocationAsReachable(location.toString());
+            }
             break;
           }
           if (stream != null) {
@@ -549,6 +583,8 @@ public static final class Builder {
 
     @VisibleForTesting Optional<String> catalog = Optional.empty();
 
+    @VisibleForTesting @Nullable FlightClientCache flightClientCache;
+
     @VisibleForTesting @Nullable Duration connectTimeout;
 
     // These two middleware are for internal use within build() and should not be exposed by builder
@@ -833,11 +869,20 @@ public Builder withCatalog(@Nullable final String catalog) {
       return this;
     }
 
+    public Builder withClientCache(FlightClientCache flightClientCache) {
+      this.flightClientCache = flightClientCache;
+      return this;
+    }
+
     public Builder withConnectTimeout(Duration connectTimeout) {
       this.connectTimeout = connectTimeout;
       return this;
     }
 
+    public String getCacheKey() {
+      return getLocation().toString();
+    }
+
     /** Get the location that this client will connect to. */
     public Location getLocation() {
       if (useEncryption) {
@@ -931,7 +976,7 @@ public ArrowFlightSqlClientHandler build() throws SQLException {
                   options.toArray(new CallOption[0])));
         }
         return ArrowFlightSqlClientHandler.createNewHandler(
-            client, this, credentialOptions, catalog);
+            getCacheKey(), client, this, credentialOptions, catalog, flightClientCache);
 
       } catch (final IllegalArgumentException
           | GeneralSecurityException

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/utils/FlightClientCache.java

@@ -0,0 +1,55 @@
+/*
+ * 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 com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
+import java.time.Duration;
+import org.apache.arrow.util.VisibleForTesting;
+
+/**
+ * A cache for Flight clients.
+ *
+ * <p>The intent is to avoid constantly recreating clients to the same locations. gRPC can multiplex
+ * multiple requests over a single TCP connection, and a cache would let us take advantage of that.
+ *
+ * <p>At the time being it only tracks whether a location is reachable or not. To actually cache
+ * clients, we would need a way to incorporate other connection parameters (authentication, etc.)
+ * into the cache key.
+ */
+public final class FlightClientCache {
+  @VisibleForTesting Cache<String, ClientCacheEntry> clientCache;
+
+  public FlightClientCache() {
+    this.clientCache = Caffeine.newBuilder().expireAfterWrite(Duration.ofSeconds(600)).build();
+  }
+
+  public boolean isDud(String key) {
+    return clientCache.getIfPresent(key) != null;
+  }
+
+  public void markLocationAsDud(String key) {
+    clientCache.put(key, new ClientCacheEntry());
+  }
+
+  public void markLocationAsReachable(String key) {
+    clientCache.invalidate(key);
+  }
+
+  /** A cache entry (empty because we only track reachability, see outer class docstring). */
+  public static final class ClientCacheEntry {}
+}

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/utils/FlightLocationQueue.java

@@ -0,0 +1,71 @@
+/*
+ * 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.ArrayDeque;
+import java.util.Deque;
+import java.util.Iterator;
+import java.util.List;
+import java.util.NoSuchElementException;
+import org.apache.arrow.flight.Location;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * A queue of Flight locations to connect to for an endpoint.
+ *
+ * <p>This helper class is intended to encapsulate the retry logic in a testable manner.
+ */
+public final class FlightLocationQueue implements Iterator<Location> {
+  private final Deque<Location> locations;
+  private final Deque<Location> badLocations;
+
+  /**
+   * Create a new queue.
+   *
+   * @param flightClientCache An optional cache used to sort previously unreachable locations to the
+   *     end.
+   * @param locations The locations to try.
+   */
+  public FlightLocationQueue(
+      @Nullable FlightClientCache flightClientCache, List<Location> locations) {
+    this.locations = new ArrayDeque<>();
+    this.badLocations = new ArrayDeque<>();
+
+    for (Location location : locations) {
+      if (flightClientCache != null && flightClientCache.isDud(location.toString())) {
+        this.badLocations.add(location);
+      } else {
+        this.locations.add(location);
+      }
+    }
+  }
+
+  @Override
+  public boolean hasNext() {
+    return !locations.isEmpty() || !badLocations.isEmpty();
+  }
+
+  @Override
+  public Location next() {
+    if (!locations.isEmpty()) {
+      return locations.pop();
+    } else if (!badLocations.isEmpty()) {
+      return badLocations.pop();
+    }
+    throw new NoSuchElementException();
+  }
+}

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/ArrowFlightConnectionConfigImpl.java

@@ -174,6 +174,11 @@ public Duration getConnectTimeout() {
     return Duration.ofMillis(timeout);
   }
 
+  /** Whether to enable the client cache. */
+  public boolean useClientCache() {
+    return ArrowFlightConnectionProperty.USE_CLIENT_CACHE.getBoolean(properties);
+  }
+
   /**
    * Gets the {@link CallOption}s from this {@link ConnectionConfig}.
    *
@@ -226,6 +231,7 @@ public enum ArrowFlightConnectionProperty implements ConnectionProperty {
     RETAIN_AUTH("retainAuth", true, Type.BOOLEAN, false),
     CATALOG("catalog", null, Type.STRING, false),
     CONNECT_TIMEOUT_MILLIS("connectTimeoutMs", 10000, Type.NUMBER, false),
+    USE_CLIENT_CACHE("useClientCache", true, Type.BOOLEAN, false),
     ;
 
     private final String camelName;

flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java

@@ -699,10 +699,9 @@ public void testFallbackUnresolvableFlightServer() throws Exception {
         }
         attempt1 = System.nanoTime();
         elapsedMs = (attempt1 - start) / 1_000_000.;
-        // TODO(GH-661): this assertion should be flipped to assertTrue.
-        assertFalse(
+        assertTrue(
             elapsedMs < 5000.,
-            String.format("Expected second attempt to be the same, but %f ms elapsed", elapsedMs));
+            String.format("Expected second attempt to be faster, but %f ms elapsed", elapsedMs));
       }
     }
   }

flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandlerBuilderTest.java

@@ -147,6 +147,7 @@ public void testDefaults() {
     assertNull(builder.clientCertificatePath);
     assertNull(builder.clientKeyPath);
     assertEquals(Optional.empty(), builder.catalog);
+    assertNull(builder.flightClientCache);
     assertNull(builder.connectTimeout);
   }
 

flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/client/utils/FlightClientCacheTest.java

@@ -0,0 +1,51 @@
+/*
+ * 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 static org.junit.jupiter.api.Assertions.*;
+
+import org.apache.arrow.flight.Location;
+import org.junit.jupiter.api.Test;
+
+class FlightClientCacheTest {
+  @Test
+  void basicOperation() {
+    FlightClientCache cache = new FlightClientCache();
+
+    Location location1 = Location.forGrpcInsecure("localhost", 8080);
+    Location location2 = Location.forGrpcInsecure("localhost", 8081);
+
+    assertFalse(cache.isDud(location1.toString()));
+    assertFalse(cache.isDud(location2.toString()));
+
+    cache.markLocationAsReachable(location1.toString());
+    assertFalse(cache.isDud(location1.toString()));
+    assertFalse(cache.isDud(location2.toString()));
+
+    cache.markLocationAsDud(location1.toString());
+    assertTrue(cache.isDud(location1.toString()));
+    assertFalse(cache.isDud(location2.toString()));
+
+    cache.markLocationAsDud(location2.toString());
+    assertTrue(cache.isDud(location1.toString()));
+    assertTrue(cache.isDud(location2.toString()));
+
+    cache.markLocationAsReachable(location1.toString());
+    assertFalse(cache.isDud(location1.toString()));
+    assertTrue(cache.isDud(location2.toString()));
+  }
+}