Integrate the MutateRowsRetryingCallable

Author: igorbernstein2

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java

@@ -15,18 +15,18 @@
  */
 package com.google.cloud.bigtable.data.v2.stub;
 
-import com.google.api.core.ApiFuture;
 import com.google.api.core.InternalApi;
-import com.google.api.gax.retrying.RetrySettings;
-import com.google.api.gax.rpc.ApiCallContext;
+import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
+import com.google.api.gax.retrying.RetryAlgorithm;
+import com.google.api.gax.retrying.RetryingExecutor;
+import com.google.api.gax.retrying.ScheduledRetryingExecutor;
 import com.google.api.gax.rpc.BatchingCallSettings;
 import com.google.api.gax.rpc.Callables;
 import com.google.api.gax.rpc.ClientContext;
 import com.google.api.gax.rpc.ServerStreamingCallSettings;
 import com.google.api.gax.rpc.ServerStreamingCallable;
 import com.google.api.gax.rpc.UnaryCallable;
 import com.google.bigtable.v2.MutateRowsRequest;
-import com.google.bigtable.v2.MutateRowsResponse;
 import com.google.bigtable.v2.ReadRowsRequest;
 import com.google.bigtable.v2.SampleRowKeysRequest;
 import com.google.bigtable.v2.SampleRowKeysResponse;
@@ -40,13 +40,14 @@
 import com.google.cloud.bigtable.data.v2.models.RowAdapter;
 import com.google.cloud.bigtable.data.v2.models.RowMutation;
 import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsBatchingDescriptor;
-import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsSpoolingCallable;
+import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsRetryingCallable;
 import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsUserFacingCallable;
 import com.google.cloud.bigtable.data.v2.stub.readrows.FilterMarkerRowsCallable;
 import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsResumptionStrategy;
 import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsRetryCompletedCallable;
 import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsUserCallable;
 import com.google.cloud.bigtable.data.v2.stub.readrows.RowMergingCallable;
+import com.google.cloud.bigtable.gaxx.retrying.ApiResultRetryAlgorithm;
 import java.io.IOException;
 import java.util.List;
 import org.threeten.bp.Duration;
@@ -251,28 +252,39 @@ private UnaryCallable<RowMutation, Void> createMutateRowCallable() {
    *   <li>Convert a {@link RowMutation} into a {@link MutateRowsRequest} with a single entry.
    *   <li>Using gax's {@link com.google.api.gax.rpc.BatchingCallable} to spool the requests and
    *       aggregate the {@link MutateRowsRequest.Entry}s.
-   *   <li>Spool the streamed responses.
+   *   <li>Process the response and schedule retries until all entries have been successfully
+   *       applied, fail permanently or there are no more retry attempts left.
+   *   <li>Wrap batch failures in a {@link
+   *       com.google.cloud.bigtable.data.v2.models.MutateRowsException}.
    *   <li>Split the responses using {@link MutateRowsBatchingDescriptor}.
-   *   <li>Apply retries to individual mutations
    * </ul>
    */
   private UnaryCallable<RowMutation, Void> createMutateRowsCallable() {
-    MutateRowsSpoolingCallable spooling = new MutateRowsSpoolingCallable(stub.mutateRowsCallable());
+    RetryAlgorithm<Void> retryAlgorithm =
+        new RetryAlgorithm<>(
+            new ApiResultRetryAlgorithm<Void>(),
+            new ExponentialRetryAlgorithm(
+                settings.mutateRowsSettings().getRetrySettings(), clientContext.getClock()));
+    RetryingExecutor<Void> retryingExecutor =
+        new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor());
+
+    UnaryCallable<MutateRowsRequest, Void> retrying =
+        new MutateRowsRetryingCallable(
+            clientContext.getDefaultCallContext(),
+            stub.mutateRowsCallable(),
+            retryingExecutor,
+            settings.mutateRowsSettings().getRetryableCodes());
 
     // recreate BatchingCallSettings with the correct descriptor
-    BatchingCallSettings.Builder<MutateRowsRequest, MutateRowsResponse> batchingCallSettings =
-        BatchingCallSettings.newBuilder(
-                new MutateRowsBatchingDescriptor(settings.mutateRowsSettings().getRetryableCodes()))
+    BatchingCallSettings.Builder<MutateRowsRequest, Void> batchingCallSettings =
+        BatchingCallSettings.newBuilder(new MutateRowsBatchingDescriptor())
             .setBatchingSettings(settings.mutateRowsSettings().getBatchingSettings());
 
-    UnaryCallable<MutateRowsRequest, MutateRowsResponse> batching =
-        Callables.batching(spooling, batchingCallSettings.build(), clientContext);
-
-    UnaryCallable<MutateRowsRequest, MutateRowsResponse> retrying =
-        Callables.retrying(batching, settings.mutateRowsSettings(), clientContext);
+    UnaryCallable<MutateRowsRequest, Void> batching =
+        Callables.batching(retrying, batchingCallSettings.build(), clientContext);
 
     MutateRowsUserFacingCallable userFacing =
-        new MutateRowsUserFacingCallable(retrying, requestContext);
+        new MutateRowsUserFacingCallable(batching, requestContext);
 
     return userFacing.withDefaultCallContext(clientContext.getDefaultCallContext());
   }

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsBatchingDescriptor.java

