perf: use direct executor (googleapis#1864)

Use a direct executor by default for JDBC connections, as these do not
use the async API, and cancelling statements can be handled by directly
in the JDBC driver. This reduces the overall number of threads that is
created by the JDBC driver.

Author: olavloite

src/main/java/com/google/cloud/spanner/connection/ConnectionOptionsHelper.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed 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 com.google.cloud.spanner.connection;
+
+import com.google.api.core.InternalApi;
+import com.google.cloud.spanner.connection.StatementExecutor.StatementExecutorType;
+
+@InternalApi
+public class ConnectionOptionsHelper {
+
+  @InternalApi
+  public static ConnectionOptions.Builder useDirectExecutorIfNotUseVirtualThreads(
+      String uri, ConnectionOptions.Builder builder) {
+    ConnectionState connectionState = new ConnectionState(ConnectionProperties.parseValues(uri));
+    if (!connectionState.getValue(ConnectionProperties.USE_VIRTUAL_THREADS).getValue()) {
+      return builder.setStatementExecutorType(StatementExecutorType.DIRECT_EXECUTOR);
+    }
+    return builder;
+  }
+
+  @InternalApi
+  public static boolean usesDirectExecutor(ConnectionOptions options) {
+    return options.getStatementExecutorType() == StatementExecutorType.DIRECT_EXECUTOR;
+  }
+
+  private ConnectionOptionsHelper() {}
+}

src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcConnection.java

@@ -21,6 +21,7 @@
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.connection.AbstractStatementParser;
 import com.google.cloud.spanner.connection.ConnectionOptions;
+import com.google.cloud.spanner.connection.ConnectionOptionsHelper;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.rpc.Code;
 import java.sql.CallableStatement;
@@ -53,6 +54,7 @@ abstract class AbstractJdbcConnection extends AbstractJdbcWrapper
   private final ConnectionOptions options;
   private final com.google.cloud.spanner.connection.Connection spanner;
   private final Properties clientInfo;
+  private final boolean usesDirectExecutor;
   private AbstractStatementParser parser;
 
   private SQLWarning firstWarning = null;
@@ -63,6 +65,7 @@ abstract class AbstractJdbcConnection extends AbstractJdbcWrapper
     this.options = options;
     this.spanner = options.getConnection();
     this.clientInfo = new Properties(JdbcDatabaseMetaData.getDefaultClientInfoProperties());
+    this.usesDirectExecutor = ConnectionOptionsHelper.usesDirectExecutor(options);
   }
 
   /** Return the corresponding {@link com.google.cloud.spanner.connection.Connection} */
@@ -83,6 +86,10 @@ Spanner getSpanner() {
     return this.spanner.getSpanner();
   }
 
