refactor: move multiplexed session handling to separate class

README.md

@@ -50,20 +50,20 @@ If you are using Maven without the BOM, add this to your dependencies:
 If you are using Gradle 5.x or later, add this to your dependencies:
 
 ```Groovy
-implementation platform('com.google.cloud:libraries-bom:26.37.0')
+implementation platform('com.google.cloud:libraries-bom:26.38.0')
 
 implementation 'com.google.cloud:google-cloud-spanner'
 ```
 If you are using Gradle without BOM, add this to your dependencies:
 
 ```Groovy
-implementation 'com.google.cloud:google-cloud-spanner:6.65.0'
+implementation 'com.google.cloud:google-cloud-spanner:6.65.1'
 ```
 
 If you are using SBT, add this to your dependencies:
 
 ```Scala
-libraryDependencies += "com.google.cloud" % "google-cloud-spanner" % "6.65.0"
+libraryDependencies += "com.google.cloud" % "google-cloud-spanner" % "6.65.1"
 ```
 <!-- {x-version-update-end} -->
 
@@ -651,7 +651,7 @@ Java is a registered trademark of Oracle and/or its affiliates.
 [kokoro-badge-link-5]: http://storage.googleapis.com/cloud-devrel-public/java/badges/java-spanner/java11.html
 [stability-image]: https://img.shields.io/badge/stability-stable-green
 [maven-version-image]: https://img.shields.io/maven-central/v/com.google.cloud/google-cloud-spanner.svg
-[maven-version-link]: https://central.sonatype.com/artifact/com.google.cloud/google-cloud-spanner/6.65.0
+[maven-version-link]: https://central.sonatype.com/artifact/com.google.cloud/google-cloud-spanner/6.65.1
 [authentication]: https://github.com/googleapis/google-cloud-java#authentication
 [auth-scopes]: https://developers.google.com/identity/protocols/oauth2/scopes
 [predefined-iam-roles]: https://cloud.google.com/iam/docs/understanding-roles#predefined_roles

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractMultiplexedSessionDatabaseClient.java

