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.

Author: olavloite

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.MILLISECONDS), TimeUnit.MILLISECONDS);
       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.MILLISECONDS);
+                    if (remainingTimeout <= 0) {
+                      remainingTimeout = 1;
+                    }
                     return GrpcCallContext.createDefault()
-                        .withTimeoutDuration(statementTimeout.asDuration());
+                        .withTimeoutDuration(Duration.ofMillis(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();
   }

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionAsyncApiTest.java

@@ -19,6 +19,7 @@
 import static com.google.cloud.spanner.SpannerApiFutures.get;
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.fail;
@@ -535,8 +536,16 @@ private void testExecuteQuery(Function<Connection, Void> connectionConfigurator)
           }
           assertThat(timeout).isFalse();
         } catch (SpannerException e) {
-          assertThat(timeout).isTrue();
-          assertThat(e.getErrorCode()).isEqualTo(ErrorCode.DEADLINE_EXCEEDED);
+          if (e.getErrorCode() == ErrorCode.FAILED_PRECONDITION
+              && e.getMessage()
+                  .contains(
+                      "This transaction has been invalidated by a later transaction in the same session")) {
+            // Ignore for regular sessions.
+            assertFalse(isMultiplexedSessionsEnabledForRW(connection.getSpanner()));
+          } else {
+            assertThat(e.getErrorCode()).isEqualTo(ErrorCode.DEADLINE_EXCEEDED);
+            assertThat(timeout).isTrue();
+          }
           // Start a new transaction if a timeout occurred on a read/write transaction, as that will
           // invalidate that transaction.
           if (!connection.isReadOnly() && connection.isInTransaction()) {
@@ -577,8 +586,16 @@ private void testExecuteUpdateAsync(
           assertThat(connection.getCommitTimestamp()).isNotNull();
           assertThat(timeout).isFalse();
         } catch (SpannerException e) {
-          assertThat(timeout).isTrue();
-          assertThat(e.getErrorCode()).isEqualTo(ErrorCode.DEADLINE_EXCEEDED);
+          if (e.getErrorCode() == ErrorCode.FAILED_PRECONDITION
+              && e.getMessage()
+                  .contains(
+                      "This transaction has been invalidated by a later transaction in the same session")) {
+            // Ignore for regular sessions.
+            assertFalse(isMultiplexedSessionsEnabledForRW(connection.getSpanner()));
+          } else {
+            assertThat(e.getErrorCode()).isEqualTo(ErrorCode.DEADLINE_EXCEEDED);
+            assertThat(timeout).isTrue();
+          }
           // Start a new transaction if a timeout occurred on a read/write transaction, as that will
           // invalidate that transaction.
           if (!connection.isReadOnly() && connection.isInTransaction()) {