From 5de65b90a68ccd908a06c9f8246bb1720509e2e8 Mon Sep 17 00:00:00 2001 From: Rahul Yadav Date: Wed, 24 Dec 2025 13:08:54 +0530 Subject: [PATCH 1/3] fix: Implement retry-on-access for multiplexed session creation failures --- .../MultiplexedSessionDatabaseClient.java | 351 ++++++++++++++---- ...edSessionDatabaseClientMockServerTest.java | 231 ++++++++++++ 2 files changed, 519 insertions(+), 63 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java index 11e83add517..81aa22e004d 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java @@ -42,17 +42,18 @@ import java.util.BitSet; import java.util.HashMap; import java.util.Map; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.locks.ReentrantLock; /** * {@link TransactionRunner} that automatically handles "UNIMPLEMENTED" errors with the message @@ -327,6 +328,31 @@ public void close() { */ @VisibleForTesting final AtomicBoolean unimplementedForPartitionedOps = new AtomicBoolean(false); + /** + * Lock to synchronize session creation attempts. This ensures that only one thread attempts to + * create a session at a time, while other threads wait for the result. + */ + private final ReentrantLock creationLock = new ReentrantLock(); + + /** + * Flag to indicate if session creation is currently in progress. Used to allow waiting threads to + * know if they should wait or trigger a new creation attempt. + */ + private final AtomicBoolean creationInProgress = new AtomicBoolean(false); + + /** + * Latch that waiting threads block on while session creation is in progress. This is replaced + * with a new latch after each creation attempt completes. + */ + private volatile CountDownLatch creationLatch = new CountDownLatch(1); + + /** + * The last error that occurred during session creation. This is stored temporarily and cleared + * when a session is successfully created. Unlike the previous implementation, this error is not + * cached forever - subsequent requests will retry session creation. + */ + @VisibleForTesting final AtomicReference lastCreationError = new AtomicReference<>(); + private SessionPool pool; MultiplexedSessionDatabaseClient(SessionClient sessionClient) { @@ -354,18 +380,16 @@ public void close() { this.sessionClient = sessionClient; this.maintainer = new MultiplexedSessionMaintainer(clock); this.tracer = sessionClient.getSpanner().getTracer(); - final SettableApiFuture initialSessionReferenceFuture = - SettableApiFuture.create(); this.readWriteBeginTransactionReferenceFuture = SettableApiFuture.create(); - this.multiplexedSessionReference = new AtomicReference<>(initialSessionReferenceFuture); + // Initialize with null - session will be set when created successfully + this.multiplexedSessionReference = new AtomicReference<>(null); + // Mark creation as in progress for the initial attempt + this.creationInProgress.set(true); this.sessionClient.asyncCreateMultiplexedSession( new SessionConsumer() { @Override public void onSessionReady(SessionImpl session) { - initialSessionReferenceFuture.set(session.getSessionReference()); - // only start the maintainer if we actually managed to create a session in the first - // place. - maintainer.start(); + onSessionCreatedSuccessfully(session); // initiate a begin transaction request to verify if read-write transactions are // supported using multiplexed sessions. @@ -392,38 +416,219 @@ public void onSessionReady(SessionImpl session) { @Override public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount) { - // Mark multiplexes sessions as unimplemented and fall back to regular sessions if - // UNIMPLEMENTED is returned. - maybeMarkUnimplemented(t); - initialSessionReferenceFuture.setException(t); + onSessionCreationFailed(t); } }); - maybeWaitForSessionCreation( - sessionClient.getSpanner().getOptions().getSessionPoolOptions(), - initialSessionReferenceFuture); + maybeWaitForInitialSessionCreation( + sessionClient.getSpanner().getOptions().getSessionPoolOptions()); } void setPool(SessionPool pool) { this.pool = pool; } - private static void maybeWaitForSessionCreation( - SessionPoolOptions sessionPoolOptions, ApiFuture future) { + /** + * Called when a multiplexed session is successfully created. This method stores the session + * reference, clears any previous error, starts the maintainer, and notifies waiting threads. + */ + private void onSessionCreatedSuccessfully(SessionImpl session) { + creationLock.lock(); + try { + multiplexedSessionReference.set(ApiFutures.immediateFuture(session.getSessionReference())); + lastCreationError.set(null); + expirationDate.set(Instant.now().plus(sessionExpirationDuration)); + creationInProgress.set(false); + // Start the maintainer if not already started + if (!maintainer.isStarted()) { + maintainer.start(); + } + // Notify all waiting threads + creationLatch.countDown(); + creationLatch = new CountDownLatch(1); + } finally { + creationLock.unlock(); + } + } + + /** + * Called when multiplexed session creation fails. This method stores the error temporarily, + * notifies waiting threads, and starts the maintainer for retry (unless UNIMPLEMENTED). + */ + private void onSessionCreationFailed(Throwable t) { + creationLock.lock(); + try { + // Mark multiplexed sessions as unimplemented and fall back to regular sessions if + // UNIMPLEMENTED is returned. + maybeMarkUnimplemented(t); + lastCreationError.set(t); + creationInProgress.set(false); + // Notify all waiting threads + creationLatch.countDown(); + creationLatch = new CountDownLatch(1); + // Start the maintainer even on failure (except for UNIMPLEMENTED) so it can retry + if (!unimplemented.get() && !maintainer.isStarted()) { + maintainer.start(); + } + } finally { + creationLock.unlock(); + } + } + + /** + * Waits for the initial session creation to complete if configured to do so. This method handles + * the case where the session creation is still in progress or has failed. + */ + private void maybeWaitForInitialSessionCreation(SessionPoolOptions sessionPoolOptions) { Duration waitDuration = sessionPoolOptions.getWaitForMinSessions(); if (waitDuration != null && !waitDuration.isZero()) { long timeoutMillis = waitDuration.toMillis(); try { - future.get(timeoutMillis, TimeUnit.MILLISECONDS); - } catch (ExecutionException executionException) { - throw SpannerExceptionFactory.asSpannerException(executionException.getCause()); + if (!creationLatch.await(timeoutMillis, TimeUnit.MILLISECONDS)) { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.DEADLINE_EXCEEDED, + "Timed out after waiting " + timeoutMillis + "ms for multiplexed session creation"); + } + // Check if there was an error during creation + Throwable error = lastCreationError.get(); + if (error != null && multiplexedSessionReference.get() == null) { + throw SpannerExceptionFactory.asSpannerException(error); + } } catch (InterruptedException interruptedException) { throw SpannerExceptionFactory.propagateInterrupt(interruptedException); - } catch (TimeoutException timeoutException) { + } + } + } + + /** + * Gets or creates a multiplexed session reference. This method implements retry-on-access + * semantics: if no session exists and creation is not in progress, it triggers a new creation + * attempt. If creation is in progress, it waits for the result. + * + * @return the session reference + * @throws SpannerException if session creation fails + */ + @VisibleForTesting + SessionReference getOrCreateSessionReference() { + // Fast path: session already exists + ApiFuture sessionFuture = multiplexedSessionReference.get(); + if (sessionFuture != null) { + try { + return sessionFuture.get(); + } catch (ExecutionException | InterruptedException e) { + // This shouldn't happen with immediateFuture, but handle it anyway + throw SpannerExceptionFactory.asSpannerException(e.getCause() != null ? e.getCause() : e); + } + } + + // Check if UNIMPLEMENTED - don't retry in this case + if (unimplemented.get()) { + Throwable error = lastCreationError.get(); + if (error != null) { + throw SpannerExceptionFactory.asSpannerException(error); + } + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.UNIMPLEMENTED, "Multiplexed sessions are not supported"); + } + + // Try to acquire the lock for creation + if (creationLock.tryLock()) { + try { + // Double-check after acquiring lock + sessionFuture = multiplexedSessionReference.get(); + if (sessionFuture != null) { + try { + return sessionFuture.get(); + } catch (ExecutionException | InterruptedException e) { + throw SpannerExceptionFactory.asSpannerException( + e.getCause() != null ? e.getCause() : e); + } + } + + // Check if creation is already in progress + if (creationInProgress.get()) { + // Wait for the ongoing creation to complete + creationLock.unlock(); + return waitForSessionCreation(); + } + + // Start a new creation attempt + creationInProgress.set(true); + CountDownLatch currentLatch = creationLatch; + creationLock.unlock(); + + // Trigger async session creation + sessionClient.asyncCreateMultiplexedSession( + new SessionConsumer() { + @Override + public void onSessionReady(SessionImpl session) { + onSessionCreatedSuccessfully(session); + } + + @Override + public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount) { + onSessionCreationFailed(t); + } + }); + + // Wait for creation to complete + try { + currentLatch.await(); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } + + // Check result + sessionFuture = multiplexedSessionReference.get(); + if (sessionFuture != null) { + try { + return sessionFuture.get(); + } catch (ExecutionException | InterruptedException e) { + throw SpannerExceptionFactory.asSpannerException( + e.getCause() != null ? e.getCause() : e); + } + } + + // Creation failed + Throwable error = lastCreationError.get(); + if (error != null) { + throw SpannerExceptionFactory.asSpannerException(error); + } throw SpannerExceptionFactory.newSpannerException( - ErrorCode.DEADLINE_EXCEEDED, - "Timed out after waiting " + timeoutMillis + "ms for multiplexed session creation"); + ErrorCode.INTERNAL, "Failed to create multiplexed session"); + } finally { + if (creationLock.isHeldByCurrentThread()) { + creationLock.unlock(); + } + } + } else { + // Another thread is creating, wait for it + return waitForSessionCreation(); + } + } + + /** Waits for an ongoing session creation to complete and returns the result. */ + private SessionReference waitForSessionCreation() { + try { + creationLatch.await(); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } + + ApiFuture sessionFuture = multiplexedSessionReference.get(); + if (sessionFuture != null) { + try { + return sessionFuture.get(); + } catch (ExecutionException | InterruptedException e) { + throw SpannerExceptionFactory.asSpannerException(e.getCause() != null ? e.getCause() : e); } } + + Throwable error = lastCreationError.get(); + if (error != null) { + throw SpannerExceptionFactory.asSpannerException(error); + } + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.INTERNAL, "Failed to create multiplexed session"); } private void maybeMarkUnimplemented(Throwable t) { @@ -548,13 +753,7 @@ MultiplexedSessionMaintainer getMaintainer() { @VisibleForTesting SessionReference getCurrentSessionReference() { - try { - return this.multiplexedSessionReference.get().get(); - } catch (ExecutionException executionException) { - throw SpannerExceptionFactory.asSpannerException(executionException.getCause()); - } catch (InterruptedException interruptedException) { - throw SpannerExceptionFactory.propagateInterrupt(interruptedException); - } + return getOrCreateSessionReference(); } @VisibleForTesting @@ -575,7 +774,8 @@ Transaction getReadWriteBeginTransactionReference() { * actual statement that is being executed (e.g. the first query that is sent to Spanner). */ private boolean isMultiplexedSessionCreated() { - return multiplexedSessionReference.get().isDone(); + ApiFuture future = multiplexedSessionReference.get(); + return future != null && future.isDone(); } private DatabaseClient createMultiplexedSessionTransaction(boolean singleUse) { @@ -587,28 +787,35 @@ private DatabaseClient createMultiplexedSessionTransaction(boolean singleUse) { private MultiplexedSessionTransaction createDirectMultiplexedSessionTransaction( boolean singleUse) { - try { - return new MultiplexedSessionTransaction( - this, - tracer.getCurrentSpan(), - // Getting the result of the SettableApiFuture that contains the multiplexed session will - // also automatically propagate any error that happened during the creation of the - // session, such as for example a DatabaseNotFound exception. We therefore do not need - // any special handling of such errors. - multiplexedSessionReference.get().get(), - singleUse ? getSingleUseChannelHint() : NO_CHANNEL_HINT, - singleUse, - this.pool); - } catch (ExecutionException executionException) { - throw SpannerExceptionFactory.asSpannerException(executionException.getCause()); - } catch (InterruptedException interruptedException) { - throw SpannerExceptionFactory.propagateInterrupt(interruptedException); - } + // Use getOrCreateSessionReference() which implements retry-on-access semantics. + // This will trigger a new creation attempt if the session doesn't exist yet. + SessionReference sessionReference = getOrCreateSessionReference(); + return new MultiplexedSessionTransaction( + this, + tracer.getCurrentSpan(), + sessionReference, + singleUse ? getSingleUseChannelHint() : NO_CHANNEL_HINT, + singleUse, + this.pool); } private DelayedMultiplexedSessionTransaction createDelayedMultiplexSessionTransaction() { + // For delayed transactions, we create a future that will be resolved when + // getOrCreateSessionReference() + // is called. This allows the transaction to be created immediately while the session creation + // may still be in progress or may need to be retried. + SettableApiFuture delayedFuture = SettableApiFuture.create(); + MAINTAINER_SERVICE.submit( + () -> { + try { + SessionReference sessionRef = getOrCreateSessionReference(); + delayedFuture.set(sessionRef); + } catch (Throwable t) { + delayedFuture.setException(t); + } + }); return new DelayedMultiplexedSessionTransaction( - this, tracer.getCurrentSpan(), multiplexedSessionReference.get(), this.pool); + this, tracer.getCurrentSpan(), delayedFuture, this.pool); } private int getSingleUseChannelHint() { @@ -765,13 +972,22 @@ public long executePartitionedUpdate(Statement stmt, UpdateOption... options) { final class MultiplexedSessionMaintainer { private final Clock clock; - private ScheduledFuture scheduledFuture; + private volatile ScheduledFuture scheduledFuture; + private volatile boolean started = false; MultiplexedSessionMaintainer(Clock clock) { this.clock = clock; } + boolean isStarted() { + return started; + } + void start() { + if (started) { + return; + } + started = true; // Schedule the maintainer to run once every ten minutes (by default). long loopFrequencyMillis = MultiplexedSessionDatabaseClient.this @@ -793,27 +1009,36 @@ void stop() { } void maintain() { - if (clock.instant().isAfter(expirationDate.get())) { + // Check if session needs to be created (either null or expired) + boolean sessionNeedsCreation = + multiplexedSessionReference.get() == null + || clock.instant().isAfter(expirationDate.get()); + + if (sessionNeedsCreation && !creationInProgress.get()) { + // Only attempt creation if not already in progress + creationLock.lock(); + try { + // Double-check after acquiring lock + if (creationInProgress.get()) { + return; + } + creationInProgress.set(true); + } finally { + creationLock.unlock(); + } + sessionClient.asyncCreateMultiplexedSession( new SessionConsumer() { @Override public void onSessionReady(SessionImpl session) { - multiplexedSessionReference.set( - ApiFutures.immediateFuture(session.getSessionReference())); - expirationDate.set( - clock - .instant() - .plus(MultiplexedSessionDatabaseClient.this.sessionExpirationDuration)); + onSessionCreatedSuccessfully(session); } @Override public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount) { - // ignore any errors during re-creation of the multiplexed session. This means that - // we continue to use the session that has passed its expiration date for now, and - // that a new attempt at creating a new session will be done in 10 minutes from now. - // The only exception to this rule is if the server returns UNIMPLEMENTED. In that - // case we invalidate the client and fall back to regular sessions. - maybeMarkUnimplemented(t); + // Store the error but allow future retries. + // The only exception is UNIMPLEMENTED which falls back to regular sessions. + onSessionCreationFailed(t); } }); } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientMockServerTest.java index 0448656475a..2d1fa0c30c5 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientMockServerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientMockServerTest.java @@ -25,6 +25,7 @@ import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotEquals; import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.junit.Assume.assumeFalse; @@ -2454,4 +2455,234 @@ private void waitForSessionToBeReplaced(DatabaseClientImpl client) { Thread.yield(); } } + + /** + * Test that when initial session creation fails with a sticky DEADLINE_EXCEEDED error, multiple + * RPCs all receive errors, but subsequent RPCs after clearing the error succeed. + */ + @Test + public void testStickyError_allRPCsFail_thenRetrySucceeds() throws Exception { + Spanner testSpanner = + SpannerOptions.newBuilder() + .setProjectId("test-project") + .setChannelProvider(channelProvider) + .setCredentials(NoCredentials.getInstance()) + .setSessionPoolOption( + SessionPoolOptions.newBuilder() + .setUseMultiplexedSession(true) + .setUseMultiplexedSessionForRW(true) + .setMultiplexedSessionMaintenanceLoopFrequency(Duration.ofMinutes(10)) + .setMultiplexedSessionMaintenanceDuration(Duration.ofDays(7)) + .setFailOnSessionLeak() + .build()) + .build() + .getService(); + + try { + // Set up the mock to return sticky DEADLINE_EXCEEDED + mockSpanner.setCreateSessionExecutionTime( + SimulatedExecutionTime.ofStickyException( + Status.DEADLINE_EXCEEDED + .withDescription("Deadline exceeded during session creation") + .asRuntimeException())); + + // Get a database client - this will start the async session creation + DatabaseClientImpl client = + (DatabaseClientImpl) testSpanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + + // Wait for initial creation to fail + Thread.sleep(200); + + // First query should fail + SpannerException error1 = + assertThrows( + SpannerException.class, + () -> { + try (ResultSet rs = client.singleUse().executeQuery(STATEMENT)) { + rs.next(); + } + }); + assertEquals(ErrorCode.DEADLINE_EXCEEDED, error1.getErrorCode()); + + // Second query should also fail (sticky error) + SpannerException error2 = + assertThrows( + SpannerException.class, + () -> { + try (ResultSet rs = client.singleUse().executeQuery(STATEMENT)) { + rs.next(); + } + }); + assertEquals(ErrorCode.DEADLINE_EXCEEDED, error2.getErrorCode()); + + // Now clear the error + mockSpanner.setCreateSessionExecutionTime(SimulatedExecutionTime.none()); + + // Next query should trigger retry and succeed + try (ResultSet resultSet = client.singleUse().executeQuery(STATEMENT)) { + assertTrue(resultSet.next()); + } + + // Verify session was created + assertNotNull(client.multiplexedSessionDatabaseClient.getCurrentSessionReference()); + } finally { + testSpanner.close(); + } + } + + /** Test that an RPC arriving after initial creation failure triggers a retry and can succeed. */ + @Test + public void testRPCAfterFailure_triggersRetryAndSucceeds() throws Exception { + // Create a spanner instance with slower maintainer to control retry timing + Spanner testSpanner = + SpannerOptions.newBuilder() + .setProjectId("test-project") + .setChannelProvider(channelProvider) + .setCredentials(NoCredentials.getInstance()) + .setSessionPoolOption( + SessionPoolOptions.newBuilder() + .setUseMultiplexedSession(true) + .setUseMultiplexedSessionForRW(true) + // Slow maintainer so it doesn't interfere with test timing + .setMultiplexedSessionMaintenanceLoopFrequency(Duration.ofMinutes(10)) + .setMultiplexedSessionMaintenanceDuration(Duration.ofDays(7)) + .setFailOnSessionLeak() + .build()) + .build() + .getService(); + + try { + // First call to CreateSession fails with DEADLINE_EXCEEDED (not sticky, so only first call) + mockSpanner.setCreateSessionExecutionTime( + SimulatedExecutionTime.ofException( + Status.DEADLINE_EXCEEDED.withDescription("Deadline exceeded").asRuntimeException())); + + // Get a database client - initial creation will fail + DatabaseClientImpl client = + (DatabaseClientImpl) testSpanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + + // Wait a bit for the initial creation to complete (and fail) + Thread.sleep(200); + + // Verify the error is stored + assertNotNull(client.multiplexedSessionDatabaseClient); + assertNotNull(client.multiplexedSessionDatabaseClient.lastCreationError.get()); + + // The non-sticky exception was already consumed, so subsequent calls succeed + // Now execute a query - this should trigger a retry and succeed + try (ResultSet resultSet = client.singleUse().executeQuery(STATEMENT)) { + assertTrue(resultSet.next()); + } + + // Verify the session was created on retry + assertNull(client.multiplexedSessionDatabaseClient.lastCreationError.get()); + assertNotNull(client.multiplexedSessionDatabaseClient.getCurrentSessionReference()); + + // Verify we used a multiplexed session + List requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + assertEquals(1, requests.size()); + Session session = mockSpanner.getSession(requests.get(0).getSession()); + assertNotNull(session); + assertTrue(session.getMultiplexed()); + } finally { + testSpanner.close(); + } + } + + /** Test that UNIMPLEMENTED errors fall back to regular sessions without retrying. */ + @Test + public void testUnimplementedError_fallsBackToRegularSessions() throws Exception { + Spanner testSpanner = + SpannerOptions.newBuilder() + .setProjectId("test-project") + .setChannelProvider(channelProvider) + .setCredentials(NoCredentials.getInstance()) + .setSessionPoolOption( + SessionPoolOptions.newBuilder() + .setUseMultiplexedSession(true) + .setMultiplexedSessionMaintenanceLoopFrequency(Duration.ofMinutes(10)) + .setMultiplexedSessionMaintenanceDuration(Duration.ofDays(7)) + .setFailOnSessionLeak() + .build()) + .build() + .getService(); + + try { + // Set sticky UNIMPLEMENTED error + mockSpanner.setCreateSessionExecutionTime( + SimulatedExecutionTime.ofStickyException( + Status.UNIMPLEMENTED + .withDescription("Multiplexed sessions not supported") + .asRuntimeException())); + + DatabaseClientImpl client = + (DatabaseClientImpl) testSpanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + + // Wait for creation to fail + Thread.sleep(200); + + // Verify unimplemented flag is set + assertNotNull(client.multiplexedSessionDatabaseClient); + assertFalse(client.multiplexedSessionDatabaseClient.isMultiplexedSessionsSupported()); + + // Query should fall back to regular session + try (ResultSet resultSet = client.singleUse().executeQuery(STATEMENT)) { + assertTrue(resultSet.next()); + } + + // Verify regular session was used + List requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + assertEquals(1, requests.size()); + Session session = mockSpanner.getSession(requests.get(0).getSession()); + assertFalse(session.getMultiplexed()); + } finally { + testSpanner.close(); + } + } + + /** Test that the maintainer retries session creation after initial non-UNIMPLEMENTED failure. */ + @Test + public void testMaintainer_retriesAfterInitialFailure() throws Exception { + Spanner testSpanner = + SpannerOptions.newBuilder() + .setProjectId("test-project") + .setChannelProvider(channelProvider) + .setCredentials(NoCredentials.getInstance()) + .setSessionPoolOption( + SessionPoolOptions.newBuilder() + .setUseMultiplexedSession(true) + // Fast maintainer for test + .setMultiplexedSessionMaintenanceLoopFrequency(Duration.ofMillis(50L)) + .setMultiplexedSessionMaintenanceDuration(Duration.ofMillis(1L)) + .setFailOnSessionLeak() + .build()) + .build() + .getService(); + + try { + // Initial creation fails (non-sticky, consumed after first call) + mockSpanner.setCreateSessionExecutionTime( + SimulatedExecutionTime.ofException( + Status.DEADLINE_EXCEEDED.withDescription("Deadline exceeded").asRuntimeException())); + + DatabaseClientImpl client = + (DatabaseClientImpl) testSpanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + + // Wait for initial creation to fail and maintainer to retry (exception is consumed) + Thread.sleep(200); + + // Session should now be created (either by retry on access or by maintainer) + try (ResultSet resultSet = client.singleUse().executeQuery(STATEMENT)) { + assertTrue(resultSet.next()); + } + + // Verify multiplexed session was used + List requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + assertEquals(1, requests.size()); + Session session = mockSpanner.getSession(requests.get(0).getSession()); + assertTrue(session.getMultiplexed()); + } finally { + testSpanner.close(); + } + } } From 3235d25d47d1db5d3c4857bf533ea282afc465e1 Mon Sep 17 00:00:00 2001 From: Rahul Yadav Date: Wed, 24 Dec 2025 13:22:17 +0530 Subject: [PATCH 2/3] fix tests --- .../MultiplexedSessionDatabaseClient.java | 166 +++++++++--------- ...OpenTelemetryBuiltInMetricsTracerTest.java | 14 +- 2 files changed, 97 insertions(+), 83 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java index 81aa22e004d..06b6c24b753 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java @@ -348,8 +348,7 @@ public void close() { /** * The last error that occurred during session creation. This is stored temporarily and cleared - * when a session is successfully created. Unlike the previous implementation, this error is not - * cached forever - subsequent requests will retry session creation. + * when a session is successfully created. */ @VisibleForTesting final AtomicReference lastCreationError = new AtomicReference<>(); @@ -385,6 +384,7 @@ public void close() { this.multiplexedSessionReference = new AtomicReference<>(null); // Mark creation as in progress for the initial attempt this.creationInProgress.set(true); + final CountDownLatch initialCreationLatch = this.creationLatch; this.sessionClient.asyncCreateMultiplexedSession( new SessionConsumer() { @Override @@ -420,7 +420,7 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount } }); maybeWaitForInitialSessionCreation( - sessionClient.getSpanner().getOptions().getSessionPoolOptions()); + sessionClient.getSpanner().getOptions().getSessionPoolOptions(), initialCreationLatch); } void setPool(SessionPool pool) { @@ -452,7 +452,21 @@ private void onSessionCreatedSuccessfully(SessionImpl session) { /** * Called when multiplexed session creation fails. This method stores the error temporarily, - * notifies waiting threads, and starts the maintainer for retry (unless UNIMPLEMENTED). + * notifies waiting threads, and starts the maintainer for retry (unless it's a permanent error). + * + *

Permanent errors that should NOT be retried: + * + *

    + *
  • UNIMPLEMENTED - multiplexed sessions are not supported + *
  • DatabaseNotFoundException - the database doesn't exist + *
  • InstanceNotFoundException - the instance doesn't exist + *
+ * + *

Note: We do NOT set {@link #resourceNotFoundException} here because that field is used by + * {@link #isValid()} to determine if the client should be recreated. Setting it during session + * creation would cause {@link SpannerImpl#getDatabaseClient} to create a new client and retry, + * which we don't want for permanent errors. Instead, we check if the error is a + * ResourceNotFoundException when deciding whether to start the maintainer. */ private void onSessionCreationFailed(Throwable t) { creationLock.lock(); @@ -465,8 +479,10 @@ private void onSessionCreationFailed(Throwable t) { // Notify all waiting threads creationLatch.countDown(); creationLatch = new CountDownLatch(1); - // Start the maintainer even on failure (except for UNIMPLEMENTED) so it can retry - if (!unimplemented.get() && !maintainer.isStarted()) { + // Start the maintainer even on failure so it can retry, but NOT for permanent errors: + // - UNIMPLEMENTED: multiplexed sessions are not supported + // - ResourceNotFoundException: database or instance doesn't exist + if (!unimplemented.get() && !isResourceNotFoundException(t) && !maintainer.isStarted()) { maintainer.start(); } } finally { @@ -474,16 +490,27 @@ private void onSessionCreationFailed(Throwable t) { } } + /** + * Checks if the throwable is a {@link DatabaseNotFoundException} or {@link + * InstanceNotFoundException}. + */ + private boolean isResourceNotFoundException(Throwable t) { + SpannerException spannerException = SpannerExceptionFactory.asSpannerException(t); + return spannerException instanceof DatabaseNotFoundException + || spannerException instanceof InstanceNotFoundException; + } + /** * Waits for the initial session creation to complete if configured to do so. This method handles * the case where the session creation is still in progress or has failed. */ - private void maybeWaitForInitialSessionCreation(SessionPoolOptions sessionPoolOptions) { + private void maybeWaitForInitialSessionCreation( + SessionPoolOptions sessionPoolOptions, CountDownLatch latchToWaitOn) { Duration waitDuration = sessionPoolOptions.getWaitForMinSessions(); if (waitDuration != null && !waitDuration.isZero()) { long timeoutMillis = waitDuration.toMillis(); try { - if (!creationLatch.await(timeoutMillis, TimeUnit.MILLISECONDS)) { + if (!latchToWaitOn.await(timeoutMillis, TimeUnit.MILLISECONDS)) { throw SpannerExceptionFactory.newSpannerException( ErrorCode.DEADLINE_EXCEEDED, "Timed out after waiting " + timeoutMillis + "ms for multiplexed session creation"); @@ -504,6 +531,10 @@ private void maybeWaitForInitialSessionCreation(SessionPoolOptions sessionPoolOp * semantics: if no session exists and creation is not in progress, it triggers a new creation * attempt. If creation is in progress, it waits for the result. * + *

This method uses a blocking lock (not tryLock) to ensure that the creationLatch is always + * read while holding the lock, avoiding a race condition where a waiting thread could read an old + * latch that has already been counted down and replaced. + * * @return the session reference * @throws SpannerException if session creation fails */ @@ -530,91 +561,66 @@ SessionReference getOrCreateSessionReference() { ErrorCode.UNIMPLEMENTED, "Multiplexed sessions are not supported"); } - // Try to acquire the lock for creation - if (creationLock.tryLock()) { - try { - // Double-check after acquiring lock - sessionFuture = multiplexedSessionReference.get(); - if (sessionFuture != null) { - try { - return sessionFuture.get(); - } catch (ExecutionException | InterruptedException e) { - throw SpannerExceptionFactory.asSpannerException( - e.getCause() != null ? e.getCause() : e); - } - } - - // Check if creation is already in progress - if (creationInProgress.get()) { - // Wait for the ongoing creation to complete - creationLock.unlock(); - return waitForSessionCreation(); - } - - // Start a new creation attempt - creationInProgress.set(true); - CountDownLatch currentLatch = creationLatch; - creationLock.unlock(); - - // Trigger async session creation - sessionClient.asyncCreateMultiplexedSession( - new SessionConsumer() { - @Override - public void onSessionReady(SessionImpl session) { - onSessionCreatedSuccessfully(session); - } - - @Override - public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount) { - onSessionCreationFailed(t); - } - }); + // Check if resource not found (database or instance) - don't retry in this case + Throwable lastError = lastCreationError.get(); + if (lastError != null && isResourceNotFoundException(lastError)) { + throw SpannerExceptionFactory.asSpannerException(lastError); + } - // Wait for creation to complete + // Use blocking lock to avoid race condition with latch replacement. + // The latch must be read while holding the lock to ensure we wait on the correct latch. + CountDownLatch latchToWaitOn; + boolean amTheCreator = false; + creationLock.lock(); + try { + // Re-check state after acquiring lock + sessionFuture = multiplexedSessionReference.get(); + if (sessionFuture != null) { try { - currentLatch.await(); - } catch (InterruptedException e) { - throw SpannerExceptionFactory.propagateInterrupt(e); + return sessionFuture.get(); + } catch (ExecutionException | InterruptedException e) { + throw SpannerExceptionFactory.asSpannerException(e.getCause() != null ? e.getCause() : e); } + } - // Check result - sessionFuture = multiplexedSessionReference.get(); - if (sessionFuture != null) { - try { - return sessionFuture.get(); - } catch (ExecutionException | InterruptedException e) { - throw SpannerExceptionFactory.asSpannerException( - e.getCause() != null ? e.getCause() : e); - } - } + // Capture the current latch while holding the lock + latchToWaitOn = creationLatch; - // Creation failed - Throwable error = lastCreationError.get(); - if (error != null) { - throw SpannerExceptionFactory.asSpannerException(error); - } - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.INTERNAL, "Failed to create multiplexed session"); - } finally { - if (creationLock.isHeldByCurrentThread()) { - creationLock.unlock(); - } + if (!creationInProgress.get()) { + // We are the creator + amTheCreator = true; + creationInProgress.set(true); } - } else { - // Another thread is creating, wait for it - return waitForSessionCreation(); + // If creationInProgress is true, we are a waiter + } finally { + creationLock.unlock(); } - } - /** Waits for an ongoing session creation to complete and returns the result. */ - private SessionReference waitForSessionCreation() { + if (amTheCreator) { + // Trigger async session creation + sessionClient.asyncCreateMultiplexedSession( + new SessionConsumer() { + @Override + public void onSessionReady(SessionImpl session) { + onSessionCreatedSuccessfully(session); + } + + @Override + public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount) { + onSessionCreationFailed(t); + } + }); + } + + // Wait for creation to complete (both creator and waiters wait on the same latch) try { - creationLatch.await(); + latchToWaitOn.await(); } catch (InterruptedException e) { throw SpannerExceptionFactory.propagateInterrupt(e); } - ApiFuture sessionFuture = multiplexedSessionReference.get(); + // Check result + sessionFuture = multiplexedSessionReference.get(); if (sessionFuture != null) { try { return sessionFuture.get(); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetryBuiltInMetricsTracerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetryBuiltInMetricsTracerTest.java index ce1bb87ed4a..e37858eec20 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetryBuiltInMetricsTracerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetryBuiltInMetricsTracerTest.java @@ -264,7 +264,7 @@ public void testMetricsWithGaxRetryUnaryRpc() { } @Test - public void testNoNetworkConnection() { + public void testNoNetworkConnection() throws InterruptedException { assumeFalse(TestHelper.isMultiplexSessionDisabled()); // Create a Spanner instance that tries to connect to a server that does not exist. // This simulates a bad network connection. @@ -308,6 +308,11 @@ public void testNoNetworkConnection() { String instance = "i"; DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("test-project", instance, "d")); + // Wait for the initial async session creation to complete (fail). + // This ensures deterministic behavior - the first creation will have failed by the time + // we execute the query, so the query will trigger a retry attempt. + Thread.sleep(100); + // Using this client will return UNAVAILABLE, as the server is not reachable and we have // disabled retries. SpannerException exception = @@ -337,9 +342,12 @@ public void testNoNetworkConnection() { getMetricData(metricReader, BuiltInMetricsConstant.ATTEMPT_COUNT_NAME); assertNotNull(attemptCountMetricData); - // Attempt count should have a failed metric point for CreateSession. + // Attempt count should have failed metric points for CreateSession. + // With retry-on-access behavior, we expect 2 attempts: + // 1. Initial async CreateSession during client construction + // 2. Retry attempt when executeQuery().next() is called assertEquals( - 1, getAggregatedValue(attemptCountMetricData, expectedAttributesCreateSessionFailed), 0); + 2, getAggregatedValue(attemptCountMetricData, expectedAttributesCreateSessionFailed), 0); assertTrue( checkIfMetricExists(metricReader, BuiltInMetricsConstant.GFE_CONNECTIVITY_ERROR_NAME)); assertTrue( From 07259cd93b2f2145970050a777a8e5a41cb86c83 Mon Sep 17 00:00:00 2001 From: Rahul Yadav Date: Wed, 24 Dec 2025 14:00:55 +0530 Subject: [PATCH 3/3] fix flakiness --- .../src/test/java/com/google/cloud/spanner/SpanTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpanTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpanTest.java index 38aa31ad0b1..8cbb0d87135 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpanTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpanTest.java @@ -261,8 +261,8 @@ public void tearDown() { @Test public void singleUseNonRetryableErrorOnNext() { + mockSpanner.addException(FAILED_PRECONDITION); try (ResultSet rs = client.singleUse().executeQuery(SELECT1)) { - mockSpanner.addException(FAILED_PRECONDITION); SpannerException e = assertThrows(SpannerException.class, () -> rs.next()); assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); }