Repository: googleapis/java-spanner-jdbc
Pull request #1673 (2024-07-03 18:13:23)

perf: use multiplexed sessions
Enable the use of multiplexed sessions by default for queries in auto-commit mode. Multiplexed sessions can handle any number of queries concurrently. This means that the JDBC driver does not need to check out a session exclusively from the internal session pool in order to execute a query. Instead, a single multiplexed session is enough for all queries that are executed by all JDBC connections that connect to the same Spanner database. This allows a higher degree of parallelism to be achieved from a single client machine.

Note that due to how the JDBC API is defined, each JDBC connection can only execute one query at a time. If you for example want to execute 1000 queries in parallel, then you also need to create 1000 JDBC connections. Spanner JDBC connection are however lightweight, as each JDBC connection internally uses a pool of gRPC channels.

A high degree of parallelism from a single client might require a higher number of gRPC channels than the default 4. You can set a higher number of gRPC channels with the JDBC connection URL option `numChannels`, e.g. `numChannels=8`.

It is also recommended to enable the use of virtual threads to achieve the highest possible degree of parallelism with the JDBC driver. This option can be set by adding `useVirtualThreads=true` to the JDBC connection URL. Note that virtual threads are only supported on Java 21 and higher. Setting this option on a JVM version 20 or lower will have no effect, and the JDBC driver will instead use platform threads.




Patch:
diff --git a/src/main/java/com/google/cloud/spanner/SessionPoolOptionsHelper.java b/src/main/java/com/google/cloud/spanner/SessionPoolOptionsHelper.java
new file mode 100644
index 0000000..71fb105
--- /dev/null
+++ b/src/main/java/com/google/cloud/spanner/SessionPoolOptionsHelper.java
@@ -0,0 +1,34 @@
+/*
+ * 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;
+
+import com.google.api.core.InternalApi;
+
+/**
+ * This class is only here to access a package-private method in the Spanner client library and will
+ * be removed in the future.
+ */
+@InternalApi
+public class SessionPoolOptionsHelper {
+  private SessionPoolOptionsHelper() {}
+
+  @InternalApi
+  public static SessionPoolOptions.Builder useMultiplexedSessions(
+      SessionPoolOptions.Builder builder) {
+    return builder.setUseMultiplexedSession(true);
+  }
+}
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/JdbcDriver.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcDriver.java
index 16538e6..d87aab0 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcDriver.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcDriver.java
@@ -18,6 +18,8 @@ package com.google.cloud.spanner.jdbc;
 
 import com.google.api.core.InternalApi;
 import com.google.auth.oauth2.GoogleCredentials;
+import com.google.cloud.spanner.SessionPoolOptions;
+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.ConnectionOptions.ConnectionProperty;
@@ -239,6 +241,9 @@ public class JdbcDriver implements Driver {
         && info.get(OPEN_TELEMETRY_PROPERTY_KEY) instanceof OpenTelemetry) {
       builder.setOpenTelemetry((OpenTelemetry) info.get(OPEN_TELEMETRY_PROPERTY_KEY));
     }
