diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/AuthorizedViewAsync.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/AuthorizedViewAsync.java index 0f2cf71a1a9d..72997953cc0e 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/AuthorizedViewAsync.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/AuthorizedViewAsync.java @@ -22,6 +22,8 @@ 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; @@ -29,6 +31,7 @@ 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 readModifyWriteRow( return f; } + public CompletableFuture mutateRows( + SessionMutateRowsRequest req, Deadline deadline) { + UnaryResponseFuture f = new UnaryResponseFuture<>(); + base.mutateRows(req, f, deadline); + return f; + } + @Override public void close() { this.base.close(); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java index b6f72659c092..f984f1da5ac3 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java +++ b/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,6 +109,7 @@ public class Client implements AutoCloseable { private final Resource channelPool; private final Resource metrics; private final Resource configManager; + @Nullable private final SessionMutateRowsRateLimiter mutateRowsRateLimiter; private final Set> sessionPools = Collections.newSetFromMap(new WeakHashMap<>()); // Guarded by sessionPools' monitor: close() sets it before snapshotting the pool set, and the @@ -113,9 +117,22 @@ public class Client implements AutoCloseable { // 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; } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/ClientSettings.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/ClientSettings.java index 6b4448945eef..427805cce5f0 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/ClientSettings.java +++ b/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() { diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/MaterializedViewAsync.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/MaterializedViewAsync.java index 856e424118f5..a8c7004fab89 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/MaterializedViewAsync.java +++ b/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, diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableAsync.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableAsync.java index 9d6a5b515ffd..d44848f02ee0 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableAsync.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableAsync.java @@ -23,6 +23,8 @@ 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; @@ -30,6 +32,7 @@ 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 readModifyWriteRow( base.readModifyWriteRow(req, f, deadline); return f; } + + public CompletableFuture mutateRows( + SessionMutateRowsRequest req, Deadline deadline) { + UnaryResponseFuture f = new UnaryResponseFuture<>(); + base.mutateRows(req, f, deadline); + return f; + } } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableBase.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableBase.java index f82861979ba6..91395e38376a 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableBase.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableBase.java @@ -21,6 +21,8 @@ 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; @@ -31,8 +33,12 @@ 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.csm.tracers.VRpcTracer; +import com.google.cloud.bigtable.data.v2.internal.middleware.MutateRowsRetryingVRpc; +import com.google.cloud.bigtable.data.v2.internal.middleware.RateLimitingVRpc; import com.google.cloud.bigtable.data.v2.internal.middleware.RetryingVRpc; +import com.google.cloud.bigtable.data.v2.internal.middleware.SessionMutateRowsRateLimiter; import com.google.cloud.bigtable.data.v2.internal.middleware.VOperationImpl; +import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc; import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc.VRpcListener; import com.google.cloud.bigtable.data.v2.internal.session.BigtableTimer; import com.google.cloud.bigtable.data.v2.internal.session.SessionPool; @@ -46,12 +52,15 @@ import io.grpc.Deadline; import io.grpc.Metadata; import java.util.concurrent.Executor; +import java.util.function.Supplier; +import javax.annotation.Nullable; class TableBase implements AutoCloseable { private final SessionPool sessionPool; private final BigtableTimer timer; private final Executor userCallbackExecutor; private final Metrics metrics; + @Nullable private final SessionMutateRowsRateLimiter mutateRowsRateLimiter; private final VRpcDescriptor readRowDescriptor; private final VRpcDescriptor mutateRowDescriptor; @@ -60,6 +69,8 @@ class TableBase implements AutoCloseable { private final VRpcDescriptor< ?, SessionReadModifyWriteRowRequest, SessionReadModifyWriteRowResponse> readModifyWriteRowDescriptor; + private final VRpcDescriptor + mutateRowsDescriptor; static TableBase createAndStart( ReqT openReq, @@ -70,6 +81,7 @@ static TableBase createAndStart( checkAndMutateRowDescriptor, VRpcDescriptor readModifyWriteRowDescriptor, + VRpcDescriptor mutateRowsDescriptor, FeatureFlags featureFlags, ClientInfo clientInfo, ClientConfigurationManager configManager, @@ -80,6 +92,48 @@ static TableBase createAndStart( BigtableTimer timer, Executor backgroundExecutor, Executor userCallbackExecutor) { + return createAndStart( + openReq, + sessionDescriptor, + readRowDescriptor, + mutateRowDescriptor, + checkAndMutateRowDescriptor, + readModifyWriteRowDescriptor, + mutateRowsDescriptor, + featureFlags, + clientInfo, + configManager, + channelPool, + callOptions, + sessionPoolName, + metrics, + timer, + backgroundExecutor, + userCallbackExecutor, + /* mutateRowsRateLimiter= */ null); + } + + static TableBase createAndStart( + ReqT openReq, + VRpcDescriptor.SessionDescriptor sessionDescriptor, + VRpcDescriptor readRowDescriptor, + VRpcDescriptor mutateRowDescriptor, + VRpcDescriptor + checkAndMutateRowDescriptor, + VRpcDescriptor + readModifyWriteRowDescriptor, + VRpcDescriptor mutateRowsDescriptor, + FeatureFlags featureFlags, + ClientInfo clientInfo, + ClientConfigurationManager configManager, + ChannelPool channelPool, + CallOptions callOptions, + String sessionPoolName, + Metrics metrics, + BigtableTimer timer, + Executor backgroundExecutor, + Executor userCallbackExecutor, + @Nullable SessionMutateRowsRateLimiter mutateRowsRateLimiter) { // Stamp the tenant key so ChannelPoolDpImpl can make tenant-aware placement decisions. CallOptions stamped = @@ -108,9 +162,11 @@ static TableBase createAndStart( mutateRowDescriptor, checkAndMutateRowDescriptor, readModifyWriteRowDescriptor, + mutateRowsDescriptor, metrics, timer, - userCallbackExecutor); + userCallbackExecutor, + mutateRowsRateLimiter); } @VisibleForTesting @@ -122,17 +178,47 @@ static TableBase createAndStart( checkAndMutateRowDescriptor, VRpcDescriptor readModifyWriteRowDescriptor, + VRpcDescriptor mutateRowsDescriptor, Metrics metrics, BigtableTimer timer, Executor userCallbackExecutor) { + this( + sessionPool, + readRowDescriptor, + mutateRowDescriptor, + checkAndMutateRowDescriptor, + readModifyWriteRowDescriptor, + mutateRowsDescriptor, + metrics, + timer, + userCallbackExecutor, + /* mutateRowsRateLimiter= */ null); + } + + @VisibleForTesting + TableBase( + SessionPool sessionPool, + VRpcDescriptor readRowDescriptor, + VRpcDescriptor mutateRowDescriptor, + VRpcDescriptor + checkAndMutateRowDescriptor, + VRpcDescriptor + readModifyWriteRowDescriptor, + VRpcDescriptor mutateRowsDescriptor, + Metrics metrics, + BigtableTimer timer, + Executor userCallbackExecutor, + @Nullable SessionMutateRowsRateLimiter mutateRowsRateLimiter) { this.sessionPool = sessionPool; this.readRowDescriptor = readRowDescriptor; this.mutateRowDescriptor = mutateRowDescriptor; this.checkAndMutateRowDescriptor = checkAndMutateRowDescriptor; this.readModifyWriteRowDescriptor = readModifyWriteRowDescriptor; + this.mutateRowsDescriptor = mutateRowsDescriptor; this.metrics = metrics; this.timer = timer; this.userCallbackExecutor = userCallbackExecutor; + this.mutateRowsRateLimiter = mutateRowsRateLimiter; } @Override @@ -202,4 +288,26 @@ public void readModifyWriteRow( new VOperationImpl<>(retry, Context.current(), userCallbackExecutor, tracer, deadline, false) .start(req, listener); } + + public void mutateRows( + SessionMutateRowsRequest req, + VRpcListener listener, + Deadline deadline) { + Supplier> attemptSupplier = + mutateRowsRateLimiter == null + ? () -> sessionPool.newCall(mutateRowsDescriptor) + : () -> + new RateLimitingVRpc( + sessionPool.newCall(mutateRowsDescriptor), mutateRowsRateLimiter, timer); + MutateRowsRetryingVRpc retry = new MutateRowsRetryingVRpc(attemptSupplier, timer); + boolean idempotent = + req.getEntriesList().stream() + .allMatch(entry -> Util.isIdempotent(entry.getMutationsList())); + VRpcTracer tracer = + metrics.newTableTracer(sessionPool.getInfo(), mutateRowsDescriptor, deadline); + + new VOperationImpl<>( + retry, Context.current(), userCallbackExecutor, tracer, deadline, idempotent) + .start(req, listener); + } } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/DisabledShim.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/DisabledShim.java index b95064df164b..bbaed487c0ff 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/DisabledShim.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/DisabledShim.java @@ -17,11 +17,13 @@ import com.google.api.gax.rpc.UnaryCallSettings; import com.google.api.gax.rpc.UnaryCallable; +import com.google.cloud.bigtable.data.v2.models.BulkMutation; import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation; import com.google.cloud.bigtable.data.v2.models.Query; import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; 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.MutateRowsAttemptResult; public class DisabledShim implements Shim { @@ -55,4 +57,11 @@ public UnaryCallable decorateReadModifyWriteRow UnaryCallSettings settings) { return classic; } + + @Override + public UnaryCallable decorateMutateRows( + UnaryCallable classic, + UnaryCallSettings settings) { + return classic; + } } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/Shim.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/Shim.java index 2fe6aade5cd0..474369bd6085 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/Shim.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/Shim.java @@ -17,11 +17,13 @@ import com.google.api.gax.rpc.UnaryCallSettings; import com.google.api.gax.rpc.UnaryCallable; +import com.google.cloud.bigtable.data.v2.models.BulkMutation; import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation; import com.google.cloud.bigtable.data.v2.models.Query; import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; 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.MutateRowsAttemptResult; public interface Shim { @@ -42,4 +44,8 @@ UnaryCallable decorateReadModifyWriteRow( UnaryCallable classic, RowAdapter rowAdapter, UnaryCallSettings settings); + + UnaryCallable decorateMutateRows( + UnaryCallable classic, + UnaryCallSettings settings); } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java index 8db91d0a9308..91a880a5fab6 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java @@ -40,6 +40,7 @@ import com.google.cloud.bigtable.data.v2.internal.compat.ops.CheckAndMutateRowShim; import com.google.cloud.bigtable.data.v2.internal.compat.ops.DivertingUnaryCallable; import com.google.cloud.bigtable.data.v2.internal.compat.ops.MutateRowShim; +import com.google.cloud.bigtable.data.v2.internal.compat.ops.MutateRowsShim; import com.google.cloud.bigtable.data.v2.internal.compat.ops.ReadModifyWriteRowShim; import com.google.cloud.bigtable.data.v2.internal.compat.ops.ReadRowShim; import com.google.cloud.bigtable.data.v2.internal.compat.ops.RowBuilderShim; @@ -51,12 +52,15 @@ import com.google.cloud.bigtable.data.v2.internal.dp.DirectAccessInvestigator; import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager; import com.google.cloud.bigtable.data.v2.internal.util.SessionScope; +import com.google.cloud.bigtable.data.v2.models.BulkMutation; import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation; import com.google.cloud.bigtable.data.v2.models.Query; import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; 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.EnhancedBigtableStubSettings; import com.google.cloud.bigtable.data.v2.stub.MetadataExtractorInterceptor; +import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsAttemptResult; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Strings; import com.google.common.collect.ImmutableMap; @@ -111,6 +115,22 @@ public class ShimImpl implements Shim { private final CheckAndMutateRowShim checkAndMutateRowShim; private final ReadModifyWriteRowShim readModifyWriteRowShim; + @VisibleForTesting + static FeatureFlags computeBaseFeatureFlags(StubSettings stubSettings) { + FeatureFlags baseFeatureFlags = Client.BASE_FEATURE_FLAGS; + if (stubSettings instanceof EnhancedBigtableStubSettings + && ((EnhancedBigtableStubSettings) stubSettings) + .bulkMutateRowsSettings() + .isServerInitiatedFlowControlEnabled()) { + baseFeatureFlags = + baseFeatureFlags.toBuilder() + .setMutateRowsRateLimit(true) + .setMutateRowsRateLimit2(true) + .build(); + } + return baseFeatureFlags; + } + public static Shim create( ClientInfo clientInfo, @Nullable Credentials callCreds, @@ -147,6 +167,8 @@ public static Shim create( transportProvider = transportProvider.withMtlsEndpoint(stubSettings.getMtlsEndpoint()); } + FeatureFlags baseFeatureFlags = computeBaseFeatureFlags(stubSettings); + // Channel provider for the client. DirectAccess with fallback by default. ChannelProvider clientChannelProvider = configureChannelProvider( @@ -155,7 +177,7 @@ public static Shim create( metrics); if (!checkDirectAccessAvailable( - Client.BASE_FEATURE_FLAGS, clientInfo, clientChannelProvider, metrics, bgExecutor)) { + baseFeatureFlags, clientInfo, clientChannelProvider, metrics, bgExecutor)) { // Directpath failed, revert to basic channel provider for the client. clientChannelProvider = configureChannelProvider( @@ -173,7 +195,7 @@ public static Shim create( ClientConfigurationManager configManager = new ClientConfigurationManager( - configChannelProvider.updateFeatureFlags(Client.BASE_FEATURE_FLAGS), + configChannelProvider.updateFeatureFlags(baseFeatureFlags), clientInfo, configChannelProvider, metrics.getDebugTagTracer(), @@ -188,7 +210,7 @@ public static Shim create( return new DisabledShim(); } - FeatureFlags featureFlags = Client.BASE_FEATURE_FLAGS; + FeatureFlags featureFlags = baseFeatureFlags; if (configManager.areSessionsRequired()) { featureFlags = featureFlags.toBuilder().setSessionsRequired(true).build(); @@ -435,4 +457,17 @@ public UnaryCallable decorateReadModifyWriteRow debugTagTracer, SessionScope.READ_MODIFY_WRITE_ROW); } + + @Override + public UnaryCallable decorateMutateRows( + UnaryCallable classic, + UnaryCallSettings settings) { + return new DivertingUnaryCallable<>( + configManager, + classic, + new MutateRowsShim(sessionMap, settings.getRetryableCodes()), + Util.extractTimeout(settings), + debugTagTracer, + SessionScope.MUTATE_ROWS); + } } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallable.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallable.java index 6500add3dee2..90accab465c1 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallable.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallable.java @@ -144,6 +144,10 @@ ApiException translateException(Throwable e) { } } + if (cause instanceof ApiException) { + return (ApiException) cause; + } + Status.Code code = Status.Code.UNKNOWN; if (cause instanceof StatusRuntimeException) { diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowsShim.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowsShim.java new file mode 100644 index 000000000000..b89a908de93c --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowsShim.java @@ -0,0 +1,226 @@ +/* + * Copyright 2026 Google LLC + * + * 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 + * + * https://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.data.v2.internal.compat.ops; + +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.StatusCode; +import com.google.bigtable.v2.SessionMutateRowsRequest; +import com.google.bigtable.v2.SessionMutateRowsResponse; +import com.google.cloud.bigtable.data.v2.internal.api.Util; +import com.google.cloud.bigtable.data.v2.internal.session.SessionPool; +import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId; +import com.google.cloud.bigtable.data.v2.models.BulkMutation; +import com.google.cloud.bigtable.data.v2.models.MutateRowsException; +import com.google.cloud.bigtable.data.v2.models.MutateRowsException.FailedMutation; +import com.google.cloud.bigtable.data.v2.models.TableId; +import com.google.cloud.bigtable.data.v2.models.TargetId; +import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsAttemptResult; +import com.google.common.collect.ImmutableSet; +import com.google.common.primitives.Ints; +import com.google.rpc.Code; +import io.grpc.Deadline; +import io.grpc.Status; +import io.grpc.StatusException; +import io.grpc.StatusRuntimeException; +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ExecutionException; + +public class MutateRowsShim implements UnaryShim { + + private static final Set DEFAULT_RETRYABLE_CODES = + ImmutableSet.of(StatusCode.Code.DEADLINE_EXCEEDED, StatusCode.Code.UNAVAILABLE); + + private final SessionPoolMap sessionMap; + private final Set retryableCodes; + + public MutateRowsShim(SessionPoolMap sessionMap) { + this(sessionMap, DEFAULT_RETRYABLE_CODES); + } + + public MutateRowsShim(SessionPoolMap sessionMap, Set retryableCodes) { + this.sessionMap = sessionMap; + this.retryableCodes = ImmutableSet.copyOf(retryableCodes); + } + + @Override + public void close() {} + + @Override + public boolean supports(BulkMutation request) { + TargetId targetId = request.getTargetId(); + SessionPool pool; + if (targetId instanceof TableId) { + pool = sessionMap.get((TableId) targetId, SessionKey.Permission.WRITE).getSessionPool(); + } else if (targetId instanceof AuthorizedViewId) { + pool = + sessionMap.get((AuthorizedViewId) targetId, SessionKey.Permission.WRITE).getSessionPool(); + } else { + return false; + } + return UnaryShim.shouldRouteToSession(pool); + } + + @Override + public CompletableFuture call(BulkMutation request, Deadline deadline) { + TargetId targetId = request.getTargetId(); + SessionMutateRowsRequest innerReq = request.toSessionProto(); + + CompletableFuture responseFuture; + if (targetId instanceof TableId) { + responseFuture = + sessionMap.apply( + (TableId) targetId, + SessionKey.Permission.WRITE, + t -> t.mutateRows(innerReq, deadline)); + } else if (targetId instanceof AuthorizedViewId) { + responseFuture = + sessionMap.apply( + (AuthorizedViewId) targetId, + SessionKey.Permission.WRITE, + v -> v.mutateRows(innerReq, deadline)); + } else { + CompletableFuture f = new CompletableFuture<>(); + f.completeExceptionally( + new UnsupportedOperationException("Unsupported targetId type: " + targetId)); + return f; + } + + return responseFuture.handle( + (response, error) -> { + if (error != null) { + throw toWholeRpcException(innerReq, error); + } + return toAttemptResult(innerReq, response); + }); + } + + private MutateRowsAttemptResult toAttemptResult( + SessionMutateRowsRequest innerReq, SessionMutateRowsResponse response) { + int entryCount = innerReq.getEntriesCount(); + boolean[] seenIndices = new boolean[entryCount]; + List failedMutations = new ArrayList<>(); + boolean anyRetryable = false; + + for (SessionMutateRowsResponse.Entry entry : response.getEntriesList()) { + int index = Ints.checkedCast(entry.getIndex()); + if (index < 0 || index >= entryCount || seenIndices[index]) { + continue; + } + seenIndices[index] = true; + + if (entry.getStatus().getCode() == Code.OK_VALUE) { + continue; + } + + Status grpcStatus = + Status.fromCodeValue(entry.getStatus().getCode()) + .withDescription(entry.getStatus().getMessage()); + StatusCode gaxStatusCode = GrpcStatusCode.of(grpcStatus.getCode()); + boolean isEntryRetryable = + retryableCodes.contains(gaxStatusCode.getCode()) + && Util.isIdempotent(innerReq.getEntries(index).getMutationsList()); + + ApiException entryError = + ApiExceptionFactory.createException( + grpcStatus.asRuntimeException(), gaxStatusCode, isEntryRetryable); + failedMutations.add(FailedMutation.create(index, entryError)); + if (isEntryRetryable) { + anyRetryable = true; + } + } + + for (int i = 0; i < entryCount; i++) { + if (!seenIndices[i]) { + ApiException missingError = + ApiExceptionFactory.createException( + "Missing entry response for entry " + i, + null, + GrpcStatusCode.of(Status.Code.INTERNAL), + false); + failedMutations.add(FailedMutation.create(i, missingError)); + } + } + + if (failedMutations.isEmpty()) { + return MutateRowsAttemptResult.success(); + } + return MutateRowsAttemptResult.create(failedMutations, anyRetryable); + } + + private RuntimeException toWholeRpcException(SessionMutateRowsRequest innerReq, Throwable error) { + Throwable cause = error; + while (cause instanceof CompletionException || cause instanceof ExecutionException) { + if (cause.getCause() != null) { + cause = cause.getCause(); + } else { + break; + } + } + + if (cause instanceof MutateRowsException) { + return (MutateRowsException) cause; + } + + ApiException rpcApiException; + if (cause instanceof ApiException) { + rpcApiException = (ApiException) cause; + } else { + Status.Code code = Status.Code.UNKNOWN; + if (cause instanceof StatusRuntimeException) { + code = ((StatusRuntimeException) cause).getStatus().getCode(); + } else if (cause instanceof StatusException) { + code = ((StatusException) cause).getStatus().getCode(); + } + StatusCode gaxStatusCode = GrpcStatusCode.of(code); + boolean rpcRetryable = retryableCodes.contains(gaxStatusCode.getCode()); + rpcApiException = + ApiExceptionFactory.createException( + cause.getMessage(), cause, gaxStatusCode, rpcRetryable); + } + + int entryCount = innerReq.getEntriesCount(); + if (entryCount == 0) { + return rpcApiException; + } + + List failedMutations = new ArrayList<>(entryCount); + boolean anyEntryRetryable = false; + for (int i = 0; i < entryCount; i++) { + boolean entryRetryable = + rpcApiException.isRetryable() + && Util.isIdempotent(innerReq.getEntries(i).getMutationsList()); + ApiException entryError = + ApiExceptionFactory.createException( + rpcApiException, + rpcApiException.getStatusCode(), + entryRetryable, + rpcApiException.getErrorDetails()); + failedMutations.add(FailedMutation.create(i, entryError)); + if (entryRetryable) { + anyEntryRetryable = true; + } + } + + return MutateRowsException.create( + rpcApiException, rpcApiException.getStatusCode(), failedMutations, anyEntryRetryable); + } +} diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/NoopMetrics.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/NoopMetrics.java index 4990987ea3bc..b43aedab4683 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/NoopMetrics.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/NoopMetrics.java @@ -109,6 +109,13 @@ public void onResponseReceived() {} @Override public void recordApplicationBlockingLatencies(Duration elapsed) {} + @Override + public void recordBatchWriteFlowControlTargetQps(double targetQps) {} + + @Override + public void recordBatchWriteFlowControlFactor( + double factor, Status.Code statusCode, boolean applied) {} + @Override public void onAttemptFinish(VRpc.VRpcResult result) {} diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/CompositeVRpcTracer.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/CompositeVRpcTracer.java index 473c8c56811c..e402ce75a03f 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/CompositeVRpcTracer.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/CompositeVRpcTracer.java @@ -18,6 +18,7 @@ import com.google.bigtable.v2.PeerInfo; import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc; +import io.grpc.Status; import java.time.Duration; import java.util.List; import java.util.logging.Level; @@ -106,6 +107,37 @@ public void recordApplicationBlockingLatencies(Duration elapsed) { }); } + @Override + public void recordBatchWriteFlowControlTargetQps(double targetQps) { + children.forEach( + c -> { + try { + c.recordBatchWriteFlowControlTargetQps(targetQps); + } catch (Throwable e) { + logger.log( + Level.WARNING, + "Failed to log batch write flow control target QPS on tracer {0}", + c.getClass().getName()); + } + }); + } + + @Override + public void recordBatchWriteFlowControlFactor( + double factor, Status.Code statusCode, boolean applied) { + children.forEach( + c -> { + try { + c.recordBatchWriteFlowControlFactor(factor, statusCode, applied); + } catch (Throwable e) { + logger.log( + Level.WARNING, + "Failed to log batch write flow control factor on tracer {0}", + c.getClass().getName()); + } + }); + } + @Override public void onAttemptFinish(VRpc.VRpcResult result) { children.forEach( diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/UserApiVRpcTracer.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/UserApiVRpcTracer.java index e794036d5aa0..5fc44e056f57 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/UserApiVRpcTracer.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/UserApiVRpcTracer.java @@ -21,6 +21,8 @@ import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc; import com.google.cloud.bigtable.data.v2.internal.session.SessionPoolInfo; import com.google.cloud.bigtable.data.v2.internal.session.VRpcDescriptor; +import com.google.cloud.bigtable.data.v2.stub.metrics.BigtableTracer; +import io.grpc.Status; import java.time.Duration; import java.util.concurrent.atomic.AtomicInteger; @@ -69,6 +71,22 @@ public void onResponseReceived() { @Override public void recordApplicationBlockingLatencies(Duration elapsed) {} + @Override + public void recordBatchWriteFlowControlTargetQps(double targetQps) { + if (delegate instanceof BigtableTracer) { + ((BigtableTracer) delegate).setBatchWriteFlowControlTargetQps(targetQps); + } + } + + @Override + public void recordBatchWriteFlowControlFactor( + double factor, Status.Code statusCode, boolean applied) { + if (delegate instanceof BigtableTracer) { + Throwable error = statusCode == Status.Code.OK ? null : statusCode.toStatus().asException(); + ((BigtableTracer) delegate).addBatchWriteFlowControlFactor(factor, error, applied); + } + } + @Override public void onAttemptFinish(VRpc.VRpcResult result) { if (result.getStatus().isOk()) { diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracer.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracer.java index 924bdd4d6006..f1b92a661dce 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracer.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracer.java @@ -18,6 +18,7 @@ import com.google.bigtable.v2.PeerInfo; import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc.VRpcResult; +import io.grpc.Status; import java.time.Duration; public interface VRpcTracer { @@ -32,6 +33,11 @@ public interface VRpcTracer { void recordApplicationBlockingLatencies(Duration elapsed); + default void recordBatchWriteFlowControlTargetQps(double targetQps) {} + + default void recordBatchWriteFlowControlFactor( + double factor, Status.Code statusCode, boolean applied) {} + void onAttemptFinish(VRpcResult result); void onOperationFinish(VRpcResult result); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerImpl.java index 0f20ae1a0214..72091e429c46 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerImpl.java @@ -26,6 +26,7 @@ import com.google.cloud.bigtable.data.v2.internal.session.SessionPoolInfo; import com.google.common.base.Stopwatch; import io.grpc.Deadline; +import io.grpc.Status; import java.time.Duration; import java.util.Optional; import java.util.concurrent.TimeUnit; @@ -133,6 +134,19 @@ public void recordApplicationBlockingLatencies(Duration elapsed) { applicationBlockingLatency = applicationBlockingLatency.plus(elapsed); } + @Override + public void recordBatchWriteFlowControlTargetQps(double targetQps) { + metricRegistry.batchWriteFlowControlTargetQps.record( + poolInfo.getClientInfo(), methodInfo, targetQps); + } + + @Override + public void recordBatchWriteFlowControlFactor( + double factor, Status.Code statusCode, boolean applied) { + metricRegistry.batchWriteFlowControlFactor.record( + poolInfo.getClientInfo(), statusCode, applied, methodInfo, factor); + } + @Override public void onAttemptFinish(VRpcResult result) { attemptTimer.stop(); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/MutateRowsRetryingVRpc.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/MutateRowsRetryingVRpc.java new file mode 100644 index 000000000000..f013faae0222 --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/MutateRowsRetryingVRpc.java @@ -0,0 +1,612 @@ +/* + * Copyright 2026 Google LLC + * + * 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 + * + * https://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.data.v2.internal.middleware; + +import static io.opentelemetry.context.Context.current; + +import com.google.bigtable.v2.SessionMutateRowsRequest; +import com.google.bigtable.v2.SessionMutateRowsResponse; +import com.google.cloud.bigtable.data.v2.internal.api.Util; +import com.google.cloud.bigtable.data.v2.internal.csm.tracers.VRpcTracer; +import com.google.cloud.bigtable.data.v2.internal.session.BigtableTimer; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Stopwatch; +import com.google.common.collect.ImmutableSet; +import com.google.common.primitives.Ints; +import com.google.protobuf.util.Durations; +import com.google.rpc.Code; +import com.google.rpc.RetryInfo; +import io.grpc.Context; +import io.grpc.Status; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.TimeUnit; +import java.util.function.Supplier; +import java.util.function.UnaryOperator; +import java.util.logging.Level; +import java.util.logging.Logger; +import javax.annotation.Nullable; + +/** + * Retrying {@link VRpc} middleware specialized for {@code SessionMutateRows}. + * + *

Handles both whole-vRPC failures (uncommitted errors, transport failures when all active + * entries are idempotent, and server {@code ErrorResponse}s carrying {@link RetryInfo}) and partial + * entry-level transient failures inside an {@code OK} {@link SessionMutateRowsResponse}. + * + *

On partial failures: + * + *

    + *
  • Entries that succeeded ({@code OK}) or failed with non-retryable statuses are locked in at + * their original request indices and excluded from subsequent attempts. + *
  • Entries that failed with transient statuses are retried only if that individual + * entry's mutations are idempotent ({@link Util#isIdempotent}). Non-idempotent entries (e.g. + * {@code AddToCell}, {@code MergeToCell}, server-assigned timestamps) are never retried after + * reaching the server. + *
  • Any entry index omitted from a server response is recorded as a permanent {@code INTERNAL} + * error. + *
  • If an earlier attempt locked in partial outcomes and a subsequent retry attempt fails at + * the whole-vRPC level, previously resolved entry outcomes are preserved and only the + * remaining unresolved entries receive the final vRPC failure status. + *
+ * + *

All mutable state is owned by the per-operation {@link OpExecutor}. + */ +public class MutateRowsRetryingVRpc + implements VRpc { + + private static final Logger LOG = Logger.getLogger(MutateRowsRetryingVRpc.class.getName()); + + static final Set DEFAULT_RETRYABLE_ENTRY_CODES = + ImmutableSet.of(Code.DEADLINE_EXCEEDED_VALUE, Code.UNAVAILABLE_VALUE); + static final Duration DEFAULT_INITIAL_BACKOFF = Duration.ofMillis(10); + static final double DEFAULT_BACKOFF_MULTIPLIER = 2.0; + static final Duration DEFAULT_MAX_BACKOFF = Duration.ofMinutes(1); + + private final Supplier> attemptFactory; + private final BigtableTimer timer; + private final Set retryableEntryCodes; + private final Duration initialBackoff; + private final double backoffMultiplier; + private final Duration maxBackoff; + private final UnaryOperator contextWrapper; + + private SessionMutateRowsRequest originalRequest; + private SessionMutateRowsRequest currentRequest; + private boolean[] entryIdempotent; + private SessionMutateRowsResponse.Entry[] resolvedEntries; + private int[] activeOriginalIndices; + private boolean hasPartialProgress = false; + @Nullable private VRpcResult lastOkServerResult; + private int partialRetryAttempt = 0; + private int consecutiveNonServerRetries = 0; + + private VRpcListener listener; + private VRpcCallContext context; + private VRpcTracer tracer; + + private State currentState; + private boolean started = false; + private boolean isCancelling = false; + + public MutateRowsRetryingVRpc( + Supplier> supplier, + BigtableTimer timer) { + this( + supplier, + timer, + DEFAULT_RETRYABLE_ENTRY_CODES, + DEFAULT_INITIAL_BACKOFF, + DEFAULT_BACKOFF_MULTIPLIER, + DEFAULT_MAX_BACKOFF); + } + + @VisibleForTesting + public MutateRowsRetryingVRpc( + Supplier> supplier, + BigtableTimer timer, + Set retryableEntryCodes, + Duration initialBackoff, + double backoffMultiplier, + Duration maxBackoff) { + this.attemptFactory = supplier; + this.timer = timer; + this.retryableEntryCodes = ImmutableSet.copyOf(retryableEntryCodes); + this.initialBackoff = initialBackoff; + this.backoffMultiplier = backoffMultiplier; + this.maxBackoff = maxBackoff; + Context grpcContext = Context.current(); + UnaryOperator otelWrap = current()::wrap; + this.contextWrapper = task -> grpcContext.wrap(otelWrap.apply(task)); + this.currentState = new Idle(); + } + + @Override + public void start( + SessionMutateRowsRequest req, + VRpcCallContext ctx, + VRpcListener listener) { + if (started) { + VRpcResult alreadyStarted = + VRpcResult.createRejectedError( + Status.FAILED_PRECONDITION.withDescription("operation is already started")); + ctx.getExecutor().execute(() -> listener.onClose(alreadyStarted)); + return; + } + + int entryCount = req.getEntriesCount(); + this.originalRequest = req; + this.currentRequest = req; + this.entryIdempotent = new boolean[entryCount]; + this.resolvedEntries = new SessionMutateRowsResponse.Entry[entryCount]; + this.activeOriginalIndices = new int[entryCount]; + for (int i = 0; i < entryCount; i++) { + this.entryIdempotent[i] = Util.isIdempotent(req.getEntries(i).getMutationsList()); + this.activeOriginalIndices[i] = i; + } + + this.listener = listener; + this.context = ctx; + this.tracer = ctx.getTracer(); + + tracer.onOperationStart(); + started = true; + + try { + currentState.onStart(); + } catch (Throwable t) { + cancel("Unexpected error in start", t); + } + } + + @Override + public void cancel(@Nullable String message, @Nullable Throwable cause) { + if (currentState.isDone() || isCancelling) { + LOG.fine("Ignoring cancel because the vRPC is already cancelled or done."); + return; + } + isCancelling = true; + Throwable finalCause = cause; + try { + currentState.onCancel(message, cause); + } catch (Throwable t) { + if (finalCause != null) { + finalCause.addSuppressed(t); + } else { + finalCause = t; + } + } + onStateChange( + new Done( + VRpcResult.createRejectedError( + Status.CANCELLED.withDescription(message).withCause(finalCause)))); + } + + @Override + public boolean isDone() { + return currentState.isDone(); + } + + @Override + public void requestNext() { + if (context != null) { + context.getExecutor().throwIfNotInThisExecutor(); + } + throw new UnsupportedOperationException("request next is not supported in unary"); + } + + private void onStateChange(State state) { + if (currentState.isDone()) { + return; + } + currentState.onExit(); + this.currentState = state; + currentState.onStart(); + } + + private boolean areAllActiveEntriesIdempotent() { + for (int origIdx : activeOriginalIndices) { + if (!entryIdempotent[origIdx]) { + return false; + } + } + return true; + } + + private Duration nextPartialRetryBackoff() { + double multiplier = Math.pow(backoffMultiplier, partialRetryAttempt); + partialRetryAttempt++; + double delayNanos = initialBackoff.toNanos() * multiplier; + if (delayNanos >= maxBackoff.toNanos()) { + return maxBackoff; + } + return Duration.ofNanos((long) delayNanos); + } + + private boolean hasRemainingDeadlineFor(Duration retryDelay) { + return context.getOperationInfo().getDeadline().timeRemaining(TimeUnit.NANOSECONDS) + - retryDelay.toNanos() + > TimeUnit.MILLISECONDS.toNanos(1); + } + + abstract static class State { + public abstract void onStart(); + + public void onExit() {} + + public void onCancel(@Nullable String reason, @Nullable Throwable throwable) {} + + public boolean isDone() { + return false; + } + } + + class Idle extends State { + @Override + public void onStart() { + onStateChange(new Active()); + } + } + + class Active extends State { + private VRpc attempt; + private boolean attemptFinished = false; + @Nullable private SessionMutateRowsResponse attemptResponse; + + private void finishAttempt(VRpcResult result) { + if (!attemptFinished) { + attemptFinished = true; + tracer.onAttemptFinish(result); + } + } + + @Override + public void onStart() { + attempt = attemptFactory.get(); + tracer.onAttemptStart(currentRequest); + attempt.start( + currentRequest, + context, + new VRpcListener() { + @Override + public void onMessage(SessionMutateRowsResponse msg) { + context.getExecutor().throwIfNotInThisExecutor(); + if (currentState != Active.this) { + LOG.log( + Level.FINE, + "Discarding response {0} because the attempt is no longer active.", + msg); + return; + } + tracer.onResponseReceived(); + attemptResponse = msg; + } + + @Override + public void onClose(VRpcResult result) { + context.getExecutor().throwIfNotInThisExecutor(); + if (currentState != Active.this) { + LOG.log( + Level.FINE, + "Discarding server close with result {0} because the attempt is no longer" + + " active.", + result); + return; + } + + if (!result.getStatus().isOk()) { + handleWholeVRpcError(result); + return; + } + + if (attemptResponse == null) { + // Protocol error: unary vRPC closed OK without delivering a response payload. + finishAttempt(result); + onStateChange(new Done(result)); + return; + } + + handleResponseAndClose(attemptResponse, result); + } + }); + } + + private void handleWholeVRpcError(VRpcResult result) { + finishAttempt(result); + if (result.getState() == VRpcResult.State.SERVER_RESULT) { + consecutiveNonServerRetries = 0; + } + if (shouldRetryWholeVRpc(result)) { + if (result.getState() != VRpcResult.State.SERVER_RESULT) { + consecutiveNonServerRetries++; + } + context = context.createForNextAttempt(); + Duration retryDelay = + Optional.ofNullable(result.getRetryInfo()) + .filter(RetryInfo::hasRetryDelay) + .map(info -> Duration.ofNanos(Durations.toNanos(info.getRetryDelay()))) + .orElse(Duration.ZERO); + if (!retryDelay.isZero() && !retryDelay.isNegative()) { + onStateChange(new Scheduled(retryDelay)); + } else { + onStateChange(new Idle()); + } + return; + } + + if (hasPartialProgress && lastOkServerResult != null) { + // Earlier attempts already locked in outcomes for some entries; synthesize per-entry + // failure statuses for the remaining unresolved entries so prior successes/permanent + // errors are not lost. + Status failureStatus = result.getStatus(); + String message = + failureStatus.getDescription() != null ? failureStatus.getDescription() : ""; + for (int origIdx : activeOriginalIndices) { + SessionMutateRowsResponse.Entry.Builder entryBuilder = + SessionMutateRowsResponse.Entry.newBuilder().setIndex(origIdx); + entryBuilder + .getStatusBuilder() + .setCode(failureStatus.getCode().value()) + .setMessage(message); + resolvedEntries[origIdx] = entryBuilder.build(); + } + SessionMutateRowsResponse merged = buildFinalResponse(null); + deliverFinalMessageAndClose(merged, lastOkServerResult); + return; + } + + onStateChange(new Done(result)); + } + + private void handleResponseAndClose(SessionMutateRowsResponse response, VRpcResult result) { + consecutiveNonServerRetries = 0; + lastOkServerResult = result; + int currentCount = activeOriginalIndices.length; + boolean[] seen = new boolean[currentCount]; + + List retryOrigIndices = new ArrayList<>(); + List pendingTransientEntries = new ArrayList<>(); + + for (SessionMutateRowsResponse.Entry respEntry : response.getEntriesList()) { + int idx = Ints.checkedCast(respEntry.getIndex()); + if (idx < 0 || idx >= currentCount || seen[idx]) { + continue; + } + seen[idx] = true; + int origIdx = activeOriginalIndices[idx]; + int code = respEntry.getStatus().getCode(); + + SessionMutateRowsResponse.Entry remapped = respEntry.toBuilder().setIndex(origIdx).build(); + + if (code == Code.OK_VALUE) { + resolvedEntries[origIdx] = remapped; + hasPartialProgress = true; + } else if (retryableEntryCodes.contains(code) && entryIdempotent[origIdx]) { + retryOrigIndices.add(origIdx); + pendingTransientEntries.add(remapped); + } else { + resolvedEntries[origIdx] = remapped; + hasPartialProgress = true; + } + } + + for (int i = 0; i < currentCount; i++) { + if (!seen[i]) { + int origIdx = activeOriginalIndices[i]; + SessionMutateRowsResponse.Entry.Builder missingBuilder = + SessionMutateRowsResponse.Entry.newBuilder().setIndex(origIdx); + missingBuilder + .getStatusBuilder() + .setCode(Code.INTERNAL_VALUE) + .setMessage("Missing entry response for entry " + origIdx); + resolvedEntries[origIdx] = missingBuilder.build(); + hasPartialProgress = true; + } + } + + if (!retryOrigIndices.isEmpty()) { + Duration backoff = nextPartialRetryBackoff(); + if (hasRemainingDeadlineFor(backoff)) { + finishAttempt(result); + SessionMutateRowsRequest.Builder nextReq = SessionMutateRowsRequest.newBuilder(); + int[] nextIndices = new int[retryOrigIndices.size()]; + for (int i = 0; i < retryOrigIndices.size(); i++) { + int origIdx = retryOrigIndices.get(i); + nextIndices[i] = origIdx; + nextReq.addEntries(originalRequest.getEntries(origIdx)); + } + currentRequest = nextReq.build(); + activeOriginalIndices = nextIndices; + context = context.createForNextAttempt(); + if (!backoff.isZero() && !backoff.isNegative()) { + onStateChange(new Scheduled(backoff)); + } else { + onStateChange(new Idle()); + } + return; + } + + // Remaining deadline is insufficient for another retry; surface the transient entry errors. + for (SessionMutateRowsResponse.Entry transientEntry : pendingTransientEntries) { + resolvedEntries[Ints.checkedCast(transientEntry.getIndex())] = transientEntry; + } + } + + SessionMutateRowsResponse finalResponse = buildFinalResponse(response); + Stopwatch appTimer = Stopwatch.createStarted(); + Throwable userThrow = null; + try { + listener.onMessage(finalResponse); + } catch (Throwable t) { + userThrow = t; + } finally { + tracer.recordApplicationBlockingLatencies(appTimer.elapsed()); + } + + if (userThrow != null) { + VRpcResult userResult = VRpcResult.createUserError(userThrow); + finishAttempt(userResult); + onStateChange(new Done(userResult)); + return; + } + + finishAttempt(result); + onStateChange(new Done(result)); + } + + private void deliverFinalMessageAndClose( + SessionMutateRowsResponse finalResponse, VRpcResult closeResult) { + Stopwatch appTimer = Stopwatch.createStarted(); + Throwable userThrow = null; + try { + listener.onMessage(finalResponse); + } catch (Throwable t) { + userThrow = t; + } finally { + tracer.recordApplicationBlockingLatencies(appTimer.elapsed()); + } + + if (userThrow != null) { + onStateChange(new Done(VRpcResult.createUserError(userThrow))); + return; + } + onStateChange(new Done(closeResult)); + } + + private SessionMutateRowsResponse buildFinalResponse( + @Nullable SessionMutateRowsResponse latestResponse) { + SessionMutateRowsResponse.Builder builder = SessionMutateRowsResponse.newBuilder(); + if (latestResponse != null && latestResponse.hasRateLimitInfo()) { + builder.setRateLimitInfo(latestResponse.getRateLimitInfo()); + } + for (SessionMutateRowsResponse.Entry entry : resolvedEntries) { + builder.addEntries(entry); + } + return builder.build(); + } + + private boolean shouldRetryWholeVRpc(VRpcResult result) { + if (!result.getStatus().isOk() + && result.getRetryInfo() != null + && result.getRetryInfo().hasRetryDelay()) { + Duration retryDelay = + Duration.ofNanos(Durations.toNanos(result.getRetryInfo().getRetryDelay())); + return hasRemainingDeadlineFor(retryDelay); + } + if (result.getRejected()) { + return false; + } + boolean isRetryable = + (result.getState() == VRpcResult.State.UNCOMMITED) + || (areAllActiveEntriesIdempotent() + && result.getState() == VRpcResult.State.TRANSPORT_FAILURE); + return isRetryable && consecutiveNonServerRetries < 3; + } + + @Override + public void onCancel(@Nullable String reason, @Nullable Throwable throwable) { + finishAttempt( + VRpcResult.createRejectedError( + Status.CANCELLED.withDescription(reason).withCause(throwable))); + if (attempt != null) { + attempt.cancel(reason, throwable); + } + } + + @Override + public void onExit() { + finishAttempt( + VRpcResult.createRejectedError( + Status.CANCELLED.withDescription("attempt abandoned during transition"))); + } + } + + class Scheduled extends State { + private final Duration retryDelay; + @Nullable private BigtableTimer.Timeout future; + @Nullable private BigtableTimer.Registration stopHook; + + Scheduled(Duration retryDelay) { + this.retryDelay = retryDelay; + } + + @Override + public void onStart() { + stopHook = timer.onStop(this::onTimerStopping); + future = + timer.newTimeout( + () -> contextWrapper.apply(() -> onStateChange(new Idle())).run(), + context.getExecutor(), + retryDelay.toNanos(), + TimeUnit.NANOSECONDS); + } + + private void onTimerStopping() { + context + .getExecutor() + .execute( + () -> { + if (currentState != Scheduled.this) { + return; + } + onStateChange( + new Done( + VRpcResult.createRejectedError( + Status.CANCELLED.withDescription( + "Client closing while retry pending")))); + }); + } + + @Override + public void onExit() { + if (stopHook != null) { + stopHook.unregister(); + stopHook = null; + } + if (future != null && !future.isCancelled()) { + future.cancel(); + } + } + } + + class Done extends State { + private final VRpcResult result; + + Done(VRpcResult result) { + this.result = result; + } + + @Override + public void onStart() { + Stopwatch appTimer = Stopwatch.createStarted(); + try { + listener.onClose(result); + } finally { + tracer.recordApplicationBlockingLatencies(appTimer.elapsed()); + tracer.onOperationFinish(result); + } + } + + @Override + public boolean isDone() { + return true; + } + } +} diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RateLimitingVRpc.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RateLimitingVRpc.java new file mode 100644 index 000000000000..8c6ad256bc96 --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RateLimitingVRpc.java @@ -0,0 +1,216 @@ +/* + * Copyright 2026 Google LLC + * + * 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 + * + * https://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.data.v2.internal.middleware; + +import static io.opentelemetry.context.Context.current; + +import com.google.bigtable.v2.SessionMutateRowsRequest; +import com.google.bigtable.v2.SessionMutateRowsResponse; +import com.google.cloud.bigtable.data.v2.internal.session.BigtableTimer; +import io.grpc.Context; +import io.grpc.Status; +import java.time.Duration; +import java.util.Optional; +import java.util.concurrent.TimeUnit; +import java.util.function.UnaryOperator; +import javax.annotation.Nullable; + +/** + * Non-blocking rate-limiting {@link VRpc} middleware for {@code SessionMutateRows}. + * + *

Reserves a permit from {@link SessionMutateRowsRateLimiter} when an attempt starts, schedules + * any required throttle delay asynchronously on {@link BigtableTimer} (never blocking {@link + * OpExecutor} or gRPC threads), and updates the limiter from server {@code RateLimitInfo} responses + * and overload statuses. + * + *

All mutable state is owned by the per-operation {@link OpExecutor}. + */ +public final class RateLimitingVRpc + implements VRpc { + + private final VRpc delegate; + private final SessionMutateRowsRateLimiter limiter; + private final BigtableTimer timer; + private final UnaryOperator contextWrapper; + + private SessionMutateRowsRequest request; + private VRpcCallContext context; + private VRpcListener listener; + + private boolean delegateStarted = false; + private boolean done = false; + @Nullable private BigtableTimer.Timeout throttleTimeout; + @Nullable private BigtableTimer.Registration stopHook; + + public RateLimitingVRpc( + VRpc delegate, + SessionMutateRowsRateLimiter limiter, + BigtableTimer timer) { + this.delegate = delegate; + this.limiter = limiter; + this.timer = timer; + Context grpcContext = Context.current(); + UnaryOperator otelWrap = current()::wrap; + this.contextWrapper = task -> grpcContext.wrap(otelWrap.apply(task)); + } + + @Override + public void start( + SessionMutateRowsRequest req, + VRpcCallContext ctx, + VRpcListener listener) { + if (this.context != null) { + VRpcResult alreadyStarted = + VRpcResult.createRejectedError( + Status.FAILED_PRECONDITION.withDescription("operation is already started")); + ctx.getExecutor().execute(() -> listener.onClose(alreadyStarted)); + return; + } + this.request = req; + this.context = ctx; + this.listener = listener; + + long remainingNanos = ctx.getOperationInfo().getDeadline().timeRemaining(TimeUnit.NANOSECONDS); + if (remainingNanos <= 0) { + failFast( + VRpcResult.createRejectedError( + Status.DEADLINE_EXCEEDED.withDescription( + "Deadline expired before acquiring batch write flow control permit"))); + return; + } + + Optional delayOpt = limiter.tryAcquireDelay(Duration.ofNanos(remainingNanos)); + if (!delayOpt.isPresent()) { + failFast( + VRpcResult.createRejectedError( + Status.DEADLINE_EXCEEDED.withDescription( + "Deadline would expire while waiting for batch write flow control permit"))); + return; + } + + Duration delay = delayOpt.get(); + if (delay.isZero()) { + startDelegate(); + return; + } + + stopHook = timer.onStop(this::onTimerStopping); + throttleTimeout = + timer.newTimeout( + () -> contextWrapper.apply(this::onThrottleComplete).run(), + ctx.getExecutor(), + delay.toNanos(), + TimeUnit.NANOSECONDS); + } + + private void failFast(VRpcResult result) { + done = true; + context.getExecutor().execute(() -> listener.onClose(result)); + } + + private void onThrottleComplete() { + context.getExecutor().throwIfNotInThisExecutor(); + if (done || delegateStarted) { + return; + } + cleanupTimerHandles(); + startDelegate(); + } + + private void onTimerStopping() { + context + .getExecutor() + .execute( + () -> { + if (done || delegateStarted) { + return; + } + cleanupTimerHandles(); + done = true; + listener.onClose( + VRpcResult.createRejectedError( + Status.CANCELLED.withDescription( + "Client closing while waiting for batch write flow control permit"))); + }); + } + + private void startDelegate() { + delegateStarted = true; + delegate.start( + request, + context, + new VRpcListener() { + @Override + public void onMessage(SessionMutateRowsResponse msg) { + context.getExecutor().throwIfNotInThisExecutor(); + limiter.onResponse(msg, context.getTracer()); + listener.onMessage(msg); + } + + @Override + public void onClose(VRpcResult result) { + context.getExecutor().throwIfNotInThisExecutor(); + done = true; + limiter.onResult(result, context.getTracer()); + listener.onClose(result); + } + }); + } + + @Override + public void cancel(@Nullable String message, @Nullable Throwable cause) { + if (done) { + return; + } + cleanupTimerHandles(); + if (delegateStarted) { + delegate.cancel(message, cause); + done = delegate.isDone(); + } else { + done = true; + if (listener != null) { + listener.onClose( + VRpcResult.createRejectedError( + Status.CANCELLED.withDescription(message).withCause(cause))); + } + } + } + + @Override + public boolean isDone() { + return done || (delegateStarted && delegate.isDone()); + } + + @Override + public void requestNext() { + if (context != null) { + context.getExecutor().throwIfNotInThisExecutor(); + } + throw new UnsupportedOperationException("request next is not supported in unary"); + } + + private void cleanupTimerHandles() { + if (stopHook != null) { + stopHook.unregister(); + stopHook = null; + } + if (throttleTimeout != null && !throttleTimeout.isCancelled()) { + throttleTimeout.cancel(); + } + throttleTimeout = null; + } +} diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/SessionMutateRowsRateLimiter.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/SessionMutateRowsRateLimiter.java new file mode 100644 index 000000000000..872e6675c5da --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/SessionMutateRowsRateLimiter.java @@ -0,0 +1,262 @@ +/* + * Copyright 2026 Google LLC + * + * 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 + * + * https://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.data.v2.internal.middleware; + +import com.google.bigtable.v2.SessionMutateRowsResponse; +import com.google.bigtable.v2.SessionMutateRowsResponse.RateLimitInfo; +import com.google.cloud.bigtable.data.v2.internal.csm.tracers.VRpcTracer; +import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc.VRpcResult; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; +import com.google.common.base.Ticker; +import com.google.common.math.LongMath; +import com.google.protobuf.util.Durations; +import io.grpc.Status; +import java.time.Duration; +import java.util.Optional; +import java.util.concurrent.TimeUnit; +import java.util.logging.Logger; + +/** + * Non-blocking conditional rate limiter for {@code SessionMutateRows}, driven by server-returned + * {@link RateLimitInfo} and overload errors. + * + *

Mirrors the behavior of {@code RateLimitingServerStreamingCallable} on the classic path, + * including smooth-bursty permit reservation (up to 1 second of stored permits), QPS bounds {@code + * [0.1, 100_000]}, factor bounds {@code [0.7, 1.3]}, period-gated updates/disabling, and {@code + * 0.7x} backoff on {@code DEADLINE_EXCEEDED}, {@code UNAVAILABLE}, and {@code RESOURCE_EXHAUSTED}, + * while returning reservation delays instead of blocking threads. + */ +public final class SessionMutateRowsRateLimiter { + private static final Logger logger = + Logger.getLogger(SessionMutateRowsRateLimiter.class.getName()); + + static final double DEFAULT_QPS = 10.0; + static final Duration DEFAULT_PERIOD = Duration.ofSeconds(10); + static final double MIN_QPS = 0.1; + static final double MAX_QPS = 100_000.0; + static final double MIN_FACTOR = 0.7; + static final double MAX_FACTOR = 1.3; + private static final double MAX_BURST_SECONDS = 1.0; + + private final Ticker ticker; + private final Object lock = new Object(); + + private boolean enabled = false; + private double rate; + private double stableIntervalNanos; + private double maxPermits; + private double storedPermits; + private long nextFreeTicketNanos; + private long nextRateUpdateNanos; + + public SessionMutateRowsRateLimiter() { + this(DEFAULT_QPS, Ticker.systemTicker()); + } + + public SessionMutateRowsRateLimiter(double defaultQps, Ticker ticker) { + Preconditions.checkArgument(defaultQps > 0.0, "defaultQps must be positive"); + this.ticker = Preconditions.checkNotNull(ticker, "ticker"); + long nowNanos = ticker.read(); + this.rate = defaultQps; + this.stableIntervalNanos = TimeUnit.SECONDS.toNanos(1L) / defaultQps; + this.maxPermits = MAX_BURST_SECONDS * defaultQps; + this.storedPermits = 0.0; + this.nextFreeTicketNanos = nowNanos; + this.nextRateUpdateNanos = nowNanos; + logger.info( + "Batch write flow control: session rate limiter initialized (disabled) with rate of " + + defaultQps + + " QPS."); + } + + /** + * Reserves a permit without blocking if the wait does not exceed {@code maxWait}. + * + * @return the non-negative delay to wait before sending the request, or {@link Optional#empty()} + * if the required wait exceeds {@code maxWait} (in which case no permit is consumed). + */ + public Optional tryAcquireDelay(Duration maxWait) { + synchronized (lock) { + if (!enabled) { + return Optional.of(Duration.ZERO); + } + long nowNanos = ticker.read(); + resyncLocked(nowNanos); + long waitNanos = Math.max(0L, nextFreeTicketNanos - nowNanos); + if (maxWait.isNegative() || waitNanos > maxWait.toNanos()) { + return Optional.empty(); + } + double storedPermitsToSpend = Math.min(1.0, storedPermits); + double freshPermits = 1.0 - storedPermitsToSpend; + long freshWaitNanos = (long) (freshPermits * stableIntervalNanos); + nextFreeTicketNanos = LongMath.saturatedAdd(nextFreeTicketNanos, freshWaitNanos); + storedPermits -= storedPermitsToSpend; + return Optional.of(Duration.ofNanos(waitNanos)); + } + } + + /** Updates limiter state and CSM metrics from a {@link SessionMutateRowsResponse}. */ + public void onResponse(SessionMutateRowsResponse response, VRpcTracer tracer) { + if (hasValidRateLimitInfo(response)) { + enable(); + RateLimitInfo info = response.getRateLimitInfo(); + updateQps( + info.getFactor(), + Duration.ofSeconds(Durations.toSeconds(info.getPeriod())), + Status.Code.OK, + tracer); + } else { + tryDisable(); + } + } + + /** Updates limiter state and CSM metrics when an attempt closes with an overload status. */ + public void onResult(VRpcResult result, VRpcTracer tracer) { + if (result.getStatus().isOk()) { + return; + } + Status.Code code = result.getStatus().getCode(); + if (code == Status.Code.DEADLINE_EXCEEDED + || code == Status.Code.UNAVAILABLE + || code == Status.Code.RESOURCE_EXHAUSTED) { + updateQps(MIN_FACTOR, DEFAULT_PERIOD, code, tracer); + } + } + + private static boolean hasValidRateLimitInfo(SessionMutateRowsResponse response) { + if (!response.hasRateLimitInfo()) { + return false; + } + RateLimitInfo info = response.getRateLimitInfo(); + return info.getFactor() > 0 && info.getPeriod().getSeconds() > 0; + } + + private void enable() { + boolean wasEnabled; + synchronized (lock) { + wasEnabled = enabled; + enabled = true; + } + if (!wasEnabled) { + logger.info("Batch write flow control: session rate limiter is enabled."); + } + } + + private void tryDisable() { + boolean wasEnabled = false; + synchronized (lock) { + long nowNanos = ticker.read(); + if (nowNanos > nextRateUpdateNanos) { + wasEnabled = enabled; + enabled = false; + } + } + if (wasEnabled) { + logger.info("Batch write flow control: session rate limiter is disabled."); + } + } + + private void updateQps( + double factor, Duration period, Status.Code statusCode, VRpcTracer tracer) { + double cappedFactor = Math.min(Math.max(factor, MIN_FACTOR), MAX_FACTOR); + double oldRate; + double newRate; + boolean applied; + + synchronized (lock) { + long nowNanos = ticker.read(); + if (nowNanos < nextRateUpdateNanos) { + oldRate = rate; + newRate = rate; + applied = false; + } else { + oldRate = rate; + newRate = Math.min(Math.max(oldRate * cappedFactor, MIN_QPS), MAX_QPS); + setRateLocked(newRate, nowNanos); + nextRateUpdateNanos = + LongMath.saturatedAdd(nowNanos, TimeUnit.SECONDS.toNanos(period.getSeconds())); + applied = true; + } + } + + if (!applied) { + tracer.recordBatchWriteFlowControlFactor(cappedFactor, statusCode, false); + return; + } + + logger.info( + "Batch write flow control: updated max rate from " + + oldRate + + " to " + + newRate + + " applied factor " + + cappedFactor + + " with period " + + period.getSeconds() + + " seconds. Status=" + + statusCode); + tracer.recordBatchWriteFlowControlTargetQps(newRate); + tracer.recordBatchWriteFlowControlFactor(cappedFactor, statusCode, true); + } + + private void resyncLocked(long nowNanos) { + if (nowNanos > nextFreeTicketNanos) { + double newPermits = (nowNanos - nextFreeTicketNanos) / stableIntervalNanos; + storedPermits = Math.min(maxPermits, storedPermits + newPermits); + nextFreeTicketNanos = nowNanos; + } + } + + private void setRateLocked(double newRate, long nowNanos) { + resyncLocked(nowNanos); + this.rate = newRate; + this.stableIntervalNanos = TimeUnit.SECONDS.toNanos(1L) / newRate; + double oldMaxPermits = this.maxPermits; + this.maxPermits = MAX_BURST_SECONDS * newRate; + this.storedPermits = + (oldMaxPermits == 0.0) ? 0.0 : this.storedPermits * this.maxPermits / oldMaxPermits; + } + + @VisibleForTesting + public boolean isEnabled() { + synchronized (lock) { + return enabled; + } + } + + @VisibleForTesting + public double getRate() { + synchronized (lock) { + return rate; + } + } + + @VisibleForTesting + void setEnabled(boolean enabled) { + synchronized (lock) { + this.enabled = enabled; + } + } + + @VisibleForTesting + void setRate(double newRate) { + synchronized (lock) { + setRateLocked(newRate, ticker.read()); + } + } +} diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcDescriptor.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcDescriptor.java index 575fb79fa26b..ad56311ed9a6 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcDescriptor.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcDescriptor.java @@ -23,6 +23,7 @@ import com.google.bigtable.v2.MaterializedViewRequest; import com.google.bigtable.v2.MaterializedViewResponse; import com.google.bigtable.v2.MutateRowRequest; +import com.google.bigtable.v2.MutateRowsRequest; import com.google.bigtable.v2.OpenAuthorizedViewRequest; import com.google.bigtable.v2.OpenMaterializedViewRequest; import com.google.bigtable.v2.OpenTableRequest; @@ -33,6 +34,8 @@ 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; @@ -189,6 +192,30 @@ public final class VRpcDescriptor { .addAllRules(req.getRulesList()) .build()); + public static final VRpcDescriptor< + OpenTableRequest, SessionMutateRowsRequest, SessionMutateRowsResponse> + MUTATE_ROWS = + new VRpcDescriptor<>( + TABLE_SESSION, + MethodInfo.of("Bigtable.MutateRows", false), + createTableEncoder(TableRequest.Builder::setMutateRows), + createTableDecoder(TableResponse::getMutateRows), + (name, appProfileId, req) -> { + MutateRowsRequest.Builder builder = + MutateRowsRequest.newBuilder().setTableName(name).setAppProfileId(appProfileId); + for (SessionMutateRowsRequest.Entry entry : req.getEntriesList()) { + MutateRowsRequest.Entry.Builder entryBuilder = + MutateRowsRequest.Entry.newBuilder() + .setRowKey(entry.getKey()) + .addAllMutations(entry.getMutationsList()); + if (entry.hasIdempotency()) { + entryBuilder.setIdempotency(entry.getIdempotency()); + } + builder.addEntries(entryBuilder); + } + return builder.build(); + }); + public static final VRpcDescriptor< OpenAuthorizedViewRequest, SessionReadRowRequest, SessionReadRowResponse> READ_ROW_AUTH_VIEW = @@ -264,6 +291,32 @@ public final class VRpcDescriptor { .addAllRules(req.getRulesList()) .build()); + public static final VRpcDescriptor< + OpenAuthorizedViewRequest, SessionMutateRowsRequest, SessionMutateRowsResponse> + MUTATE_ROWS_AUTH_VIEW = + new VRpcDescriptor<>( + AUTHORIZED_VIEW_SESSION, + MethodInfo.of("Bigtable.MutateRows", false), + createAuthViewEncoder(AuthorizedViewRequest.Builder::setMutateRows), + createAuthViewDecoder(AuthorizedViewResponse::getMutateRows), + (name, appProfileId, req) -> { + MutateRowsRequest.Builder builder = + MutateRowsRequest.newBuilder() + .setAuthorizedViewName(name) + .setAppProfileId(appProfileId); + for (SessionMutateRowsRequest.Entry entry : req.getEntriesList()) { + MutateRowsRequest.Entry.Builder entryBuilder = + MutateRowsRequest.Entry.newBuilder() + .setRowKey(entry.getKey()) + .addAllMutations(entry.getMutationsList()); + if (entry.hasIdempotency()) { + entryBuilder.setIdempotency(entry.getIdempotency()); + } + builder.addEntries(entryBuilder); + } + return builder.build(); + }); + public static final VRpcDescriptor< OpenMaterializedViewRequest, SessionReadRowRequest, SessionReadRowResponse> READ_ROW_MAT_VIEW = diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScope.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScope.java index e89892a5eba8..b7e27928b2eb 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScope.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScope.java @@ -35,7 +35,8 @@ public enum SessionScope { CHECK_AND_MUTATE_ROW( "method:Bigtable.CheckAndMutateRow", /* fallbackToGlobalSessionLoad= */ false), READ_MODIFY_WRITE_ROW( - "method:Bigtable.ReadModifyWriteRow", /* fallbackToGlobalSessionLoad= */ false); + "method:Bigtable.ReadModifyWriteRow", /* fallbackToGlobalSessionLoad= */ false), + MUTATE_ROWS("method:Bigtable.MutateRows", /* fallbackToGlobalSessionLoad= */ false); private final String key; private final boolean fallbackToGlobalSessionLoad; diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/BulkMutation.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/BulkMutation.java index fb7b04269f07..767ffc770234 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/BulkMutation.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/BulkMutation.java @@ -20,6 +20,7 @@ import com.google.api.core.BetaApi; import com.google.api.core.InternalApi; import com.google.bigtable.v2.MutateRowsRequest; +import com.google.bigtable.v2.SessionMutateRowsRequest; import com.google.cloud.bigtable.data.v2.internal.NameUtil; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.common.base.Preconditions; @@ -128,6 +129,27 @@ public int getEntryCount() { return builder.getEntriesCount(); } + @InternalApi + public TargetId getTargetId() { + return targetId; + } + + @InternalApi + public SessionMutateRowsRequest toSessionProto() { + SessionMutateRowsRequest.Builder sessionBuilder = SessionMutateRowsRequest.newBuilder(); + for (MutateRowsRequest.Entry entry : builder.getEntriesList()) { + SessionMutateRowsRequest.Entry.Builder entryBuilder = + SessionMutateRowsRequest.Entry.newBuilder() + .setKey(entry.getRowKey()) + .addAllMutations(entry.getMutationsList()); + if (entry.hasIdempotency()) { + entryBuilder.setIdempotency(entry.getIdempotency()); + } + sessionBuilder.addEntries(entryBuilder); + } + return sessionBuilder.build(); + } + @InternalApi public MutateRowsRequest toProto(RequestContext requestContext) { String resourceName = diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index cafc9e78e8b4..5dee50d8c877 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java @@ -205,9 +205,14 @@ public EnhancedBigtableStub( new MutateRowsErrorConverterUnaryCallable(bulkMutateRowsCallable); externalBulkMutateRowsCallable = new MaybePointWriteCallable( - bulkMutateRowsVoidCallable, - createPointWriteCallable(bulkMutateRowsVoidCallable), - requestContext); + bulkMutateRowsVoidCallable, + createPointWriteCallable(bulkMutateRowsVoidCallable), + requestContext) + .withDefaultCallContext( + bigtableClientContext + .getClientContext() + .getDefaultCallContext() + .withRetrySettings(perOpSettings.bulkMutateRowsSettings.getRetrySettings())); checkAndMutateRowCallable = createCheckAndMutateRowCallable(); readModifyWriteRowCallable = createReadModifyWriteRowCallable(); generateInitialChangeStreamPartitionsCallable = @@ -793,33 +798,42 @@ private UnaryCallable createMutateRowsBas UnaryCallable withCookie = new CookiesUnaryCallable<>(baseCallable); - UnaryCallable flowControlCallable = null; - if (perOpSettings.bulkMutateRowsSettings.isLatencyBasedThrottlingEnabled()) { - flowControlCallable = - new DynamicFlowControlCallable( - withCookie, - bulkMutationFlowController, - bulkMutationDynamicFlowControlStats, - perOpSettings.bulkMutateRowsSettings.getTargetRpcLatencyMs(), - FLOW_CONTROL_ADJUSTING_INTERVAL_MS); - } UnaryCallable userFacing = - new BulkMutateRowsUserFacingCallable( - flowControlCallable != null ? flowControlCallable : withCookie, requestContext); + new BulkMutateRowsUserFacingCallable(withCookie, requestContext); SpanName spanName = getSpanName("MutateRows"); UnaryCallable tracedBatcherUnaryCallable = new TracedBatcherUnaryCallable<>(userFacing); - UnaryCallable traced = - new TracedUnaryCallable<>( - tracedBatcherUnaryCallable, clientContext.getTracerFactory(), spanName); - - return traced.withDefaultCallContext( + ApiCallContext defaultCallContext = clientContext .getDefaultCallContext() - .withRetrySettings(perOpSettings.bulkMutateRowsSettings.getRetrySettings())); + .withRetrySettings(perOpSettings.bulkMutateRowsSettings.getRetrySettings()); + + UnaryCallable classic = + new TracedUnaryCallable<>( + tracedBatcherUnaryCallable, clientContext.getTracerFactory(), spanName) + .withDefaultCallContext(defaultCallContext); + + UnaryCallable diverted = + bigtableClientContext + .getSessionShim() + .decorateMutateRows(classic, perOpSettings.bulkMutateRowsSettings); + + if (perOpSettings.bulkMutateRowsSettings.isLatencyBasedThrottlingEnabled()) { + @SuppressWarnings("unchecked") + UnaryCallable flowControlled = + new DynamicFlowControlCallable( + diverted, + bulkMutationFlowController, + bulkMutationDynamicFlowControlStats, + perOpSettings.bulkMutateRowsSettings.getTargetRpcLatencyMs(), + FLOW_CONTROL_ADJUSTING_INTERVAL_MS); + diverted = flowControlled; + } + + return diverted.withDefaultCallContext(defaultCallContext); } /** diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/TableBaseTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/TableBaseTest.java index 64b32d92c27d..ca3b4f20af77 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/TableBaseTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/TableBaseTest.java @@ -17,25 +17,34 @@ package com.google.cloud.bigtable.data.v2.internal.api; import static com.google.cloud.bigtable.data.v2.internal.test_helpers.VRpcCallContextSubject.assertThat; +import static com.google.common.truth.Truth.assertThat; import com.google.bigtable.v2.CloseSessionRequest; +import com.google.bigtable.v2.FeatureFlags; import com.google.bigtable.v2.Mutation; import com.google.bigtable.v2.OpenTableRequest; import com.google.bigtable.v2.ReadModifyWriteRule; import com.google.bigtable.v2.SessionCheckAndMutateRowRequest; 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.VirtualRpcResponse; import com.google.cloud.bigtable.data.v2.internal.csm.Metrics; 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.middleware.VRpc; +import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc.VRpcResult; 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.SessionPoolInfo; import com.google.cloud.bigtable.data.v2.internal.session.VRpcDescriptor; import com.google.common.util.concurrent.MoreExecutors; import com.google.protobuf.Message; +import com.google.protobuf.util.Durations; +import com.google.rpc.Code; import io.grpc.Deadline; import io.grpc.Metadata; import java.time.Duration; @@ -80,6 +89,7 @@ public void setup() { VRpcDescriptor.MUTATE_ROW, VRpcDescriptor.CHECK_AND_MUTATE_ROW, VRpcDescriptor.READ_MODIFY_WRITE_ROW, + VRpcDescriptor.MUTATE_ROWS, noopMetrics, mockTimer, MoreExecutors.directExecutor()); @@ -200,6 +210,145 @@ public void testReadModifyWriteRowNotIdempotent() { assertThat(fakeSessionPool.lastVRpc.ctx).isNotIdempotent(); } + @Test + public void testMutateRowsIdempotent() { + table.mutateRows( + SessionMutateRowsRequest.newBuilder() + .addEntries( + SessionMutateRowsRequest.Entry.newBuilder() + .addMutations( + Mutation.newBuilder() + .setSetCell(Mutation.SetCell.getDefaultInstance()) + .build())) + .build(), + new UnaryResponseFuture<>(), + deadline); + assertThat(fakeSessionPool.lastVRpc.ctx).isIdempotent(); + } + + @Test + public void testMutateRowsNotIdempotent() { + table.mutateRows( + SessionMutateRowsRequest.newBuilder() + .addEntries( + SessionMutateRowsRequest.Entry.newBuilder() + .addMutations( + Mutation.newBuilder() + .setAddToCell(Mutation.AddToCell.getDefaultInstance()) + .build())) + .build(), + new UnaryResponseFuture<>(), + deadline); + assertThat(fakeSessionPool.lastVRpc.ctx).isNotIdempotent(); + } + + @Test + public void testMutateRowsMixedEntriesNotIdempotent() { + table.mutateRows( + SessionMutateRowsRequest.newBuilder() + .addEntries( + SessionMutateRowsRequest.Entry.newBuilder() + .addMutations( + Mutation.newBuilder() + .setSetCell(Mutation.SetCell.newBuilder().setTimestampMicros(1000L)))) + .addEntries( + SessionMutateRowsRequest.Entry.newBuilder() + .addMutations( + Mutation.newBuilder() + .setAddToCell(Mutation.AddToCell.getDefaultInstance()))) + .build(), + new UnaryResponseFuture<>(), + deadline); + assertThat(fakeSessionPool.lastVRpc.ctx).isNotIdempotent(); + } + + @Test + public void testMutateRowsWithRateLimiterUpdatesLimiterFromResponse() throws Exception { + SessionMutateRowsRateLimiter limiter = new SessionMutateRowsRateLimiter(); + TableBase rateLimitedTable = + new TableBase( + fakeSessionPool, + VRpcDescriptor.READ_ROW, + VRpcDescriptor.MUTATE_ROW, + VRpcDescriptor.CHECK_AND_MUTATE_ROW, + VRpcDescriptor.READ_MODIFY_WRITE_ROW, + VRpcDescriptor.MUTATE_ROWS, + noopMetrics, + mockTimer, + MoreExecutors.directExecutor(), + limiter); + + assertThat(limiter.isEnabled()).isFalse(); + assertThat(limiter.getRate()).isEqualTo(10.0); + + UnaryResponseFuture future = new UnaryResponseFuture<>(); + rateLimitedTable.mutateRows( + SessionMutateRowsRequest.newBuilder() + .addEntries( + SessionMutateRowsRequest.Entry.newBuilder() + .addMutations( + Mutation.newBuilder() + .setDeleteFromRow(Mutation.DeleteFromRow.getDefaultInstance()))) + .build(), + future, + deadline); + + SessionMutateRowsResponse.Entry.Builder entry = + SessionMutateRowsResponse.Entry.newBuilder().setIndex(0); + entry.getStatusBuilder().setCode(Code.OK_VALUE); + SessionMutateRowsResponse response = + SessionMutateRowsResponse.newBuilder() + .addEntries(entry) + .setRateLimitInfo( + SessionMutateRowsResponse.RateLimitInfo.newBuilder() + .setFactor(0.8) + .setPeriod(Durations.fromSeconds(10))) + .build(); + + @SuppressWarnings("unchecked") + VRpc.VRpcListener listener = + (VRpc.VRpcListener) fakeSessionPool.lastVRpc.listener; + fakeSessionPool + .lastVRpc + .ctx + .getExecutor() + .execute( + () -> { + listener.onMessage(response); + listener.onClose(VRpcResult.createServerOk(VirtualRpcResponse.getDefaultInstance())); + }); + + assertThat(future.get()).isEqualTo(response); + assertThat(limiter.isEnabled()).isTrue(); + assertThat(limiter.getRate()).isWithin(0.001).of(8.0); + } + + @Test + public void testClientComputeBaseFeatureFlagsReflectsFlowControlSetting() { + ChannelProviders.ChannelProvider channelProvider = + Mockito.mock(ChannelProviders.ChannelProvider.class); + ClientSettings defaultSettings = + ClientSettings.builder() + .setChannelProvider(channelProvider) + .setInstanceName(clientInfo.getInstanceName()) + .setAppProfileId("default") + .build(); + FeatureFlags defaultFlags = Client.computeBaseFeatureFlags(defaultSettings); + assertThat(defaultFlags.getMutateRowsRateLimit()).isFalse(); + assertThat(defaultFlags.getMutateRowsRateLimit2()).isFalse(); + + ClientSettings enabledSettings = + ClientSettings.builder() + .setChannelProvider(channelProvider) + .setInstanceName(clientInfo.getInstanceName()) + .setAppProfileId("default") + .setServerInitiatedFlowControlEnabled(true) + .build(); + FeatureFlags enabledFlags = Client.computeBaseFeatureFlags(enabledSettings); + assertThat(enabledFlags.getMutateRowsRateLimit()).isTrue(); + assertThat(enabledFlags.getMutateRowsRateLimit2()).isTrue(); + } + static class FakeSessionPool implements SessionPool { private FakeVRpc lastVRpc = null; @@ -241,10 +390,12 @@ public boolean hasSession() { static class FakeVRpc implements VRpc { private VRpcCallContext ctx; + private VRpcListener listener; @Override - public void start(Object req, VRpcCallContext ctx, VRpcListener ignored) { + public void start(ReqT req, VRpcCallContext ctx, VRpcListener listener) { this.ctx = ctx; + this.listener = listener; } @Override diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImplTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImplTest.java index 4dfab36d0694..e95a07f4adc3 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImplTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImplTest.java @@ -22,9 +22,11 @@ import com.google.api.gax.core.FixedExecutorProvider; import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider; +import com.google.bigtable.v2.FeatureFlags; import com.google.bigtable.v2.SessionReadModifyWriteRowResponse; import com.google.bigtable.v2.SessionReadRowResponse; import com.google.cloud.bigtable.data.v2.internal.api.Client.Resource; +import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings; import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executor; import java.util.concurrent.ExecutorService; @@ -116,4 +118,20 @@ public void selectUserCallbackExecutor_unset_allocatesOwnedCachedPool() throws E resource.close(); assertThat(((ExecutorService) exec).isShutdown()).isTrue(); } + + @Test + public void computeBaseFeatureFlags_reflectsServerInitiatedFlowControlSetting() { + EnhancedBigtableStubSettings defaultSettings = + EnhancedBigtableStubSettings.newBuilder().setProjectId("p").setInstanceId("i").build(); + FeatureFlags defaultFlags = ShimImpl.computeBaseFeatureFlags(defaultSettings); + assertThat(defaultFlags.getMutateRowsRateLimit()).isFalse(); + assertThat(defaultFlags.getMutateRowsRateLimit2()).isFalse(); + + EnhancedBigtableStubSettings.Builder enabledBuilder = + EnhancedBigtableStubSettings.newBuilder().setProjectId("p").setInstanceId("i"); + enabledBuilder.bulkMutateRowsSettings().setServerInitiatedFlowControl(true); + FeatureFlags enabledFlags = ShimImpl.computeBaseFeatureFlags(enabledBuilder.build()); + assertThat(enabledFlags.getMutateRowsRateLimit()).isTrue(); + assertThat(enabledFlags.getMutateRowsRateLimit2()).isTrue(); + } } diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallableTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallableTest.java index 729f378b7a07..e1de16abaabe 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallableTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallableTest.java @@ -16,13 +16,18 @@ package com.google.cloud.bigtable.data.v2.internal.compat.ops; import static com.google.common.truth.Truth.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; import com.google.api.gax.grpc.GrpcCallContext; +import com.google.api.gax.grpc.GrpcStatusCode; import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.ApiExceptionFactory; +import com.google.api.gax.rpc.StatusCode; import com.google.api.gax.rpc.UnaryCallable; import com.google.auth.Credentials; import com.google.bigtable.v2.ClientConfiguration; @@ -34,7 +39,12 @@ import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DebugTagTracer; import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager; import com.google.cloud.bigtable.data.v2.internal.util.SessionScope; +import com.google.cloud.bigtable.data.v2.models.MutateRowsException; +import com.google.cloud.bigtable.data.v2.models.MutateRowsException.FailedMutation; +import com.google.common.collect.ImmutableList; import io.grpc.Deadline; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; import java.io.IOException; import java.net.URI; import java.time.Duration; @@ -43,6 +53,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -184,6 +195,42 @@ void perScopeDiversion_nonMatchingScope_routesToClassic() throws Exception { assertThat(experimental.callCount).isEqualTo(0); } + @Test + void experimentalFailure_preservesMutateRowsException() { + ApiException entryError = + ApiExceptionFactory.createException( + "bad entry", null, GrpcStatusCode.of(Status.Code.INVALID_ARGUMENT), false); + MutateRowsException mutateRowsException = + MutateRowsException.create( + null, + GrpcStatusCode.of(Status.Code.OK), + ImmutableList.of(FailedMutation.create(1, entryError)), + false); + experimental.nextResponse = new CompletableFuture<>(); + experimental.nextResponse.completeExceptionally(mutateRowsException); + + DivertingUnaryCallable callable = callable(SessionScope.READ_ROW); + ApiFuture future = callable.futureCall("req", GrpcCallContext.createDefault()); + + ExecutionException ex = assertThrows(ExecutionException.class, future::get); + assertThat(ex.getCause()).isSameInstanceAs(mutateRowsException); + } + + @Test + void experimentalFailure_translatesStatusRuntimeExceptionToApiException() { + experimental.nextResponse = new CompletableFuture<>(); + experimental.nextResponse.completeExceptionally( + new StatusRuntimeException(Status.UNAVAILABLE.withDescription("backend unavailable"))); + + DivertingUnaryCallable callable = callable(SessionScope.READ_ROW); + ApiFuture future = callable.futureCall("req", GrpcCallContext.createDefault()); + + ExecutionException ex = assertThrows(ExecutionException.class, future::get); + assertThat(ex.getCause()).isInstanceOf(ApiException.class); + assertThat(((ApiException) ex.getCause()).getStatusCode().getCode()) + .isEqualTo(StatusCode.Code.UNAVAILABLE); + } + // --- helpers --- private DivertingUnaryCallable callable() { @@ -252,11 +299,12 @@ public ApiFuture futureCall(ReqT request, ApiCallContext context) { private static class RecordingShim implements UnaryShim { int callCount; + CompletableFuture nextResponse = new CompletableFuture<>(); @Override public CompletableFuture call(ReqT request, Deadline deadline) { callCount++; - return new CompletableFuture<>(); // never completes — test doesn't wait on it + return nextResponse; // never completes by default — test doesn't wait on it } @Override diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowsShimTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowsShimTest.java new file mode 100644 index 000000000000..024485a7036c --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowsShimTest.java @@ -0,0 +1,379 @@ +/* + * Copyright 2026 Google LLC + * + * 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 + * + * https://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.data.v2.internal.compat.ops; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.api.core.SettableApiFuture; +import com.google.api.gax.batching.BatchEntry; +import com.google.api.gax.grpc.GrpcCallContext; +import com.google.api.gax.rpc.ApiException; +import com.google.api.gax.rpc.StatusCode; +import com.google.api.gax.rpc.UnaryCallable; +import com.google.bigtable.v2.ClientConfiguration; +import com.google.bigtable.v2.SessionClientConfiguration; +import com.google.bigtable.v2.SessionDiversionConfiguration; +import com.google.bigtable.v2.SessionMutateRowsResponse; +import com.google.bigtable.v2.SessionScopeDiversionConfiguration; +import com.google.cloud.bigtable.data.v2.internal.api.AuthorizedViewAsync; +import com.google.cloud.bigtable.data.v2.internal.api.Client; +import com.google.cloud.bigtable.data.v2.internal.api.TableAsync; +import com.google.cloud.bigtable.data.v2.internal.api.VRpcException; +import com.google.cloud.bigtable.data.v2.internal.csm.NoopMetrics; +import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc.VRpcResult; +import com.google.cloud.bigtable.data.v2.internal.session.SessionPool; +import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager; +import com.google.cloud.bigtable.data.v2.internal.util.SessionScope; +import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId; +import com.google.cloud.bigtable.data.v2.models.BulkMutation; +import com.google.cloud.bigtable.data.v2.models.MutateRowsException; +import com.google.cloud.bigtable.data.v2.models.MutateRowsException.FailedMutation; +import com.google.cloud.bigtable.data.v2.models.Mutation; +import com.google.cloud.bigtable.data.v2.models.RowMutationEntry; +import com.google.cloud.bigtable.data.v2.models.TableId; +import com.google.cloud.bigtable.data.v2.models.TargetId; +import com.google.cloud.bigtable.data.v2.stub.MutateRowsErrorConverterUnaryCallable; +import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsAttemptResult; +import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsBatchingDescriptor; +import com.google.common.collect.ImmutableList; +import com.google.rpc.Code; +import io.grpc.Deadline; +import io.grpc.Status; +import java.time.Duration; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class MutateRowsShimTest { + + private Client client; + private TableAsync tableAsync; + private AuthorizedViewAsync authViewAsync; + private SessionPool tablePool; + private SessionPool authViewPool; + private MutateRowsShim shim; + + @BeforeEach + @SuppressWarnings("unchecked") + void setUp() { + client = mock(Client.class); + tableAsync = mock(TableAsync.class); + authViewAsync = mock(AuthorizedViewAsync.class); + tablePool = mock(SessionPool.class); + authViewPool = mock(SessionPool.class); + + when(client.openTableAsync(any(), any())).thenReturn(tableAsync); + when(client.openAuthorizedViewAsync(any(), any(), any())).thenReturn(authViewAsync); + doReturn(tablePool).when(tableAsync).getSessionPool(); + doReturn(authViewPool).when(authViewAsync).getSessionPool(); + + shim = new MutateRowsShim(new SessionPoolMap(client)); + } + + @Test + void supports_tableId_belowThreshold_returnsTrue() { + when(tablePool.getConsecutiveUnimplementedFailures()) + .thenReturn(UnaryShim.MAX_CONSECUTIVE_UNIMPLEMENTED_FAILURES - 1); + + BulkMutation request = BulkMutation.create(TableId.of("table1")); + assertThat(shim.supports(request)).isTrue(); + } + + @Test + void supports_tableId_atThresholdNoSession_returnsFalse() { + when(tablePool.getConsecutiveUnimplementedFailures()) + .thenReturn(UnaryShim.MAX_CONSECUTIVE_UNIMPLEMENTED_FAILURES); + when(tablePool.hasSession()).thenReturn(false); + + BulkMutation request = BulkMutation.create(TableId.of("table1")); + assertThat(shim.supports(request)).isFalse(); + } + + @Test + void supports_tableId_atThresholdWithSession_returnsTrue() { + when(tablePool.getConsecutiveUnimplementedFailures()) + .thenReturn(UnaryShim.MAX_CONSECUTIVE_UNIMPLEMENTED_FAILURES); + when(tablePool.hasSession()).thenReturn(true); + + BulkMutation request = BulkMutation.create(TableId.of("table1")); + assertThat(shim.supports(request)).isTrue(); + } + + @Test + void supports_authorizedViewId_belowThreshold_returnsTrue() { + when(authViewPool.getConsecutiveUnimplementedFailures()) + .thenReturn(UnaryShim.MAX_CONSECUTIVE_UNIMPLEMENTED_FAILURES - 1); + + BulkMutation request = BulkMutation.create(AuthorizedViewId.of("table1", "view1")); + assertThat(shim.supports(request)).isTrue(); + } + + @Test + void supports_unsupportedTargetId_returnsFalse() { + TargetId unsupported = mock(TargetId.class); + BulkMutation request = BulkMutation.create(unsupported); + assertThat(shim.supports(request)).isFalse(); + } + + @Test + void call_tableId_allOk_returnsSuccess() throws Exception { + SessionMutateRowsResponse response = + SessionMutateRowsResponse.newBuilder() + .addEntries(entry(0, Code.OK_VALUE, "")) + .addEntries(entry(1, Code.OK_VALUE, "")) + .build(); + when(tableAsync.mutateRows(any(), any())) + .thenReturn(CompletableFuture.completedFuture(response)); + + BulkMutation request = + BulkMutation.create(TableId.of("table1")) + .add("k0", Mutation.create().setCell("cf", "q", 1000L, "v0")) + .add("k1", Mutation.create().setCell("cf", "q", 1000L, "v1")); + Deadline deadline = Deadline.after(1, TimeUnit.SECONDS); + + MutateRowsAttemptResult result = shim.call(request, deadline).get(); + assertThat(result.getFailedMutations()).isEmpty(); + assertThat(result.getIsRetryable()).isFalse(); + } + + @Test + void call_authorizedViewId_allOk_returnsSuccess() throws Exception { + SessionMutateRowsResponse response = + SessionMutateRowsResponse.newBuilder().addEntries(entry(0, Code.OK_VALUE, "")).build(); + when(authViewAsync.mutateRows(any(), any())) + .thenReturn(CompletableFuture.completedFuture(response)); + + BulkMutation request = + BulkMutation.create(AuthorizedViewId.of("table1", "view1")) + .add("k0", Mutation.create().deleteRow()); + Deadline deadline = Deadline.after(1, TimeUnit.SECONDS); + + MutateRowsAttemptResult result = shim.call(request, deadline).get(); + assertThat(result.getFailedMutations()).isEmpty(); + assertThat(result.getIsRetryable()).isFalse(); + } + + @Test + void call_partialEntryFailures_mapsErrorsAndRetryabilityPerEntry() throws Exception { + SessionMutateRowsResponse response = + SessionMutateRowsResponse.newBuilder() + .addEntries(entry(0, Code.OK_VALUE, "")) + .addEntries(entry(1, Code.UNAVAILABLE_VALUE, "tablet unavailable")) + .addEntries(entry(2, Code.INVALID_ARGUMENT_VALUE, "bad family")) + .addEntries(entry(3, Code.UNAVAILABLE_VALUE, "non-idempotent unavailable")) + .build(); + when(tableAsync.mutateRows(any(), any())) + .thenReturn(CompletableFuture.completedFuture(response)); + + BulkMutation request = + BulkMutation.create(TableId.of("table1")) + .add("k0", Mutation.create().setCell("cf", "q", 1000L, "v0")) + .add("k1", Mutation.create().setCell("cf", "q", 1000L, "v1")) + .add("k2", Mutation.create().setCell("cf", "q", 1000L, "v2")) + // setCell without explicit timestamp uses server timestamp (-1), which is + // non-idempotent + .add("k3", Mutation.createUnsafe().setCell("cf", "q", -1L, "v3")); + Deadline deadline = Deadline.after(1, TimeUnit.SECONDS); + + MutateRowsAttemptResult result = shim.call(request, deadline).get(); + assertThat(result.getIsRetryable()).isTrue(); + assertThat(result.getFailedMutations()).hasSize(3); + + FailedMutation f1 = result.getFailedMutations().get(0); + assertThat(f1.getIndex()).isEqualTo(1); + assertThat(f1.getError().getStatusCode().getCode()).isEqualTo(StatusCode.Code.UNAVAILABLE); + assertThat(f1.getError().isRetryable()).isTrue(); + + FailedMutation f2 = result.getFailedMutations().get(1); + assertThat(f2.getIndex()).isEqualTo(2); + assertThat(f2.getError().getStatusCode().getCode()).isEqualTo(StatusCode.Code.INVALID_ARGUMENT); + assertThat(f2.getError().isRetryable()).isFalse(); + + FailedMutation f3 = result.getFailedMutations().get(2); + assertThat(f3.getIndex()).isEqualTo(3); + assertThat(f3.getError().getStatusCode().getCode()).isEqualTo(StatusCode.Code.UNAVAILABLE); + assertThat(f3.getError().isRetryable()).isFalse(); + } + + @Test + void call_missingEntryInResponse_recordsPermanentInternalFailure() throws Exception { + SessionMutateRowsResponse response = + SessionMutateRowsResponse.newBuilder() + .addEntries(entry(0, Code.OK_VALUE, "")) + .addEntries(entry(2, Code.OK_VALUE, "")) + .build(); + when(tableAsync.mutateRows(any(), any())) + .thenReturn(CompletableFuture.completedFuture(response)); + + BulkMutation request = + BulkMutation.create(TableId.of("table1")) + .add("k0", Mutation.create().deleteRow()) + .add("k1", Mutation.create().deleteRow()) + .add("k2", Mutation.create().deleteRow()); + Deadline deadline = Deadline.after(1, TimeUnit.SECONDS); + + MutateRowsAttemptResult result = shim.call(request, deadline).get(); + assertThat(result.getIsRetryable()).isFalse(); + assertThat(result.getFailedMutations()).hasSize(1); + FailedMutation missing = result.getFailedMutations().get(0); + assertThat(missing.getIndex()).isEqualTo(1); + assertThat(missing.getError().getStatusCode().getCode()).isEqualTo(StatusCode.Code.INTERNAL); + assertThat(missing.getError().isRetryable()).isFalse(); + assertThat(missing.getError().getMessage()).contains("Missing entry response for entry 1"); + } + + @Test + void call_wholeVRpcFailure_failsWithMutateRowsExceptionForAllEntries() { + CompletableFuture failedFuture = new CompletableFuture<>(); + failedFuture.completeExceptionally( + VRpcException.create( + VRpcResult.createLocalTransportError( + Status.UNAVAILABLE.withDescription("session disconnected")))); + when(tableAsync.mutateRows(any(), any())).thenReturn(failedFuture); + + BulkMutation request = + BulkMutation.create(TableId.of("table1")) + .add("k0", Mutation.create().setCell("cf", "q", 1000L, "v0")) + .add("k1", Mutation.createUnsafe().setCell("cf", "q", -1L, "v1")); + Deadline deadline = Deadline.after(1, TimeUnit.SECONDS); + + ExecutionException ex = + assertThrows(ExecutionException.class, () -> shim.call(request, deadline).get()); + assertThat(ex.getCause()).isInstanceOf(MutateRowsException.class); + MutateRowsException mrEx = (MutateRowsException) ex.getCause(); + assertThat(mrEx.getStatusCode().getCode()).isEqualTo(StatusCode.Code.UNAVAILABLE); + assertThat(mrEx.isRetryable()).isTrue(); + assertThat(mrEx.getFailedMutations()).hasSize(2); + assertThat(mrEx.getFailedMutations().get(0).getIndex()).isEqualTo(0); + assertThat(mrEx.getFailedMutations().get(0).getError().isRetryable()).isTrue(); + assertThat(mrEx.getFailedMutations().get(1).getIndex()).isEqualTo(1); + assertThat(mrEx.getFailedMutations().get(1).getError().isRetryable()).isFalse(); + } + + @Test + void call_wholeVRpcFailureEmptyRequest_failsWithApiException() { + CompletableFuture failedFuture = new CompletableFuture<>(); + failedFuture.completeExceptionally( + VRpcException.create( + VRpcResult.createRejectedError( + Status.PERMISSION_DENIED.withDescription("access denied")))); + when(tableAsync.mutateRows(any(), any())).thenReturn(failedFuture); + + BulkMutation emptyRequest = BulkMutation.create(TableId.of("table1")); + Deadline deadline = Deadline.after(1, TimeUnit.SECONDS); + + ExecutionException ex = + assertThrows(ExecutionException.class, () -> shim.call(emptyRequest, deadline).get()); + assertThat(ex.getCause()).isInstanceOf(ApiException.class); + assertThat(ex.getCause()).isNotInstanceOf(MutateRowsException.class); + ApiException apiEx = (ApiException) ex.getCause(); + assertThat(apiEx.getStatusCode().getCode()).isEqualTo(StatusCode.Code.PERMISSION_DENIED); + } + + @Test + void divertingCallable_integratesWithErrorConverterAndBatchingDescriptor() throws Exception { + ClientConfigurationManager configManager = mock(ClientConfigurationManager.class); + when(configManager.getClientConfiguration()) + .thenReturn( + ClientConfiguration.newBuilder() + .setSessionConfiguration( + SessionClientConfiguration.newBuilder() + .setSessionDiversionConfiguration( + SessionDiversionConfiguration.newBuilder() + .setGlobalDiversion( + SessionScopeDiversionConfiguration.newBuilder() + .setSessionLoad(1.0f)))) + .build()); + + @SuppressWarnings("unchecked") + UnaryCallable unusedClassic = mock(UnaryCallable.class); + DivertingUnaryCallable diverted = + new DivertingUnaryCallable<>( + configManager, + unusedClassic, + shim, + Duration.ofSeconds(10), + new NoopMetrics().getDebugTagTracer(), + SessionScope.MUTATE_ROWS); + MutateRowsErrorConverterUnaryCallable bulkVoidCallable = + new MutateRowsErrorConverterUnaryCallable(diverted); + + SessionMutateRowsResponse partialFailureResponse = + SessionMutateRowsResponse.newBuilder() + .addEntries(entry(0, Code.OK_VALUE, "")) + .addEntries(entry(1, Code.INVALID_ARGUMENT_VALUE, "bad col")) + .build(); + when(tableAsync.mutateRows(any(), any())) + .thenReturn(CompletableFuture.completedFuture(partialFailureResponse)); + + RowMutationEntry row0 = RowMutationEntry.create("k0").deleteRow(); + RowMutationEntry row1 = RowMutationEntry.create("k1").deleteRow(); + BulkMutation request = BulkMutation.create(TableId.of("table1")).add(row0).add(row1); + + // 1. bulkMutateRows path surfaces MutateRowsException with OK RPC status and entry 1 failure. + MutateRowsException thrown = + assertThrows( + MutateRowsException.class, + () -> bulkVoidCallable.call(request, GrpcCallContext.createDefault())); + assertThat(thrown.getStatusCode().getCode()).isEqualTo(StatusCode.Code.OK); + assertThat(thrown.getFailedMutations()).hasSize(1); + assertThat(thrown.getFailedMutations().get(0).getIndex()).isEqualTo(1); + assertThat(thrown.getFailedMutations().get(0).getError().getStatusCode().getCode()) + .isEqualTo(StatusCode.Code.INVALID_ARGUMENT); + + // 2. Batcher path splits MutateRowsAttemptResult so row 0 succeeds and only row 1 fails. + MutateRowsAttemptResult attemptResult = diverted.call(request, GrpcCallContext.createDefault()); + SettableApiFuture f0 = SettableApiFuture.create(); + SettableApiFuture f1 = SettableApiFuture.create(); + List> batchEntries = + ImmutableList.of(BatchEntry.create(row0, f0), BatchEntry.create(row1, f1)); + new MutateRowsBatchingDescriptor().splitResponse(attemptResult, batchEntries); + + assertThat(f0.get()).isNull(); + ExecutionException f1Ex = assertThrows(ExecutionException.class, f1::get); + assertThat(f1Ex.getCause()).isInstanceOf(ApiException.class); + assertThat(((ApiException) f1Ex.getCause()).getStatusCode().getCode()) + .isEqualTo(StatusCode.Code.INVALID_ARGUMENT); + } + + @Test + void call_unsupportedTargetId_futureFailsWithUnsupportedOperationException() { + TargetId unsupported = mock(TargetId.class); + BulkMutation request = BulkMutation.create(unsupported); + Deadline deadline = Deadline.after(1, TimeUnit.SECONDS); + + CompletableFuture future = shim.call(request, deadline); + + assertThat(future.isCompletedExceptionally()).isTrue(); + ExecutionException ex = assertThrows(ExecutionException.class, future::get); + assertThat(ex.getCause()).isInstanceOf(UnsupportedOperationException.class); + } + + private static SessionMutateRowsResponse.Entry entry(int index, int code, String message) { + SessionMutateRowsResponse.Entry.Builder builder = + SessionMutateRowsResponse.Entry.newBuilder().setIndex(index); + builder.getStatusBuilder().setCode(code).setMessage(message); + return builder.build(); + } +} diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/CompositeVRpcTracerTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/CompositeVRpcTracerTest.java index 12f041fa851c..c0f1d0d8e68b 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/CompositeVRpcTracerTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/CompositeVRpcTracerTest.java @@ -89,4 +89,18 @@ void onOperationFinishTest() { verify(child1, times(1)).onOperationFinish(result); verify(child2, times(1)).onOperationFinish(result); } + + @Test + void recordBatchWriteFlowControlTargetQpsTest() { + compositeVRpcTracer.recordBatchWriteFlowControlTargetQps(12.5); + verify(child1, times(1)).recordBatchWriteFlowControlTargetQps(12.5); + verify(child2, times(1)).recordBatchWriteFlowControlTargetQps(12.5); + } + + @Test + void recordBatchWriteFlowControlFactorTest() { + compositeVRpcTracer.recordBatchWriteFlowControlFactor(0.8, Status.Code.OK, true); + verify(child1, times(1)).recordBatchWriteFlowControlFactor(0.8, Status.Code.OK, true); + verify(child2, times(1)).recordBatchWriteFlowControlFactor(0.8, Status.Code.OK, true); + } } diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java index 54685b907700..2120ae5ed5ca 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java @@ -69,11 +69,13 @@ import io.grpc.InsecureChannelCredentials; import io.grpc.Metadata; import io.grpc.Server; +import io.grpc.Status; import io.opencensus.stats.Stats; import io.opencensus.tags.Tags; import io.opentelemetry.api.common.Attributes; import io.opentelemetry.sdk.OpenTelemetrySdk; import io.opentelemetry.sdk.metrics.SdkMeterProvider; +import io.opentelemetry.sdk.metrics.data.DoublePointData; import io.opentelemetry.sdk.metrics.data.HistogramPointData; import io.opentelemetry.sdk.metrics.data.LongPointData; import io.opentelemetry.sdk.metrics.data.MetricData; @@ -410,6 +412,36 @@ public void onAttemptFinish(VRpc.VRpcResult result) { assertThat(clientBlockingLatency).isAtLeast(sessionDelay); } + @Test + public void batchWriteFlowControlMetricsTest() { + MethodInfo methodInfo = + MethodInfo.builder().setName("Bigtable.MutateRows").setStreaming(false).build(); + VRpcTracer tracer = + new VRpcTracerImpl( + recorderRegistry, poolInfo, methodInfo, Deadline.after(1, TimeUnit.MINUTES)); + + tracer.recordBatchWriteFlowControlTargetQps(13.0); + tracer.recordBatchWriteFlowControlFactor(1.3, Status.Code.OK, true); + + double targetQps = + getDoubleMetricValue( + metricReader, + "bigtable.googleapis.com/internal/client/batch_write_flow_control_target_qps", + Attributes.builder().put("method", "Bigtable.MutateRows").build()); + double factor = + getDoubleMetricValue( + metricReader, + "bigtable.googleapis.com/internal/client/batch_write_flow_control_factor", + Attributes.builder() + .put("method", "Bigtable.MutateRows") + .put("status", "OK") + .put("applied", true) + .build()); + + assertThat(targetQps).isWithin(0.001).of(13.0); + assertThat(factor).isWithin(0.001).of(1.3); + } + private Action createErrorResponse( int errorCode, com.google.protobuf.Duration retryDelay, long rpcId) { com.google.rpc.Status status = @@ -484,6 +516,33 @@ private static long getMetricValue( "MetricData was missing then appeared, this should never happen"); } + private static double getDoubleMetricValue( + InMemoryMetricReader reader, String metricName, Attributes attributes) { + Collection allMetricData = reader.collectAllMetrics(); + List matchingMetadata = + allMetricData.stream() + .filter(md -> METRIC_DATA_BY_NAME.compare(md, metricName)) + .collect(Collectors.toList()); + assertWithMessage( + "Found multiple MetricData with the same name: %s, in: %s", + metricName, matchingMetadata) + .that(matchingMetadata.size()) + .isEqualTo(1); + + MetricData metricData = matchingMetadata.get(0); + List points = + metricData.getDoubleGaugeData().getPoints().stream() + .filter( + pd -> + pd.getAttributes() + .asMap() + .entrySet() + .containsAll(attributes.asMap().entrySet())) + .collect(Collectors.toList()); + assertWithMessage("ambiguous labels matched points").that(points).hasSize(1); + return points.get(0).getValue(); + } + static class DelayedVRpc implements VRpc { private final Supplier> supplier; diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/MutateRowsRetryingVRpcTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/MutateRowsRetryingVRpcTest.java new file mode 100644 index 000000000000..654823457353 --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/MutateRowsRetryingVRpcTest.java @@ -0,0 +1,655 @@ +/* + * Copyright 2026 Google LLC + * + * 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 + * + * https://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.data.v2.internal.middleware; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; + +import com.google.bigtable.v2.ErrorResponse; +import com.google.bigtable.v2.Mutation; +import com.google.bigtable.v2.SessionMutateRowsRequest; +import com.google.bigtable.v2.SessionMutateRowsResponse; +import com.google.bigtable.v2.SessionMutateRowsResponse.RateLimitInfo; +import com.google.bigtable.v2.VirtualRpcResponse; +import com.google.cloud.bigtable.data.v2.internal.api.UnaryResponseFuture; +import com.google.cloud.bigtable.data.v2.internal.api.VRpcException; +import com.google.cloud.bigtable.data.v2.internal.csm.NoopMetrics; +import com.google.cloud.bigtable.data.v2.internal.csm.tracers.VRpcTracer; +import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc.VRpcCallContext; +import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc.VRpcListener; +import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc.VRpcResult; +import com.google.cloud.bigtable.data.v2.internal.session.BigtableTimer; +import com.google.protobuf.ByteString; +import com.google.protobuf.util.Durations; +import com.google.rpc.Code; +import com.google.rpc.RetryInfo; +import io.grpc.Deadline; +import io.grpc.Status; +import java.time.Duration; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Deque; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; +import java.util.concurrent.TimeUnit; +import java.util.function.BiConsumer; +import javax.annotation.Nullable; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +@Timeout(30) +class MutateRowsRetryingVRpcTest { + + private ManualTimer timer; + private ScriptedAttemptFactory attemptFactory; + private VRpcTracer noopTracer; + + @BeforeEach + void setUp() { + timer = new ManualTimer(); + attemptFactory = new ScriptedAttemptFactory(); + noopTracer = NoopMetrics.NoopVrpcTracer.INSTANCE; + } + + @Test + void allEntriesSucceed_noRetryAndBalancedTracer() throws Exception { + VRpcTracer tracer = mock(VRpcTracer.class); + SessionMutateRowsResponse okResponse = + SessionMutateRowsResponse.newBuilder() + .addEntries(respEntry(0, Code.OK_VALUE, "")) + .addEntries(respEntry(1, Code.OK_VALUE, "")) + .build(); + attemptFactory.enqueueResponse(okResponse); + + MutateRowsRetryingVRpc retrying = new MutateRowsRetryingVRpc(attemptFactory::next, timer); + UnaryResponseFuture future = new UnaryResponseFuture<>(); + SessionMutateRowsRequest req = + SessionMutateRowsRequest.newBuilder() + .addEntries(idempotentReqEntry("k0")) + .addEntries(idempotentReqEntry("k1")) + .build(); + + retrying.start( + req, VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), true, tracer), future); + + assertThat(future.get()).isEqualTo(okResponse); + assertThat(attemptFactory.requests).containsExactly(req); + assertThat(timer.pendingCount()).isEqualTo(0); + + verify(tracer, times(1)).onOperationStart(); + verify(tracer, times(1)).onAttemptStart(req); + verify(tracer, times(1)).onResponseReceived(); + verify(tracer, times(1)).onAttemptFinish(any()); + verify(tracer, times(1)).onOperationFinish(any()); + } + + @Test + void partialTransientFailure_retriesOnlyFailedIdempotentEntriesAndRemapsOriginalIndices() + throws Exception { + RateLimitInfo latestRateLimit = + RateLimitInfo.newBuilder().setFactor(0.9).setPeriod(Durations.fromSeconds(10)).build(); + + // Attempt 1 (k0, k1, k2, k3): + // 0 -> OK, 1 -> UNAVAILABLE (retry), 2 -> INVALID_ARGUMENT (permanent), 3 -> DEADLINE_EXCEEDED + attemptFactory.enqueueResponse( + SessionMutateRowsResponse.newBuilder() + .addEntries(respEntry(0, Code.OK_VALUE, "")) + .addEntries(respEntry(1, Code.UNAVAILABLE_VALUE, "transient-1")) + .addEntries(respEntry(2, Code.INVALID_ARGUMENT_VALUE, "bad-family")) + .addEntries(respEntry(3, Code.DEADLINE_EXCEEDED_VALUE, "transient-3")) + .build()); + + // Attempt 2 (k1 at index 0, k3 at index 1): + // 0 (k1) -> DEADLINE_EXCEEDED (retry again), 1 (k3) -> OK + attemptFactory.enqueueResponse( + SessionMutateRowsResponse.newBuilder() + .addEntries(respEntry(0, Code.DEADLINE_EXCEEDED_VALUE, "still-slow")) + .addEntries(respEntry(1, Code.OK_VALUE, "")) + .build()); + + // Attempt 3 (k1 at index 0): + // 0 (k1) -> OK, with RateLimitInfo + attemptFactory.enqueueResponse( + SessionMutateRowsResponse.newBuilder() + .addEntries(respEntry(0, Code.OK_VALUE, "")) + .setRateLimitInfo(latestRateLimit) + .build()); + + MutateRowsRetryingVRpc retrying = + new MutateRowsRetryingVRpc( + attemptFactory::next, + timer, + MutateRowsRetryingVRpc.DEFAULT_RETRYABLE_ENTRY_CODES, + Duration.ofMillis(10), + 2.0, + Duration.ofSeconds(1)); + UnaryResponseFuture future = new UnaryResponseFuture<>(); + SessionMutateRowsRequest req = + SessionMutateRowsRequest.newBuilder() + .addEntries(idempotentReqEntry("k0")) + .addEntries(idempotentReqEntry("k1")) + .addEntries(idempotentReqEntry("k2")) + .addEntries(idempotentReqEntry("k3")) + .build(); + + retrying.start( + req, VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), true, noopTracer), future); + + assertThat(future.isDone()).isFalse(); + assertThat(attemptFactory.requests).hasSize(1); + assertThat(timer.lastScheduledDelayNanos()).isEqualTo(Duration.ofMillis(10).toNanos()); + + // Fire backoff for Attempt 2 + timer.fireNext(); + assertThat(future.isDone()).isFalse(); + assertThat(attemptFactory.requests).hasSize(2); + assertThat(rowKeysOf(attemptFactory.requests.get(1))).containsExactly("k1", "k3").inOrder(); + assertThat(timer.lastScheduledDelayNanos()).isEqualTo(Duration.ofMillis(20).toNanos()); + + // Fire backoff for Attempt 3 + timer.fireNext(); + assertThat(future.isDone()).isTrue(); + assertThat(attemptFactory.requests).hasSize(3); + assertThat(rowKeysOf(attemptFactory.requests.get(2))).containsExactly("k1"); + + SessionMutateRowsResponse merged = future.get(); + assertThat(merged.getRateLimitInfo()).isEqualTo(latestRateLimit); + assertThat(merged.getEntriesList()) + .containsExactly( + respEntry(0, Code.OK_VALUE, ""), + respEntry(1, Code.OK_VALUE, ""), + respEntry(2, Code.INVALID_ARGUMENT_VALUE, "bad-family"), + respEntry(3, Code.OK_VALUE, "")) + .inOrder(); + } + + @Test + void nonIdempotentEntries_notRetriedOnTransientEntryFailure() throws Exception { + // 0: idempotent, 1: AddToCell (non-idempotent), 2: SetCell ts=-1 (non-idempotent) + attemptFactory.enqueueResponse( + SessionMutateRowsResponse.newBuilder() + .addEntries(respEntry(0, Code.UNAVAILABLE_VALUE, "transient-0")) + .addEntries(respEntry(1, Code.UNAVAILABLE_VALUE, "transient-add")) + .addEntries(respEntry(2, Code.UNAVAILABLE_VALUE, "transient-server-ts")) + .build()); + attemptFactory.enqueueResponse( + SessionMutateRowsResponse.newBuilder().addEntries(respEntry(0, Code.OK_VALUE, "")).build()); + + MutateRowsRetryingVRpc retrying = new MutateRowsRetryingVRpc(attemptFactory::next, timer); + UnaryResponseFuture future = new UnaryResponseFuture<>(); + SessionMutateRowsRequest req = + SessionMutateRowsRequest.newBuilder() + .addEntries(idempotentReqEntry("k0")) + .addEntries(nonIdempotentAddReqEntry("k1")) + .addEntries(nonIdempotentServerTsReqEntry("k2")) + .build(); + + retrying.start( + req, + VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), false, noopTracer), + future); + + timer.fireNext(); + + assertThat(attemptFactory.requests).hasSize(2); + assertThat(rowKeysOf(attemptFactory.requests.get(1))).containsExactly("k0"); + + SessionMutateRowsResponse merged = future.get(); + assertThat(merged.getEntriesList()) + .containsExactly( + respEntry(0, Code.OK_VALUE, ""), + respEntry(1, Code.UNAVAILABLE_VALUE, "transient-add"), + respEntry(2, Code.UNAVAILABLE_VALUE, "transient-server-ts")) + .inOrder(); + } + + @Test + void missingEntryInResponse_recordedAsPermanentInternalError() throws Exception { + // Entry 1 is omitted by server; entry 2 has transient error and is retried. + attemptFactory.enqueueResponse( + SessionMutateRowsResponse.newBuilder() + .addEntries(respEntry(0, Code.OK_VALUE, "")) + .addEntries(respEntry(2, Code.UNAVAILABLE_VALUE, "retry-2")) + .build()); + attemptFactory.enqueueResponse( + SessionMutateRowsResponse.newBuilder().addEntries(respEntry(0, Code.OK_VALUE, "")).build()); + + MutateRowsRetryingVRpc retrying = new MutateRowsRetryingVRpc(attemptFactory::next, timer); + UnaryResponseFuture future = new UnaryResponseFuture<>(); + SessionMutateRowsRequest req = + SessionMutateRowsRequest.newBuilder() + .addEntries(idempotentReqEntry("k0")) + .addEntries(idempotentReqEntry("k1")) + .addEntries(idempotentReqEntry("k2")) + .build(); + + retrying.start( + req, VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), true, noopTracer), future); + + timer.fireNext(); + assertThat(rowKeysOf(attemptFactory.requests.get(1))).containsExactly("k2"); + + SessionMutateRowsResponse merged = future.get(); + assertThat(merged.getEntriesList()) + .containsExactly( + respEntry(0, Code.OK_VALUE, ""), + respEntry(1, Code.INTERNAL_VALUE, "Missing entry response for entry 1"), + respEntry(2, Code.OK_VALUE, "")) + .inOrder(); + } + + @Test + void wholeVRpcFailureAfterPartialProgress_preservesPriorOutcomes() throws Exception { + // Attempt 1 resolves entry 0 (OK) and entry 1 (INVALID_ARGUMENT); entry 2 needs retry. + attemptFactory.enqueueResponse( + SessionMutateRowsResponse.newBuilder() + .addEntries(respEntry(0, Code.OK_VALUE, "")) + .addEntries(respEntry(1, Code.INVALID_ARGUMENT_VALUE, "bad-arg")) + .addEntries(respEntry(2, Code.UNAVAILABLE_VALUE, "retry-me")) + .build()); + // Attempt 2 fails at the whole-vRPC level with non-retryable PERMISSION_DENIED. + ErrorResponse serverError = errorResponse(Code.PERMISSION_DENIED_VALUE, "token revoked", null); + attemptFactory.enqueueClose(VRpcResult.createServerError(serverError)); + + MutateRowsRetryingVRpc retrying = new MutateRowsRetryingVRpc(attemptFactory::next, timer); + UnaryResponseFuture future = new UnaryResponseFuture<>(); + SessionMutateRowsRequest req = + SessionMutateRowsRequest.newBuilder() + .addEntries(idempotentReqEntry("k0")) + .addEntries(idempotentReqEntry("k1")) + .addEntries(idempotentReqEntry("k2")) + .build(); + + retrying.start( + req, VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), true, noopTracer), future); + timer.fireNext(); + + SessionMutateRowsResponse merged = future.get(); + assertThat(merged.getEntriesList()) + .containsExactly( + respEntry(0, Code.OK_VALUE, ""), + respEntry(1, Code.INVALID_ARGUMENT_VALUE, "bad-arg"), + respEntry(2, Code.PERMISSION_DENIED_VALUE, "token revoked")) + .inOrder(); + } + + @Test + void wholeVRpcFailureOnInitialAttempt_retriesRetryInfoAndPropagatesNonRetryableError() + throws Exception { + ErrorResponse retryableServerError = + errorResponse( + Code.UNAVAILABLE_VALUE, + "temporarily unavailable", + RetryInfo.newBuilder().setRetryDelay(Durations.fromMillis(15)).build()); + attemptFactory.enqueueClose(VRpcResult.createServerError(retryableServerError)); + SessionMutateRowsResponse okResponse = + SessionMutateRowsResponse.newBuilder().addEntries(respEntry(0, Code.OK_VALUE, "")).build(); + attemptFactory.enqueueResponse(okResponse); + + MutateRowsRetryingVRpc retrying = new MutateRowsRetryingVRpc(attemptFactory::next, timer); + UnaryResponseFuture future = new UnaryResponseFuture<>(); + SessionMutateRowsRequest req = + SessionMutateRowsRequest.newBuilder().addEntries(idempotentReqEntry("k0")).build(); + + retrying.start( + req, VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), true, noopTracer), future); + assertThat(timer.lastScheduledDelayNanos()).isEqualTo(Duration.ofMillis(15).toNanos()); + timer.fireNext(); + assertThat(future.get()).isEqualTo(okResponse); + + // Without RetryInfo on initial attempt, server error is not retried and propagates + // VRpcException + ScriptedAttemptFactory nonRetryFactory = new ScriptedAttemptFactory(); + ErrorResponse terminalError = errorResponse(Code.PERMISSION_DENIED_VALUE, "denied", null); + nonRetryFactory.enqueueClose(VRpcResult.createServerError(terminalError)); + + MutateRowsRetryingVRpc nonRetrying = new MutateRowsRetryingVRpc(nonRetryFactory::next, timer); + UnaryResponseFuture failFuture = new UnaryResponseFuture<>(); + nonRetrying.start( + req, + VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), true, noopTracer), + failFuture); + + ExecutionException ex = assertThrows(ExecutionException.class, failFuture::get); + assertThat(ex.getCause()).isInstanceOf(VRpcException.class); + assertThat(((VRpcException) ex.getCause()).getStatus().getCode()) + .isEqualTo(Status.Code.PERMISSION_DENIED); + } + + @Test + void transportFailure_checksActiveEntriesIdempotency() throws Exception { + SessionMutateRowsRequest mixedReq = + SessionMutateRowsRequest.newBuilder() + .addEntries(idempotentReqEntry("k0")) + .addEntries(nonIdempotentAddReqEntry("k1")) + .build(); + + // Case 1: Initial attempt has a non-idempotent entry active -> TRANSPORT_FAILURE is NOT retried + ScriptedAttemptFactory factory1 = new ScriptedAttemptFactory(); + factory1.enqueueClose( + VRpcResult.createLocalTransportError(Status.UNAVAILABLE.withDescription("rst"))); + MutateRowsRetryingVRpc retrying1 = new MutateRowsRetryingVRpc(factory1::next, timer); + UnaryResponseFuture f1 = new UnaryResponseFuture<>(); + retrying1.start( + mixedReq, + VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), false, noopTracer), + f1); + assertThrows(ExecutionException.class, f1::get); + assertThat(factory1.requests).hasSize(1); + + // Case 2: Attempt 1 finishes non-idempotent k1 (OK) and leaves only idempotent k0 active; + // Attempt 2 hits TRANSPORT_FAILURE -> safely retried on Attempt 3! + attemptFactory.enqueueResponse( + SessionMutateRowsResponse.newBuilder() + .addEntries(respEntry(0, Code.UNAVAILABLE_VALUE, "retry-k0")) + .addEntries(respEntry(1, Code.OK_VALUE, "")) + .build()); + attemptFactory.enqueueClose( + VRpcResult.createLocalTransportError(Status.UNAVAILABLE.withDescription("rst"))); + attemptFactory.enqueueResponse( + SessionMutateRowsResponse.newBuilder().addEntries(respEntry(0, Code.OK_VALUE, "")).build()); + + MutateRowsRetryingVRpc retrying2 = new MutateRowsRetryingVRpc(attemptFactory::next, timer); + UnaryResponseFuture f2 = new UnaryResponseFuture<>(); + retrying2.start( + mixedReq, + VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), false, noopTracer), + f2); + timer.fireNext(); + + assertThat(attemptFactory.requests).hasSize(3); + assertThat(f2.get().getEntriesList()) + .containsExactly(respEntry(0, Code.OK_VALUE, ""), respEntry(1, Code.OK_VALUE, "")) + .inOrder(); + } + + @Test + void deadlineExhausted_stopsPartialRetriesAndSurfacesTransientErrors() throws Exception { + attemptFactory.enqueueResponse( + SessionMutateRowsResponse.newBuilder() + .addEntries(respEntry(0, Code.OK_VALUE, "")) + .addEntries(respEntry(1, Code.UNAVAILABLE_VALUE, "no-time-left")) + .build()); + + // Initial backoff (5s) exceeds 100ms deadline, so no retry is scheduled. + MutateRowsRetryingVRpc retrying = + new MutateRowsRetryingVRpc( + attemptFactory::next, + timer, + MutateRowsRetryingVRpc.DEFAULT_RETRYABLE_ENTRY_CODES, + Duration.ofSeconds(5), + 2.0, + Duration.ofSeconds(10)); + UnaryResponseFuture future = new UnaryResponseFuture<>(); + SessionMutateRowsRequest req = + SessionMutateRowsRequest.newBuilder() + .addEntries(idempotentReqEntry("k0")) + .addEntries(idempotentReqEntry("k1")) + .build(); + + retrying.start( + req, + VRpcCallContext.create(Deadline.after(100, TimeUnit.MILLISECONDS), true, noopTracer), + future); + + assertThat(attemptFactory.requests).hasSize(1); + assertThat(timer.pendingCount()).isEqualTo(0); + assertThat(future.get().getEntriesList()) + .containsExactly( + respEntry(0, Code.OK_VALUE, ""), respEntry(1, Code.UNAVAILABLE_VALUE, "no-time-left")) + .inOrder(); + } + + @Test + void cancelAndTimerStopInScheduledState_transitionToCancelled() { + attemptFactory.enqueueResponse( + SessionMutateRowsResponse.newBuilder() + .addEntries(respEntry(0, Code.UNAVAILABLE_VALUE, "retry")) + .build()); + + MutateRowsRetryingVRpc retrying = new MutateRowsRetryingVRpc(attemptFactory::next, timer); + UnaryResponseFuture future = new UnaryResponseFuture<>(); + SessionMutateRowsRequest req = + SessionMutateRowsRequest.newBuilder().addEntries(idempotentReqEntry("k0")).build(); + + retrying.start( + req, VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), true, noopTracer), future); + assertThat(timer.pendingCount()).isEqualTo(1); + + retrying.cancel("caller cancelled", null); + assertThat(retrying.isDone()).isTrue(); + ExecutionException ex = assertThrows(ExecutionException.class, future::get); + assertThat(((VRpcException) ex.getCause()).getStatus().getCode()) + .isEqualTo(Status.Code.CANCELLED); + + // Also verify timer.stop() drives a pending Scheduled retry to CANCELLED + ManualTimer stopTimer = new ManualTimer(); + ScriptedAttemptFactory stopFactory = new ScriptedAttemptFactory(); + stopFactory.enqueueResponse( + SessionMutateRowsResponse.newBuilder() + .addEntries(respEntry(0, Code.UNAVAILABLE_VALUE, "retry")) + .build()); + MutateRowsRetryingVRpc stopRetrying = new MutateRowsRetryingVRpc(stopFactory::next, stopTimer); + UnaryResponseFuture stopFuture = new UnaryResponseFuture<>(); + stopRetrying.start( + req, + VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), true, noopTracer), + stopFuture); + + stopTimer.stop(); + assertThat(stopRetrying.isDone()).isTrue(); + ExecutionException stopEx = assertThrows(ExecutionException.class, stopFuture::get); + assertThat(((VRpcException) stopEx.getCause()).getStatus().getCode()) + .isEqualTo(Status.Code.CANCELLED); + } + + private static SessionMutateRowsRequest.Entry idempotentReqEntry(String rowKey) { + return SessionMutateRowsRequest.Entry.newBuilder() + .setKey(ByteString.copyFromUtf8(rowKey)) + .addMutations( + Mutation.newBuilder() + .setSetCell(Mutation.SetCell.newBuilder().setTimestampMicros(1000L))) + .build(); + } + + private static SessionMutateRowsRequest.Entry nonIdempotentAddReqEntry(String rowKey) { + return SessionMutateRowsRequest.Entry.newBuilder() + .setKey(ByteString.copyFromUtf8(rowKey)) + .addMutations(Mutation.newBuilder().setAddToCell(Mutation.AddToCell.getDefaultInstance())) + .build(); + } + + private static SessionMutateRowsRequest.Entry nonIdempotentServerTsReqEntry(String rowKey) { + return SessionMutateRowsRequest.Entry.newBuilder() + .setKey(ByteString.copyFromUtf8(rowKey)) + .addMutations( + Mutation.newBuilder().setSetCell(Mutation.SetCell.newBuilder().setTimestampMicros(-1L))) + .build(); + } + + private static SessionMutateRowsResponse.Entry respEntry(int index, int code, String message) { + SessionMutateRowsResponse.Entry.Builder builder = + SessionMutateRowsResponse.Entry.newBuilder().setIndex(index); + builder.getStatusBuilder().setCode(code).setMessage(message); + return builder.build(); + } + + private static ErrorResponse errorResponse( + int code, String message, @Nullable RetryInfo retryInfo) { + ErrorResponse.Builder builder = ErrorResponse.newBuilder(); + builder.getStatusBuilder().setCode(code).setMessage(message); + if (retryInfo != null) { + builder.setRetryInfo(retryInfo); + } + return builder.build(); + } + + private static List rowKeysOf(SessionMutateRowsRequest req) { + List keys = new ArrayList<>(req.getEntriesCount()); + for (SessionMutateRowsRequest.Entry entry : req.getEntriesList()) { + keys.add(entry.getKey().toStringUtf8()); + } + return keys; + } + + static final class ScriptedAttemptFactory { + final List requests = new ArrayList<>(); + private final Deque< + BiConsumer>> + actions = new ArrayDeque<>(); + + void enqueueResponse(SessionMutateRowsResponse response) { + actions.addLast( + (req, listener) -> { + listener.onMessage(response); + listener.onClose(VRpcResult.createServerOk(VirtualRpcResponse.getDefaultInstance())); + }); + } + + void enqueueClose(VRpcResult result) { + actions.addLast((req, listener) -> listener.onClose(result)); + } + + VRpc next() { + BiConsumer> action = + actions.removeFirst(); + return new VRpc() { + private boolean done = false; + + @Override + public void start( + SessionMutateRowsRequest req, + VRpcCallContext ctx, + VRpcListener listener) { + requests.add(req); + done = true; + ctx.getExecutor().execute(() -> action.accept(req, listener)); + } + + @Override + public void cancel(@Nullable String message, @Nullable Throwable cause) { + done = true; + } + + @Override + public boolean isDone() { + return done; + } + + @Override + public void requestNext() { + throw new UnsupportedOperationException(); + } + }; + } + } + + static final class ManualTimer implements BigtableTimer { + private final Deque scheduled = new ArrayDeque<>(); + private final Set stopHooks = new LinkedHashSet<>(); + private boolean stopped = false; + private long lastDelayNanos = 0L; + + @Override + public BigtableTimer.Timeout newTimeout( + Runnable task, Executor executor, long delay, TimeUnit unit) { + if (stopped) { + throw new IllegalStateException("Timer is stopped"); + } + lastDelayNanos = unit.toNanos(delay); + ManualTimeout timeout = new ManualTimeout(task, executor); + scheduled.addLast(timeout); + return timeout; + } + + @Override + public void stop() { + if (stopped) { + return; + } + stopped = true; + scheduled.clear(); + List hooks = new ArrayList<>(stopHooks); + stopHooks.clear(); + for (Runnable hook : hooks) { + hook.run(); + } + } + + @Override + public Registration onStop(Runnable hook) { + if (stopped) { + hook.run(); + return () -> {}; + } + stopHooks.add(hook); + return () -> stopHooks.remove(hook); + } + + int pendingCount() { + return (int) scheduled.stream().filter(t -> !t.cancelled).count(); + } + + long lastScheduledDelayNanos() { + return lastDelayNanos; + } + + void fireNext() { + while (!scheduled.isEmpty()) { + ManualTimeout next = scheduled.removeFirst(); + if (!next.cancelled) { + next.executor.execute(next.task); + return; + } + } + throw new IllegalStateException("No pending timeout to fire"); + } + + private static final class ManualTimeout implements BigtableTimer.Timeout { + private final Runnable task; + private final Executor executor; + private boolean cancelled = false; + + ManualTimeout(Runnable task, Executor executor) { + this.task = task; + this.executor = executor; + } + + @Override + public boolean cancel() { + if (cancelled) { + return false; + } + cancelled = true; + return true; + } + + @Override + public boolean isCancelled() { + return cancelled; + } + } + } +} diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/SessionMutateRowsRateLimiterTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/SessionMutateRowsRateLimiterTest.java new file mode 100644 index 000000000000..49038390e08c --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/SessionMutateRowsRateLimiterTest.java @@ -0,0 +1,386 @@ +/* + * Copyright 2026 Google LLC + * + * 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 + * + * https://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.data.v2.internal.middleware; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; + +import com.google.bigtable.v2.ErrorResponse; +import com.google.bigtable.v2.SessionMutateRowsRequest; +import com.google.bigtable.v2.SessionMutateRowsResponse; +import com.google.bigtable.v2.SessionMutateRowsResponse.RateLimitInfo; +import com.google.bigtable.v2.VirtualRpcResponse; +import com.google.cloud.bigtable.data.v2.internal.api.UnaryResponseFuture; +import com.google.cloud.bigtable.data.v2.internal.api.VRpcException; +import com.google.cloud.bigtable.data.v2.internal.csm.NoopMetrics; +import com.google.cloud.bigtable.data.v2.internal.csm.tracers.VRpcTracer; +import com.google.cloud.bigtable.data.v2.internal.middleware.MutateRowsRetryingVRpcTest.ManualTimer; +import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc.VRpcCallContext; +import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc.VRpcListener; +import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc.VRpcResult; +import com.google.common.base.Ticker; +import com.google.protobuf.util.Durations; +import com.google.rpc.Code; +import io.grpc.Deadline; +import io.grpc.Status; +import java.time.Duration; +import java.util.Optional; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import javax.annotation.Nullable; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +@Timeout(30) +class SessionMutateRowsRateLimiterTest { + + private FakeTicker ticker; + private SessionMutateRowsRateLimiter limiter; + private VRpcTracer tracer; + + @BeforeEach + void setUp() { + ticker = new FakeTicker(); + limiter = new SessionMutateRowsRateLimiter(SessionMutateRowsRateLimiter.DEFAULT_QPS, ticker); + tracer = mock(VRpcTracer.class); + } + + @Test + void initialState_startsDisabledAtDefaultQpsWithZeroDelay() { + assertThat(limiter.isEnabled()).isFalse(); + assertThat(limiter.getRate()).isEqualTo(10.0); + assertThat(limiter.tryAcquireDelay(Duration.ofMillis(1))).hasValue(Duration.ZERO); + } + + @Test + void validRateLimitInfo_enablesLimiterUpdatesRateAndRecordsMetrics() { + SessionMutateRowsResponse response = + responseWithRateLimitInfo(/* factor= */ 0.8, /* periodSeconds= */ 10); + + limiter.onResponse(response, tracer); + + assertThat(limiter.isEnabled()).isTrue(); + assertThat(limiter.getRate()).isWithin(0.001).of(8.0); + verify(tracer).recordBatchWriteFlowControlTargetQps(8.0); + verify(tracer).recordBatchWriteFlowControlFactor(0.8, Status.Code.OK, true); + } + + @Test + void factorAndQpsBounds_areClamped() { + // Factor 0.1 is clamped to MIN_FACTOR (0.7) -> 10.0 * 0.7 = 7.0 + limiter.onResponse(responseWithRateLimitInfo(0.1, 10), tracer); + assertThat(limiter.getRate()).isWithin(0.001).of(7.0); + verify(tracer).recordBatchWriteFlowControlFactor(0.7, Status.Code.OK, true); + + // Advance past period; factor 5.0 is clamped to MAX_FACTOR (1.3) -> 7.0 * 1.3 = 9.1 + ticker.advance(Duration.ofSeconds(11)); + limiter.onResponse(responseWithRateLimitInfo(5.0, 10), tracer); + assertThat(limiter.getRate()).isWithin(0.001).of(9.1); + verify(tracer).recordBatchWriteFlowControlFactor(1.3, Status.Code.OK, true); + + // Clamp to MIN_QPS (0.1) + limiter.setRate(0.12); + ticker.advance(Duration.ofSeconds(11)); + limiter.onResponse(responseWithRateLimitInfo(0.7, 10), tracer); + assertThat(limiter.getRate()).isWithin(0.0001).of(SessionMutateRowsRateLimiter.MIN_QPS); + + // Clamp to MAX_QPS (100_000) + limiter.setRate(90_000.0); + ticker.advance(Duration.ofSeconds(11)); + limiter.onResponse(responseWithRateLimitInfo(1.3, 10), tracer); + assertThat(limiter.getRate()).isWithin(0.1).of(SessionMutateRowsRateLimiter.MAX_QPS); + } + + @Test + void periodCooldown_blocksRateUpdateAndDisableWithinPeriodAndAllowsEnable() { + // Apply initial update with 10s period: 10.0 * 0.8 = 8.0 + limiter.onResponse(responseWithRateLimitInfo(0.8, 10), tracer); + assertThat(limiter.isEnabled()).isTrue(); + assertThat(limiter.getRate()).isWithin(0.001).of(8.0); + + // Within 10s period: another RateLimitInfo does NOT update QPS, records applied=false + ticker.advance(Duration.ofSeconds(5)); + VRpcTracer secondTracer = mock(VRpcTracer.class); + limiter.onResponse(responseWithRateLimitInfo(0.7, 10), secondTracer); + assertThat(limiter.getRate()).isWithin(0.001).of(8.0); + verify(secondTracer).recordBatchWriteFlowControlFactor(0.7, Status.Code.OK, false); + verify(secondTracer, never()).recordBatchWriteFlowControlTargetQps(8.0); + + // Within 10s period: missing RateLimitInfo does NOT disable limiter + limiter.onResponse(SessionMutateRowsResponse.getDefaultInstance(), secondTracer); + assertThat(limiter.isEnabled()).isTrue(); + + // If manually disabled while still within period, valid RateLimitInfo re-enables immediately + // without updating QPS + limiter.setEnabled(false); + limiter.onResponse(responseWithRateLimitInfo(0.7, 10), secondTracer); + assertThat(limiter.isEnabled()).isTrue(); + assertThat(limiter.getRate()).isWithin(0.001).of(8.0); + + // After 10s period expires (> nextRateUpdateNanos), missing RateLimitInfo disables limiter + ticker.advance(Duration.ofSeconds(6)); + limiter.onResponse(SessionMutateRowsResponse.getDefaultInstance(), secondTracer); + assertThat(limiter.isEnabled()).isFalse(); + assertThat(limiter.getRate()).isWithin(0.001).of(8.0); + } + + @Test + void invalidRateLimitInfo_doesNotEnableOrUpdateRate() { + // Factor == 0 -> invalid + limiter.onResponse(responseWithRateLimitInfo(0.0, 10), tracer); + assertThat(limiter.isEnabled()).isFalse(); + assertThat(limiter.getRate()).isEqualTo(10.0); + + // Period == 0 -> invalid + limiter.onResponse(responseWithRateLimitInfo(0.8, 0), tracer); + assertThat(limiter.isEnabled()).isFalse(); + assertThat(limiter.getRate()).isEqualTo(10.0); + } + + @Test + void overloadResults_reduceQpsByMinFactor() { + // UNAVAILABLE -> 10.0 * 0.7 = 7.0 + limiter.onResult(VRpcResult.createLocalTransportError(Status.UNAVAILABLE), tracer); + assertThat(limiter.getRate()).isWithin(0.001).of(7.0); + verify(tracer).recordBatchWriteFlowControlFactor(0.7, Status.Code.UNAVAILABLE, true); + + // DEADLINE_EXCEEDED after period -> 7.0 * 0.7 = 4.9 + ticker.advance(Duration.ofSeconds(11)); + limiter.onResult(VRpcResult.createRejectedError(Status.DEADLINE_EXCEEDED), tracer); + assertThat(limiter.getRate()).isWithin(0.001).of(4.9); + verify(tracer).recordBatchWriteFlowControlFactor(0.7, Status.Code.DEADLINE_EXCEEDED, true); + + // RESOURCE_EXHAUSTED after period -> 4.9 * 0.7 = 3.43 + ticker.advance(Duration.ofSeconds(11)); + ErrorResponse.Builder exhausted = ErrorResponse.newBuilder(); + exhausted.getStatusBuilder().setCode(Code.RESOURCE_EXHAUSTED_VALUE); + limiter.onResult(VRpcResult.createServerError(exhausted.build()), tracer); + assertThat(limiter.getRate()).isWithin(0.001).of(3.43); + verify(tracer).recordBatchWriteFlowControlFactor(0.7, Status.Code.RESOURCE_EXHAUSTED, true); + + // Non-overload error (INVALID_ARGUMENT) does not change QPS + ticker.advance(Duration.ofSeconds(11)); + ErrorResponse.Builder invalidArg = ErrorResponse.newBuilder(); + invalidArg.getStatusBuilder().setCode(Code.INVALID_ARGUMENT_VALUE); + limiter.onResult(VRpcResult.createServerError(invalidArg.build()), tracer); + assertThat(limiter.getRate()).isWithin(0.001).of(3.43); + } + + @Test + void tryAcquireDelay_reservesPermitsSmoothlyAndDoesNotConsumePermitWhenExceedingMaxWait() { + SessionMutateRowsRateLimiter twoQpsLimiter = new SessionMutateRowsRateLimiter(2.0, ticker); + twoQpsLimiter.setEnabled(true); + + // First call acquires immediately (0ms wait) and pushes nextFreeTicketNanos +500ms + assertThat(twoQpsLimiter.tryAcquireDelay(Duration.ofSeconds(5))).hasValue(Duration.ZERO); + + // Second call immediately after gets 500ms delay and pushes nextFreeTicketNanos to +1000ms + assertThat(twoQpsLimiter.tryAcquireDelay(Duration.ofSeconds(5))) + .hasValue(Duration.ofMillis(500)); + + // Third call with maxWait=400ms cannot wait 1000ms -> returns empty WITHOUT consuming a permit + Optional rejected = twoQpsLimiter.tryAcquireDelay(Duration.ofMillis(400)); + assertThat(rejected).isEmpty(); + + // Subsequent call with maxWait=2s still sees 1000ms (proving rejected call burned no permit) + assertThat(twoQpsLimiter.tryAcquireDelay(Duration.ofSeconds(2))) + .hasValue(Duration.ofSeconds(1)); + } + + @Test + void rateLimitingVRpc_schedulesThrottleOnTimerWithoutBlockingAndUpdatesLimiter() + throws Exception { + ManualTimer timer = new ManualTimer(); + SessionMutateRowsRateLimiter twoQpsLimiter = new SessionMutateRowsRateLimiter(2.0, ticker); + twoQpsLimiter.setEnabled(true); + // Reserve the immediate permit so the next call requires a 500ms throttle wait + assertThat(twoQpsLimiter.tryAcquireDelay(Duration.ofSeconds(5))).hasValue(Duration.ZERO); + + AtomicBoolean delegateStarted = new AtomicBoolean(false); + SessionMutateRowsResponse serverResp = + responseWithRateLimitInfo(/* factor= */ 1.2, /* periodSeconds= */ 10); + VRpc delegate = + new VRpc() { + @Override + public void start( + SessionMutateRowsRequest req, + VRpcCallContext ctx, + VRpcListener listener) { + delegateStarted.set(true); + ctx.getExecutor() + .execute( + () -> { + listener.onMessage(serverResp); + listener.onClose( + VRpcResult.createServerOk(VirtualRpcResponse.getDefaultInstance())); + }); + } + + @Override + public void cancel(@Nullable String message, @Nullable Throwable cause) {} + + @Override + public boolean isDone() { + return delegateStarted.get(); + } + + @Override + public void requestNext() {} + }; + + RateLimitingVRpc rateLimitingVRpc = new RateLimitingVRpc(delegate, twoQpsLimiter, timer); + UnaryResponseFuture future = new UnaryResponseFuture<>(); + rateLimitingVRpc.start( + SessionMutateRowsRequest.getDefaultInstance(), + VRpcCallContext.create(Deadline.after(10, TimeUnit.SECONDS), true, tracer), + future); + + // Must not start delegate until timer fires 500ms throttle timeout + assertThat(delegateStarted.get()).isFalse(); + assertThat(timer.pendingCount()).isEqualTo(1); + assertThat(timer.lastScheduledDelayNanos()).isEqualTo(Duration.ofMillis(500).toNanos()); + + timer.fireNext(); + + assertThat(delegateStarted.get()).isTrue(); + assertThat(future.get()).isEqualTo(serverResp); + assertThat(twoQpsLimiter.getRate()).isWithin(0.001).of(2.4); + } + + @Test + void rateLimitingVRpc_failsFastWithDeadlineExceededWhenWaitExceedsRemainingDeadline() { + ManualTimer timer = new ManualTimer(); + SessionMutateRowsRateLimiter oneQpsLimiter = new SessionMutateRowsRateLimiter(1.0, ticker); + oneQpsLimiter.setEnabled(true); + // Consume initial permit so next call needs 1000ms wait + assertThat(oneQpsLimiter.tryAcquireDelay(Duration.ofSeconds(5))).hasValue(Duration.ZERO); + + AtomicBoolean delegateStarted = new AtomicBoolean(false); + VRpc delegate = + noopDelegate(delegateStarted); + + RateLimitingVRpc rateLimitingVRpc = new RateLimitingVRpc(delegate, oneQpsLimiter, timer); + UnaryResponseFuture future = new UnaryResponseFuture<>(); + rateLimitingVRpc.start( + SessionMutateRowsRequest.getDefaultInstance(), + VRpcCallContext.create( + Deadline.after(100, TimeUnit.MILLISECONDS), true, NoopMetrics.NoopVrpcTracer.INSTANCE), + future); + + assertThat(delegateStarted.get()).isFalse(); + assertThat(timer.pendingCount()).isEqualTo(0); + ExecutionException ex = assertThrows(ExecutionException.class, future::get); + assertThat(ex.getCause()).isInstanceOf(VRpcException.class); + assertThat(((VRpcException) ex.getCause()).getStatus().getCode()) + .isEqualTo(Status.Code.DEADLINE_EXCEEDED); + } + + @Test + void rateLimitingVRpc_cancelAndTimerStopWhileThrottled_closesCancelledWithoutStartingDelegate() { + ManualTimer timer = new ManualTimer(); + SessionMutateRowsRateLimiter oneQpsLimiter = new SessionMutateRowsRateLimiter(1.0, ticker); + oneQpsLimiter.setEnabled(true); + assertThat(oneQpsLimiter.tryAcquireDelay(Duration.ofSeconds(5))).hasValue(Duration.ZERO); + + AtomicBoolean delegateStarted = new AtomicBoolean(false); + RateLimitingVRpc vRpc1 = + new RateLimitingVRpc(noopDelegate(delegateStarted), oneQpsLimiter, timer); + UnaryResponseFuture f1 = new UnaryResponseFuture<>(); + vRpc1.start( + SessionMutateRowsRequest.getDefaultInstance(), + VRpcCallContext.create( + Deadline.after(10, TimeUnit.SECONDS), true, NoopMetrics.NoopVrpcTracer.INSTANCE), + f1); + assertThat(timer.pendingCount()).isEqualTo(1); + + vRpc1.cancel("cancelled while throttled", null); + assertThat(delegateStarted.get()).isFalse(); + assertThat(timer.pendingCount()).isEqualTo(0); + ExecutionException ex1 = assertThrows(ExecutionException.class, f1::get); + assertThat(((VRpcException) ex1.getCause()).getStatus().getCode()) + .isEqualTo(Status.Code.CANCELLED); + + // Also test timer.stop() while throttled + RateLimitingVRpc vRpc2 = + new RateLimitingVRpc(noopDelegate(delegateStarted), oneQpsLimiter, timer); + UnaryResponseFuture f2 = new UnaryResponseFuture<>(); + vRpc2.start( + SessionMutateRowsRequest.getDefaultInstance(), + VRpcCallContext.create( + Deadline.after(10, TimeUnit.SECONDS), true, NoopMetrics.NoopVrpcTracer.INSTANCE), + f2); + + timer.stop(); + assertThat(delegateStarted.get()).isFalse(); + ExecutionException ex2 = assertThrows(ExecutionException.class, f2::get); + assertThat(((VRpcException) ex2.getCause()).getStatus().getCode()) + .isEqualTo(Status.Code.CANCELLED); + } + + private static SessionMutateRowsResponse responseWithRateLimitInfo( + double factor, long periodSeconds) { + return SessionMutateRowsResponse.newBuilder() + .setRateLimitInfo( + RateLimitInfo.newBuilder() + .setFactor(factor) + .setPeriod(Durations.fromSeconds(periodSeconds))) + .build(); + } + + private static VRpc noopDelegate( + AtomicBoolean started) { + return new VRpc() { + @Override + public void start( + SessionMutateRowsRequest req, + VRpcCallContext ctx, + VRpcListener listener) { + started.set(true); + } + + @Override + public void cancel(@Nullable String message, @Nullable Throwable cause) {} + + @Override + public boolean isDone() { + return false; + } + + @Override + public void requestNext() {} + }; + } + + private static final class FakeTicker extends Ticker { + private final AtomicLong nanos = new AtomicLong(1_000_000_000L); + + @Override + public long read() { + return nanos.get(); + } + + void advance(Duration duration) { + nanos.addAndGet(duration.toNanos()); + } + } +} diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcDescriptorTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcDescriptorTest.java index fa86d4dcd4b1..93bbfdb33e6b 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcDescriptorTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcDescriptorTest.java @@ -24,7 +24,9 @@ import com.google.bigtable.v2.CheckAndMutateRowRequest; import com.google.bigtable.v2.Column; import com.google.bigtable.v2.Family; +import com.google.bigtable.v2.Idempotency; import com.google.bigtable.v2.MutateRowRequest; +import com.google.bigtable.v2.MutateRowsRequest; import com.google.bigtable.v2.Mutation; import com.google.bigtable.v2.ReadModifyWriteRowRequest; import com.google.bigtable.v2.ReadModifyWriteRule; @@ -34,6 +36,8 @@ import com.google.bigtable.v2.SessionCheckAndMutateRowRequest; import com.google.bigtable.v2.SessionCheckAndMutateRowResponse; import com.google.bigtable.v2.SessionMutateRowRequest; +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; @@ -42,6 +46,7 @@ import com.google.bigtable.v2.TableResponse; import com.google.protobuf.ByteString; import com.google.protobuf.InvalidProtocolBufferException; +import com.google.rpc.Status; import org.junit.jupiter.api.Test; class VRpcDescriptorTest { @@ -302,4 +307,104 @@ private static SessionReadModifyWriteRowRequest sampleReadModifyWriteRowRequest( .setIncrementAmount(1)) .build(); } + + @Test + void testMutateRows() throws InvalidProtocolBufferException { + assertThat(VRpcDescriptor.MUTATE_ROWS.getSessionDescriptor()) + .isEqualTo(VRpcDescriptor.TABLE_SESSION); + + TableResponse tableResp = + TableResponse.newBuilder() + .setMutateRows( + SessionMutateRowsResponse.newBuilder() + .addEntries( + SessionMutateRowsResponse.Entry.newBuilder() + .setIndex(0) + .setStatus(Status.newBuilder().setCode(0).build()))) + .build(); + assertThat(VRpcDescriptor.MUTATE_ROWS.decode(tableResp.toByteString())) + .isEqualTo(tableResp.getMutateRows()); + + SessionMutateRowsRequest req = sampleMutateRowsRequest(); + assertThat(TableRequest.parseFrom(VRpcDescriptor.MUTATE_ROWS.encode(req))) + .isEqualTo(TableRequest.newBuilder().setMutateRows(req).build()); + } + + @Test + void testMutateRowsAuthView() throws InvalidProtocolBufferException { + assertThat(VRpcDescriptor.MUTATE_ROWS_AUTH_VIEW.getSessionDescriptor()) + .isEqualTo(VRpcDescriptor.AUTHORIZED_VIEW_SESSION); + + AuthorizedViewResponse authViewResp = + AuthorizedViewResponse.newBuilder() + .setMutateRows( + SessionMutateRowsResponse.newBuilder() + .addEntries( + SessionMutateRowsResponse.Entry.newBuilder() + .setIndex(0) + .setStatus(Status.newBuilder().setCode(0).build()))) + .build(); + assertThat(VRpcDescriptor.MUTATE_ROWS_AUTH_VIEW.decode(authViewResp.toByteString())) + .isEqualTo(authViewResp.getMutateRows()); + + SessionMutateRowsRequest req = sampleMutateRowsRequest(); + assertThat(AuthorizedViewRequest.parseFrom(VRpcDescriptor.MUTATE_ROWS_AUTH_VIEW.encode(req))) + .isEqualTo(AuthorizedViewRequest.newBuilder().setMutateRows(req).build()); + } + + @Test + void testToLegacyProtoMutateRows() { + SessionMutateRowsRequest req = sampleMutateRowsRequest(); + MutateRowsRequest legacyReq = + (MutateRowsRequest) VRpcDescriptor.MUTATE_ROWS.toLegacyProto("table1", "app1", req); + assertThat(legacyReq.getTableName()).isEqualTo("table1"); + assertThat(legacyReq.getAuthorizedViewName()).isEmpty(); + assertThat(legacyReq.getAppProfileId()).isEqualTo("app1"); + assertThat(legacyReq.getEntriesCount()).isEqualTo(1); + assertThat(legacyReq.getEntries(0).getRowKey()).isEqualTo(req.getEntries(0).getKey()); + assertThat(legacyReq.getEntries(0).getMutationsList()) + .isEqualTo(req.getEntries(0).getMutationsList()); + } + + @Test + void testToLegacyProtoMutateRowsAuthView() { + SessionMutateRowsRequest req = sampleMutateRowsRequest(); + MutateRowsRequest legacyReq = + (MutateRowsRequest) + VRpcDescriptor.MUTATE_ROWS_AUTH_VIEW.toLegacyProto("view1", "app1", req); + assertThat(legacyReq.getAuthorizedViewName()).isEqualTo("view1"); + assertThat(legacyReq.getTableName()).isEmpty(); + assertThat(legacyReq.getAppProfileId()).isEqualTo("app1"); + assertThat(legacyReq.getEntriesCount()).isEqualTo(1); + assertThat(legacyReq.getEntries(0).getRowKey()).isEqualTo(req.getEntries(0).getKey()); + assertThat(legacyReq.getEntries(0).getMutationsList()) + .isEqualTo(req.getEntries(0).getMutationsList()); + } + + @Test + void testToLegacyProtoMutateRowsWithIdempotency() { + Idempotency idempotency = + Idempotency.newBuilder().setToken(ByteString.copyFromUtf8("token123")).build(); + SessionMutateRowsRequest req = + SessionMutateRowsRequest.newBuilder() + .addEntries( + SessionMutateRowsRequest.Entry.newBuilder() + .setKey(ByteString.copyFromUtf8("rowkey1")) + .addMutations(sampleSetCell()) + .setIdempotency(idempotency)) + .build(); + MutateRowsRequest legacyReq = + (MutateRowsRequest) VRpcDescriptor.MUTATE_ROWS.toLegacyProto("table1", "app1", req); + assertThat(legacyReq.getEntries(0).hasIdempotency()).isTrue(); + assertThat(legacyReq.getEntries(0).getIdempotency()).isEqualTo(idempotency); + } + + private static SessionMutateRowsRequest sampleMutateRowsRequest() { + return SessionMutateRowsRequest.newBuilder() + .addEntries( + SessionMutateRowsRequest.Entry.newBuilder() + .setKey(ByteString.copyFromUtf8("rowkey1")) + .addMutations(sampleSetCell())) + .build(); + } } diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScopeTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScopeTest.java index 190a03e59528..9bf93e3f4dab 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScopeTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScopeTest.java @@ -40,6 +40,7 @@ void noDiversionConfig_nonFallbackScopesIgnoreFlatSessionLoad() { assertThat(SessionScope.CHECK_AND_MUTATE_ROW.resolveSessionLoad(config)).isEqualTo(0f); assertThat(SessionScope.READ_MODIFY_WRITE_ROW.resolveSessionLoad(config)).isEqualTo(0f); + assertThat(SessionScope.MUTATE_ROWS.resolveSessionLoad(config)).isEqualTo(0f); } @Test @@ -94,7 +95,7 @@ void globalDiversionAppliesToEveryScope() { void unservedScope_doesNotCountAsDivertingTraffic() { SessionClientConfiguration config = SessionClientConfiguration.newBuilder() - .setSessionDiversionConfiguration(perScope("method:Bigtable.MutateRows", 1.0f)) + .setSessionDiversionConfiguration(perScope("method:Bigtable.ExecuteQuery", 1.0f)) .build(); for (SessionScope scope : SessionScope.values()) { diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/BulkMutationTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/BulkMutationTest.java index 498e3ee034c5..7c8591691790 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/BulkMutationTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/BulkMutationTest.java @@ -17,7 +17,9 @@ import static com.google.common.truth.Truth.assertThat; +import com.google.bigtable.v2.Idempotency; import com.google.bigtable.v2.MutateRowsRequest; +import com.google.bigtable.v2.SessionMutateRowsRequest; import com.google.cloud.bigtable.data.v2.internal.NameUtil; import com.google.cloud.bigtable.data.v2.internal.RequestContext; import com.google.protobuf.ByteString; @@ -314,4 +316,54 @@ public void testManyMutations() { } bulkMutation.add("key", mutation); } + + @Test + public void testGetTargetId() { + BulkMutation mutation = BulkMutation.create(TABLE_ID); + assertThat(mutation.getTargetId()).isEqualTo(TABLE_ID); + + AuthorizedViewId viewId = AuthorizedViewId.of(TABLE_ID, AUTHORIZED_VIEW_ID); + mutation = BulkMutation.create(viewId); + assertThat(mutation.getTargetId()).isEqualTo(viewId); + } + + @Test + public void testToSessionProto() { + BulkMutation m = + BulkMutation.create(TABLE_ID) + .add( + "key-a", + Mutation.create() + .setCell("fake-family1", "fake-qualifier1", 1_000, "fake-value1") + .setCell("fake-family2", "fake-qualifier2", 2_000, "fake-value2")) + .add( + ByteString.copyFromUtf8("key-b"), + Mutation.create().setCell("fake-family3", "fake-qualifier3", 3_000, "fake-value3")); + + SessionMutateRowsRequest sessionReq = m.toSessionProto(); + assertThat(sessionReq.getEntriesCount()).isEqualTo(2); + assertThat(sessionReq.getEntries(0).getKey()).isEqualTo(ByteString.copyFromUtf8("key-a")); + assertThat(sessionReq.getEntries(0).getMutationsCount()).isEqualTo(2); + assertThat(sessionReq.getEntries(1).getKey()).isEqualTo(ByteString.copyFromUtf8("key-b")); + assertThat(sessionReq.getEntries(1).getMutationsCount()).isEqualTo(1); + } + + @Test + public void testToSessionProtoWithIdempotency() { + Idempotency idempotency = + Idempotency.newBuilder().setToken(ByteString.copyFromUtf8("token123")).build(); + MutateRowsRequest proto = + MutateRowsRequest.newBuilder() + .setTableName("projects/p/instances/i/tables/t") + .addEntries( + MutateRowsRequest.Entry.newBuilder() + .setRowKey(ByteString.copyFromUtf8("key-a")) + .setIdempotency(idempotency)) + .build(); + BulkMutation m = BulkMutation.fromProto(proto); + SessionMutateRowsRequest sessionReq = m.toSessionProto(); + assertThat(sessionReq.getEntriesCount()).isEqualTo(1); + assertThat(sessionReq.getEntries(0).hasIdempotency()).isTrue(); + assertThat(sessionReq.getEntries(0).getIdempotency()).isEqualTo(idempotency); + } }