fix: handling of totalTimeout on sql plan refresh (#2541)

jackdingilian · web-flow · commit bf49cf93f9a3 · 2025-03-24T22:58:25.000Z
Change-Id: I76d99d11c014b5602c28495c6f82b515542860e0 Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://togithub.com/googleapis/java-bigtable/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) - [ ] Rollback plan is reviewed and LGTMed - [ ] All new data plane features have a completed end to end testing plan Fixes #<issue_number_goes_here> ☕️ If you write sample code, please follow the [samples format]( https://togithub.com/GoogleCloudPlatform/java-docs-samples/blob/main/SAMPLE_FORMAT.md).
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallable.java
@@ -31,6 +31,7 @@
 import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatementRefreshTimeoutException;
 import com.google.cloud.bigtable.data.v2.models.sql.ResultSetMetadata;
 import com.google.cloud.bigtable.data.v2.stub.SafeResponseObserver;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.rpc.PreconditionFailure;
 import com.google.rpc.PreconditionFailure.Violation;
 import io.grpc.Deadline;
@@ -104,7 +105,8 @@ public void call(
 
   // Checks for an attempt timeout first, then a total timeout. If found, converts the timeout
   // to an absolute deadline. Adjusts totalTimeout based on the time since startTimeOfOverallRequest
-  private static @Nullable Deadline getDeadline(
+  @VisibleForTesting
+  static @Nullable Deadline getDeadline(
       GrpcCallContext grpcCallContext, Instant startTimeOfOverallRequest) {
     Optional<Deadline> attemptDeadline =
         Optional.ofNullable(grpcCallContext)
@@ -123,7 +125,7 @@ public void call(
               Duration elapsedTime = Duration.between(startTimeOfOverallRequest, Instant.now());
               Duration remaining = d.minus(elapsedTime);
               // zero is treated as no deadline, so if full deadline is elapsed pass 1 nano
-              long adjusted = Math.max(remaining.getNano(), 1);
+              long adjusted = Math.max(remaining.toNanos(), 1);
               return Deadline.after(adjusted, TimeUnit.NANOSECONDS);
             })
         .orElse(null);
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/PlanRefreshingCallableTest.java
@@ -60,6 +60,7 @@
 import io.grpc.Deadline;
 import io.grpc.Status.Code;
 import java.time.Duration;
+import java.time.Instant;
 import java.util.Collections;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
@@ -282,4 +283,47 @@ public void planRefreshDelayIsFactoredIntoExecuteTimeout() throws InterruptedExc
     Deadline executeDeadline = grpcCallContext.getCallOptions().getDeadline();
     assertThat(executeDeadline.isBefore(paddedDeadlineAtStartOfCall)).isTrue();
   }
+
+  @Test
+  public void testGetDeadlineWithAttemptTimeout() {
+    GrpcCallContext callContext =
+        GrpcCallContext.createDefault().withTimeoutDuration(Duration.ofMinutes(1));
+    // startTimeOfOverallRequest doesn't matter here
+    Deadline deadline = PlanRefreshingCallable.getDeadline(callContext, Instant.now());
+    long millisRemaining = deadline.timeRemaining(TimeUnit.MILLISECONDS);
+    assertThat(millisRemaining).isLessThan((60 * 1000) + 1);
+    // Give some padding in case tests are very slow
+    assertThat(millisRemaining).isGreaterThan(58 * 1000);
+  }
+
+  @Test
+  public void testGetDeadlineWithTotalTimeout() {
+    GrpcCallContext callContext =
+        GrpcCallContext.createDefault()
+            .withRetrySettings(
+                RetrySettings.newBuilder()
+                    .setTotalTimeout(org.threeten.bp.Duration.ofMinutes(1))
+                    .build());
+    Deadline deadline = PlanRefreshingCallable.getDeadline(callContext, Instant.now());
+    long millisRemaining = deadline.timeRemaining(TimeUnit.MILLISECONDS);
+    assertThat(millisRemaining).isLessThan((60 * 1000) + 1);
+    // Give some padding in case tests are very slow
+    assertThat(millisRemaining).isGreaterThan(58 * 1000);
+  }
+
+  @Test
+  public void testAttemptTimeoutUsedOverTotalTimeout() {
+    GrpcCallContext callContext =
+        GrpcCallContext.createDefault()
+            .withTimeoutDuration(Duration.ofMinutes(1))
+            .withRetrySettings(
+                RetrySettings.newBuilder()
+                    .setTotalTimeout(org.threeten.bp.Duration.ofHours(1))
+                    .build());
+    Deadline deadline = PlanRefreshingCallable.getDeadline(callContext, Instant.now());
+    long millisRemaining = deadline.timeRemaining(TimeUnit.MILLISECONDS);
+    assertThat(millisRemaining).isLessThan((60 * 1000) + 1);
+    // Give some padding in case tests are very slow
+    assertThat(millisRemaining).isGreaterThan(58 * 1000);
+  }
 }
