fix: Retry creation of multiplexed session

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

@@ -40,6 +40,7 @@
 import java.time.Duration;
 import java.time.Instant;
 import java.util.BitSet;
+import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
@@ -262,6 +263,9 @@ public void close() {
    */
   private static final Map<SpannerImpl, BitSet> CHANNEL_USAGE = new HashMap<>();
 
+  private static final EnumSet<ErrorCode> RETRYABLE_ERROR_CODES =
+      EnumSet.of(ErrorCode.DEADLINE_EXCEEDED, ErrorCode.RESOURCE_EXHAUSTED, ErrorCode.UNAVAILABLE);
+
   private final BitSet channelUsage;
 
   private final int numChannels;
@@ -358,11 +362,19 @@ public void close() {
         SettableApiFuture.create();
     this.readWriteBeginTransactionReferenceFuture = SettableApiFuture.create();
     this.multiplexedSessionReference = new AtomicReference<>(initialSessionReferenceFuture);
+    asyncCreateMultiplexedSession(initialSessionReferenceFuture);
+    maybeWaitForSessionCreation(
+        sessionClient.getSpanner().getOptions().getSessionPoolOptions(),
+        initialSessionReferenceFuture);
+  }
+
+  private void asyncCreateMultiplexedSession(
+      SettableApiFuture<SessionReference> sessionReferenceFuture) {
     this.sessionClient.asyncCreateMultiplexedSession(
         new SessionConsumer() {
           @Override
           public void onSessionReady(SessionImpl session) {
-            initialSessionReferenceFuture.set(session.getSessionReference());
+            sessionReferenceFuture.set(session.getSessionReference());
             // only start the maintainer if we actually managed to create a session in the first
             // place.
             maintainer.start();
@@ -395,33 +407,62 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount
             // Mark multiplexes sessions as unimplemented and fall back to regular sessions if
             // UNIMPLEMENTED is returned.
             maybeMarkUnimplemented(t);
-            initialSessionReferenceFuture.setException(t);
+            sessionReferenceFuture.setException(t);
           }
         });
-    maybeWaitForSessionCreation(
-        sessionClient.getSpanner().getOptions().getSessionPoolOptions(),
-        initialSessionReferenceFuture);
   }
 
   void setPool(SessionPool pool) {
     this.pool = pool;
   }
 
-  private static void maybeWaitForSessionCreation(
-      SessionPoolOptions sessionPoolOptions, ApiFuture<SessionReference> future) {
+  private void maybeWaitForSessionCreation(
+      SessionPoolOptions sessionPoolOptions,
+      SettableApiFuture<SessionReference> initialSessionReferenceFuture) {
     Duration waitDuration = sessionPoolOptions.getWaitForMinSessions();
     if (waitDuration != null && !waitDuration.isZero()) {
-      long timeoutMillis = waitDuration.toMillis();
-      try {
-        future.get(timeoutMillis, TimeUnit.MILLISECONDS);
-      } catch (ExecutionException executionException) {
-        throw SpannerExceptionFactory.asSpannerException(executionException.getCause());
-      } catch (InterruptedException interruptedException) {
-        throw SpannerExceptionFactory.propagateInterrupt(interruptedException);
-      } catch (TimeoutException timeoutException) {
-        throw SpannerExceptionFactory.newSpannerException(
-            ErrorCode.DEADLINE_EXCEEDED,
-            "Timed out after waiting " + timeoutMillis + "ms for multiplexed session creation");
+
+      SpannerException lastException = null;
+      SettableApiFuture<SessionReference> sessionReferenceFuture = initialSessionReferenceFuture;
+      Duration remainingTime;
+
+      Instant endTime = Instant.now().plus(waitDuration);
+      while ((remainingTime = Duration.between(Instant.now(), endTime)).toMillis() > 0) {
+        // If any exception is thrown, then retry the multiplexed session creation
+        if (sessionReferenceFuture == null) {
+          sessionReferenceFuture = SettableApiFuture.create();
+          asyncCreateMultiplexedSession(sessionReferenceFuture);
+          this.multiplexedSessionReference.set(sessionReferenceFuture);
+        }
+        try {
+          sessionReferenceFuture.get(remainingTime.toMillis(), TimeUnit.MILLISECONDS);
+          lastException = null;
+          break;
+        } catch (ExecutionException executionException) {
+          lastException = SpannerExceptionFactory.asSpannerException(executionException.getCause());
+        } catch (InterruptedException interruptedException) {
+          lastException = SpannerExceptionFactory.propagateInterrupt(interruptedException);
+        } catch (TimeoutException timeoutException) {
+          lastException =
+              SpannerExceptionFactory.newSpannerException(
+                  ErrorCode.DEADLINE_EXCEEDED,
+                  "Timed out after waiting "
+                      + waitDuration.toMillis()
+                      + "ms for multiplexed session creation");
+        }
+        // if any exception is thrown, then set the session reference to null to retry the
+        // multiplexed session creation only if the error code is DEADLINE EXCEEDED, UNAVAILABLE or
+        // RESOURCE_EXHAUSTED
+        if (RETRYABLE_ERROR_CODES.contains(lastException.getErrorCode())) {
+          sessionReferenceFuture = null;
+        } else {
+          break;
+        }
+      }
+      // if the wait time elapsed and multiplexed session fetch failed then throw the last exception
+      // that we have received
+      if (lastException != null) {
+        throw lastException;
       }
     }
   }

google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java

@@ -542,16 +542,16 @@ void simulateExecutionTime(
         boolean stickyGlobalExceptions,
         CountDownLatch freezeLock) {
       Uninterruptibles.awaitUninterruptibly(freezeLock);
-      checkException(globalExceptions, stickyGlobalExceptions);
-      if (streamIndices.isEmpty()) {
-        checkException(this.exceptions, stickyException);
-      }
       if (minimumExecutionTime > 0 || randomExecutionTime > 0) {
         Uninterruptibles.sleepUninterruptibly(
             (randomExecutionTime == 0 ? 0 : RANDOM.nextInt(randomExecutionTime))
                 + minimumExecutionTime,
             TimeUnit.MILLISECONDS);
       }
+      checkException(globalExceptions, stickyGlobalExceptions);
+      if (streamIndices.isEmpty()) {
+        checkException(this.exceptions, stickyException);
+      }
     }
 
     private static void checkException(Queue<Exception> exceptions, boolean keepException) {