+    // Enable multiplexed sessions by default for the JDBC driver.
+    builder.setSessionPoolOptions(
+        SessionPoolOptionsHelper.useMultiplexedSessions(SessionPoolOptions.newBuilder()).build());
     return builder.build();
   }
 


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/MockServerHelper.java b/src/test/java/com/google/cloud/spanner/MockServerHelper.java
new file mode 100644
index 0000000..95bdf1e
--- /dev/null
+++ b/src/test/java/com/google/cloud/spanner/MockServerHelper.java
@@ -0,0 +1,28 @@
+/*
+ * 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;
+
+import com.google.spanner.v1.Session;
+
+public class MockServerHelper {
+
+  private MockServerHelper() {}
+
+  public static Session getSession(MockSpannerServiceImpl server, String sessionName) {
+    return server.getSession(sessionName);
+  }
+}
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/JdbcConnectionUrlTest.java b/src/test/java/com/google/cloud/spanner/jdbc/JdbcConnectionUrlTest.java
index c0f9473..bcf8256 100644
--- a/src/test/java/com/google/cloud/spanner/jdbc/JdbcConnectionUrlTest.java
+++ b/src/test/java/com/google/cloud/spanner/jdbc/JdbcConnectionUrlTest.java
@@ -68,7 +68,7 @@ public class JdbcConnectionUrlTest {
     }
 
     protected String getBaseUrl() {
-      return super.getBaseUrl() + ";maxSessions=1";
+      return super.getBaseUrl() + ";minSessions=0;maxSessions=1";
     }
 
     @Test
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/MultiplexedSessionsMockServerTest.java b/src/test/java/com/google/cloud/spanner/jdbc/MultiplexedSessionsMockServerTest.java
new file mode 100644
index 0000000..7688999
--- /dev/null
+++ b/src/test/java/com/google/cloud/spanner/jdbc/MultiplexedSessionsMockServerTest.java
@@ -0,0 +1,217 @@
+/*
+ * Copyright 2023 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.jdbc;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import com.google.cloud.spanner.Dialect;
+import com.google.cloud.spanner.MockServerHelper;
+import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import com.google.cloud.spanner.connection.AbstractMockServerTest;
+import com.google.cloud.spanner.connection.SpannerPool;
+import com.google.spanner.v1.CreateSessionRequest;
+import com.google.spanner.v1.ExecuteSqlRequest;
+import com.google.spanner.v1.Session;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+import org.junit.runners.Parameterized.Parameters;
+
+@RunWith(Parameterized.class)
+public class MultiplexedSessionsMockServerTest extends AbstractMockServerTest {
+  private static final String SELECT_RANDOM_SQL = SELECT_RANDOM_STATEMENT.getSql();
+
+  private static final String INSERT_SQL = INSERT_STATEMENT.getSql();
+
+  @Parameter public Dialect dialect;
+
+  private Dialect currentDialect;
+
+  @Parameters(name = "dialect = {0}")
+  public static Object[] data() {
+    return Dialect.values();
+  }
+
+  @Before
+  public void setupDialect() {
+    if (this.dialect != currentDialect) {
+      mockSpanner.putStatementResult(StatementResult.detectDialectResult(this.dialect));
+      this.currentDialect = dialect;
+    }
+  }
+
+  @After
+  public void clearRequests() {
+    mockSpanner.clearRequests();
+    SpannerPool.closeSpannerPool();
+  }
+
+  private String createUrl() {
+    return String.format(
+        "jdbc:cloudspanner://localhost:%d/projects/%s/instances/%s/databases/%s?usePlainText=true",
+        getPort(), "proj", "inst", "db" + (dialect == Dialect.POSTGRESQL ? "pg" : ""));
+  }
+
+  private Connection createConnection() throws SQLException {
+    return DriverManager.getConnection(createUrl());
+  }
+
+  @Test
+  public void testUsesMultiplexedSessionForQueryInAutoCommit() throws SQLException {
+    try (Connection connection = createConnection()) {
+      assertTrue(connection.getAutoCommit());
+      try (ResultSet resultSet = connection.createStatement().executeQuery(SELECT_RANDOM_SQL)) {
+        //noinspection StatementWithEmptyBody
+        while (resultSet.next()) {
+          // Just consume the results
+        }
+      }
+    }
+    // Verify that one multiplexed session was created and used.
+    assertEquals(1, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+    CreateSessionRequest request = mockSpanner.getRequestsOfType(CreateSessionRequest.class).get(0);
+    assertTrue(request.getSession().getMultiplexed());
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    String sessionId = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0).getSession();
+    Session session = MockServerHelper.getSession(mockSpanner, sessionId);
+    assertNotNull(session);
+    assertTrue(session.getMultiplexed());
+  }
+
+  @Test
+  public void testUsesMultiplexedSessionForQueryInReadOnlyTransaction() throws SQLException {
+    int numQueries = 2;
+    try (Connection connection = createConnection()) {
+      connection.setReadOnly(true);
+      connection.setAutoCommit(false);
+
+      for (int ignore = 0; ignore < numQueries; ignore++) {
+        try (ResultSet resultSet = connection.createStatement().executeQuery(SELECT_RANDOM_SQL)) {
+          //noinspection StatementWithEmptyBody
+          while (resultSet.next()) {
+            // Just consume the results
+          }
+        }
+      }
+    }
+    // Verify that one multiplexed session was created and used.
+    assertEquals(1, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+    CreateSessionRequest request = mockSpanner.getRequestsOfType(CreateSessionRequest.class).get(0);
+    assertTrue(request.getSession().getMultiplexed());
+
+    // Verify that both queries used the multiplexed session.
+    assertEquals(numQueries, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    for (int index = 0; index < numQueries; index++) {
+      String sessionId =
+          mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(index).getSession();
+      Session session = MockServerHelper.getSession(mockSpanner, sessionId);
+      assertNotNull(session);
+      assertTrue(session.getMultiplexed());
+    }
+  }
+
+  @Test
+  public void testUsesRegularSessionForDmlInAutoCommit() throws SQLException {
+    try (Connection connection = createConnection()) {
+      assertTrue(connection.getAutoCommit());
+      assertEquals(1, connection.createStatement().executeUpdate(INSERT_SQL));
+    }
+    // The JDBC connection creates a multiplexed session by default, because it executes a query to
+    // check what dialect the database uses. This query is executed using a multiplexed session.
+    assertEquals(1, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+    CreateSessionRequest request = mockSpanner.getRequestsOfType(CreateSessionRequest.class).get(0);
+    assertTrue(request.getSession().getMultiplexed());
+    // Verify that a regular session was used for the insert statement.
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    assertEquals(
+        INSERT_SQL, mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0).getSql());
+    String sessionId = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0).getSession();
+    Session session = MockServerHelper.getSession(mockSpanner, sessionId);
+    assertNotNull(session);
+    assertFalse(session.getMultiplexed());
+  }
+
+  @Test
+  public void testUsesRegularSessionForQueryInTransaction() throws SQLException {
+    try (Connection connection = createConnection()) {
+      connection.setAutoCommit(false);
+      assertFalse(connection.getAutoCommit());
+
+      try (ResultSet resultSet = connection.createStatement().executeQuery(SELECT_RANDOM_SQL)) {
+        //noinspection StatementWithEmptyBody
+        while (resultSet.next()) {
+          // Just consume the results
+        }
+      }
+      connection.commit();
+    }
+    // The JDBC connection creates a multiplexed session by default, because it executes a query to
+    // check what dialect the database uses. This query is executed using a multiplexed session.
+    assertEquals(1, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+    CreateSessionRequest request = mockSpanner.getRequestsOfType(CreateSessionRequest.class).get(0);
+    assertTrue(request.getSession().getMultiplexed());
+    // Verify that a regular session was used for the select statement.
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    assertEquals(
+        SELECT_RANDOM_SQL, mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0).getSql());
+    String sessionId = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0).getSession();
+    Session session = MockServerHelper.getSession(mockSpanner, sessionId);
+    assertNotNull(session);
+    assertFalse(session.getMultiplexed());
+  }
+
+  @Test
+  public void testUsesMultiplexedSessionInCombinationWithSessionPoolOptions() throws SQLException {
+    // Create a connection that uses a session pool with MinSessions=0.
+    // This should stop any regular sessions from being created.
+    // TODO: Modify this test once https://github.com/googleapis/java-spanner/pull/3197 has been
+    //       released.
+    try (Connection connection = DriverManager.getConnection(createUrl() + ";minSessions=0")) {
+      assertTrue(connection.getAutoCommit());
+      try (ResultSet resultSet = connection.createStatement().executeQuery(SELECT_RANDOM_SQL)) {
+        //noinspection StatementWithEmptyBody
+        while (resultSet.next()) {
+          // Just consume the results
+        }
+      }
+    }
+    // TODO: Remove this line once https://github.com/googleapis/java-spanner/pull/3197 has been
+    //       released.
+    // Adding 'minSessions=X' or 'maxSessions=x' to the connection URL currently disables the use of
+    // multiplexed sessions due to a bug in the Spanner Java client.
+    assertEquals(0, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+
+    // Verify that one multiplexed session was created and used.
+    // TODO: Uncomment
+    //    assertEquals(1, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+    //    CreateSessionRequest request =
+    // mockSpanner.getRequestsOfType(CreateSessionRequest.class).get(0);
+    //    assertTrue(request.getSession().getMultiplexed());
+    //    // There should be no regular sessions in use.
+    //    assertEquals(0, mockSpanner.countRequestsOfType(BatchCreateSessionsRequest.class));
+  }
+}