test(spanner): unflake testCreateSessionDeadlineExceededWithNoSessionCreateWaitTime

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

@@ -324,16 +324,37 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount
               MultiplexedSessionDatabaseClient.this.resourceNotFoundException.set(
                   (ResourceNotFoundException) spannerException);
             }
-            // Set the exception to trigger an error for all waiters.
-            // Then retry the session creation if the error is (potentially) transient and the
-            // client has not been closed in the meantime.
-            sessionReferenceFuture.setException(t);
+            SettableApiFuture<SessionReference> retryFuture = null;
             if (remainingAttempts > 1
                 && !isClientClosed()
                 && RETRYABLE_ERROR_CODES.contains(spannerException.getErrorCode())) {
-              final SettableApiFuture<SessionReference> future = SettableApiFuture.create();
-              MultiplexedSessionDatabaseClient.this.multiplexedSessionReference.set(future);
-              asyncCreateMultiplexedSession(future, remainingAttempts - 1);
+              final SettableApiFuture<SessionReference> nextFuture = SettableApiFuture.create();
+              if (MultiplexedSessionDatabaseClient.this.multiplexedSessionReference.compareAndSet(
+                  sessionReferenceFuture, nextFuture)) {
+                retryFuture = nextFuture;
+              }
+            }
+
+            // 1. Unblock all waiters of the failed attempt first.
+            // This is executed OUTSIDE any lock to avoid deadlocks in foreign listener callbacks.
+            // Because the retryFuture was installed into multiplexedSessionReference above,
+            // any waking waiter that retries immediately is guaranteed to observe the retry future.
+            sessionReferenceFuture.setException(t);
+
+            // 2. If a retry was scheduled, initiate it:
+            if (retryFuture != null) {
+              // Guard against concurrent client close that happened while waking waiters.
+              if (isClientClosed()) {
+                retryFuture.setException(newClosedException());
+                return;
+              }
+              try {
+                asyncCreateMultiplexedSession(retryFuture, remainingAttempts - 1);
+              } catch (Throwable e) {
+                // Failsafe: ensure waiting threads never hang if RPC scheduling throws
+                // synchronously.
+                retryFuture.setException(e);
+              }
             }
           }
         });
