fix(bigtable): check remaining deadline on client-initiated retries in RetryingVRpc

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpc.java

@@ -293,25 +293,25 @@ boolean shouldRetry(VRpcResult result) {
       if (result.getRejected()) {
         return false;
       }
+      long timeRemainingMs =
+          context.getOperationInfo().getDeadline().timeRemaining(TimeUnit.MILLISECONDS);
+
       // If the error has RetryInfo, it means it comes from the server and should
-      // be retried.
+      // be retried as long as there's time remaining after the backoff.
       if (result.getRetryInfo() != null && result.getRetryInfo().hasRetryDelay()) {
-        long retryDelay = Durations.toMillis(result.getRetryInfo().getRetryDelay());
-        // only schedule retry if there's still a chance for the request to succeed
-        return context.getOperationInfo().getDeadline().timeRemaining(TimeUnit.MILLISECONDS)
-                - retryDelay
-            > 1;
+        long retryDelayMs = Durations.toMillis(result.getRetryInfo().getRetryDelay());
+        return timeRemainingMs - retryDelayMs > 1;
       }
-      // If the error didn't leave the client or failed in transport and is idempotent, we
-      // can retry up to 3 times.
-      boolean isRetryable =
-          (result.getState() == VRpcResult.State.UNCOMMITED)
-              || (context.getOperationInfo().isIdempotent()
-                  && result.getState() == VRpcResult.State.TRANSPORT_FAILURE);
-      if (isRetryable && context.getOperationInfo().getAttemptNumber() < 3) {
-        return true;
+
+      // Client-initiated retries require remaining deadline and are capped at 3 attempts.
+      if (timeRemainingMs <= 1 || context.getOperationInfo().getAttemptNumber() >= 3) {
+        return false;
       }
-      return false;
+
+      // Safe to retry if the error never left the client, or failed in transport and is idempotent.
+      return result.getState() == VRpcResult.State.UNCOMMITED
+          || (result.getState() == VRpcResult.State.TRANSPORT_FAILURE
+              && context.getOperationInfo().isIdempotent());
     }
   }
 

java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpcTest.java

@@ -471,4 +471,54 @@ public void requestNext() {}
     assertThat(vRpcException.getStatus().getCode()).isEqualTo(Status.Code.UNAVAILABLE);
     assertThat(attempts.get()).isEqualTo(1);
   }
+
+  @Test
+  void clientRetryRespectsDeadline() {
+    AtomicInteger attempts = new AtomicInteger();
+    RetryingVRpc<SessionFakeScriptedRequest, SessionFakeScriptedResponse> retrying =
+        new RetryingVRpc<>(
+            () ->
+                new VRpc<SessionFakeScriptedRequest, SessionFakeScriptedResponse>() {
+                  private boolean done;
+
+                  @Override
+                  public void start(
+                      SessionFakeScriptedRequest req,
+                      VRpcCallContext ctx,
+                      VRpcListener<SessionFakeScriptedResponse> listener) {
+                    attempts.incrementAndGet();
+                    done = true;
+                    ctx.getExecutor()
+                        .execute(
+                            () ->
+                                listener.onClose(
+                                    VRpcResult.createUncommitedError(
+                                        Status.UNAVAILABLE.withDescription("uncommitted error"))));
+                  }
+
+                  @Override
+                  public void cancel(String message, Throwable cause) {}
+
+                  @Override
+                  public boolean isDone() {
+                    return done;
+                  }
+
+                  @Override
+                  public void requestNext() {}
+                },
+            timer);
+
+    UnaryResponseFuture<SessionFakeScriptedResponse> f = new UnaryResponseFuture<>();
+    retrying.start(
+        SessionFakeScriptedRequest.getDefaultInstance(),
+        VRpc.VRpcCallContext.create(Deadline.after(0, TimeUnit.MILLISECONDS), true, tracer),
+        f);
+
+    ExecutionException executionException = assertThrows(ExecutionException.class, f::get);
+    assertThat(executionException.getCause()).isInstanceOf(VRpcException.class);
+    VRpcException vRpcException = (VRpcException) executionException.getCause();
+    assertThat(vRpcException.getStatus().getCode()).isEqualTo(Status.Code.UNAVAILABLE);
+    assertThat(attempts.get()).isEqualTo(1);
+  }
 }