feat: Session support for mutateRows

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/AuthorizedViewAsync.java

@@ -22,13 +22,16 @@
 import com.google.bigtable.v2.SessionCheckAndMutateRowResponse;
 import com.google.bigtable.v2.SessionMutateRowRequest;
 import com.google.bigtable.v2.SessionMutateRowResponse;
+import com.google.bigtable.v2.SessionMutateRowsRequest;
+import com.google.bigtable.v2.SessionMutateRowsResponse;
 import com.google.bigtable.v2.SessionReadModifyWriteRowRequest;
 import com.google.bigtable.v2.SessionReadModifyWriteRowResponse;
 import com.google.bigtable.v2.SessionReadRowRequest;
 import com.google.bigtable.v2.SessionReadRowResponse;
 import com.google.cloud.bigtable.data.v2.internal.channels.ChannelPool;
 import com.google.cloud.bigtable.data.v2.internal.csm.Metrics;
 import com.google.cloud.bigtable.data.v2.internal.csm.attributes.ClientInfo;
+import com.google.cloud.bigtable.data.v2.internal.middleware.SessionMutateRowsRateLimiter;
 import com.google.cloud.bigtable.data.v2.internal.session.BigtableTimer;
 import com.google.cloud.bigtable.data.v2.internal.session.SessionPool;
 import com.google.cloud.bigtable.data.v2.internal.session.VRpcDescriptor;
@@ -38,6 +41,7 @@
 import java.io.Closeable;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
