perf: optimize isValid implementation (googleapis#1444)

olavloite · web-flow · commit a1042f8a21ec · 2023-12-19T07:30:49.000+01:00
Cloud Spanner JDBC connections do not maintain a physical connection to Cloud Spanner, but are merely a wrapper around the underlying Java client library. This again uses a pool of gRPC channels to communicate with Cloud Spanner. This means that a single JDBC connection will never lose its network connection with Cloud Spanner, and checking whether it is valid or not by executing a query every time is not useful. Instead, the check should: 1. Verify that a connection has successfully been established with Cloud Spanner. The result should be cached for all JDBC connections using the same Cloud Spanner client library instance. 2. Verify that the connection has not been closed. The above can be achieved by checking that the dialect of the database that the connection is connected to has been successfully fetched. This result is cached in the client library, and being able to get that means that there has been a valid connection. This means that the isValid method will still return true if the network connectivity has been lost completely between the client and Cloud Spanner. However, this check is mostly used by connection pools to determine whether a connection is safe to be handed out to an application, and when all network connectivity has been lost, this will apply to all JDBC connections and not just one, meaning that the check is void. The original isValid check can be enabled by setting the System property spanner.jdbc.use_legacy_is_valid_check to true or setting the Environment variable SPANNER_JDBC_USE_LEGACY_IS_VALID_CHECK to true. Fixes googleapis#1443
diff --git a/java-spanner-jdbc/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java b/java-spanner-jdbc/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java
@@ -29,6 +29,7 @@
 import com.google.cloud.spanner.connection.ConnectionOptions;
 import com.google.cloud.spanner.connection.SavepointSupport;
 import com.google.cloud.spanner.connection.TransactionMode;
+import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterators;
 import java.sql.Array;
@@ -57,14 +58,42 @@ class JdbcConnection extends AbstractJdbcConnection {
       "Only result sets with concurrency CONCUR_READ_ONLY are supported";
   private static final String ONLY_CLOSE_CURSORS_AT_COMMIT =
       "Only result sets with holdability CLOSE_CURSORS_AT_COMMIT are supported";
-  static final String IS_VALID_QUERY = "SELECT 1";
+
+  /**
+   * This query is used to check the aliveness of the connection if legacy alive check has been
+   * enabled. As Cloud Spanner JDBC connections do not maintain a physical or logical connection to
+   * Cloud Spanner, there is also no point in repeatedly executing a simple query to check whether a
+   * connection is alive. Instead, we rely on the result from the initial query to Spanner that
+   * determines the dialect to determine whether the connection is alive or not. This result is
+   * cached for all JDBC connections using the same {@link com.google.cloud.spanner.Spanner}
+   * instance.
+   *
+   * <p>The legacy {@link #isValid(int)} check using a SELECT 1 statement can be enabled by setting
+   * the System property spanner.jdbc.use_legacy_is_valid_check to true or setting the environment
+   * variable SPANNER_JDBC_USE_LEGACY_IS_VALID_CHECK to true.
+   */
+  static final String LEGACY_IS_VALID_QUERY = "SELECT 1";
 
   static final ImmutableList<String> NO_GENERATED_KEY_COLUMNS = ImmutableList.of();
 
   private Map<String, Class<?>> typeMap = new HashMap<>();
 
+  private final boolean useLegacyIsValidCheck;
+
   JdbcConnection(String connectionUrl, ConnectionOptions options) throws SQLException {
     super(connectionUrl, options);
+    this.useLegacyIsValidCheck = useLegacyValidCheck();
+  }
+
+  static boolean useLegacyValidCheck() {
+    String value = System.getProperty("spanner.jdbc.use_legacy_is_valid_check");
+    if (Strings.isNullOrEmpty(value)) {
+      value = System.getenv("SPANNER_JDBC_USE_LEGACY_IS_VALID_CHECK");
+    }
+    if (!Strings.isNullOrEmpty(value)) {
+      return Boolean.parseBoolean(value);
+    }
+    return false;
   }
 
   @Override
@@ -347,23 +376,38 @@ public void setTypeMap(Map<String, Class<?>> map) throws SQLException {
     this.typeMap = new HashMap<>(map);
   }
 
+  boolean isUseLegacyIsValidCheck() {
+    return useLegacyIsValidCheck;
+  }
+
   @Override
   public boolean isValid(int timeout) throws SQLException {
     JdbcPreconditions.checkArgument(timeout >= 0, "timeout must be >= 0");
     if (!isClosed()) {
+      if (isUseLegacyIsValidCheck()) {
+        return legacyIsValid(timeout);
+      }
       try {
-        Statement statement = createStatement();
-        statement.setQueryTimeout(timeout);
-        try (ResultSet rs = statement.executeQuery(IS_VALID_QUERY)) {
-          if (rs.next()) {
-            if (rs.getLong(1) == 1L) {
-              return true;
-            }
+        return getDialect() != null;
+      } catch (Exception ignore) {
+        // ignore and fall through.
+      }
+    }
+    return false;
+  }
+
+  private boolean legacyIsValid(int timeout) throws SQLException {
+    try (Statement statement = createStatement()) {
+      statement.setQueryTimeout(timeout);
+      try (ResultSet rs = statement.executeQuery(LEGACY_IS_VALID_QUERY)) {
+        if (rs.next()) {
+          if (rs.getLong(1) == 1L) {
+            return true;
           }
         }
-      } catch (SQLException e) {
-        // ignore
       }
+    } catch (SQLException e) {
+      // ignore and fall through.
     }
     return false;
   }
diff --git a/java-spanner-jdbc/src/test/java/com/google/cloud/spanner/jdbc/JdbcConnectionTest.java b/java-spanner-jdbc/src/test/java/com/google/cloud/spanner/jdbc/JdbcConnectionTest.java
@@ -502,7 +502,7 @@ public void testIsValid() throws SQLException {
         mock(com.google.cloud.spanner.connection.Connection.class);
     when(spannerConnection.getDialect()).thenReturn(dialect);
     when(options.getConnection()).thenReturn(spannerConnection);
-    Statement statement = Statement.of(JdbcConnection.IS_VALID_QUERY);
+    Statement statement = Statement.of(JdbcConnection.LEGACY_IS_VALID_QUERY);
 
     // Verify that an opened connection that returns a result set is valid.
     try (JdbcConnection connection = new JdbcConnection("url", options)) {
@@ -517,7 +517,7 @@ public void testIsValid() throws SQLException {
       }
 
       // Now let the query return an error. isValid should now return false.
-      when(spannerConnection.executeQuery(statement))
+      when(spannerConnection.getDialect())
           .thenThrow(
               SpannerExceptionFactory.newSpannerException(
                   ErrorCode.ABORTED, "the current transaction has been aborted"));
