Repository: GoogleCloudPlatform/cloud-bigtable-client
Pull request #1570 (2017-12-19 13:32:08)

Adding a MutateRowsRequestManager
This new class encapsulates the logic of handling the MutateRowsRequest/Response interaction across retries.  This change is very similar to the structure of ReadRowsRequestManager / RetryingReadRowsOperation.

sduskis: I improved the testing, and added volatile where necessary.  @igorbernstein2, PTAL.

Patch:
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/MutateRowsRequestManager.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/MutateRowsRequestManager.java
new file mode 100644
index 000000000..34e138820
--- /dev/null
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/MutateRowsRequestManager.java
@@ -0,0 +1,195 @@
+/*
+ * Copyright 2017 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.bigtable.grpc.async;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import com.google.bigtable.v2.MutateRowsRequest;
+import com.google.bigtable.v2.MutateRowsResponse;
+import com.google.bigtable.v2.MutateRowsResponse.Entry;
+import com.google.cloud.bigtable.config.RetryOptions;
+import com.google.cloud.bigtable.grpc.BigtableDataClient;
+import com.google.rpc.Status;
+
+import io.grpc.Status.Code;
+
+/**
+ * Performs retries for {@link BigtableDataClient#mutateRows(MutateRowsRequest)} operations.
+ */
+public class MutateRowsRequestManager  {
+  private final static Status STATUS_INTERNAL =
+      Status.newBuilder().setCode(io.grpc.Status.Code.INTERNAL.value()).build();
+
+  private static Code getGrpcCode(Status status) {
+    return status == null ? null : io.grpc.Status.fromCodeValue(status.getCode()).getCode();
+  }
+
+  private static Entry createEntry(int i, Status status) {
+    return MutateRowsResponse.Entry.newBuilder().setIndex(i).setStatus(status).build();
+  }
+
+  public enum ProcessingStatus {
+    /* All responses produced OK */
+    SUCCESS,
+    /* All responses produced OK or a retryable code */
+    RETRYABLE,
+    /* Some responses a non-retryable code */
+    NOT_RETRYABLE,
+    /* The response was invalid - missing indexes and etc. */
+    INVALID;
+  }
+
+  /**
+   * The current request to send. This starts as the original request. If retries occur, this
+   * request will contain the subset of Mutations that need to be retried.
+   */
+  private volatile MutateRowsRequest currentRequest;
+
+  /**
+   * When doing retries, the retry sends a partial set of the original mutations that failed with a
+   * retryable status. This array contains a mapping of indices from the {@link #currentRequest}
+   * to {@link #originalRequest}.
+   */
+  private int[] mapToOriginalIndex;
+
+  /**
+   * This array tracks the cumulative set of results across all RPC requests.
+   */
+  private final Status[] results;
+  private final RetryOptions retryOptions;
+  private final MutateRowsRequest originalRequest;
+
+  private boolean messageIsInvalid = false;
+
+  public MutateRowsRequestManager(RetryOptions retryOptions, MutateRowsRequest originalRequest) {
+    this.currentRequest = originalRequest;
+    this.originalRequest = originalRequest;
+    this.retryOptions = retryOptions;
+    results = new Status[originalRequest.getEntriesCount()];
+
+    // This map should is a map between currentRequest and originalRquest. For now, currentRequest
+    // == originalRquest, but they could diverge if a retry occurs.
+    mapToOriginalIndex = new int[originalRequest.getEntriesCount()];
+    for (int i = 0; i < mapToOriginalIndex.length; i++) {
+      mapToOriginalIndex[i] = i;
+    }
+  }
+
+  /**
+   * Adds the content of the message to the {@link #results}.
+   */
+  public void onMessage(MutateRowsResponse message) {
+    for (Entry entry : message.getEntriesList()) {
+      int index = (int) entry.getIndex();
+
+      // Sanity check to make sure that the index returned from the server is valid.
+      if (index >= mapToOriginalIndex.length || index < 0) {
+        messageIsInvalid = true;
+        break;
+      }
+
+      // Set the result.
+      results[mapToOriginalIndex[index]] = entry.getStatus();
+    }
+  }
+
+  /**
+   * This is called when all calls to {@link #onMessage(MutateRowsResponse)} are complete.
+   *
+   * @return {@link ProcessingStatus} of the accumulated responses - success, invalid, retrable,
+   *         non-retryable.
+   */
+  public ProcessingStatus onOK() {
+    // Sanity check to make sure that every mutation received a response.
+    if (!messageIsInvalid) {
+      for (int i = 0; i < results.length; i++) {
+        if (results[i] == null) {
+          messageIsInvalid = true;
+          break;
+        }
+      }
+    }
+
+    // There was a problem in the data found in onMessage(), so fail the RPC.
+    if (messageIsInvalid) {
+      return ProcessingStatus.INVALID;
+    }
+
+    List<Integer> toRetry = new ArrayList<>();
+    ProcessingStatus processingStatus = ProcessingStatus.SUCCESS;
+
+    // Check the current state to determine the state of the results.
+    // There are three states: OK, Fail, or Partial Retry.
+    for (int i = 0; i < results.length; i++) {
+      Status status = results[i];
+      if (status.getCode() == io.grpc.Status.Code.OK.value()) {
+        continue;
+      } else if (retryOptions.isRetryable(getGrpcCode(status))) {
+        // An individual mutation failed with a retryable code, usually DEADLINE_EXCEEDED.
+        toRetry.add(i);
+        if (processingStatus == ProcessingStatus.SUCCESS) {
+          processingStatus = ProcessingStatus.RETRYABLE;
+        }
+      } else {
+        // Don't retry if even a single response is not retryable.
+        processingStatus = ProcessingStatus.NOT_RETRYABLE;
+        break;
+      }
+    }
+
+    if (!toRetry.isEmpty()) {
+      currentRequest = createRetryRequest(toRetry);
+    }
+
+    return processingStatus;
+  }
+
+  public MutateRowsRequest getRetryRequest() {
+    return currentRequest;
+  }
+
+  /**
+   * Creates a new {@link MutateRowsRequest} that's a subset of the original request that
+   * corresponds to a set of indices.
+   *
+   * @param indiciesToRetry
+   * @return the new {@link MutateRowsRequest}.
+   */
+  private MutateRowsRequest createRetryRequest(List<Integer> indiciesToRetry) {
+    MutateRowsRequest.Builder updatedRequest = MutateRowsRequest.newBuilder()
+        .setTableName(originalRequest.getTableName());
+    mapToOriginalIndex = new int[indiciesToRetry.size()];
+    for (int i = 0; i < indiciesToRetry.size(); i++) {
+      mapToOriginalIndex[i] = indiciesToRetry.get(i);
+      updatedRequest.addEntries(originalRequest.getEntries(indiciesToRetry.get(i)));
+    }
+    return updatedRequest.build();
+  }
+
+  /**
+   * @return a {@link MutateRowsResponse} built from the accumulation of all calls to
+   *         onMessage/onOK.
+   */
+  public MutateRowsResponse buildResponse() {
+    List<MutateRowsResponse.Entry> entries = new ArrayList<>();
+    for (int i = 0; i < results.length; i++) {
+      Status status = (results[i] == null) ? STATUS_INTERNAL : results[i];
+      entries.add(createEntry(i, status));
+    }
+    return MutateRowsResponse.newBuilder().addAllEntries(entries).build();
+  }
+}
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/RetryingMutateRowsOperation.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/RetryingMutateRowsOperation.java
index 93f634f4d..62af2464e 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/RetryingMutateRowsOperation.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/RetryingMutateRowsOperation.java
@@ -15,7 +15,6 @@
  */
 package com.google.cloud.bigtable.grpc.async;
 
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.ScheduledExecutorService;
@@ -23,15 +22,13 @@ import java.util.concurrent.ScheduledExecutorService;
 import com.google.api.client.util.BackOff;
 import com.google.bigtable.v2.MutateRowsRequest;
 import com.google.bigtable.v2.MutateRowsResponse;