+  boolean usesDirectExecutor() {
+    return this.usesDirectExecutor;
+  }
+
   @Override
   public Dialect getDialect() {
     return spanner.getDialect();

src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java

@@ -34,6 +34,8 @@
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
@@ -43,6 +45,8 @@ abstract class AbstractJdbcStatement extends AbstractJdbcWrapper implements Stat
   private static final String CURSORS_NOT_SUPPORTED = "Cursors are not supported";
   private static final String ONLY_FETCH_FORWARD_SUPPORTED = "Only fetch_forward is supported";
   final AbstractStatementParser parser;
+  private final Lock executingLock;
+  private volatile Thread executingThread;
   private boolean closed;
   private boolean closeOnCompletion;
   private boolean poolable;
@@ -52,6 +56,11 @@ abstract class AbstractJdbcStatement extends AbstractJdbcWrapper implements Stat
   AbstractJdbcStatement(JdbcConnection connection) throws SQLException {
     this.connection = connection;
     this.parser = connection.getParser();
+    if (connection.usesDirectExecutor()) {
+      this.executingLock = new ReentrantLock();
+    } else {
+      this.executingLock = null;
+    }
   }
 
   @Override
@@ -239,6 +248,10 @@ private <T> T doWithStatementTimeout(
       Supplier<T> runnable, Function<T, Boolean> shouldResetTimeout) throws SQLException {
     StatementTimeout originalTimeout = setTemporaryStatementTimeout();
     T result = null;
+    if (this.executingLock != null) {
+      this.executingLock.lock();
+      this.executingThread = Thread.currentThread();
+    }
     try {
       Stopwatch stopwatch = Stopwatch.createStarted();
       result = runnable.get();
@@ -248,6 +261,10 @@ private <T> T doWithStatementTimeout(
     } catch (SpannerException spannerException) {
       throw JdbcSqlExceptionFactory.of(spannerException);
     } finally {
+      if (this.executingLock != null) {
+        this.executingThread = null;
+        this.executingLock.unlock();
+      }
       if (shouldResetTimeout.apply(result)) {
         resetStatementTimeout(originalTimeout);
       }
@@ -353,7 +370,16 @@ void setQueryTimeout(@Nonnull Duration duration) throws SQLException {
   @Override
   public void cancel() throws SQLException {
     checkClosed();
-    connection.getSpannerConnection().cancel();
+    if (this.executingThread != null) {
+      // This is a best-effort operation. It could be that the executing thread is set to null
+      // between the if-check and the actual execution. Just ignore if that happens.
+      try {
+        this.executingThread.interrupt();
+      } catch (NullPointerException ignore) {
+      }
+    } else {
+      connection.getSpannerConnection().cancel();
+    }
   }
 
   @Override

src/main/java/com/google/cloud/spanner/jdbc/JdbcDriver.java

@@ -22,6 +22,7 @@
 import com.google.cloud.spanner.SessionPoolOptionsHelper;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.connection.ConnectionOptions;
+import com.google.cloud.spanner.connection.ConnectionOptionsHelper;
 import com.google.cloud.spanner.connection.ConnectionPropertiesHelper;
 import com.google.cloud.spanner.connection.ConnectionProperty;
 import com.google.rpc.Code;
@@ -245,6 +246,9 @@ private ConnectionOptions buildConnectionOptions(String connectionUrl, Propertie
     // Enable multiplexed sessions by default for the JDBC driver.
     builder.setSessionPoolOptions(
         SessionPoolOptionsHelper.useMultiplexedSessions(SessionPoolOptions.newBuilder()).build());
+    // Enable direct executor for JDBC, as we don't use the async API.
+    builder =
+        ConnectionOptionsHelper.useDirectExecutorIfNotUseVirtualThreads(connectionUrl, builder);
     return builder.build();
   }
 

src/test/java/com/google/cloud/spanner/jdbc/JdbcStatementTimeoutTest.java

@@ -19,24 +19,45 @@
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 
 import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
 import com.google.cloud.spanner.connection.AbstractMockServerTest;
+import com.google.cloud.spanner.jdbc.JdbcSqlExceptionFactory.JdbcSqlExceptionImpl;
 import com.google.cloud.spanner.jdbc.JdbcSqlExceptionFactory.JdbcSqlTimeoutException;
+import com.google.rpc.Code;
+import com.google.spanner.v1.ExecuteSqlRequest;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.time.Duration;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import org.junit.After;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+import org.junit.runners.Parameterized.Parameters;
 
 /** Tests setting a statement timeout. */
-@RunWith(JUnit4.class)
+@RunWith(Parameterized.class)
 public class JdbcStatementTimeoutTest extends AbstractMockServerTest {
 
+  @Parameter public boolean useVirtualThreads;
+
+  @Parameters(name = "useVirtualThreads = {0}")
+  public static Object[] data() {
+    return new Boolean[] {false, true};
+  }
+
+  @Override
+  protected String getBaseUrl() {
+    return super.getBaseUrl() + ";useVirtualThreads=" + this.useVirtualThreads;
+  }
+
   @After
   public void resetExecutionTimes() {
     mockSpanner.removeAllExecutionTimes();
@@ -122,4 +143,35 @@ public void testExecuteBatchTimeout() throws SQLException {
       }
     }
   }
+
+  @Test
+  public void testCancel() throws Exception {
+    ExecutorService service = Executors.newSingleThreadExecutor();
+    String sql = INSERT_STATEMENT.getSql();
+
+    try (java.sql.Connection connection = createJdbcConnection();
+        Statement statement = connection.createStatement()) {
+      mockSpanner.freeze();
+      Future<Void> future =
+          service.submit(
+              () -> {
+                // Wait until the request has landed on the server and then cancel the statement.
+                mockSpanner.waitForRequestsToContain(
+                    message ->
+                        message instanceof ExecuteSqlRequest
+                            && ((ExecuteSqlRequest) message).getSql().equals(sql),
+                    5000L);
+                System.out.println("Cancelling statement");
+                statement.cancel();
+                return null;
+              });
+      JdbcSqlExceptionImpl exception =
+          assertThrows(JdbcSqlExceptionImpl.class, () -> statement.execute(sql));
+      assertEquals(Code.CANCELLED, exception.getCode());
+      assertNull(future.get());
+    } finally {
+      mockSpanner.unfreeze();
+      service.shutdown();
+    }
+  }
 }