fix test

Author: sakthivelmanii

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

@@ -284,6 +284,9 @@ public void close() {
   /** The current multiplexed session that is used by this client. */
   private final AtomicReference<ApiFuture<SessionReference>> multiplexedSessionReference;
 
+  /** The in progress multiplexed session reference that is used by this client. */
+  private final AtomicReference<ApiFuture<SessionReference>> inProgressMultiplexedSessionReference;
+
   /**
    * The Transaction response returned by the BeginTransaction request with read-write when a
    * multiplexed session is created during client initialization.
@@ -320,7 +323,7 @@ public void close() {
    * This flag is set to true if create session RPC is in progress. This flag prevents application
    * from firing two requests concurrently
    */
-  private final AtomicBoolean retryingSessionCreation = new AtomicBoolean(false);
+  private final AtomicBoolean sessionCreationInProgress = new AtomicBoolean(true);
 
   /**
    * This lock is used to prevent two threads from retrying createSession RPC requests in
@@ -371,6 +374,7 @@ public void close() {
         SettableApiFuture.create();
     this.readWriteBeginTransactionReferenceFuture = SettableApiFuture.create();
     this.multiplexedSessionReference = new AtomicReference<>(initialSessionReferenceFuture);
+    this.inProgressMultiplexedSessionReference = new AtomicReference<>(initialSessionReferenceFuture);
     asyncCreateMultiplexedSession(initialSessionReferenceFuture);
     maybeWaitForSessionCreation(
         sessionClient.getSpanner().getOptions().getSessionPoolOptions(),
@@ -383,7 +387,8 @@ private void asyncCreateMultiplexedSession(
         new SessionConsumer() {
           @Override
           public void onSessionReady(SessionImpl session) {
-            retryingSessionCreation.set(false);
+            sessionCreationInProgress.set(false);
+            multiplexedSessionReference.set(sessionReferenceFuture);
             sessionReferenceFuture.set(session.getSessionReference());
             // only start the maintainer if we actually managed to create a session in the first
             // place.
@@ -416,7 +421,8 @@ public void onSessionReady(SessionImpl session) {
           public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount) {
             // Mark multiplexes sessions as unimplemented and fall back to regular sessions if
             // UNIMPLEMENTED is returned.
-            retryingSessionCreation.set(false);
+            sessionCreationInProgress.set(false);
+            multiplexedSessionReference.set(sessionReferenceFuture);
             maybeMarkUnimplemented(t);
             sessionReferenceFuture.setException(t);
           }
@@ -567,13 +573,12 @@ MultiplexedSessionMaintainer getMaintainer() {
   }
 
   ApiFuture<SessionReference> getCurrentSessionReferenceFuture() {
-    System.out.println("Accessing Multiplexed Session " + Instant.now());
     return ApiFutures.catchingAsync(
         this.multiplexedSessionReference.get(),
         Throwable.class,
         (throwable) -> {
           maybeRetrySessionCreation();
-          return this.multiplexedSessionReference.get();
+          return this.inProgressMultiplexedSessionReference.get();
         },
         MoreExecutors.directExecutor());
   }
@@ -583,11 +588,10 @@ private void maybeRetrySessionCreation() {
     try {
       if (isValid()
           && isMultiplexedSessionsSupported()
-          && retryingSessionCreation.compareAndSet(false, true)) {
-        System.out.println("Retrying Multiplexed Session " + Instant.now());
+          && sessionCreationInProgress.compareAndSet(false, true)) {
         SettableApiFuture<SessionReference> settableApiFuture = SettableApiFuture.create();
         asyncCreateMultiplexedSession(settableApiFuture);
-        multiplexedSessionReference.set(settableApiFuture);
+        inProgressMultiplexedSessionReference.set(settableApiFuture);
       }
     } finally {
       sessionCreationLock.unlock();
@@ -840,6 +844,7 @@ void maintain() {
             new SessionConsumer() {
               @Override
               public void onSessionReady(SessionImpl session) {
+                System.out.println("Session ready");
                 multiplexedSessionReference.set(
                     ApiFutures.immediateFuture(session.getSessionReference()));
                 expirationDate.set(