-import com.google.bigtable.v2.MutateRowsResponse.Entry;
 import com.google.cloud.bigtable.config.RetryOptions;
 import com.google.cloud.bigtable.grpc.BigtableDataClient;
+import com.google.cloud.bigtable.grpc.async.MutateRowsRequestManager.ProcessingStatus;
 import com.google.common.collect.ImmutableMap;
-import com.google.rpc.Status;
 
 import io.grpc.CallOptions;
 import io.grpc.Metadata;
-import io.grpc.Status.Code;
 import io.opencensus.trace.AttributeValue;
 
 /**
@@ -41,180 +38,60 @@ public class RetryingMutateRowsOperation extends
     AbstractRetryingOperation<MutateRowsRequest, MutateRowsResponse, List<MutateRowsResponse>> {
   private final static io.grpc.Status INVALID_RESPONSE =
       io.grpc.Status.INTERNAL.withDescription("The server returned an invalid response");
-  private final static Status STATUS_INTERNAL =
-      Status.newBuilder().setCode(io.grpc.Status.Code.INTERNAL.value()).build();
 
-  private static Code getGrpcCode(Status status) {
-    return status == null ? null : io.grpc.Status.fromCodeValue(status.getCode()).getCode();
-  }
-
-  private static Entry createEntry(int i, Status status) {
-    return MutateRowsResponse.Entry.newBuilder().setIndex(i).setStatus(status).build();
-  }
-
-
-  /**
-   * The current request to send. This starts as the original request. If retries occur, this
-   * request will contain the subset of Mutations that need to be retried.
-   */
-  private MutateRowsRequest currentRequest;
-
-  /**
-   * When doing retries, the retry sends a partial set of the original mutations that failed with a
-   * retryable status. This array contains a mapping of indices from the {@link #getRetryRequest()}
-   * to {@link #getOriginalRequest()}.
-   */
-  private int[] mapToOriginalIndex;
-
-  /**
-   * This array tracks the cumulative set of results across all RPC requests.
-   */
-  private final Status[] results;
-  private boolean messageIsInvalid;
+  private final MutateRowsRequestManager requestManager;
 
   public RetryingMutateRowsOperation(RetryOptions retryOptions, MutateRowsRequest originalRquest,
       BigtableAsyncRpc<MutateRowsRequest, MutateRowsResponse> retryableRpc, CallOptions callOptions,
       ScheduledExecutorService retryExecutorService, Metadata originalMetadata) {
     super(retryOptions, originalRquest, retryableRpc, callOptions, retryExecutorService, originalMetadata);
-    this.currentRequest = originalRquest;
-    results = new Status[originalRquest.getEntriesCount()];
-
-    // This map should is a map between currentRequest and originalRquest. For now, currentRequest
-    // == originalRquest, but they could diverge if a retry occurs.
-    mapToOriginalIndex = new int[originalRquest.getEntriesCount()];
-    for (int i = 0; i < mapToOriginalIndex.length; i++) {
-      mapToOriginalIndex[i] = i;
-    }
+    requestManager = new MutateRowsRequestManager(retryOptions, originalRquest);
     operationSpan.addAnnotation("MutationCount", ImmutableMap.of("count",
       AttributeValue.longAttributeValue(originalRquest.getEntriesCount())));
   }
 
