fix

Change-Id: I572ff9eb1569f83a87ad67c59fbf5f023ef45a4c

Author: mutianf

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/LargeRowException.java

@@ -16,36 +16,27 @@
 package com.google.cloud.bigtable.data.v2.models;
 
 import com.google.api.core.InternalApi;
+import com.google.api.gax.grpc.GrpcStatusCode;
 import com.google.api.gax.rpc.ApiException;
-import com.google.api.gax.rpc.StatusCode;
 import com.google.protobuf.ByteString;
+import io.grpc.Status;
 import java.util.List;
 import javax.annotation.Nonnull;
-import javax.annotation.Nullable;
 
-/** Thrown by the ReadRows when at least one row was too large to be read. */
+/** Exception that wraps all the large row keys encountered in a ReadRows request. */
 public final class LargeRowException extends ApiException {
   private final List<ByteString> largeRowKeys;
 
-  /**
-   * This constructor is considered an internal implementation detail and not meant to be used by
-   * applications.
-   */
   @InternalApi
-  public LargeRowException(
-      @Nullable Throwable cause,
-      @Nonnull StatusCode statusCode,
-      boolean retryable,
-      @Nonnull List<ByteString> largeRowKeys) {
-    super(cause, statusCode, retryable);
+  public LargeRowException(@Nonnull List<ByteString> largeRowKeys) {
+    super(
+        "Encountered large row keys",
+        null,
+        GrpcStatusCode.of(Status.Code.FAILED_PRECONDITION),
+        false);
     this.largeRowKeys = largeRowKeys;
   }
 
-  @Override
-  public String getMessage() {
-    return "Large rows encountered: " + largeRowKeys;
-  }
-
   /** Retrieve the keys of the rows that were too large to be read. */
   @Nonnull
   public List<ByteString> getLargeRowKeys() {

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

@@ -451,11 +451,10 @@ public <ReqT, RowT> ServerStreamingCallable<Query, RowT> createSkipLargeRowsCall
     // ReadRowsRequest -> ReadRowsResponse callable).
     // We override the resumption strategy to use LargeReadRowsResumptionStrategy here (which skips
     // the large rows) instead of ReadRowResumptionStrategy
-    LargeReadRowsResumptionStrategy<RowT> strategy =
-        new LargeReadRowsResumptionStrategy<>(rowAdapter, failOnLargeRows);
     ServerStreamingCallSettings<ReadRowsRequest, RowT> innerSettings =
         ServerStreamingCallSettings.<ReadRowsRequest, RowT>newBuilder()
-            .setResumptionStrategy(strategy)
+            .setResumptionStrategy(
+                new LargeReadRowsResumptionStrategy<>(rowAdapter, failOnLargeRows))
             .setRetryableCodes(readRowsSettings.getRetryableCodes())
             .setRetrySettings(readRowsSettings.getRetrySettings())
             .setIdleTimeout(readRowsSettings.getIdleTimeout())

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

@@ -683,6 +683,7 @@ private Builder(EnhancedBigtableStubSettings settings) {
       jwtAudience = settings.jwtAudience;
       this.directPathConfig = settings.getDirectPathConfig();
       sessionsEnabled = settings.sessionsEnabled;
+      failOnLargeRows = settings.failOnLargeRows;
 
       this.perOpSettings = new ClientOperationSettings.Builder(settings.perOpSettings);
 
@@ -1072,6 +1073,7 @@ public String toString() {
         .add("jwtAudience", jwtAudience)
         .add("directPathConfig", getDirectPathConfig().toString())
         .add("sessionsEnabled", sessionsEnabled)
+        .add("failOnLargeRows", failOnLargeRows)
         .add("parent", super.toString())
         .toString();
   }

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

@@ -60,24 +60,12 @@ public class LargeReadRowsResumptionStrategy<RowT>
   // row-keys
   private RowSet previousFailedRequestRowset = null;
 
-  private final List<ByteString> encounteredLargeRowKeys;
+  private final List<ByteString> largeRowKeys = new ArrayList<>();
   private final boolean failOnLargeRows;
 
   public LargeReadRowsResumptionStrategy(RowAdapter<RowT> rowAdapter, boolean failOnLargeRows) {
-    this(rowAdapter, failOnLargeRows, new ArrayList<ByteString>());
-  }
-
-  private LargeReadRowsResumptionStrategy(
-      RowAdapter<RowT> rowAdapter,
-      boolean failOnLargeRows,
-      List<ByteString> encounteredLargeRowKeys) {
     this.rowAdapter = rowAdapter;
     this.failOnLargeRows = failOnLargeRows;
-    this.encounteredLargeRowKeys = encounteredLargeRowKeys;
-  }
-
-  public List<ByteString> getEncounteredLargeRowKeys() {
-    return encounteredLargeRowKeys;
   }
 
   @Override
@@ -87,8 +75,7 @@ public boolean canResume() {
 
   @Override
   public StreamResumptionStrategy<ReadRowsRequest, RowT> createNew() {
-    return new LargeReadRowsResumptionStrategy<>(
-        rowAdapter, failOnLargeRows, encounteredLargeRowKeys);
+    return new LargeReadRowsResumptionStrategy<>(rowAdapter, failOnLargeRows);
   }
 
   @Override
@@ -113,24 +100,26 @@ public Throwable processError(Throwable throwable) {
       this.largeRowKey = rowKeyExtracted;
       numProcessed = numProcessed + 1;
       if (failOnLargeRows) {
-        encounteredLargeRowKeys.add(rowKeyExtracted);
+        largeRowKeys.add(rowKeyExtracted);
       } else {
         LOGGER.warning("skipping large row " + rowKeyExtracted);
       }
     }
     return throwable;
   }
 
+  public List<ByteString> getLargeRowKeys() {
+    return largeRowKeys;
+  }
+
   private ByteString extractLargeRowKey(Throwable t) {
-    if (t instanceof ApiException) {
-      ApiException ae = (ApiException) t;
-      if (ae.getReason() != null && ae.getReason().equals("LargeRowReadError")) {
-        String rowKey = ae.getMetadata().get("rowKeyBase64Encoded");
-        if (rowKey != null) {
-          byte[] decodedBytes = Base64.getDecoder().decode(rowKey);
-          return ByteString.copyFrom(decodedBytes);
-        }
-      }
+    if (t instanceof ApiException
+        && ((ApiException) t).getReason() != null
+        && ((ApiException) t).getReason().equals("LargeRowReadError")) {
+      String rowKey = ((ApiException) t).getMetadata().get("rowKeyBase64Encoded");
+
+      byte[] decodedBytes = Base64.getDecoder().decode(rowKey);
+      return ByteString.copyFrom(decodedBytes);
     }
     return null;
   }

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/LargeRowServerStreamingAttemptCallable.java

@@ -20,44 +20,34 @@
 import com.google.api.core.ApiFutureCallback;
 import com.google.api.core.ApiFutures;
 import com.google.api.core.InternalApi;
-import com.google.api.gax.grpc.GrpcStatusCode;
 import com.google.api.gax.retrying.RetryingFuture;
 import com.google.api.gax.retrying.ScheduledRetryingExecutor;
 import com.google.api.gax.retrying.ServerStreamingAttemptException;
 import com.google.api.gax.rpc.ApiCallContext;
-import com.google.api.gax.rpc.ApiException;
 import com.google.api.gax.rpc.ResponseObserver;
 import com.google.api.gax.rpc.ServerStreamingCallable;
-import com.google.api.gax.rpc.StatusCode;
 import com.google.bigtable.v2.ReadRowsRequest;
 import com.google.cloud.bigtable.data.v2.models.LargeRowException;
 import com.google.cloud.bigtable.data.v2.stub.readrows.LargeReadRowsResumptionStrategy;
 import com.google.protobuf.ByteString;
-import io.grpc.Status;
-import io.grpc.StatusRuntimeException;
 import java.util.List;
 
-/**
- * A ServerStreamingCallable that implements resumable retries for large rows.
- *
- * <p>This class is a specialized version of {@link RetryingServerStreamingCallable} that uses
- * {@link LargeRowServerStreamingAttemptCallable} and {@link LargeReadRowsResumptionStrategy}.
- */
+/** A ServerStreamingCallable that throws large row keys at the end of the stream. */
 @InternalApi
 public final class RetryingLargeRowServerStreamingCallable<ResponseT>
     extends ServerStreamingCallable<ReadRowsRequest, ResponseT> {
 
   private final ServerStreamingCallable<ReadRowsRequest, ResponseT> innerCallable;
   private final ScheduledRetryingExecutor<Void> executor;
-  private final LargeReadRowsResumptionStrategy<ResponseT> resumptionStrategyPrototype;
+  private final LargeReadRowsResumptionStrategy<ResponseT> resumptionStrategy;
 
   public RetryingLargeRowServerStreamingCallable(
       ServerStreamingCallable<ReadRowsRequest, ResponseT> innerCallable,
       ScheduledRetryingExecutor<Void> executor,
-      LargeReadRowsResumptionStrategy<ResponseT> resumptionStrategyPrototype) {
+      LargeReadRowsResumptionStrategy<ResponseT> resumptionStrategy) {
     this.innerCallable = innerCallable;
     this.executor = executor;
-    this.resumptionStrategyPrototype = resumptionStrategyPrototype;
+    this.resumptionStrategy = resumptionStrategy;
   }
 
   @Override
@@ -66,13 +56,12 @@ public void call(
       final ResponseObserver<ResponseT> responseObserver,
       ApiCallContext context) {
 
-    LargeRowServerStreamingAttemptCallable<ResponseT> attemptCallable =
-        new LargeRowServerStreamingAttemptCallable<>(
-            innerCallable,
-            (LargeReadRowsResumptionStrategy<ResponseT>) resumptionStrategyPrototype.createNew(),
-            request,
-            context,
-            responseObserver);
+    final LargeReadRowsResumptionStrategy<ResponseT> strategy =
+        (LargeReadRowsResumptionStrategy<ResponseT>) resumptionStrategy.createNew();
+
+    ServerStreamingAttemptCallable<ReadRowsRequest, ResponseT> attemptCallable =
+        new ServerStreamingAttemptCallable<>(
+            innerCallable, strategy, request, context, responseObserver);
 
     RetryingFuture<Void> retryingFuture = executor.createFuture(attemptCallable, context);
     attemptCallable.setExternalFuture(retryingFuture);
@@ -88,35 +77,18 @@ public void onFailure(Throwable throwable) {
             if (throwable instanceof ServerStreamingAttemptException) {
               throwable = throwable.getCause();
             }
-            List<ByteString> encounteredKeys =
-                resumptionStrategyPrototype.getEncounteredLargeRowKeys();
+            List<ByteString> encounteredKeys = strategy.getLargeRowKeys();
             if (!encounteredKeys.isEmpty()) {
-              StatusCode statusCode = GrpcStatusCode.of(Status.Code.FAILED_PRECONDITION);
-              boolean isRetryable = false;
-              if (throwable instanceof ApiException) {
-                statusCode = ((ApiException) throwable).getStatusCode();
-                isRetryable = ((ApiException) throwable).isRetryable();
-              } else if (throwable instanceof StatusRuntimeException) {
-                statusCode =
-                    GrpcStatusCode.of(((StatusRuntimeException) throwable).getStatus().getCode());
-              }
-              throwable =
-                  new LargeRowException(throwable, statusCode, isRetryable, encounteredKeys);
+              throwable.addSuppressed(new LargeRowException(encounteredKeys));
             }
             responseObserver.onError(throwable);
           }
 
           @Override
           public void onSuccess(Void ignored) {
-            List<ByteString> encounteredKeys =
-                resumptionStrategyPrototype.getEncounteredLargeRowKeys();
+            List<ByteString> encounteredKeys = strategy.getLargeRowKeys();
             if (!encounteredKeys.isEmpty()) {
-              responseObserver.onError(
-                  new LargeRowException(
-                      null,
-                      GrpcStatusCode.of(Status.Code.FAILED_PRECONDITION),
-                      false,
-                      encounteredKeys));
+              responseObserver.onError(new LargeRowException(encounteredKeys));
             } else {
               responseObserver.onComplete();
             }