diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractLazyInitializer.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractLazyInitializer.java index 734368526020..710c2d0b4130 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractLazyInitializer.java +++ b/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 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()}. + * + *

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. diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java index ecd82500a9c4..66eea4a882ef 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java +++ b/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() { @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 getDialectAsync() { - return multiplexedSessionDatabaseClient.getDialectAsync(); - } - boolean isValid() { return multiplexedSessionDatabaseClient.isValid(); } diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java index afe498219139..03558aded997 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java @@ -182,6 +182,13 @@ private SharedChannelUsage(int numChannels) { private static final EnumSet 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 diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index 8acb0a7b725b..02a3b4b80172 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/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 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 = diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/LazySpannerInitializerTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/LazySpannerInitializerTest.java index aa4879875e32..c7bb78d537fa 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/LazySpannerInitializerTest.java +++ b/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> 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 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(); + } }