fix(bigtable): check isOk and getRejected first in RetryingVRpc.shouldRetry

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

@@ -285,21 +285,23 @@ public void onExit() {
     }
 
     boolean shouldRetry(VRpcResult result) {
+      // Short-circuit successful RPCs.
+      if (result.getStatus().isOk()) {
+        return false;
+      }
+      // Do not retry results that are explicitly rejected by the client (ie. client shutdown).
+      if (result.getRejected()) {
+        return false;
+      }
       // If the error has RetryInfo, it means it comes from the server and should
       // be retried.
-      if (!result.getStatus().isOk()
-          && result.getRetryInfo() != null
-          && result.getRetryInfo().hasRetryDelay()) {
+      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;
       }
-      // Do not retry result that is explicitly rejected
-      if (result.getRejected()) {
-        return false;
-      }
       // If the error didn't leave the client or failed in transport and is idempotent, we
       // can retry up to 3 times.
       boolean isRetryable =

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

@@ -48,6 +48,7 @@
 import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeSessionService;
 import com.google.cloud.bigtable.data.v2.internal.session.fake.PeerInfoInterceptor;
 import com.google.common.base.Suppliers;
+import com.google.common.collect.ImmutableList;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.Duration;
 import com.google.protobuf.util.Durations;
@@ -65,6 +66,7 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -414,4 +416,59 @@ void cancelInScheduledState() throws Exception {
     VRpcException vRpcException = (VRpcException) executionException.getCause();
     assertThat(vRpcException.getStatus().getCode()).isEqualTo(Status.Code.CANCELLED);
   }
+
+  @Test
+  void rejectedResultNotRetriedEvenWithRetryInfo() {
+    AtomicInteger attempts = new AtomicInteger();
+    VRpc.VRpcResult rejectedWithRetryInfo =
+        new AutoValue_VRpc_VRpcResult(
+            VRpc.VRpcResult.State.UNCOMMITED,
+            Status.UNAVAILABLE.withDescription("rejected by client"),
+            ImmutableList.of(),
+            null,
+            java.time.Duration.ZERO,
+            RetryInfo.newBuilder().setRetryDelay(Durations.fromMillis(1)).build(),
+            /* rejected= */ true);
+
+    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(rejectedWithRetryInfo));
+                  }
+
+                  @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(1, TimeUnit.MINUTES), 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);
+  }
 }