@@ -0,0 +1,96 @@
+/*
+ * 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.gax.rpc.ServerStream;
+import com.google.cloud.Timestamp;
+import com.google.cloud.spanner.Options.TransactionOption;
+import com.google.cloud.spanner.Options.UpdateOption;
+import com.google.spanner.v1.BatchWriteResponse;
+
+/**
+ * Base class for the Multiplexed Session {@link DatabaseClient} implementation. Throws {@link
+ * UnsupportedOperationException} for all methods that are currently not supported for multiplexed
+ * sessions. The concrete implementation implements the methods that are supported with multiplexed
+ * sessions.
+ */
+abstract class AbstractMultiplexedSessionDatabaseClient implements DatabaseClient {
+
+  @Override
+  public Dialect getDialect() {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public String getDatabaseRole() {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public Timestamp write(Iterable<Mutation> mutations) throws SpannerException {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public CommitResponse writeWithOptions(Iterable<Mutation> mutations, TransactionOption... options)
+      throws SpannerException {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public Timestamp writeAtLeastOnce(Iterable<Mutation> mutations) throws SpannerException {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public CommitResponse writeAtLeastOnceWithOptions(
+      Iterable<Mutation> mutations, TransactionOption... options) throws SpannerException {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
+      Iterable<MutationGroup> mutationGroups, TransactionOption... options)
+      throws SpannerException {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public TransactionRunner readWriteTransaction(TransactionOption... options) {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public TransactionManager transactionManager(TransactionOption... options) {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public AsyncRunner runAsync(TransactionOption... options) {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public AsyncTransactionManager transactionManagerAsync(TransactionOption... options) {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public long executePartitionedUpdate(Statement stmt, UpdateOption... options) {
+    throw new UnsupportedOperationException();
+  }
+}

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java

@@ -844,11 +844,14 @@ public void onTransactionMetadata(Transaction transaction, boolean shouldInclude
 
   @Override
   public SpannerException onError(SpannerException e, boolean withBeginTransaction) {
+    this.session.onError(e);
     return e;
   }
 
   @Override
-  public void onDone(boolean withBeginTransaction) {}
+  public void onDone(boolean withBeginTransaction) {
+    this.session.onReadDone();
+  }
 
   private ResultSet readInternal(
       String table,

google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java

@@ -21,7 +21,6 @@
 import com.google.cloud.spanner.Options.TransactionOption;
 import com.google.cloud.spanner.Options.UpdateOption;
 import com.google.cloud.spanner.SessionPool.PooledSessionFuture;
-import com.google.cloud.spanner.SessionPool.SessionFutureWrapper;
 import com.google.cloud.spanner.SpannerImpl.ClosedException;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Function;
@@ -36,15 +35,26 @@ class DatabaseClientImpl implements DatabaseClient {
   private final TraceWrapper tracer;
   @VisibleForTesting final String clientId;
   @VisibleForTesting final SessionPool pool;
+  @VisibleForTesting final MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient;
 
   @VisibleForTesting
   DatabaseClientImpl(SessionPool pool, TraceWrapper tracer) {
-    this("", pool, tracer);
+    this("", pool, /* multiplexedSessionDatabaseClient = */ null, tracer);
   }
 
+  @VisibleForTesting
   DatabaseClientImpl(String clientId, SessionPool pool, TraceWrapper tracer) {
+    this(clientId, pool, /* multiplexedSessionDatabaseClient = */ null, tracer);
+  }
+
+  DatabaseClientImpl(
+      String clientId,
+      SessionPool pool,
+      @Nullable MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient,
+      TraceWrapper tracer) {
     this.clientId = clientId;
     this.pool = pool;
+    this.multiplexedSessionDatabaseClient = multiplexedSessionDatabaseClient;
     this.tracer = tracer;
   }
 
@@ -54,7 +64,11 @@ PooledSessionFuture getSession() {
   }
 
   @VisibleForTesting
-  SessionFutureWrapper getMultiplexedSession() {
+  DatabaseClient getMultiplexedSession() {
+    if (this.multiplexedSessionDatabaseClient != null
+        && this.multiplexedSessionDatabaseClient.isMultiplexedSessionsSupported()) {
+      return this.multiplexedSessionDatabaseClient;
+    }
     return pool.getMultiplexedSessionWithFallback();
   }
 
@@ -270,7 +284,18 @@ private <T> T runWithSessionRetry(Function<Session, T> callable) {
     }
   }
 
+  boolean isValid() {
+    return pool.isValid()
+        && (multiplexedSessionDatabaseClient == null
+            || multiplexedSessionDatabaseClient.isValid()
+            || !multiplexedSessionDatabaseClient.isMultiplexedSessionsSupported());
+  }
+
   ListenableFuture<Void> closeAsync(ClosedException closedException) {
+    if (this.multiplexedSessionDatabaseClient != null) {
+      // This method is non-blocking.
+      this.multiplexedSessionDatabaseClient.close();
+    }
     return pool.closeAsync(closedException);
   }
 }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedMultiplexedSessionTransaction.java

@@ -0,0 +1,122 @@
+/*
+ * 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 static com.google.cloud.spanner.SessionImpl.NO_CHANNEL_HINT;
+
+import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutures;
+import com.google.cloud.spanner.DelayedReadContext.DelayedReadOnlyTransaction;
+import com.google.cloud.spanner.MultiplexedSessionDatabaseClient.MultiplexedSessionTransaction;
+import com.google.common.util.concurrent.MoreExecutors;
+
+/**
+ * Represents a delayed execution of a transaction on a multiplexed session. The execution is
+ * delayed because the multiplexed session is not yet ready. This class is only used during client
+ * creation before the multiplexed session has been created. The use of this class while the
+ * multiplexed session is still being created ensures that the creation of a {@link DatabaseClient}
+ * is non-blocking.
+ */
+class DelayedMultiplexedSessionTransaction extends AbstractMultiplexedSessionDatabaseClient {
+
+  private final MultiplexedSessionDatabaseClient client;
+
+  private final ISpan span;
+
+  private final ApiFuture<SessionReference> sessionFuture;
+
+  DelayedMultiplexedSessionTransaction(
+      MultiplexedSessionDatabaseClient client,
+      ISpan span,
+      ApiFuture<SessionReference> sessionFuture) {
+    this.client = client;
+    this.span = span;
+    this.sessionFuture = sessionFuture;
+  }
+
+  @Override
+  public ReadContext singleUse() {
+    return new DelayedReadContext<>(
+        ApiFutures.transform(
+            this.sessionFuture,
+            sessionReference ->
+                new MultiplexedSessionTransaction(
+                        client, span, sessionReference, NO_CHANNEL_HINT, true)
+                    .singleUse(),
+            MoreExecutors.directExecutor()));
+  }
+
+  @Override
+  public ReadContext singleUse(TimestampBound bound) {
+    return new DelayedReadContext<>(
+        ApiFutures.transform(
+            this.sessionFuture,
+            sessionReference ->
+                new MultiplexedSessionTransaction(
+                        client, span, sessionReference, NO_CHANNEL_HINT, true)
+                    .singleUse(bound),
+            MoreExecutors.directExecutor()));
+  }
+
+  @Override
+  public ReadOnlyTransaction singleUseReadOnlyTransaction() {
+    return new DelayedReadOnlyTransaction(
+        ApiFutures.transform(
+            this.sessionFuture,
+            sessionReference ->
+                new MultiplexedSessionTransaction(
+                        client, span, sessionReference, NO_CHANNEL_HINT, true)
+                    .singleUseReadOnlyTransaction(),
+            MoreExecutors.directExecutor()));
+  }
+
+  @Override
+  public ReadOnlyTransaction singleUseReadOnlyTransaction(TimestampBound bound) {
+    return new DelayedReadOnlyTransaction(
+        ApiFutures.transform(
+            this.sessionFuture,
+            sessionReference ->
+                new MultiplexedSessionTransaction(
+                        client, span, sessionReference, NO_CHANNEL_HINT, true)
+                    .singleUseReadOnlyTransaction(bound),
+            MoreExecutors.directExecutor()));
+  }
+
+  @Override
+  public ReadOnlyTransaction readOnlyTransaction() {
+    return new DelayedReadOnlyTransaction(
+        ApiFutures.transform(
+            this.sessionFuture,
+            sessionReference ->
+                new MultiplexedSessionTransaction(
+                        client, span, sessionReference, NO_CHANNEL_HINT, false)
+                    .readOnlyTransaction(),
+            MoreExecutors.directExecutor()));
+  }
+
+  @Override
+  public ReadOnlyTransaction readOnlyTransaction(TimestampBound bound) {
+    return new DelayedReadOnlyTransaction(
+        ApiFutures.transform(
+            this.sessionFuture,
+            sessionReference ->
+                new MultiplexedSessionTransaction(
+                        client, span, sessionReference, NO_CHANNEL_HINT, false)
+                    .readOnlyTransaction(bound),
+            MoreExecutors.directExecutor()));
+  }
+}