From ab941f696085147d5a57423356e3d552fcc979d2 Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Fri, 30 Oct 2026 10:16:11 +0000 Subject: [PATCH 1/6] test(spanner): unflake testCreateSessionDeadlineExceededWithNoSessionCreateWaitTime The initial future was completed with setException(t) before multiplexedSessionReference was updated to the retry future. Waiters waking up and polling the reference encountered the stale completed-with-error future. This fix installs the retry future before completing the initial future exceptionally to eliminate the race. --- .../MultiplexedSessionDatabaseClient.java | 14 +++- .../MultiplexedSessionDatabaseClientTest.java | 83 +++++++++++++++++++ 2 files changed, 93 insertions(+), 4 deletions(-) diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java index 03558aded997..e05ca7055f9b 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java @@ -324,16 +324,22 @@ 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); + // Retry the session creation if the error is (potentially) transient and the client + // has not been closed in the meantime. The future for the retry is installed before + // the exception is set, so that transactions that are started after the waiters of + // this attempt have received the error, wait for the retry instead of failing with + // the error of this attempt. if (remainingAttempts > 1 && !isClientClosed() && RETRYABLE_ERROR_CODES.contains(spannerException.getErrorCode())) { final SettableApiFuture future = SettableApiFuture.create(); MultiplexedSessionDatabaseClient.this.multiplexedSessionReference.set(future); + // Set the exception to trigger an error for all waiters of this attempt. + sessionReferenceFuture.setException(t); asyncCreateMultiplexedSession(future, remainingAttempts - 1); + } else { + // Set the exception to trigger an error for all waiters. + sessionReferenceFuture.setException(t); } } }); diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java index dfbe267bfdcd..45c36d273496 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java +++ b/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,88 @@ public void testClosedClientIgnoresRefreshedSessionThatArrivesAfterClose() { assertEquals(sessionReference1, client.getCurrentSessionReference()); } + @Test + public void testRetryOfInitialSessionIsInstalledBeforeFailureIsPropagated() throws Exception { + 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 initialConsumer = new AtomicReference<>(); + AtomicReference 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); + assertNotNull(initialConsumer.get()); + + AtomicReference> currentReference = + getMultiplexedSessionReference(client); + ApiFuture initialFuture = currentReference.get(); + // 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> futureSeenByNewTransaction = + new AtomicReference<>(); + initialFuture.addListener( + () -> futureSeenByNewTransaction.set(currentReference.get()), + MoreExecutors.directExecutor()); + + initialConsumer + .get() + .onSessionCreateFailure( + SpannerExceptionFactory.newSpannerException(ErrorCode.DEADLINE_EXCEEDED, "test"), 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()); + // New transactions that started after the failure was propagated wait for the retry. + assertNotNull(futureSeenByNewTransaction.get()); + assertFalse(futureSeenByNewTransaction.get().isDone()); + assertNotNull(retryConsumer.get()); + retryConsumer.get().onSessionReady(session); + assertEquals(sessionReference, futureSeenByNewTransaction.get().get()); + assertEquals(sessionReference, client.getCurrentSessionReference()); + client.close(); + } + + @SuppressWarnings("unchecked") + private static AtomicReference> getMultiplexedSessionReference( + MultiplexedSessionDatabaseClient client) throws Exception { + Field field = + MultiplexedSessionDatabaseClient.class.getDeclaredField("multiplexedSessionReference"); + field.setAccessible(true); + return (AtomicReference>) field.get(client); + } + @Test public void testDisableMultiplexedSessionEnvVar() throws Exception { assumeTrue(isJava8() && !isWindows()); From a506f9eeb1a7dd92a75d6d0a7f47d0c65dd610b8 Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Fri, 30 Oct 2026 12:35:07 +0000 Subject: [PATCH 2/6] Addressed comments --- .../MultiplexedSessionDatabaseClient.java | 16 +++-- .../MultiplexedSessionDatabaseClientTest.java | 65 ++++++++++--------- 2 files changed, 43 insertions(+), 38 deletions(-) diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java index e05ca7055f9b..79bdbf2a50d3 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java @@ -333,14 +333,16 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount && !isClientClosed() && RETRYABLE_ERROR_CODES.contains(spannerException.getErrorCode())) { final SettableApiFuture future = SettableApiFuture.create(); - MultiplexedSessionDatabaseClient.this.multiplexedSessionReference.set(future); - // Set the exception to trigger an error for all waiters of this attempt. - sessionReferenceFuture.setException(t); - asyncCreateMultiplexedSession(future, remainingAttempts - 1); - } else { - // Set the exception to trigger an error for all waiters. - sessionReferenceFuture.setException(t); + if (MultiplexedSessionDatabaseClient.this.multiplexedSessionReference.compareAndSet( + sessionReferenceFuture, future)) { + // Set the exception to trigger an error for all waiters of this attempt. + sessionReferenceFuture.setException(t); + asyncCreateMultiplexedSession(future, remainingAttempts - 1); + return; + } } + // Set the exception to trigger an error for all waiters. + sessionReferenceFuture.setException(t); } }); } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java index 45c36d273496..20a09c194141 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java @@ -288,37 +288,40 @@ public void testRetryOfInitialSessionIsInstalledBeforeFailureIsPropagated() thro .asyncCreateMultiplexedSession(any(SessionConsumer.class)); MultiplexedSessionDatabaseClient client = new MultiplexedSessionDatabaseClient(sessionClient, clock); - assertNotNull(initialConsumer.get()); - - AtomicReference> currentReference = - getMultiplexedSessionReference(client); - ApiFuture initialFuture = currentReference.get(); - // 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> futureSeenByNewTransaction = - new AtomicReference<>(); - initialFuture.addListener( - () -> futureSeenByNewTransaction.set(currentReference.get()), - MoreExecutors.directExecutor()); - - initialConsumer - .get() - .onSessionCreateFailure( - SpannerExceptionFactory.newSpannerException(ErrorCode.DEADLINE_EXCEEDED, "test"), 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()); - // New transactions that started after the failure was propagated wait for the retry. - assertNotNull(futureSeenByNewTransaction.get()); - assertFalse(futureSeenByNewTransaction.get().isDone()); - assertNotNull(retryConsumer.get()); - retryConsumer.get().onSessionReady(session); - assertEquals(sessionReference, futureSeenByNewTransaction.get().get()); - assertEquals(sessionReference, client.getCurrentSessionReference()); - client.close(); + try { + assertNotNull(initialConsumer.get()); + + AtomicReference> currentReference = + getMultiplexedSessionReference(client); + ApiFuture initialFuture = currentReference.get(); + // 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> futureSeenByNewTransaction = + new AtomicReference<>(); + initialFuture.addListener( + () -> futureSeenByNewTransaction.set(currentReference.get()), + MoreExecutors.directExecutor()); + + initialConsumer + .get() + .onSessionCreateFailure( + SpannerExceptionFactory.newSpannerException(ErrorCode.DEADLINE_EXCEEDED, "test"), 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()); + // New transactions that started after the failure was propagated wait for the retry. + assertNotNull(futureSeenByNewTransaction.get()); + assertFalse(futureSeenByNewTransaction.get().isDone()); + assertNotNull(retryConsumer.get()); + retryConsumer.get().onSessionReady(session); + assertEquals(sessionReference, futureSeenByNewTransaction.get().get()); + assertEquals(sessionReference, client.getCurrentSessionReference()); + } finally { + client.close(); + } } @SuppressWarnings("unchecked") From 210c240b0044fc95fd739b798d0b7d21876bc91d Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Fri, 30 Oct 2026 12:47:56 +0000 Subject: [PATCH 3/6] Addressed comments --- .../cloud/spanner/MultiplexedSessionDatabaseClient.java | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java index 79bdbf2a50d3..4b4c783b7d23 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java @@ -337,7 +337,11 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount sessionReferenceFuture, future)) { // Set the exception to trigger an error for all waiters of this attempt. sessionReferenceFuture.setException(t); - asyncCreateMultiplexedSession(future, remainingAttempts - 1); + try { + asyncCreateMultiplexedSession(future, remainingAttempts - 1); + } catch (Throwable e) { + future.setException(e); + } return; } } From 53c1676020f964d33ef5bc7046c2d8d417e73bfb Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Thu, 5 Nov 2026 19:45:35 +0000 Subject: [PATCH 4/6] Addressed comments --- .../MultiplexedSessionDatabaseClient.java | 25 ++++++++++++++----- .../MultiplexedSessionDatabaseClientTest.java | 16 +++--------- 2 files changed, 23 insertions(+), 18 deletions(-) diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java index 4b4c783b7d23..7b61b9c42dfd 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java @@ -330,20 +330,24 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount // this attempt have received the error, wait for the retry instead of failing with // the error of this attempt. if (remainingAttempts > 1 - && !isClientClosed() && RETRYABLE_ERROR_CODES.contains(spannerException.getErrorCode())) { final SettableApiFuture future = SettableApiFuture.create(); - if (MultiplexedSessionDatabaseClient.this.multiplexedSessionReference.compareAndSet( - sessionReferenceFuture, future)) { - // Set the exception to trigger an error for all waiters of this attempt. - sessionReferenceFuture.setException(t); + boolean replaced = MultiplexedSessionDatabaseClient.this.multiplexedSessionReference.compareAndSet( + sessionReferenceFuture, future); + // Set the exception to trigger an error for all waiters of this attempt. + sessionReferenceFuture.setException(t); + if (replaced) { + if (isClientClosed()) { + future.setException(newClosedException()); + return; + } try { asyncCreateMultiplexedSession(future, remainingAttempts - 1); } catch (Throwable e) { future.setException(e); } - return; } + return; } // Set the exception to trigger an error for all waiters. sessionReferenceFuture.setException(t); @@ -450,6 +454,10 @@ void close() { releaseChannelUsage = true; } } + ApiFuture future = this.multiplexedSessionReference.get(); + if (future instanceof SettableApiFuture) { + ((SettableApiFuture) future).setException(newClosedException()); + } if (releaseChannelUsage) { // The multiplexed session is no longer maintained, so this client must no longer be asked // for a session when priming dynamic channel pool channels. @@ -466,6 +474,11 @@ void close() { } } + @VisibleForTesting + AtomicReference> getMultiplexedSessionReference() { + return this.multiplexedSessionReference; + } + @VisibleForTesting MultiplexedSessionMaintainer getMaintainer() { return this.maintainer; diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java index 20a09c194141..e93cb6ed05e4 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java @@ -291,16 +291,11 @@ public void testRetryOfInitialSessionIsInstalledBeforeFailureIsPropagated() thro try { assertNotNull(initialConsumer.get()); - AtomicReference> currentReference = - getMultiplexedSessionReference(client); - ApiFuture initialFuture = currentReference.get(); - // 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. + ApiFuture initialFuture = client.getMultiplexedSessionFuture(); AtomicReference> futureSeenByNewTransaction = new AtomicReference<>(); initialFuture.addListener( - () -> futureSeenByNewTransaction.set(currentReference.get()), + () -> futureSeenByNewTransaction.set(client.getMultiplexedSessionFuture()), MoreExecutors.directExecutor()); initialConsumer @@ -326,11 +321,8 @@ public void testRetryOfInitialSessionIsInstalledBeforeFailureIsPropagated() thro @SuppressWarnings("unchecked") private static AtomicReference> getMultiplexedSessionReference( - MultiplexedSessionDatabaseClient client) throws Exception { - Field field = - MultiplexedSessionDatabaseClient.class.getDeclaredField("multiplexedSessionReference"); - field.setAccessible(true); - return (AtomicReference>) field.get(client); + MultiplexedSessionDatabaseClient client) { + return client.getMultiplexedSessionReference(); } @Test From 58d70f26d6002872e54f3cc631340c8a42b186b3 Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Thu, 5 Nov 2026 20:49:57 +0000 Subject: [PATCH 5/6] Fix test compilation --- .../spanner/MultiplexedSessionDatabaseClientTest.java | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java index e93cb6ed05e4..6be3d8e5be6c 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java @@ -291,11 +291,16 @@ public void testRetryOfInitialSessionIsInstalledBeforeFailureIsPropagated() thro try { assertNotNull(initialConsumer.get()); - ApiFuture initialFuture = client.getMultiplexedSessionFuture(); + AtomicReference> currentReference = + client.getMultiplexedSessionReference(); + ApiFuture initialFuture = currentReference.get(); + // 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> futureSeenByNewTransaction = new AtomicReference<>(); initialFuture.addListener( - () -> futureSeenByNewTransaction.set(client.getMultiplexedSessionFuture()), + () -> futureSeenByNewTransaction.set(currentReference.get()), MoreExecutors.directExecutor()); initialConsumer From 7775cba69e0d981268d9c0498bcecad4c98860f9 Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Fri, 6 Nov 2026 09:52:36 +0000 Subject: [PATCH 6/6] Addressed comments --- .../MultiplexedSessionDatabaseClient.java | 65 +++++++++++-------- .../MultiplexedSessionDatabaseClientTest.java | 20 +++--- 2 files changed, 45 insertions(+), 40 deletions(-) diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java index 7b61b9c42dfd..34a7b94eb977 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java @@ -324,33 +324,38 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount MultiplexedSessionDatabaseClient.this.resourceNotFoundException.set( (ResourceNotFoundException) spannerException); } - // Retry the session creation if the error is (potentially) transient and the client - // has not been closed in the meantime. The future for the retry is installed before - // the exception is set, so that transactions that are started after the waiters of - // this attempt have received the error, wait for the retry instead of failing with - // the error of this attempt. + SettableApiFuture retryFuture = null; if (remainingAttempts > 1 + && !isClientClosed() && RETRYABLE_ERROR_CODES.contains(spannerException.getErrorCode())) { - final SettableApiFuture future = SettableApiFuture.create(); - boolean replaced = MultiplexedSessionDatabaseClient.this.multiplexedSessionReference.compareAndSet( - sessionReferenceFuture, future); - // Set the exception to trigger an error for all waiters of this attempt. - sessionReferenceFuture.setException(t); - if (replaced) { - if (isClientClosed()) { - future.setException(newClosedException()); - return; - } - try { - asyncCreateMultiplexedSession(future, remainingAttempts - 1); - } catch (Throwable e) { - future.setException(e); - } + final SettableApiFuture nextFuture = SettableApiFuture.create(); + if (MultiplexedSessionDatabaseClient.this.multiplexedSessionReference.compareAndSet( + sessionReferenceFuture, nextFuture)) { + retryFuture = nextFuture; } - return; } - // Set the exception to trigger an error for all waiters. + + // 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); + } + } } }); } @@ -447,18 +452,22 @@ public String getChannelPrimeSessionName() { void close() { boolean releaseChannelUsage = false; + ApiFuture currentSessionFuture = null; synchronized (this) { if (!this.isClosed) { this.isClosed = true; this.maintainer.stop(); releaseChannelUsage = true; + currentSessionFuture = this.multiplexedSessionReference.get(); } } - ApiFuture future = this.multiplexedSessionReference.get(); - if (future instanceof SettableApiFuture) { - ((SettableApiFuture) future).setException(newClosedException()); - } 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) 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); @@ -475,8 +484,8 @@ void close() { } @VisibleForTesting - AtomicReference> getMultiplexedSessionReference() { - return this.multiplexedSessionReference; + ApiFuture getMultiplexedSessionFuture() { + return this.multiplexedSessionReference.get(); } @VisibleForTesting diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java index 6be3d8e5be6c..53fb3e519821 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java @@ -250,6 +250,7 @@ public void testClosedClientIgnoresRefreshedSessionThatArrivesAfterClose() { @Test public void testRetryOfInitialSessionIsInstalledBeforeFailureIsPropagated() throws Exception { + assumeTrue(isJava8()); Clock clock = mock(Clock.class); when(clock.instant()).thenReturn(Instant.now()); SessionClient sessionClient = mock(SessionClient.class); @@ -291,16 +292,14 @@ public void testRetryOfInitialSessionIsInstalledBeforeFailureIsPropagated() thro try { assertNotNull(initialConsumer.get()); - AtomicReference> currentReference = - client.getMultiplexedSessionReference(); - ApiFuture initialFuture = currentReference.get(); + ApiFuture 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> futureSeenByNewTransaction = new AtomicReference<>(); initialFuture.addListener( - () -> futureSeenByNewTransaction.set(currentReference.get()), + () -> futureSeenByNewTransaction.set(client.getMultiplexedSessionFuture()), MoreExecutors.directExecutor()); initialConsumer @@ -308,14 +307,17 @@ public void testRetryOfInitialSessionIsInstalledBeforeFailureIsPropagated() thro .onSessionCreateFailure( SpannerExceptionFactory.newSpannerException(ErrorCode.DEADLINE_EXCEEDED, "test"), 1); - // Waiters for the initial attempt receive the error of that attempt. + // 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()); - // New transactions that started after the failure was propagated wait for the retry. + + // 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()); @@ -324,12 +326,6 @@ public void testRetryOfInitialSessionIsInstalledBeforeFailureIsPropagated() thro } } - @SuppressWarnings("unchecked") - private static AtomicReference> getMultiplexedSessionReference( - MultiplexedSessionDatabaseClient client) { - return client.getMultiplexedSessionReference(); - } - @Test public void testDisableMultiplexedSessionEnvVar() throws Exception { assumeTrue(isJava8() && !isWindows());