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 5f4facf148..aaa984ad5a 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 151ac89e3a..8fec34c267 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 201651cfd5..72c160a5d9 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(); }