test(spanner): unflake testCreateSessionDeadlineExceededWithNoSessionCreateWaitTime

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

@@ -286,6 +286,10 @@ static boolean shouldCacheDialectError(Exception e) {
 
   private void asyncCreateMultiplexedSession(
       SettableApiFuture<SessionReference> sessionReferenceFuture, int remainingAttempts) {
+    if (isClosed || sessionReferenceFuture.isDone()) {
+      sessionReferenceFuture.setException(newClosedException());
+      return;
+    }
     this.sessionClient.asyncCreateMultiplexedSession(
         new SessionConsumer() {
           @Override
@@ -324,16 +328,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.isDone()) {
+                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 +456,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 +487,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/MultiplexedSessionDatabaseClientMockServerTest.java

@@ -42,6 +42,7 @@
 import com.google.cloud.spanner.Options.RpcPriority;
 import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl;
 import com.google.cloud.spanner.connection.RandomResultSetGenerator;
+import com.google.common.base.Stopwatch;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 import com.google.common.util.concurrent.MoreExecutors;
@@ -54,6 +55,7 @@
 import java.util.*;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
@@ -150,6 +152,91 @@ public void testCreateSessionDeadlineExceededWithNoSessionCreateWaitTime() throw
     testSpanner.close();
   }
 
+  @Test
+  public void testCloseUnblocksWaitingTransaction() throws Exception {
+    mockSpanner.freeze();
+
+    Spanner testSpanner =
+        SpannerOptions.newBuilder()
+            .setProjectId("test-project")
+            .setChannelProvider(channelProvider)
+            .setCredentials(NoCredentials.getInstance())
+            .build()
+            .getService();
+    try {
+      DatabaseClient client = testSpanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+
+      // Wait until the initial CreateSession request is in flight and frozen on the server.
+      mockSpanner.waitForRequestsToContain(CreateSessionRequest.class, 5000);
+
+      // Start a query that blocks waiting for the session to be created.
+      try (ReadContext readContext = client.singleUse();
+          ResultSet resultSet = readContext.executeQuery(STATEMENT)) {
+        // Close the client while the transaction is waiting for session creation.
+        testSpanner.close();
+
+        // The transaction must immediately fail with FAILED_PRECONDITION (client closed),
+        // rather than hanging indefinitely waiting for mockSpanner to unfreeze.
+        SpannerException exception = assertThrows(SpannerException.class, resultSet::next);
+        assertEquals(ErrorCode.FAILED_PRECONDITION, exception.getErrorCode());
+        assertTrue(exception.getMessage().contains("This client has been closed"));
+      }
+    } finally {
+      mockSpanner.unfreeze();
+      if (!testSpanner.isClosed()) {
+        testSpanner.close();
+      }
+    }
+  }
+
+  @Test
+  public void testCloseUnblocksTransactionWaitingOnRetry() throws Exception {
+    // Attempt 1 fails immediately with DEADLINE_EXCEEDED, and freezeAfter(1) atomically freezes
+    // the mock server as soon as the retry CreateSessionRequest (attempt 2) arrives.
+    mockSpanner.setCreateSessionExecutionTime(
+        SimulatedExecutionTime.ofException(Status.DEADLINE_EXCEEDED.asRuntimeException()));
+    mockSpanner.freezeAfter(1);
+
+    Spanner testSpanner =
+        SpannerOptions.newBuilder()
+            .setProjectId("test-project")
+            .setChannelProvider(channelProvider)
+            .setCredentials(NoCredentials.getInstance())
+            .build()
+            .getService();
+    try {
+      DatabaseClient client = testSpanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+
+      // Wait until attempt 1 has failed and attempt 2 has reached the mock server and frozen.
+      Stopwatch watch = Stopwatch.createStarted();
+      while (mockSpanner.countRequestsOfType(CreateSessionRequest.class) < 2) {
+        if (watch.elapsed(TimeUnit.MILLISECONDS) > 5000) {
+          throw new TimeoutException(
+              "Timed out waiting for retry CreateSessionRequest to reach mock server");
+        }
+        Thread.sleep(1L);
+      }
+
+      // Start a query that binds to the pending retry future (attempt 2).
+      try (ReadContext readContext = client.singleUse();
+          ResultSet resultSet = readContext.executeQuery(STATEMENT)) {
+        // Close the client while the query is waiting for attempt 2.
+        testSpanner.close();
+
+        // The transaction must immediately fail with FAILED_PRECONDITION (client closed),
+        // instead of hanging while attempt 2 is frozen.
+        SpannerException exception = assertThrows(SpannerException.class, resultSet::next);
+        assertEquals(ErrorCode.FAILED_PRECONDITION, exception.getErrorCode());
+        assertTrue(exception.getMessage().contains("This client has been closed"));
+      }
+    } finally {
+      mockSpanner.unfreeze();
+      if (!testSpanner.isClosed()) {
+        testSpanner.close();
+      }
+    }
+  }
+
   @Test
   public void testMultiUseReadOnlyTransactionUsesSameSession() {
     // Execute two queries using the same transaction. Both queries should use the same

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,85 @@ public void testClosedClientIgnoresRefreshedSessionThatArrivesAfterClose() {
     assertEquals(sessionReference1, client.getCurrentSessionReference());
   }
 
+  @Test
+  public void testRetryOfInitialSessionIsInstalledBeforeFailureIsPropagated() throws Exception {
+    // Mockito dynamic mock creation is not supported in GraalVM native-image tests.
+    assumeFalse(isNativeImage());
+    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());
@@ -591,6 +671,10 @@ private boolean isJava8() {
     return JavaVersionUtil.getJavaMajorVersion() == 8;
   }
 
+  private boolean isNativeImage() {
+    return System.getProperty("org.graalvm.nativeimage.imagecode") != null;
+  }
+
   private boolean isWindows() {
     return System.getProperty("os.name").toLowerCase().contains("windows");
   }