@@ -18,52 +18,28 @@
 import com.google.api.core.InternalApi;
 import com.google.api.gax.batching.PartitionKey;
 import com.google.api.gax.batching.RequestBuilder;
-import com.google.api.gax.grpc.GrpcStatusCode;
 import com.google.api.gax.rpc.ApiException;
-import com.google.api.gax.rpc.ApiExceptionFactory;
 import com.google.api.gax.rpc.BatchedRequestIssuer;
 import com.google.api.gax.rpc.BatchingDescriptor;
-import com.google.api.gax.rpc.StatusCode;
 import com.google.bigtable.v2.MutateRowsRequest;
-import com.google.bigtable.v2.MutateRowsResponse;
-import com.google.common.base.Preconditions;
-import com.google.common.collect.ImmutableSet;
-import com.google.common.primitives.Ints;
-import com.google.rpc.Code;
-import com.google.rpc.Status;
-import io.grpc.StatusException;
-import io.grpc.StatusRuntimeException;
+import com.google.cloud.bigtable.data.v2.models.MutateRowsException;
+import com.google.cloud.bigtable.data.v2.models.MutateRowsException.FailedMutation;
+import com.google.common.base.Function;
+import com.google.common.collect.Maps;
 import java.util.Collection;
-import java.util.Set;
+import java.util.List;
+import java.util.Map;
+import javax.annotation.Nullable;
 
 /**
- * A custom implementation of a {@link BatchingDescriptor} to split individual results of a bulk
- * MutateRowsResponse. Each individual result will be matched with its issuer. Since the embedded
- * results bypass gax's result processing chains, this class is responsible for wrapping errors in
- * {@link ApiException}s and marking each error as retryable.
+ * A custom implementation of a {@link BatchingDescriptor} to split individual results in a {@link
+ * MutateRowsException}. Each individual result will be matched with its issuer.
  *
  * <p>This class is considered an internal implementation detail and not meant to be used by
  * applications directly.
  */
 @InternalApi
