From 1a4cd0b82c73dea24f2091bf81f3eb063fa91878 Mon Sep 17 00:00:00 2001
From: Igor Bernstein <igorbernstein@google.com>
Date: Thu, 10 Jan 2019 13:05:41 -0500
Subject: [PATCH] Bigtable: fix handling of unset rpc timeouts

When the rpc timeout is zero, then it should be treated as disabled not actual 0
---
 .../mutaterows/MutateRowsAttemptCallable.java | 28 +++++++++----------
 .../MutateRowsAttemptCallableTest.java        | 24 +++++++++++++++-
 2 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/google-cloud-clients/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsAttemptCallable.java b/google-cloud-clients/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsAttemptCallable.java
index 6e228d741b35..cb66f767d98f 100644
--- a/google-cloud-clients/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsAttemptCallable.java
+++ b/google-cloud-clients/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsAttemptCallable.java
@@ -101,14 +101,14 @@ public Object getTransportCode() {
       };
 
   // Everything needed to issue an RPC
-  private final UnaryCallable<MutateRowsRequest, List<MutateRowsResponse>> innerCallable;
-  private final ApiCallContext callContext;
-  private MutateRowsRequest currentRequest;
+  @Nonnull private final UnaryCallable<MutateRowsRequest, List<MutateRowsResponse>> innerCallable;
+  @Nonnull private final ApiCallContext callContext;
+  @Nonnull private MutateRowsRequest currentRequest;
 
   // Everything needed to build a retry request
-  private List<Integer> originalIndexes;
-  private final Set<StatusCode.Code> retryableCodes;
-  private final List<FailedMutation> permanentFailures;
+  @Nullable private List<Integer> originalIndexes;
+  @Nonnull private final Set<StatusCode.Code> retryableCodes;
+  @Nullable private final List<FailedMutation> permanentFailures;
 
   // Parent controller
   private RetryingFuture<Void> externalFuture;
@@ -135,12 +135,12 @@ public List<MutateRowsResponse> apply(Throwable throwable) {
   MutateRowsAttemptCallable(
       @Nonnull UnaryCallable<MutateRowsRequest, List<MutateRowsResponse>> innerCallable,
       @Nonnull MutateRowsRequest originalRequest,
-      @Nullable ApiCallContext callContext,
+      @Nonnull ApiCallContext callContext,
       @Nonnull Set<StatusCode.Code> retryableCodes) {
-    this.innerCallable = Preconditions.checkNotNull(innerCallable);
-    this.currentRequest = Preconditions.checkNotNull(originalRequest);
-    this.callContext = callContext;
-    this.retryableCodes = Preconditions.checkNotNull(retryableCodes);
+    this.innerCallable = Preconditions.checkNotNull(innerCallable, "innerCallable");
+    this.currentRequest = Preconditions.checkNotNull(originalRequest, "currentRequest");
+    this.callContext = Preconditions.checkNotNull(callContext, "callContext");
+    this.retryableCodes = Preconditions.checkNotNull(retryableCodes, "retryableCodes");
 
     permanentFailures = Lists.newArrayList();
   }
@@ -167,10 +167,10 @@ public Void call() {
           currentRequest.getEntriesCount() > 0, "Request doesn't have any mutations to send");
 
       // Configure the deadline
-      ApiCallContext currentCallContext = null;
-      if (callContext != null) {
+      ApiCallContext currentCallContext = callContext;
+      if (!externalFuture.getAttemptSettings().getRpcTimeout().isZero()) {
         currentCallContext =
-            callContext.withTimeout(externalFuture.getAttemptSettings().getRpcTimeout());
+            currentCallContext.withTimeout(externalFuture.getAttemptSettings().getRpcTimeout());
       }
 
       // Handle concurrent cancellation
diff --git a/google-cloud-clients/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsAttemptCallableTest.java b/google-cloud-clients/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsAttemptCallableTest.java
index ae4397c4db48..2df2aaf2b4a9 100644
--- a/google-cloud-clients/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsAttemptCallableTest.java
+++ b/google-cloud-clients/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsAttemptCallableTest.java
@@ -92,6 +92,28 @@ public void singleEntrySuccessTest() throws Exception {
     assertThat(innerCallable.lastRequest).isEqualTo(request);
   }
 
+  @Test
+  public void testNoRpcTimeout() {
+    parentFuture.timedAttemptSettings =
+        parentFuture.timedAttemptSettings.toBuilder().setRpcTimeout(Duration.ZERO).build();
+
+    MutateRowsRequest request =
+        MutateRowsRequest.newBuilder().addEntries(Entry.getDefaultInstance()).build();
+
+    innerCallable.response.add(
+        MutateRowsResponse.newBuilder()
+            .addEntries(
+                MutateRowsResponse.Entry.newBuilder().setIndex(0).setStatus(OK_STATUS_PROTO))
+            .build());
+
+    MutateRowsAttemptCallable attemptCallable =
+        new MutateRowsAttemptCallable(innerCallable, request, callContext, retryCodes);
+    attemptCallable.setExternalFuture(parentFuture);
+    attemptCallable.call();
+
+    assertThat(innerCallable.lastContext.getTimeout()).isNull();
+  }
+
   @Test
   public void mixedTest() {
     // Setup the request & response
@@ -340,7 +362,7 @@ public ApiFuture<List<MutateRowsResponse>> futureCall(
   static class MockRetryingFuture extends AbstractApiFuture<Void> implements RetryingFuture<Void> {
     ApiFuture<Void> attemptFuture;
 
-    final TimedAttemptSettings timedAttemptSettings;
+    TimedAttemptSettings timedAttemptSettings;
 
     MockRetryingFuture() {
       this(Duration.ofSeconds(5));