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..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 @@ -324,17 +324,29 @@ 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); - asyncCreateMultiplexedSession(future, remainingAttempts - 1); + if (MultiplexedSessionDatabaseClient.this.multiplexedSessionReference.compareAndSet( + sessionReferenceFuture, future)) { + // Set the exception to trigger an error for all waiters of this attempt. + sessionReferenceFuture.setException(t); + try { + asyncCreateMultiplexedSession(future, remainingAttempts - 1); + } catch (Throwable e) { + future.setException(e); + } + 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 dfbe267bfdcd..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 @@ -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,91 @@ 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); + 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") + 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());