chore(spanner): fallback for TransactionRunner

Author: harshachinta

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

@@ -27,6 +27,8 @@
 import com.google.cloud.spanner.Options.TransactionOption;
 import com.google.cloud.spanner.Options.UpdateOption;
 import com.google.cloud.spanner.SessionClient.SessionConsumer;
+import com.google.cloud.spanner.SessionPool.PooledSessionReplacementHandler;
+import com.google.cloud.spanner.SessionPool.SessionPoolTransactionRunner;
 import com.google.cloud.spanner.SpannerException.ResourceNotFoundException;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
@@ -52,6 +54,56 @@
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 
+class MultiplexedSessionTransactionRunner implements TransactionRunner {
+  private final SessionPool sessionPool;
+  private final TransactionRunnerImpl transactionRunner;
+  private final TransactionOption[] options;
+
+  public MultiplexedSessionTransactionRunner(SessionImpl session, SessionPool sessionPool, TransactionOption... options) {
+    this.sessionPool = sessionPool;
+    this.transactionRunner = new TransactionRunnerImpl(session, options); // Uses multiplexed session initially
+    session.setActive(this.transactionRunner);
+    this.options = options;
+  }
+
+  @Override
+  public <T> T run(TransactionCallable<T> callable) {
+    boolean useRegularSession = false;
+    while (true) {
+      try {
+        if (useRegularSession) {
+          TransactionRunner runner = new SessionPoolTransactionRunner(sessionPool.getSession(), null, options);
+          return runner.run(callable);
+        } else {
+          return transactionRunner.run(callable); // Run using multiplexed session
+        }
+      } catch (SpannerException e) {
+        if (e.getErrorCode() == ErrorCode.UNIMPLEMENTED) {
+          useRegularSession = true; // Switch to regular session
+        } else {
+          throw e; // Other errors propagate
+        }
+      }
+    }
+  }
+
+  @Override
+  public Timestamp getCommitTimestamp() {
+    return null;
+  }
+
+  @Override
+  public CommitResponse getCommitResponse() {
+    return null;
+  }
+
+  @Override
+  public TransactionRunner allowNestedTransaction() {
+    return null;
+  }
+}
+
+
 /**
  * {@link DatabaseClient} implementation that uses a single multiplexed session to execute
  * transactions.
@@ -75,18 +127,30 @@ static class MultiplexedSessionTransaction extends SessionImpl {
     private final int singleUseChannelHint;
 
     private boolean done;
+    private SessionPool pool;
 
     MultiplexedSessionTransaction(
         MultiplexedSessionDatabaseClient client,
         ISpan span,
         SessionReference sessionReference,
         int singleUseChannelHint,
         boolean singleUse) {
+      this(client, span, sessionReference, singleUseChannelHint, singleUse, null);
+    }
+
+    MultiplexedSessionTransaction(
+        MultiplexedSessionDatabaseClient client,
+        ISpan span,
+        SessionReference sessionReference,
+        int singleUseChannelHint,
+        boolean singleUse,
+        SessionPool pool) {
       super(client.sessionClient.getSpanner(), sessionReference, singleUseChannelHint);
       this.client = client;
       this.singleUse = singleUse;
       this.singleUseChannelHint = singleUseChannelHint;
       this.client.numSessionsAcquired.incrementAndGet();
+      this.pool = pool;
       setCurrentSpan(span);
     }
 
@@ -134,6 +198,12 @@ public CommitResponse writeAtLeastOnceWithOptions(
       return response;
     }
 
+    @Override
+    public TransactionRunner readWriteTransaction(TransactionOption... options) {
+      return new MultiplexedSessionTransactionRunner(this, pool, options);
+    }
+
+
     @Override
     void onTransactionDone() {
       boolean markedDone = false;
@@ -225,6 +295,8 @@ public void close() {
    */
   @VisibleForTesting final AtomicBoolean unimplementedForPartitionedOps = new AtomicBoolean(false);
 
+  private SessionPool pool;
+
   MultiplexedSessionDatabaseClient(SessionClient sessionClient) {
     this(sessionClient, Clock.systemUTC());
   }
@@ -299,6 +371,10 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount
         initialSessionReferenceFuture);
   }
 
+  void setPool(SessionPool pool) {
+    this.pool = pool;
+  }
+
   private static void maybeWaitForSessionCreation(
       SessionPoolOptions sessionPoolOptions, ApiFuture<SessionReference> future) {
     Duration waitDuration = sessionPoolOptions.getWaitForMinSessions();
@@ -489,7 +565,8 @@ private MultiplexedSessionTransaction createDirectMultiplexedSessionTransaction(
           // any special handling of such errors.
           multiplexedSessionReference.get().get(),
           singleUse ? getSingleUseChannelHint() : NO_CHANNEL_HINT,
-          singleUse);
+          singleUse,
+          this.pool);
     } catch (ExecutionException executionException) {
       throw SpannerExceptionFactory.asSpannerException(executionException.getCause());
     } catch (InterruptedException interruptedException) {

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

@@ -1004,15 +1004,15 @@ public TransactionState getState() {
    * {@link TransactionRunner} that automatically handles {@link SessionNotFoundException}s by
    * replacing the underlying session and then restarts the transaction.
    */
-  private static final class SessionPoolTransactionRunner<I extends SessionFuture>
+  static final class SessionPoolTransactionRunner<I extends SessionFuture>
       implements TransactionRunner {
 
     private I session;
     private final SessionReplacementHandler<I> sessionReplacementHandler;
     private final TransactionOption[] options;
     private TransactionRunner runner;
 
-    private SessionPoolTransactionRunner(
+    SessionPoolTransactionRunner(
         I session,
         SessionReplacementHandler<I> sessionReplacementHandler,
         TransactionOption... options) {

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

@@ -332,6 +332,9 @@ DatabaseClientImpl createDatabaseClient(
       boolean useMultiplexedSessionPartitionedOps,
       boolean useMultiplexedSessionForRW,
       Attributes commonAttributes) {
+    if (multiplexedSessionClient != null) {
+      multiplexedSessionClient.setPool(pool);
+    }
     return new DatabaseClientImpl(
         clientId,
         pool,