perf: use multiplexed sessions (#1673)

* 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.

It is 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.

* test: add test for multi-use read-only transaction

Author: olavloite

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);
+  }
+}

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

@@ -18,6 +18,8 @@
 
 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 @@ private ConnectionOptions buildConnectionOptions(String connectionUrl, Propertie
         && 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();
   }
 

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);
+  }
+}

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

@@ -68,7 +68,7 @@ public static void reset() {
     }
 
     protected String getBaseUrl() {
-      return super.getBaseUrl() + ";maxSessions=1";
+      return super.getBaseUrl() + ";minSessions=0;maxSessions=1";
     }
 
     @Test

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));
+  }
+}