-public class MutateRowsBatchingDescriptor
-    implements BatchingDescriptor<MutateRowsRequest, MutateRowsResponse> {
-
-  // Shared response to notify individual issuers of a successful mutation.
-  private static final MutateRowsResponse OK_RESPONSE =
-      MutateRowsResponse.newBuilder()
-          .addEntries(
-              MutateRowsResponse.Entry.newBuilder()
-                  .setIndex(0)
-                  .setStatus(Status.newBuilder().setCode(Code.OK_VALUE)))
-          .build();
-
-  private final ImmutableSet<StatusCode.Code> retryableCodes;
-
-  public MutateRowsBatchingDescriptor(Set<StatusCode.Code> retryableCodes) {
-    this.retryableCodes = ImmutableSet.copyOf(retryableCodes);
-  }
-
+public class MutateRowsBatchingDescriptor implements BatchingDescriptor<MutateRowsRequest, Void> {
   /** Return the target table name. This will be used to combine batcheable requests */
   @Override
   public PartitionKey getBatchPartitionKey(MutateRowsRequest request) {
@@ -76,44 +52,49 @@ public RequestBuilder<MutateRowsRequest> getRequestBuilder() {
     return new MyRequestBuilder();
   }
 
-  /** {@inheritDoc} */
   @Override
   public void splitResponse(
-      MutateRowsResponse batchResponse,
-      Collection<? extends BatchedRequestIssuer<MutateRowsResponse>> batch) {
+      Void batchResponse, Collection<? extends BatchedRequestIssuer<Void>> batch) {
 
-    // Sort the result entries by index.
-    Status[] sortedEntries = new Status[batchResponse.getEntriesCount()];
-
-    for (MutateRowsResponse.Entry entry : batchResponse.getEntriesList()) {
-      int index = Ints.checkedCast(entry.getIndex());
-      Preconditions.checkState(
-          sortedEntries[index] == null, "Got multiple results for the same sub-mutation");
-      sortedEntries[index] = entry.getStatus();
+    for (BatchedRequestIssuer<Void> issuer : batch) {
+      issuer.setResponse(null);
     }
+  }
 
-    // Notify all of issuers of the corresponding result.
-    int i = 0;
-    for (BatchedRequestIssuer<MutateRowsResponse> issuer : batch) {
-      Status entry = sortedEntries[i++];
-      Preconditions.checkState(entry != null, "Missing result for entry");
+  @Override
+  public void splitException(
+      Throwable throwable, Collection<? extends BatchedRequestIssuer<Void>> batch) {
 
-      if (entry.getCode() == Code.OK_VALUE) {
-        issuer.setResponse(OK_RESPONSE);
-      } else {
-        issuer.setException(createElementException(entry));
+    if (!(throwable instanceof MutateRowsException)) {
+      for (BatchedRequestIssuer<Void> issuer : batch) {
+        issuer.setException(throwable);
       }
+      return;
     }
-  }
 
-  /** {@inheritDoc} */
-  @Override
-  public void splitException(
-      Throwable throwable, Collection<? extends BatchedRequestIssuer<MutateRowsResponse>> batch) {
-    throwable = createElementException(throwable);
+    List<FailedMutation> failedMutations = ((MutateRowsException) throwable).getFailedMutations();
+
+    Map<Integer, FailedMutation> errorsByIndex =
+        Maps.uniqueIndex(
+            failedMutations,
+            new Function<FailedMutation, Integer>() {
+              @Nullable
+              @Override
+              public Integer apply(@Nullable FailedMutation input) {
+                return input.getIndex();
+              }
+            });
 
-    for (BatchedRequestIssuer<MutateRowsResponse> responder : batch) {
-      responder.setException(throwable);
+    int i = 0;
+    for (BatchedRequestIssuer<Void> issuer : batch) {
+      for (int j = 0; j < issuer.getMessageCount(); j++) {
+        FailedMutation failure = errorsByIndex.get(i++);
+        if (failure == null) {
+          issuer.setResponse(null);
+        } else {
+          issuer.setException(failure.getError());
+        }
+      }
     }
   }
 
@@ -129,38 +110,6 @@ public long countBytes(MutateRowsRequest request) {
     return request.getSerializedSize();
   }
 
-  /** Convert an element error Status into an ApiException */
-  private ApiException createElementException(Status protoStatus) {
-    Preconditions.checkArgument(protoStatus.getCode() != Code.OK_VALUE, "OK is not an error");
-
-    StatusRuntimeException throwable =
-        io.grpc.Status.fromCodeValue(protoStatus.getCode())
-            .withDescription(protoStatus.getMessage())
-            .asRuntimeException();
-
-    return createElementException(throwable);
-  }
-
-  /** Convert a Throwable into an ApiException, marking it as retryable when appropriate. */
-  private ApiException createElementException(Throwable throwable) {
-    final io.grpc.Status.Code code;
-
-    if (throwable instanceof ApiException) {
-      return (ApiException) throwable;
-    } else if (throwable instanceof StatusRuntimeException) {
-      code = ((StatusRuntimeException) throwable).getStatus().getCode();
-    } else if (throwable instanceof StatusException) {
-      code = ((StatusException) throwable).getStatus().getCode();
-    } else {
-      code = io.grpc.Status.Code.UNKNOWN;
-    }
-
-    GrpcStatusCode gaxStatusCode = GrpcStatusCode.of(code);
-    boolean isRetryable = retryableCodes.contains(gaxStatusCode.getCode());
-
-    return ApiExceptionFactory.createException(throwable, gaxStatusCode, isRetryable);
-  }
-
   /** A {@link com.google.api.gax.batching.RequestBuilder} that can aggregate MutateRowsRequest */
   static class MyRequestBuilder implements RequestBuilder<MutateRowsRequest> {
     private MutateRowsRequest.Builder builder;

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsSpoolingCallable.java

@@ -15,14 +15,11 @@
  */
 package com.google.cloud.bigtable.data.v2.stub.mutaterows;
 
-import com.google.api.core.ApiFunction;
 import com.google.api.core.ApiFuture;
-import com.google.api.core.ApiFutures;
 import com.google.api.core.InternalApi;
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.UnaryCallable;
 import com.google.bigtable.v2.MutateRowsRequest;
-import com.google.bigtable.v2.MutateRowsResponse;
 import com.google.cloud.bigtable.data.v2.internal.RequestContext;
 import com.google.cloud.bigtable.data.v2.models.RowMutation;
 
@@ -34,28 +31,18 @@
  */
 @InternalApi
 public class MutateRowsUserFacingCallable extends UnaryCallable<RowMutation, Void> {
-  private final UnaryCallable<MutateRowsRequest, MutateRowsResponse> inner;
+  private final UnaryCallable<MutateRowsRequest, Void> inner;
   private final RequestContext requestContext;
 
   public MutateRowsUserFacingCallable(
-      UnaryCallable<MutateRowsRequest, MutateRowsResponse> inner, RequestContext requestContext) {
+      UnaryCallable<MutateRowsRequest, Void> inner, RequestContext requestContext) {
 
     this.inner = inner;
     this.requestContext = requestContext;
   }
 
   @Override
   public ApiFuture<Void> futureCall(RowMutation request, ApiCallContext context) {
-    ApiFuture<MutateRowsResponse> rawResponse =
-        inner.futureCall(request.toBulkProto(requestContext), context);
-
-    return ApiFutures.transform(
-        rawResponse,
-        new ApiFunction<MutateRowsResponse, Void>() {
-          @Override
-          public Void apply(MutateRowsResponse mutateRowsResponse) {
-            return null;
-          }
-        });
+    return inner.futureCall(request.toBulkProto(requestContext), context);
   }
 }