-  /**
-   * Adds the content of the message to the {@link #results}.
-   */
   @Override
   public void onMessage(MutateRowsResponse message) {
-    for (Entry entry : message.getEntriesList()) {
-      int index = (int) entry.getIndex();
-
-      // Sanity check to make sure that the index returned from the server is valid.
-      if (index >= mapToOriginalIndex.length || index < 0) {
-        messageIsInvalid = true;
-        break;
-      }
-
-      // Set the result.
-      results[mapToOriginalIndex[index]] = entry.getStatus();
-    }
+    requestManager.onMessage(message);
   }
 
   @Override
   protected MutateRowsRequest getRetryRequest() {
-    return currentRequest;
+    return requestManager.getRetryRequest();
   }
 
   @Override
   protected boolean onOK(Metadata trailers) {
-    // Sanity check to make sure that every mutation received a response.
-    for (int i = 0; i < results.length; i++) {
-      if (results[i] == null) {
-        messageIsInvalid = true;
-        break;
-      }
-    }
+    ProcessingStatus status = requestManager.onOK();
 
-    // There was a problem in the data found in onMessage(), so fail the RPC.
-    if (messageIsInvalid) {
+    if (status == ProcessingStatus.INVALID) {
+      // Set an exception.
       onError(INVALID_RESPONSE, trailers);
-      return false;
-    }
-
-    boolean fail = false;
-    List<Integer> toRetry = new ArrayList<>();
-
-    // Check the current state to determine the state of the results.
-    // There are three states: OK, Fail, or Partial Retry.
-    for (int i = 0; i < results.length; i++) {
-      Status status = results[i];
-      if (status.getCode() == io.grpc.Status.Code.OK.value()) {
-        continue;
-      }
-
-      if (retryOptions.isRetryable(getGrpcCode(status))) {
-        // An individual mutation failed with a retryable code, usually DEADLINE_EXCEEDED.
-        toRetry.add(i);
-      } else {
-        // Don't retry if even a single response is not retryable.
-        fail = true;
-      }
-    }
-
-    if (fail) {
-      rpc.getRpcMetrics().markFailure();
-    }
-
-    if (!toRetry.isEmpty()) {
-      operationSpan.addAnnotation("MutationCount", ImmutableMap.of("retryCount",
-        AttributeValue.longAttributeValue(toRetry.size())));
+      return true;
     }
 
-    // OK or Fail should set the future, and end the operation.
-    // The caller should handle the processing of the failed mutations.
-    if (toRetry.isEmpty() || fail) {
-      completionFuture.set(Arrays.asList(buildResponse()));
+    // There was a problem in the data found in onMessage(), so fail the RPC.
+    if (status == ProcessingStatus.SUCCESS || status == ProcessingStatus.NOT_RETRYABLE) {
+      // Set the response, with either success, or non-retryable responses.
+      completionFuture.set(Arrays.asList(requestManager.buildResponse()));
       return true;
     }
 
     // Perform a partial retry, if the backoff policy allows it.
     long nextBackOff = getNextBackoff();
     if (nextBackOff == BackOff.STOP) {
+      // Return the response as is, and don't retry;
       rpc.getRpcMetrics().markRetriesExhasted();
-      completionFuture.set(Arrays.asList(buildResponse()));
+      completionFuture.set(Arrays.asList(requestManager.buildResponse()));
+      operationSpan.addAnnotation("MutationCount", ImmutableMap.of("failureCount",
+        AttributeValue.longAttributeValue(requestManager.getRetryRequest().getEntriesCount())));
       return true;
-    } else {
-      currentRequest = createRetryRequest(toRetry);
-      mapToOriginalIndex = toIntArray(toRetry);
-      performRetry(nextBackOff);
-      return false;
-    }
-  }
-
-  /**
-   * Creates a new {@link MutateRowsRequest} that's a subset of the original request that
-   * corresponds to a set of indices.
-   *
-   * @param indiciesToRetry
-   * @return the new {@link MutateRowsRequest}.
-   */
-  private MutateRowsRequest createRetryRequest(List<Integer> indiciesToRetry) {
-    MutateRowsRequest originalRequest = getOriginalRequest();
-    MutateRowsRequest.Builder updatedRequest = MutateRowsRequest.newBuilder()
-        .setTableName(originalRequest.getTableName());
-
-    for (int i = 0; i < indiciesToRetry.size(); i++) {
-      updatedRequest.addEntries(originalRequest.getEntries(indiciesToRetry.get(i)));
     }
 
-    return updatedRequest.build();
-  }
-
-  /**
-   * Converts the List<Integer> to int[].
-   */
-  private int[] toIntArray(List<Integer> list) {
-    int[] array = new int[list.size()];
-    for (int i = 0; i < list.size(); i++) {
-      array[i] = list.get(i);
-    }
-    return array;
+    performRetry(nextBackOff);
+    operationSpan.addAnnotation("MutationCount", ImmutableMap.of("retryCount",
+      AttributeValue.longAttributeValue(requestManager.getRetryRequest().getEntriesCount())));
+    return false;
   }
 
-  private MutateRowsResponse buildResponse() {
-    List<MutateRowsResponse.Entry> entries = new ArrayList<>();
-    for (int i = 0; i < results.length; i++) {
-      Status status = (results[i] == null) ? STATUS_INTERNAL : results[i];
-      entries.add(createEntry(i, status));
-    }
-    return MutateRowsResponse.newBuilder().addAllEntries(entries).build();
-  }
 }
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/scanner/ReadRowsRequestManager.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/scanner/ReadRowsRequestManager.java
index 8e92ae188..5d06e98df 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/scanner/ReadRowsRequestManager.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/scanner/ReadRowsRequestManager.java
@@ -38,9 +38,9 @@ class ReadRowsRequestManager {
   private final ReadRowsRequest originalRequest;
 
   // The number of rows read so far.
-  private long rowCount = 0;
+  private volatile long rowCount = 0;
 
-  private ByteString lastFoundKey;
+  private volatile ByteString lastFoundKey;
 
   /**
    * <p>
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/scanner/RetryingReadRowsOperation.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/scanner/RetryingReadRowsOperation.java
index 4169c929e..12c81c572 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/scanner/RetryingReadRowsOperation.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/scanner/RetryingReadRowsOperation.java
@@ -118,6 +118,7 @@ public class RetryingReadRowsOperation extends
   private CallToStreamObserverAdapter<ReadRowsRequest> adapter;
   private StreamObserver<ReadRowsResponse> resultObserver;
   private int totalRowsProcessed = 0;
+  private volatile ReadRowsRequest nextRequest;
 
   public RetryingReadRowsOperation(
       StreamObserver<FlatRow> observer,
@@ -130,6 +131,7 @@ public class RetryingReadRowsOperation extends
     super(retryOptions, request, retryableRpc, callOptions, retryExecutorService, originalMetadata);
     this.rowObserver = observer;
     this.requestManager = new ReadRowsRequestManager(request);
+    this.nextRequest = request;
   }
 
   // This observer will be notified after ReadRowsResponse have been fully processed.
@@ -142,7 +144,7 @@ public class RetryingReadRowsOperation extends
    */
   @Override
   protected ReadRowsRequest getRetryRequest() {
-    return requestManager.buildUpdatedRequest();
+    return nextRequest;
   }
 
   @SuppressWarnings("unchecked")
@@ -306,6 +308,17 @@ public class RetryingReadRowsOperation extends
     }
   }
 
