From 0bd58d47cf98a341e47218100d3f3fa8061cf8bc Mon Sep 17 00:00:00 2001
From: Mattie Fu <mattiefu@google.com>
Date: Wed, 13 Nov 2024 09:31:56 -0500
Subject: [PATCH] update calculation and test

---
 .../data/v2/stub/metrics/BuiltinMetricsTracer.java  | 13 ++++++-------
 .../v2/stub/metrics/BuiltinMetricsTracerTest.java   | 13 ++++++++-----
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracer.java
index 806d1c1c00..e639ea5627 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracer.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracer.java
@@ -303,8 +303,8 @@ public void grpcMessageSent() {
   public void setTotalTimeoutDuration(java.time.Duration totalTimeoutDuration) {
     // This method is called by BigtableTracerStreamingCallable and
     // BigtableTracerUnaryCallable which is called per attempt. We only set
-    // the operationDeadline on the first attempt.
-    if (operationDeadline == null) {
+    // the operationDeadline on the first attempt and when totalTimeout is set.
+    if (operationDeadline == null && !totalTimeoutDuration.isZero()) {
       this.operationDeadline =
           Deadline.after(totalTimeoutDuration.toMillis(), TimeUnit.MILLISECONDS);
       this.remainingDeadlineAtAttemptStart = totalTimeoutDuration.toMillis();
@@ -405,11 +405,10 @@ private void recordAttemptCompletion(@Nullable Throwable status) {
     attemptLatenciesHistogram.record(
         convertToMs(attemptTimer.elapsed(TimeUnit.NANOSECONDS)), attributes);
 
-    // Total timeout is set to 0 when there's no timeout. In this case remaining deadline will be
-    // <= 0. Skip recording the value because it doesn't represent how much time is left to
-    // retry the request.
-    if (remainingDeadlineAtAttemptStart > 0) {
-      remainingDeadlineHistogram.record(remainingDeadlineAtAttemptStart, attributes);
+    // When operationDeadline is set, it's possible that the deadline is passed by the time we send
+    // a new attempt. In this case we'll record 0.
+    if (operationDeadline != null) {
+      remainingDeadlineHistogram.record(Math.max(0, remainingDeadlineAtAttemptStart), attributes);
     }
 
     if (serverLatencies != null) {
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerTest.java
index e902f0271f..62e4f37a86 100644
--- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerTest.java
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerTest.java
@@ -134,6 +134,7 @@ public class BuiltinMetricsTracerTest {
   private static final long APPLICATION_LATENCY = 200;
   private static final long SLEEP_VARIABILITY = 15;
   private static final String CLIENT_NAME = "java-bigtable/" + Version.VERSION;
+  private static final long READ_ROWS_TOTAL_TIMEOUT = 1000;
 
   private static final long CHANNEL_BLOCKING_LATENCY = 200;
 
@@ -220,10 +221,10 @@ public void sendHeaders(Metadata headers) {
     stubSettingsBuilder
         .readRowsSettings()
         .retrySettings()
-        .setTotalTimeoutDuration(Duration.ofMillis(9000))
-        .setMaxRpcTimeoutDuration(Duration.ofMillis(6000))
+        .setTotalTimeoutDuration(Duration.ofMillis(READ_ROWS_TOTAL_TIMEOUT))
+        .setMaxRpcTimeoutDuration(Duration.ofMillis(READ_ROWS_TOTAL_TIMEOUT))
         .setRpcTimeoutMultiplier(1)
-        .setInitialRpcTimeoutDuration(Duration.ofMillis(6000))
+        .setInitialRpcTimeoutDuration(Duration.ofMillis(800))
         .setInitialRetryDelayDuration(Duration.ofMillis(10))
         .setRetryDelayMultiplier(1)
         .setMaxRetryDelayDuration(Duration.ofMillis(10));
@@ -789,7 +790,7 @@ public void testRemainingDeadline() {
 
     double retryRemainingDeadline = retryHistogramPointData.getSum();
     // The retry remaining deadline should be equivalent to the original timeout.
-    assertThat(retryRemainingDeadline).isEqualTo(9000);
+    assertThat(retryRemainingDeadline).isEqualTo((double) READ_ROWS_TOTAL_TIMEOUT);
 
     Attributes okAttributes =
         baseAttributes
@@ -809,7 +810,9 @@ public void testRemainingDeadline() {
             .get(0);
 
     double okRemainingDeadline = okHistogramPointData.getSum();
-    assertThat(okRemainingDeadline).isWithin(200).of(8500);
+    // first attempt latency + retry delay
+    double expected = READ_ROWS_TOTAL_TIMEOUT - SERVER_LATENCY - CHANNEL_BLOCKING_LATENCY - 10;
+    assertThat(okRemainingDeadline).isWithin(10).of(expected);
   }
 
   private static class FakeService extends BigtableGrpc.BigtableImplBase {