+import javax.annotation.Nullable;
 
 public class AuthorizedViewAsync implements AutoCloseable, Closeable {
 
@@ -56,6 +60,36 @@ static AuthorizedViewAsync createAndStart(
       BigtableTimer timer,
       Executor backgroundExecutor,
       Executor userCallbackExecutor) {
+    return createAndStart(
+        featureFlags,
+        clientInfo,
+        configManager,
+        channelPool,
+        callOptions,
+        tableId,
+        viewId,
+        permission,
+        metrics,
+        timer,
+        backgroundExecutor,
+        userCallbackExecutor,
+        /* mutateRowsRateLimiter= */ null);
+  }
+
+  static AuthorizedViewAsync createAndStart(
+      FeatureFlags featureFlags,
+      ClientInfo clientInfo,
+      ClientConfigurationManager configManager,
+      ChannelPool channelPool,
+      CallOptions callOptions,
+      String tableId,
+      String viewId,
+      Permission permission,
+      Metrics metrics,
+      BigtableTimer timer,
+      Executor backgroundExecutor,
+      Executor userCallbackExecutor,
+      @Nullable SessionMutateRowsRateLimiter mutateRowsRateLimiter) {
 
     AuthorizedViewName viewName =
         AuthorizedViewName.builder()
@@ -80,6 +114,7 @@ static AuthorizedViewAsync createAndStart(
             VRpcDescriptor.MUTATE_ROW_AUTH_VIEW,
             VRpcDescriptor.CHECK_AND_MUTATE_ROW_AUTH_VIEW,
             VRpcDescriptor.READ_MODIFY_WRITE_ROW_AUTH_VIEW,
+            VRpcDescriptor.MUTATE_ROWS_AUTH_VIEW,
             featureFlags,
             clientInfo,
             configManager,
@@ -89,7 +124,8 @@ static AuthorizedViewAsync createAndStart(
             metrics,
             timer,
             backgroundExecutor,
-            userCallbackExecutor);
+            userCallbackExecutor,
+            mutateRowsRateLimiter);
 
     return new AuthorizedViewAsync(base);
   }
@@ -130,6 +166,13 @@ public CompletableFuture<SessionReadModifyWriteRowResponse> readModifyWriteRow(
     return f;
   }
 
+  public CompletableFuture<SessionMutateRowsResponse> mutateRows(
+      SessionMutateRowsRequest req, Deadline deadline) {
+    UnaryResponseFuture<SessionMutateRowsResponse> f = new UnaryResponseFuture<>();
+    base.mutateRows(req, f, deadline);
+    return f;
+  }
+
   @Override
   public void close() {
     this.base.close();

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java

@@ -33,10 +33,12 @@
 import com.google.cloud.bigtable.data.v2.internal.csm.MetricsImpl;
 import com.google.cloud.bigtable.data.v2.internal.csm.NoopMetrics;
 import com.google.cloud.bigtable.data.v2.internal.csm.attributes.ClientInfo;
+import com.google.cloud.bigtable.data.v2.internal.middleware.SessionMutateRowsRateLimiter;
 import com.google.cloud.bigtable.data.v2.internal.session.BigtableTimer;
 import com.google.cloud.bigtable.data.v2.internal.session.HashedWheelTimer;
 import com.google.cloud.bigtable.data.v2.internal.session.SessionPool;
 import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import io.grpc.CallOptions;
 import io.opencensus.stats.Stats;
@@ -56,6 +58,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Level;
 import java.util.logging.Logger;
+import javax.annotation.Nullable;
 
 public class Client implements AutoCloseable {
   private static final Logger logger = Logger.getLogger(Client.class.getName());
@@ -106,16 +109,30 @@ public class Client implements AutoCloseable {
   private final Resource<ChannelPool> channelPool;
   private final Resource<Metrics> metrics;
   private final Resource<ClientConfigurationManager> configManager;
+  @Nullable private final SessionMutateRowsRateLimiter mutateRowsRateLimiter;
 
   private final Set<SessionPool<?>> sessionPools = Collections.newSetFromMap(new WeakHashMap<>());
   // Guarded by sessionPools' monitor: close() sets it before snapshotting the pool set, and the
   // open* methods check it before adding a new pool, so a racing open cannot insert a pool that
   // close() has already missed in its snapshot.
   private boolean closed = false;
 
+  @VisibleForTesting
+  static FeatureFlags computeBaseFeatureFlags(ClientSettings settings) {
+    FeatureFlags baseFeatureFlags = BASE_FEATURE_FLAGS;
+    if (settings.isServerInitiatedFlowControlEnabled()) {
+      baseFeatureFlags =
+          baseFeatureFlags.toBuilder()
+              .setMutateRowsRateLimit(true)
+              .setMutateRowsRateLimit2(true)
+              .build();
+    }
+    return baseFeatureFlags;
+  }
+
   public static Client create(ClientSettings settings) throws IOException {
     FeatureFlags featureFlags =
-        settings.getChannelProvider().updateFeatureFlags(BASE_FEATURE_FLAGS);
+        settings.getChannelProvider().updateFeatureFlags(computeBaseFeatureFlags(settings));
 
     ClientInfo clientInfo =
         ClientInfo.builder()
@@ -216,6 +233,8 @@ public Client(
     this.backgroundExecutor = bgExecutor;
     this.userCallbackExecutor = userCallbackExecutor;
     this.sessionTimer = new HashedWheelTimer("bigtable-session-timer");
+    this.mutateRowsRateLimiter =
+        this.featureFlags.getMutateRowsRateLimit() ? new SessionMutateRowsRateLimiter() : null;
 
     defaultCallOptions = CallOptions.DEFAULT;
 
@@ -262,6 +281,8 @@ public Client(
     this.userCallbackExecutor = userCallbackExecutor;
     this.channelPool = sharedChannelPool;
     this.sessionTimer = new HashedWheelTimer("bigtable-session-timer");
+    this.mutateRowsRateLimiter =
+        this.featureFlags.getMutateRowsRateLimit() ? new SessionMutateRowsRateLimiter() : null;
     defaultCallOptions = CallOptions.DEFAULT;
   }
 
@@ -368,7 +389,8 @@ public TableAsync openTableAsync(String tableId, Permission permission) {
               metrics.get(),
               sessionTimer,
               backgroundExecutor.get(),
-              userCallbackExecutor.get());
+              userCallbackExecutor.get(),
+              mutateRowsRateLimiter);
       sessionPools.add(tableAsync.getSessionPool());
       return tableAsync;
     }
@@ -393,7 +415,8 @@ public AuthorizedViewAsync openAuthorizedViewAsync(
               metrics.get(),
               sessionTimer,
               backgroundExecutor.get(),
-              userCallbackExecutor.get());
+              userCallbackExecutor.get(),
+              mutateRowsRateLimiter);
       sessionPools.add(viewAsync.getSessionPool());
       return viewAsync;
     }

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/ClientSettings.java

@@ -32,8 +32,10 @@ public abstract class ClientSettings {
 
   public abstract String getAppProfileId();
 
+  public abstract boolean isServerInitiatedFlowControlEnabled();
+
   public static Builder builder() {
-    return new AutoValue_ClientSettings.Builder();
+    return new AutoValue_ClientSettings.Builder().setServerInitiatedFlowControlEnabled(false);
   }
 
   @AutoValue.Builder
@@ -47,6 +49,8 @@ public abstract static class Builder {
 
     public abstract Builder setAppProfileId(String appProfileId);
 
+    public abstract Builder setServerInitiatedFlowControlEnabled(boolean enabled);
+
     abstract ClientSettings autoBuild(); // not public
 
     public ClientSettings build() {

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/MaterializedViewAsync.java

@@ -71,6 +71,7 @@ public static MaterializedViewAsync createAndStart(
             /* mutateRowDescriptor= */ null, // Materialized views are read-only
             /* checkAndMutateRowDescriptor= */ null,
             /* readModifyWriteRowDescriptor= */ null,
+            /* mutateRowsDescriptor= */ null,
             featureFlags,
             clientInfo,
             configManager,

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableAsync.java

@@ -23,13 +23,16 @@
 import com.google.bigtable.v2.SessionCheckAndMutateRowResponse;
 import com.google.bigtable.v2.SessionMutateRowRequest;
 import com.google.bigtable.v2.SessionMutateRowResponse;
+import com.google.bigtable.v2.SessionMutateRowsRequest;
+import com.google.bigtable.v2.SessionMutateRowsResponse;
 import com.google.bigtable.v2.SessionReadModifyWriteRowRequest;
 import com.google.bigtable.v2.SessionReadModifyWriteRowResponse;
 import com.google.bigtable.v2.SessionReadRowRequest;
 import com.google.bigtable.v2.SessionReadRowResponse;
 import com.google.cloud.bigtable.data.v2.internal.channels.ChannelPool;
 import com.google.cloud.bigtable.data.v2.internal.csm.Metrics;
 import com.google.cloud.bigtable.data.v2.internal.csm.attributes.ClientInfo;
+import com.google.cloud.bigtable.data.v2.internal.middleware.SessionMutateRowsRateLimiter;
 import com.google.cloud.bigtable.data.v2.internal.session.BigtableTimer;
 import com.google.cloud.bigtable.data.v2.internal.session.SessionPool;
 import com.google.cloud.bigtable.data.v2.internal.session.VRpcDescriptor;
@@ -39,6 +42,7 @@
 import java.io.Closeable;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
+import javax.annotation.Nullable;
 
 public class TableAsync implements AutoCloseable, Closeable {
   private final TableBase base;
@@ -55,6 +59,34 @@ public static TableAsync createAndStart(
       BigtableTimer timer,
       Executor backgroundExecutor,
       Executor userCallbackExecutor) {
+    return createAndStart(
+        featureFlags,
+        clientInfo,
+        configManager,
+        channelPool,
+        callOptions,
+        tableId,
+        permission,
+        metrics,
+        timer,
+        backgroundExecutor,
+        userCallbackExecutor,
+        /* mutateRowsRateLimiter= */ null);
+  }
+
+  public static TableAsync createAndStart(
+      FeatureFlags featureFlags,
+      ClientInfo clientInfo,
+      ClientConfigurationManager configManager,
+      ChannelPool channelPool,
+      CallOptions callOptions,
+      String tableId,
+      Permission permission,
+      Metrics metrics,
+      BigtableTimer timer,
+      Executor backgroundExecutor,
+      Executor userCallbackExecutor,
+      @Nullable SessionMutateRowsRateLimiter mutateRowsRateLimiter) {
 
     TableName tableName =
         TableName.builder()
@@ -78,6 +110,7 @@ public static TableAsync createAndStart(
             VRpcDescriptor.MUTATE_ROW,
             VRpcDescriptor.CHECK_AND_MUTATE_ROW,
             VRpcDescriptor.READ_MODIFY_WRITE_ROW,
+            VRpcDescriptor.MUTATE_ROWS,
             featureFlags,
             clientInfo,
             configManager,
@@ -87,7 +120,8 @@ public static TableAsync createAndStart(
             metrics,
             timer,
             backgroundExecutor,
-            userCallbackExecutor);
+            userCallbackExecutor,
+            mutateRowsRateLimiter);
 
     return new TableAsync(base);
   }
@@ -139,4 +173,11 @@ public CompletableFuture<SessionReadModifyWriteRowResponse> readModifyWriteRow(
     base.readModifyWriteRow(req, f, deadline);
     return f;
   }
+
+  public CompletableFuture<SessionMutateRowsResponse> mutateRows(
+      SessionMutateRowsRequest req, Deadline deadline) {
+    UnaryResponseFuture<SessionMutateRowsResponse> f = new UnaryResponseFuture<>();
+    base.mutateRows(req, f, deadline);
+    return f;
+  }
 }