+  @Override
+  protected void performRetry(long nextBackOff) {
+    buildUpdatedRequst();
+    super.performRetry(nextBackOff);
+  }
+
+  @VisibleForTesting
+  ReadRowsRequest buildUpdatedRequst() {
+    return nextRequest = requestManager.buildUpdatedRequest();
+  }
+
   @VisibleForTesting
   int getTimeoutRetryCount() {
     return timeoutRetryCount;


Test patch:
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestMutateRowsRequestManager.java b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestMutateRowsRequestManager.java
new file mode 100644
index 000000000..5988ce175
--- /dev/null
+++ b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestMutateRowsRequestManager.java
@@ -0,0 +1,242 @@
+/*
+ * Copyright 2017 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.bigtable.grpc.async;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicLong;
+
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.MockitoAnnotations;
+
+import com.google.api.client.util.NanoClock;
+import com.google.bigtable.v2.MutateRowsRequest;
+import com.google.bigtable.v2.MutateRowsRequest.Entry;
+import com.google.bigtable.v2.MutateRowsResponse;
+import com.google.bigtable.v2.Mutation;
+import com.google.bigtable.v2.Mutation.SetCell;
+import com.google.cloud.bigtable.config.RetryOptions;
+import com.google.cloud.bigtable.config.RetryOptionsUtil;
+import com.google.cloud.bigtable.grpc.async.MutateRowsRequestManager.ProcessingStatus;
+import com.google.rpc.Status;
+
+import io.grpc.Status.Code;
+
+/**
+ * Tests for {@link MutateRowsRequestManager}
+ */
+@RunWith(JUnit4.class)
+public class TestMutateRowsRequestManager {
+
+  private static Status OK = statusOf(io.grpc.Status.Code.OK);
+  private static Status DEADLINE_EXCEEDED = statusOf(io.grpc.Status.Code.DEADLINE_EXCEEDED);
+  private static Status NOT_FOUND = statusOf(io.grpc.Status.Code.NOT_FOUND);
+
+  private static MutateRowsRequest createRequest(int entryCount) {
+    MutateRowsRequest.Builder builder = MutateRowsRequest.newBuilder();
+    for (int i = 0; i < entryCount; i++) {
+      Mutation mutation = Mutation.newBuilder()
+          .setSetCell(SetCell.newBuilder().setFamilyName("Family" + i).build()).build();
+      builder.addEntries(Entry.newBuilder().addMutations(mutation));
+    }
+    return builder.build();
+  }
+
+  private static MutateRowsRequest createRequest(MutateRowsRequest original, int ... indices) {
+    MutateRowsRequest.Builder builder = MutateRowsRequest.newBuilder();
+    for (int i : indices) {
+      builder.addEntries(original.getEntries(i));
+    }
+    return builder.build();
+  }
+
+  private static MutateRowsResponse createResponse(Status... statuses) {
+    MutateRowsResponse.Builder builder = MutateRowsResponse.newBuilder();
+    for (int i = 0; i < statuses.length; i++) {
+      builder.addEntries(toEntry(i, statuses[i]));
+    }
+    return builder.build();
+  }
+
+  private static com.google.bigtable.v2.MutateRowsResponse.Entry toEntry(int i, Status status) {
+    return MutateRowsResponse.Entry.newBuilder().setIndex(i).setStatus(status).build();
+  }
+
+  private static Status statusOf(Code code) {
+    return Status.newBuilder().setCode(code.value()).build();
+  }
+
+  private static ProcessingStatus send(MutateRowsRequestManager underTest,
+      MutateRowsResponse response) {
+    underTest.onMessage(response);
+    return underTest.onOK();
+  }
+
+  private AtomicLong time = new AtomicLong();
+  private NanoClock nanoClock = new NanoClock() {
+    @Override
+    public long nanoTime() {
+      return time.get();
+    }
+  };
+
+  private RetryOptions retryOptions;
+
+  @Before
+  public void setup() {
+    MockitoAnnotations.initMocks(this);
+    retryOptions = RetryOptionsUtil.createTestRetryOptions(nanoClock);
+  }
+
+  @Test
+  /**
+   * An empty request should return an empty response
+   */
+  public void testEmptySuccess() {
+    MutateRowsRequestManager underTest =
+        new MutateRowsRequestManager(retryOptions, createRequest(0));
+    send(underTest, createResponse());
+    Assert.assertEquals(createResponse(), underTest.buildResponse());
+  }
+
+  @Test
+  /**
+   * A single successful entry should work.
+   */
+  public void testSingleSuccess() {
+    MutateRowsRequestManager underTest =
+        new MutateRowsRequestManager(retryOptions, createRequest(1));
+    send(underTest, createResponse(OK));
+    Assert.assertEquals(createResponse(OK), underTest.buildResponse());
+  }
+
+  @Test
+  /**
+   * Two individual calls with one retry should work.
+   */
+  public void testTwoTrySuccessOneFailure() {
+    MutateRowsRequest originalRequest = createRequest(3);
+    MutateRowsRequestManager underTest =
+        new MutateRowsRequestManager(retryOptions, originalRequest);
+    send(underTest, createResponse(OK, DEADLINE_EXCEEDED, OK));
+    Assert.assertEquals(createRequest(originalRequest, 1), underTest.getRetryRequest());
+    Assert.assertEquals(createResponse(OK, DEADLINE_EXCEEDED, OK), underTest.buildResponse());
+    send(underTest, createResponse(OK));
+    Assert.assertEquals(createResponse(OK, OK, OK), underTest.buildResponse());
+  }
+
+  @Test
+  /**
+   * Two individual calls in a more complicated case with one retry should work.
+   */
+  public void testMultiSuccess() {
+    MutateRowsRequest originalRequest = createRequest(10);
+    MutateRowsRequestManager underTest =
+        new MutateRowsRequestManager(retryOptions, originalRequest);
+
+    // 5 mutations succeed, 5 mutations are retryable.
+    MutateRowsResponse firstResponse = createResponse(OK, DEADLINE_EXCEEDED, OK, DEADLINE_EXCEEDED,
+      OK, DEADLINE_EXCEEDED, OK, DEADLINE_EXCEEDED, OK, DEADLINE_EXCEEDED);
+    send(underTest, firstResponse);
+    Assert.assertEquals(createRequest(originalRequest, 1, 3, 5, 7, 9), underTest.getRetryRequest());
+    Assert.assertEquals(firstResponse, underTest.buildResponse());
+
+    // 3 mutations succeed, 2 mutations are retryable.
+    send(underTest, createResponse(OK, DEADLINE_EXCEEDED, OK, OK, DEADLINE_EXCEEDED));
+    Assert.assertEquals(createRequest(originalRequest, 3, 9), underTest.getRetryRequest());
+    MutateRowsResponse secondResponse = createResponse(OK, OK, OK, DEADLINE_EXCEEDED,
+      OK, OK, OK, OK, OK, DEADLINE_EXCEEDED);
+    Assert.assertEquals(secondResponse, underTest.buildResponse());
+
+    // The final 2 mutations are OK
+    send(underTest, createResponse(OK, OK));
+    Assert.assertEquals(createResponse(OK, OK, OK, OK, OK, OK, OK, OK, OK, OK),
+      underTest.buildResponse());
+  }
+
+  @Test
+  /**
+   * Multiple attempts at retries should work as expected. 10 mutations are added, and 1 gets an OK
+   * status for 9 rounds until 1 mutation is left.  Each success shows up in a random location.
+   */
+  public void testMultiAttempt() {
+    MutateRowsRequest originalRequest = createRequest(10);
+    MutateRowsRequestManager underTest =
+        new MutateRowsRequestManager(retryOptions, originalRequest);
+
+    // At the beginning, all mutations are outstanding
+    List<Integer> remaining = new ArrayList<>();
+    for (int i = 0; i < 10; i++) {
+      remaining.add(i);
+    }
+
+    for (int i = 0; i < 9; i++) {
+      int remainingMutationCount = remaining.size();
+      // Randomly choose a successful mutation
+      int successIndex = (int) (Math.random() * remainingMutationCount);
+
+      // Create a Status array filled with DEADLINE_EXCEEDED except for the chose one.
+      Status statuses[] = new Status[remainingMutationCount];
+      Arrays.fill(statuses, DEADLINE_EXCEEDED);
+      statuses[successIndex] = OK;
+
+      // The successful status can now be removed.
+      remaining.remove(successIndex);
+
+      // Make sure that the request is retryable, and that the retry request looks reasonable.
+      Assert.assertEquals(ProcessingStatus.RETRYABLE, send(underTest, createResponse(statuses)));
+      Assert.assertEquals(createRequest(originalRequest, toIntArray(remaining)),
+        underTest.getRetryRequest());
+    }
+
+    // Only one Mutation should be outstanding at this point. Create a response that has all OKs,
+    // with the exception of the remaining statuses.
+    Status[] statuses = new Status[10];
+    Arrays.fill(statuses, OK);
+    statuses[remaining.get(0)] = DEADLINE_EXCEEDED;
+    Assert.assertEquals(createResponse(statuses), underTest.buildResponse());
+  }
+
+  private final int[] toIntArray(List<Integer> integers) {
+    int indicesToRetry[] = new int[integers.size()];
+    for (int i = 0; i < integers.size(); i++) {
+      indicesToRetry[i] = integers.get(i);
+    }
+    return indicesToRetry;
+  }
+
+  @Test
+  public void testNotRetryable() {
+    MutateRowsRequestManager underTest =
+        new MutateRowsRequestManager(retryOptions, createRequest(3));
+    Assert.assertEquals(ProcessingStatus.NOT_RETRYABLE,
+      send(underTest, createResponse(OK, OK, NOT_FOUND)));
+    Assert.assertEquals(createResponse(OK, OK, NOT_FOUND), underTest.buildResponse());
+  }
+
+  @Test
+  public void testInvalid() {
+    // Create 3 muntations, but only 2 OKs.  That should be invalid
+    MutateRowsRequestManager underTest =
+        new MutateRowsRequestManager(retryOptions, createRequest(3));
+    Assert.assertEquals(ProcessingStatus.INVALID, send(underTest, createResponse(OK, OK)));
+  }
+}
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/scanner/RetryingReadRowsOperationTest.java b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/scanner/RetryingReadRowsOperationTest.java
index 73c9184ae..1286338eb 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/scanner/RetryingReadRowsOperationTest.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/scanner/RetryingReadRowsOperationTest.java
@@ -381,8 +381,9 @@ public class RetryingReadRowsOperationTest {
     verify(mockRpcTimerContext, times(expectedRetryCount + 1)).close();
   }
 
-  private static void checkRetryRequest(RetryingReadRowsOperation underTest, ByteString key, int rowCount) {
-    ReadRowsRequest request = underTest.getRetryRequest();
+  private static void checkRetryRequest(RetryingReadRowsOperation underTest, ByteString key,
+      int rowCount) {
+    ReadRowsRequest request = underTest.buildUpdatedRequst();
     Assert.assertEquals(key, request.getRows().getRowRanges(0).getStartKeyOpen());
     Assert.assertEquals(rowCount, request.getRowsLimit());
   }