fix(spanner): do not cache retryable errors in AbstractLazyInitializer

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractLazyInitializer.java

@@ -18,7 +18,8 @@
 
 /**
  * Generic {@link AbstractLazyInitializer} for any heavy-weight object that might throw an exception
- * during initialization. The underlying object is initialized at most once.
+ * during initialization. The underlying object is initialized at most once unless an exception is
+ * thrown and {@link #shouldCacheError(Exception)} returns false.
  *
  * @param <T> Object which is to be initialized lazily
  */
@@ -36,10 +37,16 @@ public T get() throws Exception {
         if (!initialized) {
           try {
             object = initialize();
+            error = null;
+            initialized = true;
           } catch (Exception e) {
-            error = e;
+            if (shouldCacheError(e)) {
+              error = e;
+              initialized = true;
+            } else {
+              throw e;
+            }
           }
-          initialized = true;
         }
       }
     }
@@ -49,6 +56,22 @@ public T get() throws Exception {
     return object;
   }
 
+  /**
+   * Determines whether an exception thrown by {@link #initialize()} should be cached. If cached,
+   * subsequent calls to {@link #get()} will immediately rethrow the cached exception. If not
+   * cached, the exception is rethrown without marking this initializer as initialized, allowing
+   * subsequent calls to retry {@link #initialize()}.
+   *
+   * <p>Defaults to {@code true}. Subclasses can override this method to avoid caching transient or
+   * retryable errors.
+   *
+   * @param e the exception thrown during initialization
+   * @return {@code true} if the exception should be cached, {@code false} otherwise
+   */
+  protected boolean shouldCacheError(Exception e) {
+    return true;
+  }
+
   /**
    * Initializes the actual object that should be returned. Is called once the first time an
    * instance of T is required.

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java

@@ -30,10 +30,6 @@
 import io.opentelemetry.api.common.Attributes;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.Future;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicInteger;
 import javax.annotation.Nullable;
 
@@ -88,14 +84,15 @@ public Dialect getDialect() {
       new AbstractLazyInitializer<StatementFactory>() {
         @Override
         protected StatementFactory initialize() {
-          try {
-            Dialect dialect = getDialectAsync().get(30, TimeUnit.SECONDS);
-            return new StatementFactory(dialect);
-          } catch (ExecutionException | TimeoutException e) {
-            throw SpannerExceptionFactory.asSpannerException(e);
-          } catch (InterruptedException e) {
-            throw SpannerExceptionFactory.propagateInterrupt(e);
-          }
+          return new StatementFactory(getDialect());
+        }
+
+        @Override
+        protected boolean shouldCacheError(Exception e) {
+          // getDialect() already decides whether to cache or retry dialect errors.
+          // Never cache errors here so that statementFactorySupplier always delegates
+          // the retry/caching lifecycle entirely to getDialect().
+          return false;
         }
       };
 
@@ -298,10 +295,6 @@ public long executePartitionedUpdate(final Statement stmt, final UpdateOption...
     return multiplexedSessionDatabaseClient.executePartitionedUpdate(stmt, options);
   }
 
-  private Future<Dialect> getDialectAsync() {
-    return multiplexedSessionDatabaseClient.getDialectAsync();
-  }
-
   boolean isValid() {
     return multiplexedSessionDatabaseClient.isValid();
   }

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java

@@ -182,6 +182,13 @@ private SharedChannelUsage(int numChannels) {
   private static final EnumSet<ErrorCode> RETRYABLE_ERROR_CODES =
       EnumSet.of(ErrorCode.DEADLINE_EXCEEDED, ErrorCode.RESOURCE_EXHAUSTED, ErrorCode.UNAVAILABLE);
 
+  static boolean shouldCacheDialectError(Exception e) {
+    if (e instanceof SpannerException) {
+      return !RETRYABLE_ERROR_CODES.contains(((SpannerException) e).getErrorCode());
+    }
+    return true;
+  }
+
   private final BitSet channelUsage;
 
   private final int numChannels;
@@ -539,6 +546,11 @@ protected Dialect initialize() {
           // This should not really happen, but it is the safest fallback value.
           return Dialect.GOOGLE_STANDARD_SQL;
         }
+
+        @Override
+        protected boolean shouldCacheError(Exception e) {
+          return shouldCacheDialectError(e);
+        }
       };
 
   @Override

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java

@@ -29,6 +29,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
@@ -115,6 +116,7 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -3145,6 +3147,122 @@ public void testGetDialectDefaultPreloaded_FailsDirectlyIfDatabaseNotFound() {
     }
   }
 
+  @Test
+  public void testGetDialect_RecoversAfterTransientCreateSessionFailure() throws Exception {
+    mockSpanner.putStatementResult(StatementResult.detectDialectResult(Dialect.POSTGRESQL));
+    mockSpanner.setCreateSessionExecutionTime(
+        SimulatedExecutionTime.ofException(
+            Status.DEADLINE_EXCEEDED
+                .withDescription("Simulated transient outage on createSession call #1")
+                .asRuntimeException()));
+    mockSpanner.freeze();
+    try {
+      DatabaseClient client =
+          spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+      AtomicReference<Throwable> firstCallError = new AtomicReference<>();
+      Thread firstCaller =
+          new Thread(
+              () -> {
+                try {
+                  client.getDialect();
+                } catch (Throwable t) {
+                  firstCallError.set(t);
+                }
+              });
+      firstCaller.start();
+      // Wait (up to 10s) until firstCaller is blocked waiting on the initial multiplexed session.
+      Stopwatch watch = Stopwatch.createStarted();
+      while (firstCaller.getState() != Thread.State.WAITING
+          && firstCaller.getState() != Thread.State.TIMED_WAITING
+          && firstCaller.getState() != Thread.State.TERMINATED
+          && watch.elapsed(TimeUnit.SECONDS) < 10) {
+        Thread.yield();
+      }
+      mockSpanner.unfreeze();
+      firstCaller.join(10_000L);
+      assertFalse(firstCaller.isAlive());
+
+      assertThat(firstCallError.get()).isInstanceOf(SpannerException.class);
+      SpannerException firstException = (SpannerException) firstCallError.get();
+      assertEquals(ErrorCode.DEADLINE_EXCEEDED, firstException.getErrorCode());
+
+      // Wait for the retry CreateSession request to reach the server so multiplexedSessionReference
+      // has been updated to the new session future.
+      mockSpanner.waitForRequestsToContain(
+          message -> message instanceof CreateSessionRequest, 10_000L);
+      try (ResultSet rs = client.singleUse().executeQuery(SELECT1)) {
+        assertTrue(rs.next());
+      }
+
+      assertEquals(Dialect.POSTGRESQL, client.getDialect());
+      assertEquals(Dialect.POSTGRESQL, client.getDialect());
+      assertEquals(
+          Statement.newBuilder("select id from test where b=$1").bind("p1").to(true).build(),
+          client
+              .getStatementFactory()
+              .withUnnamedParameters("select id from test where b=?", true));
+    } finally {
+      mockSpanner.unfreeze();
+      mockSpanner.putStatementResult(
+          StatementResult.detectDialectResult(Dialect.GOOGLE_STANDARD_SQL));
+    }
+  }
+
+  @Test
+  public void testGetDialectAndStatementFactory_RecoverAfterTransientQueryFailure() {
+    mockSpanner.putStatementResult(
+        StatementResult.exception(
+            MultiplexedSessionDatabaseClient.DETERMINE_DIALECT_STATEMENT,
+            Status.DEADLINE_EXCEEDED
+                .withDescription("Transient dialect query timeout")
+                .asRuntimeException()));
+    try {
+      DatabaseClient client =
+          spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+
+      SpannerException dialectException = assertThrows(SpannerException.class, client::getDialect);
+      assertEquals(ErrorCode.DEADLINE_EXCEEDED, dialectException.getErrorCode());
+
+      SpannerException factoryException =
+          assertThrows(SpannerException.class, client::getStatementFactory);
+      assertEquals(ErrorCode.DEADLINE_EXCEEDED, factoryException.getErrorCode());
+
+      mockSpanner.putStatementResult(StatementResult.detectDialectResult(Dialect.POSTGRESQL));
+
+      assertEquals(Dialect.POSTGRESQL, client.getDialect());
+      assertEquals(
+          Statement.newBuilder("select id from test where b=$1").bind("p1").to(true).build(),
+          client
+              .getStatementFactory()
+              .withUnnamedParameters("select id from test where b=?", true));
+    } finally {
+      mockSpanner.putStatementResult(
+          StatementResult.detectDialectResult(Dialect.GOOGLE_STANDARD_SQL));
+    }
+  }
+
+  @Test
+  public void testGetDialect_CachesNonRetryableQueryFailure() {
+    mockSpanner.putStatementResult(
+        StatementResult.exception(
+            MultiplexedSessionDatabaseClient.DETERMINE_DIALECT_STATEMENT,
+            Status.PERMISSION_DENIED
+                .withDescription("Permission denied on information_schema")
+                .asRuntimeException()));
+    try {
+      DatabaseClient client =
+          spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, "non-retryable-db"));
+
+      SpannerException e1 = assertThrows(SpannerException.class, client::getDialect);
+      assertEquals(ErrorCode.PERMISSION_DENIED, e1.getErrorCode());
+      SpannerException e2 = assertThrows(SpannerException.class, client::getDialect);
+      assertSame(e1, e2);
+    } finally {
+      mockSpanner.putStatementResult(
+          StatementResult.detectDialectResult(Dialect.GOOGLE_STANDARD_SQL));
+    }
+  }
+
   @Test
   public void testUntypedNullParameters() {
     Statement statement =

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/LazySpannerInitializerTest.java

@@ -100,4 +100,94 @@ public Spanner initialize() {
     assertThat(count.get()).isEqualTo(1);
     executor.shutdown();
   }
+
+  @Test
+  public void testGet_shouldRetryOnRetryableSpannerException() throws Exception {
+    final AtomicInteger attempts = new AtomicInteger();
+    final Spanner expectedSpanner = mock(Spanner.class);
+    final LazySpannerInitializer initializer =
+        new LazySpannerInitializer() {
+          @Override
+          public Spanner initialize() {
+            int attempt = attempts.incrementAndGet();
+            if (attempt == 1) {
+              throw SpannerExceptionFactory.newSpannerException(
+                  ErrorCode.UNAVAILABLE, "Transient unavailable");
+            }
+            if (attempt == 2) {
+              throw SpannerExceptionFactory.newSpannerException(
+                  ErrorCode.DEADLINE_EXCEEDED, "Transient deadline exceeded");
+            }
+            return expectedSpanner;
+          }
+
+          @Override
+          protected boolean shouldCacheError(Exception e) {
+            return MultiplexedSessionDatabaseClient.shouldCacheDialectError(e);
+          }
+        };
+    SpannerException e1 = assertThrows(SpannerException.class, initializer::get);
+    assertThat(e1.getErrorCode()).isEqualTo(ErrorCode.UNAVAILABLE);
+    SpannerException e2 = assertThrows(SpannerException.class, initializer::get);
+    assertThat(e2.getErrorCode()).isEqualTo(ErrorCode.DEADLINE_EXCEEDED);
+
+    Spanner s1 = initializer.get();
+    Spanner s2 = initializer.get();
+    assertThat(s1).isSameInstanceAs(expectedSpanner);
+    assertThat(s2).isSameInstanceAs(expectedSpanner);
+    assertThat(attempts.get()).isEqualTo(3);
+  }
+
+  @Test
+  public void testGet_concurrentRetryAfterRetryableFailure() throws InterruptedException {
+    final AtomicInteger count = new AtomicInteger();
+    final Spanner expectedSpanner = mock(Spanner.class);
+    final LazySpannerInitializer initializer =
+        new LazySpannerInitializer() {
+          @Override
+          public Spanner initialize() {
+            if (count.incrementAndGet() == 1) {
+              throw SpannerExceptionFactory.newSpannerException(
+                  ErrorCode.UNAVAILABLE, "Transient error");
+            }
+            return expectedSpanner;
+          }
+
+          @Override
+          protected boolean shouldCacheError(Exception e) {
+            return MultiplexedSessionDatabaseClient.shouldCacheDialectError(e);
+          }
+        };
+    final int threads = 16;
+    final CountDownLatch latch = new CountDownLatch(threads);
+    ListeningExecutorService executor =
+        MoreExecutors.listeningDecorator(Executors.newFixedThreadPool(threads));
+    List<ListenableFuture<Spanner>> futures = new ArrayList<>(threads);
+    for (int i = 0; i < threads; i++) {
+      futures.add(
+          executor.submit(
+              () -> {
+                latch.countDown();
+                latch.await(10L, TimeUnit.SECONDS);
+                return initializer.get();
+              }));
+    }
+    int failureCount = 0;
+    int successCount = 0;
+    for (ListenableFuture<Spanner> future : futures) {
+      try {
+        assertThat(future.get()).isSameInstanceAs(expectedSpanner);
+        successCount++;
+      } catch (ExecutionException e) {
+        assertThat(e.getCause()).isInstanceOf(SpannerException.class);
+        assertThat(((SpannerException) e.getCause()).getErrorCode())
+            .isEqualTo(ErrorCode.UNAVAILABLE);
+        failureCount++;
+      }
+    }
+    assertThat(failureCount).isEqualTo(1);
+    assertThat(successCount).isEqualTo(threads - 1);
+    assertThat(count.get()).isEqualTo(2);
+    executor.shutdown();
+  }
 }