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..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 @@ -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,30 @@ 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. + */ + @VisibleForTesting final AtomicReference lastCreationError = new AtomicReference<>(); + private SessionPool pool; MultiplexedSessionDatabaseClient(SessionClient sessionClient) { @@ -354,18 +379,17 @@ 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); + final CountDownLatch initialCreationLatch = this.creationLatch; 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,40 +416,227 @@ 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(), initialCreationLatch); } 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 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(); + 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 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 { + creationLock.unlock(); + } + } + + /** + * 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, CountDownLatch latchToWaitOn) { 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 (!latchToWaitOn.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) { - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.DEADLINE_EXCEEDED, - "Timed out after waiting " + timeoutMillis + "ms for multiplexed session creation"); } } } + /** + * 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. + * + *

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 + */ + @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"); + } + + // 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); + } + + // 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 { + 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; + + if (!creationInProgress.get()) { + // We are the creator + amTheCreator = true; + creationInProgress.set(true); + } + // If creationInProgress is true, we are a waiter + } finally { + creationLock.unlock(); + } + + 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 { + latchToWaitOn.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); + } + } + + 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) { SpannerException spannerException = SpannerExceptionFactory.asSpannerException(t); if (spannerException.getErrorCode() == ErrorCode.UNIMPLEMENTED) { @@ -548,13 +759,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 +780,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 +793,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 +978,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 +1015,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(); + } + } } 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( 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()); }