fix: recalculate remaining statement timeout after retry

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 @@ <T> ApiFuture<T> 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 <ReqT, RespT> ApiCallContext configure(
                     ApiCallContext context, ReqT request, MethodDescriptor<ReqT, RespT> 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;
                 }

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()

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();
   }