@@ -431,14 +452,22 @@ public String getChannelPrimeSessionName() {
 
   void close() {
     boolean releaseChannelUsage = false;
+    ApiFuture<SessionReference> currentSessionFuture = null;
     synchronized (this) {
       if (!this.isClosed) {
         this.isClosed = true;
         this.maintainer.stop();
         releaseChannelUsage = true;
+        currentSessionFuture = this.multiplexedSessionReference.get();
       }
     }
     if (releaseChannelUsage) {
+      // If a session creation or retry is currently in flight, fail it immediately so
+      // transactions waiting on the closed client do not hang.
+      if (currentSessionFuture instanceof SettableApiFuture && !currentSessionFuture.isDone()) {
+        ((SettableApiFuture<SessionReference>) currentSessionFuture)
+            .setException(newClosedException());
+      }
       // The multiplexed session is no longer maintained, so this client must no longer be asked
       // for a session when priming dynamic channel pool channels.
       spanner.getRpc().unregisterChannelPrimeSessionSource(this);
@@ -454,6 +483,11 @@ void close() {
     }
   }
 
+  @VisibleForTesting
+  ApiFuture<SessionReference> getMultiplexedSessionFuture() {
+    return this.multiplexedSessionReference.get();
+  }
+
   @VisibleForTesting
   MultiplexedSessionMaintainer getMaintainer() {
     return this.maintainer;

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

@@ -39,6 +39,7 @@
 import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory;
 import com.google.cloud.spanner.SessionClient.SessionConsumer;
 import com.google.cloud.spanner.spi.v1.SpannerRpc;
+import com.google.common.util.concurrent.MoreExecutors;
 import java.io.PrintWriter;
 import java.io.StringWriter;
 import java.lang.reflect.Field;
@@ -247,6 +248,84 @@ public void testClosedClientIgnoresRefreshedSessionThatArrivesAfterClose() {
     assertEquals(sessionReference1, client.getCurrentSessionReference());
   }
 
+  @Test
+  public void testRetryOfInitialSessionIsInstalledBeforeFailureIsPropagated() throws Exception {
+    assumeTrue(isJava8());
+    Clock clock = mock(Clock.class);
+    when(clock.instant()).thenReturn(Instant.now());
+    SessionClient sessionClient = mock(SessionClient.class);
+    SpannerImpl spanner = mock(SpannerImpl.class);
+    SpannerOptions spannerOptions = mock(SpannerOptions.class);
+    SessionPoolOptions sessionPoolOptions = mock(SessionPoolOptions.class);
+    when(sessionClient.getSpanner()).thenReturn(spanner);
+    when(sessionClient.getDatabaseId()).thenReturn(TEST_DATABASE_ID);
+    when(spanner.getRpc()).thenReturn(mock(SpannerRpc.class));
+    when(spanner.getOptions()).thenReturn(spannerOptions);
+    when(spannerOptions.getSessionPoolOptions()).thenReturn(sessionPoolOptions);
+    when(sessionPoolOptions.getMultiplexedSessionMaintenanceDuration())
+        .thenReturn(Duration.ofDays(7));
+    when(sessionPoolOptions.getMultiplexedSessionMaintenanceLoopFrequency())
+        .thenReturn(Duration.ofMinutes(10));
+    SessionImpl session = mock(SessionImpl.class);
+    SessionReference sessionReference = mock(SessionReference.class);
+    when(session.getSessionReference()).thenReturn(sessionReference);
+
+    // Capture the consumers of the initial attempt and of the retry.
+    AtomicReference<SessionConsumer> initialConsumer = new AtomicReference<>();
+    AtomicReference<SessionConsumer> retryConsumer = new AtomicReference<>();
+    doAnswer(
+            (Answer<?>)
+                invocationOnMock -> {
+                  initialConsumer.set(invocationOnMock.getArgument(0));
+                  return null;
+                })
+        .doAnswer(
+            (Answer<?>)
+                invocationOnMock -> {
+                  retryConsumer.set(invocationOnMock.getArgument(0));
+                  return null;
+                })
+        .when(sessionClient)
+        .asyncCreateMultiplexedSession(any(SessionConsumer.class));
+    MultiplexedSessionDatabaseClient client =
+        new MultiplexedSessionDatabaseClient(sessionClient, clock);
+    try {
+      assertNotNull(initialConsumer.get());
+
+      ApiFuture<SessionReference> initialFuture = client.getMultiplexedSessionFuture();
+      // A transaction that waits for the initial attempt is woken up as soon as that attempt fails.
+      // At that moment, a new transaction must already see the pending retry instead of the failed
+      // initial attempt, as it would otherwise fail with the (retryable) error of that attempt.
+      AtomicReference<ApiFuture<SessionReference>> futureSeenByNewTransaction =
+          new AtomicReference<>();
+      initialFuture.addListener(
+          () -> futureSeenByNewTransaction.set(client.getMultiplexedSessionFuture()),
+          MoreExecutors.directExecutor());
+
+      initialConsumer
+          .get()
+          .onSessionCreateFailure(
+              SpannerExceptionFactory.newSpannerException(ErrorCode.DEADLINE_EXCEEDED, "test"), 1);
+
+      // 1. Waiters for the initial attempt receive the error of that attempt.
+      SpannerException exception =
+          assertThrows(SpannerException.class, () -> SpannerApiFutures.get(initialFuture));
+      assertEquals(ErrorCode.DEADLINE_EXCEEDED, exception.getErrorCode());
+
+      // 2. New transactions arriving after the failure observe the retry future.
+      assertNotNull(futureSeenByNewTransaction.get());
+      assertFalse(futureSeenByNewTransaction.get().isDone());
+      assertNotNull(retryConsumer.get());
+
+      // 3. Completing the retry succeeds the pending future.
+      retryConsumer.get().onSessionReady(session);
+      assertEquals(sessionReference, futureSeenByNewTransaction.get().get());
+      assertEquals(sessionReference, client.getCurrentSessionReference());
+    } finally {
+      client.close();
+    }
+  }
+
   @Test
   public void testDisableMultiplexedSessionEnvVar() throws Exception {
     assumeTrue(isJava8() && !isWindows());