fix tests

Author: rahul2393

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<Throwable> 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).
+   *
+   * <p>Permanent errors that should NOT be retried:
+   *
+   * <ul>
+   *   <li>UNIMPLEMENTED - multiplexed sessions are not supported
+   *   <li>DatabaseNotFoundException - the database doesn't exist
+   *   <li>InstanceNotFoundException - the instance doesn't exist
+   * </ul>
+   *
+   * <p>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,25 +479,38 @@ 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 {
       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) {
+  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.
    *
+   * <p>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<SessionReference> sessionFuture = multiplexedSessionReference.get();
+    // Check result
+    sessionFuture = multiplexedSessionReference.get();
     if (sessionFuture != null) {
       try {
         return sessionFuture.get();

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(