From 7ae9eed407638e3ff20fa31587aedffcab4bbd95 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Mon, 22 Sep 2025 15:15:55 +0200 Subject: [PATCH] fix: recalculate remaining statement timeout after retry If a connection has a statement timeout and a statement is executed in a read/write transaction, the statement timeout would be reset to the original value during each retry, instead of being reduced after each attempt. This could cause multiple transaction retries to make the execution of a single statement take much longer than the configured timeout value. --- .../connection/AbstractBaseUnitOfWork.java | 14 +++++++++-- .../cloud/spanner/connection/AbortedTest.java | 25 +++++++++++++++++++ .../connection/AbstractMockServerTest.java | 1 + 3 files changed, 38 insertions(+), 2 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java index 5f4facf1489..aaa984ad5a1 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java @@ -39,17 +39,18 @@ import com.google.cloud.spanner.Struct; import com.google.cloud.spanner.Type.StructField; import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement; -import com.google.cloud.spanner.connection.ReadWriteTransaction.Builder; import com.google.cloud.spanner.connection.StatementExecutor.StatementTimeout; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; import com.google.common.util.concurrent.MoreExecutors; import io.grpc.Context; +import io.grpc.Deadline; import io.grpc.MethodDescriptor; import io.grpc.Status; import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.Span; import io.opentelemetry.context.Scope; +import java.time.Duration; import java.util.Collection; import java.util.Collections; import java.util.HashSet; @@ -358,6 +359,9 @@ ApiFuture executeStatementAsync( } Context context = Context.current(); if (statementTimeout.hasTimeout() && !applyStatementTimeoutToMethods.isEmpty()) { + Deadline deadline = + Deadline.after( + statementTimeout.getTimeoutValue(TimeUnit.NANOSECONDS), TimeUnit.NANOSECONDS); context = context.withValue( SpannerOptions.CALL_CONTEXT_CONFIGURATOR_KEY, @@ -367,8 +371,14 @@ public ApiCallContext configure( ApiCallContext context, ReqT request, MethodDescriptor method) { if (statementTimeout.hasTimeout() && applyStatementTimeoutToMethods.contains(method)) { + // Calculate the remaining timeout. This method could be called multiple times + // if the transaction is retried. + long remainingTimeout = deadline.timeRemaining(TimeUnit.NANOSECONDS); + if (remainingTimeout <= 0) { + remainingTimeout = 1; + } return GrpcCallContext.createDefault() - .withTimeoutDuration(statementTimeout.asDuration()); + .withTimeoutDuration(Duration.ofNanos(remainingTimeout)); } return null; } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbortedTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbortedTest.java index 151ac89e3a0..8fec34c267e 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbortedTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbortedTest.java @@ -27,6 +27,7 @@ import com.google.cloud.Timestamp; import com.google.cloud.spanner.AbortedDueToConcurrentModificationException; import com.google.cloud.spanner.ErrorCode; +import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode; import com.google.cloud.spanner.ResultSet; @@ -52,6 +53,7 @@ import io.grpc.StatusRuntimeException; import java.util.Collections; import java.util.List; +import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; import java.util.stream.LongStream; import org.junit.Test; @@ -583,6 +585,29 @@ public void testAbortedWithBitReversedSequence() { } } + @Test + public void testTimeoutWithRetries() { + // Verifies that even though a single execution of a statement does not exceed the deadline, + // repeated retries of the statement does cause the deadline to be exceeded. + try (ITConnection connection = createConnection()) { + for (boolean autoCommit : new boolean[] {true, false}) { + connection.setAutocommit(autoCommit); + mockSpanner.setAbortProbability(1.0); + mockSpanner.setExecuteSqlExecutionTime(SimulatedExecutionTime.ofMinimumAndRandomTime(1, 0)); + + connection.setStatementTimeout(10, TimeUnit.MILLISECONDS); + SpannerException exception = + assertThrows(SpannerException.class, () -> connection.execute(INSERT_STATEMENT)); + assertEquals(ErrorCode.DEADLINE_EXCEEDED, exception.getErrorCode()); + if (!autoCommit) { + connection.rollback(); + } + } + } finally { + mockSpanner.setAbortProbability(0.0); + } + } + static com.google.spanner.v1.ResultSet createBitReversedSequenceResultSet( long startValue, long endValue) { return com.google.spanner.v1.ResultSet.newBuilder() diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java index 201651cfd57..72c160a5d91 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java @@ -249,6 +249,7 @@ public static void stopServer() { @Before public void setupResults() { mockSpanner.clearRequests(); + mockSpanner.removeAllExecutionTimes(); mockDatabaseAdmin.getRequests().clear(); mockInstanceAdmin.getRequests().clear(); }