From e2111eeecd0dce3cb1a9e366fd068cc6baf052d9 Mon Sep 17 00:00:00 2001 From: Alex O'Neill Date: Tue, 13 Oct 2026 10:49:00 -0400 Subject: [PATCH 1/3] feat(bigtable): support SessionMutateRows on the session client --- .../v2/internal/api/AuthorizedViewAsync.java | 10 ++ .../internal/api/MaterializedViewAsync.java | 1 + .../data/v2/internal/api/TableAsync.java | 10 ++ .../data/v2/internal/api/TableBase.java | 31 ++++ .../data/v2/internal/compat/DisabledShim.java | 7 + .../data/v2/internal/compat/Shim.java | 4 + .../data/v2/internal/compat/ShimImpl.java | 16 ++ .../internal/compat/ops/MutateRowsShim.java | 77 +++++++++ .../v2/internal/session/VRpcDescriptor.java | 53 +++++++ .../data/v2/internal/util/SessionScope.java | 3 +- .../bigtable/data/v2/models/BulkMutation.java | 22 +++ .../data/v2/stub/EnhancedBigtableStub.java | 11 +- .../data/v2/internal/api/TableBaseTest.java | 34 ++++ .../compat/ops/MutateRowsShimTest.java | 149 ++++++++++++++++++ .../internal/session/VRpcDescriptorTest.java | 105 ++++++++++++ .../v2/internal/util/SessionScopeTest.java | 3 +- .../data/v2/models/BulkMutationTest.java | 52 ++++++ 17 files changed, 585 insertions(+), 3 deletions(-) create mode 100644 java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowsShim.java create mode 100644 java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowsShimTest.java 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..0b3ef63e30b8 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; @@ -80,6 +82,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, @@ -130,6 +133,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/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..086dd79bdf07 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; @@ -78,6 +80,7 @@ public static TableAsync createAndStart( VRpcDescriptor.MUTATE_ROW, VRpcDescriptor.CHECK_AND_MUTATE_ROW, VRpcDescriptor.READ_MODIFY_WRITE_ROW, + VRpcDescriptor.MUTATE_ROWS, featureFlags, clientInfo, configManager, @@ -139,4 +142,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..f8d6ec880927 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 @@ -15,12 +15,17 @@ */ package com.google.cloud.bigtable.data.v2.internal.api; +import static com.google.common.collect.ImmutableList.toImmutableList; + import com.google.bigtable.v2.CloseSessionRequest; import com.google.bigtable.v2.FeatureFlags; +import com.google.bigtable.v2.Mutation; import com.google.bigtable.v2.SessionCheckAndMutateRowRequest; 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; @@ -40,6 +45,7 @@ import com.google.cloud.bigtable.data.v2.internal.session.VRpcDescriptor; import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager; import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableList; import com.google.protobuf.Message; import io.grpc.CallOptions; import io.grpc.Context; @@ -60,6 +66,8 @@ class TableBase implements AutoCloseable { private final VRpcDescriptor< ?, SessionReadModifyWriteRowRequest, SessionReadModifyWriteRowResponse> readModifyWriteRowDescriptor; + private final VRpcDescriptor + mutateRowsDescriptor; static TableBase createAndStart( ReqT openReq, @@ -70,6 +78,7 @@ static TableBase createAndStart( checkAndMutateRowDescriptor, VRpcDescriptor readModifyWriteRowDescriptor, + VRpcDescriptor mutateRowsDescriptor, FeatureFlags featureFlags, ClientInfo clientInfo, ClientConfigurationManager configManager, @@ -108,6 +117,7 @@ static TableBase createAndStart( mutateRowDescriptor, checkAndMutateRowDescriptor, readModifyWriteRowDescriptor, + mutateRowsDescriptor, metrics, timer, userCallbackExecutor); @@ -122,6 +132,7 @@ static TableBase createAndStart( checkAndMutateRowDescriptor, VRpcDescriptor readModifyWriteRowDescriptor, + VRpcDescriptor mutateRowsDescriptor, Metrics metrics, BigtableTimer timer, Executor userCallbackExecutor) { @@ -130,6 +141,7 @@ static TableBase createAndStart( this.mutateRowDescriptor = mutateRowDescriptor; this.checkAndMutateRowDescriptor = checkAndMutateRowDescriptor; this.readModifyWriteRowDescriptor = readModifyWriteRowDescriptor; + this.mutateRowsDescriptor = mutateRowsDescriptor; this.metrics = metrics; this.timer = timer; this.userCallbackExecutor = userCallbackExecutor; @@ -202,4 +214,23 @@ public void readModifyWriteRow( new VOperationImpl<>(retry, Context.current(), userCallbackExecutor, tracer, deadline, false) .start(req, listener); } + + public void mutateRows( + SessionMutateRowsRequest req, + VRpcListener listener, + Deadline deadline) { + RetryingVRpc retry = + new RetryingVRpc<>(() -> sessionPool.newCall(mutateRowsDescriptor), timer); + ImmutableList mutations = + req.getEntriesList().stream() + .flatMap(entry -> entry.getMutationsList().stream()) + .collect(toImmutableList()); + boolean idempotent = Util.isIdempotent(mutations); + 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..beabeb1d69b4 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,6 +17,7 @@ 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; @@ -55,4 +56,10 @@ 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..2fbfde7a1e6e 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,6 +17,7 @@ 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; @@ -42,4 +43,7 @@ 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..5742f2906565 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,6 +52,7 @@ 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; @@ -110,6 +112,7 @@ public class ShimImpl implements Shim { private final MutateRowShim mutateRowShim; private final CheckAndMutateRowShim checkAndMutateRowShim; private final ReadModifyWriteRowShim readModifyWriteRowShim; + private final MutateRowsShim mutateRowsShim; public static Shim create( ClientInfo clientInfo, @@ -228,6 +231,7 @@ public ShimImpl( this.mutateRowShim = new MutateRowShim(sessionMap); this.checkAndMutateRowShim = new CheckAndMutateRowShim(sessionMap); this.readModifyWriteRowShim = new ReadModifyWriteRowShim(sessionMap); + this.mutateRowsShim = new MutateRowsShim(sessionMap); } /** @@ -435,4 +439,16 @@ public UnaryCallable decorateReadModifyWriteRow debugTagTracer, SessionScope.READ_MODIFY_WRITE_ROW); } + + @Override + public UnaryCallable decorateMutateRows( + UnaryCallable classic, UnaryCallSettings settings) { + return new DivertingUnaryCallable<>( + configManager, + classic, + mutateRowsShim, + 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/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..ad843a02c584 --- /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,77 @@ +/* + * 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.bigtable.v2.SessionMutateRowsRequest; +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.TableId; +import com.google.cloud.bigtable.data.v2.models.TargetId; +import io.grpc.Deadline; +import java.util.concurrent.CompletableFuture; + +public class MutateRowsShim implements UnaryShim { + + private final SessionPoolMap sessionMap; + + public MutateRowsShim(SessionPoolMap sessionMap) { + this.sessionMap = sessionMap; + } + + @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(); + + if (targetId instanceof TableId) { + return sessionMap.apply( + (TableId) targetId, + SessionKey.Permission.WRITE, + t -> t.mutateRows(innerReq, deadline).thenApply(r -> null)); + } + if (targetId instanceof AuthorizedViewId) { + return sessionMap.apply( + (AuthorizedViewId) targetId, + SessionKey.Permission.WRITE, + v -> v.mutateRows(innerReq, deadline).thenApply(r -> null)); + } + + CompletableFuture f = new CompletableFuture<>(); + f.completeExceptionally( + new UnsupportedOperationException("Unsupported targetId type: " + targetId)); + return f; + } +} 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..11f7cb18576d 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 @@ -203,11 +203,20 @@ public EnhancedBigtableStub( bulkMutateRowsCallable = createMutateRowsBaseCallable(); UnaryCallable bulkMutateRowsVoidCallable = new MutateRowsErrorConverterUnaryCallable(bulkMutateRowsCallable); - externalBulkMutateRowsCallable = + UnaryCallable maybePointWriteCallable = new MaybePointWriteCallable( bulkMutateRowsVoidCallable, createPointWriteCallable(bulkMutateRowsVoidCallable), requestContext); + externalBulkMutateRowsCallable = + bigtableClientContext + .getSessionShim() + .decorateMutateRows(maybePointWriteCallable, perOpSettings.bulkMutateRowsSettings) + .withDefaultCallContext( + bigtableClientContext + .getClientContext() + .getDefaultCallContext() + .withRetrySettings(perOpSettings.bulkMutateRowsSettings.getRetrySettings())); checkAndMutateRowCallable = createCheckAndMutateRowCallable(); readModifyWriteRowCallable = createReadModifyWriteRowCallable(); generateInitialChangeStreamPartitionsCallable = 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..1207a47da84e 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 @@ -25,6 +25,7 @@ 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.SessionReadModifyWriteRowRequest; import com.google.cloud.bigtable.data.v2.internal.csm.Metrics; import com.google.cloud.bigtable.data.v2.internal.csm.NoopMetrics; @@ -80,6 +81,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 +202,38 @@ 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(); + } + static class FakeSessionPool implements SessionPool { private FakeVRpc lastVRpc = null; 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..17ea5e6c7301 --- /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,149 @@ +/* + * 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.assertNull; +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.bigtable.v2.SessionMutateRowsResponse; +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.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.TableId; +import com.google.cloud.bigtable.data.v2.models.TargetId; +import io.grpc.Deadline; +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_returnsNull() throws Exception { + SessionMutateRowsResponse response = SessionMutateRowsResponse.getDefaultInstance(); + when(tableAsync.mutateRows(any(), any())) + .thenReturn(CompletableFuture.completedFuture(response)); + + BulkMutation request = BulkMutation.create(TableId.of("table1")); + Deadline deadline = Deadline.after(1, TimeUnit.SECONDS); + + assertNull(shim.call(request, deadline).get()); + } + + @Test + void call_authorizedViewId_returnsNull() throws Exception { + SessionMutateRowsResponse response = SessionMutateRowsResponse.getDefaultInstance(); + when(authViewAsync.mutateRows(any(), any())) + .thenReturn(CompletableFuture.completedFuture(response)); + + BulkMutation request = BulkMutation.create(AuthorizedViewId.of("table1", "view1")); + Deadline deadline = Deadline.after(1, TimeUnit.SECONDS); + + assertNull(shim.call(request, deadline).get()); + } + + @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); + } +} 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); + } } From 76335925c121a986154c2c3e2dfb8f33b998d871 Mon Sep 17 00:00:00 2001 From: Alex O'Neill Date: Tue, 27 Oct 2026 23:28:02 -0400 Subject: [PATCH 2/3] fix(bigtable): handle entry errors, partial retries, and RateLimitInfo in SessionMutateRows --- .../v2/internal/api/AuthorizedViewAsync.java | 35 +- .../bigtable/data/v2/internal/api/Client.java | 29 +- .../data/v2/internal/api/ClientSettings.java | 6 +- .../data/v2/internal/api/TableAsync.java | 33 +- .../data/v2/internal/api/TableBase.java | 99 ++- .../data/v2/internal/compat/DisabledShim.java | 6 +- .../data/v2/internal/compat/Shim.java | 6 +- .../data/v2/internal/compat/ShimImpl.java | 35 +- .../compat/ops/DivertingUnaryCallable.java | 4 + .../internal/compat/ops/MutateRowsShim.java | 181 ++++- .../data/v2/internal/csm/NoopMetrics.java | 7 + .../csm/tracers/CompositeVRpcTracer.java | 32 + .../csm/tracers/UserApiVRpcTracer.java | 18 + .../v2/internal/csm/tracers/VRpcTracer.java | 6 + .../internal/csm/tracers/VRpcTracerImpl.java | 14 + .../middleware/MutateRowsRetryingVRpc.java | 612 ++++++++++++++++ .../internal/middleware/RateLimitingVRpc.java | 216 ++++++ .../SessionMutateRowsRateLimiter.java | 262 +++++++ .../data/v2/stub/EnhancedBigtableStub.java | 57 +- .../data/v2/internal/api/TableBaseTest.java | 119 +++- .../data/v2/internal/compat/ShimImplTest.java | 18 + .../ops/DivertingUnaryCallableTest.java | 50 +- .../compat/ops/MutateRowsShimTest.java | 250 ++++++- .../csm/tracers/CompositeVRpcTracerTest.java | 14 + .../internal/csm/tracers/VRpcTracerTest.java | 59 ++ .../MutateRowsRetryingVRpcTest.java | 655 ++++++++++++++++++ .../SessionMutateRowsRateLimiterTest.java | 386 +++++++++++ 27 files changed, 3126 insertions(+), 83 deletions(-) create mode 100644 java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/MutateRowsRetryingVRpc.java create mode 100644 java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RateLimitingVRpc.java create mode 100644 java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/SessionMutateRowsRateLimiter.java create mode 100644 java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/MutateRowsRetryingVRpcTest.java create mode 100644 java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/SessionMutateRowsRateLimiterTest.java 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 0b3ef63e30b8..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 @@ -31,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; @@ -40,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 { @@ -58,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() @@ -92,7 +124,8 @@ static AuthorizedViewAsync createAndStart( metrics, timer, backgroundExecutor, - userCallbackExecutor); + userCallbackExecutor, + mutateRowsRateLimiter); return new AuthorizedViewAsync(base); } 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/TableAsync.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableAsync.java index 086dd79bdf07..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 @@ -32,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; @@ -41,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; @@ -57,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() @@ -90,7 +120,8 @@ public static TableAsync createAndStart( metrics, timer, backgroundExecutor, - userCallbackExecutor); + userCallbackExecutor, + mutateRowsRateLimiter); return new TableAsync(base); } 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 f8d6ec880927..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 @@ -15,11 +15,8 @@ */ package com.google.cloud.bigtable.data.v2.internal.api; -import static com.google.common.collect.ImmutableList.toImmutableList; - import com.google.bigtable.v2.CloseSessionRequest; import com.google.bigtable.v2.FeatureFlags; -import com.google.bigtable.v2.Mutation; import com.google.bigtable.v2.SessionCheckAndMutateRowRequest; import com.google.bigtable.v2.SessionCheckAndMutateRowResponse; import com.google.bigtable.v2.SessionMutateRowRequest; @@ -36,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; @@ -45,19 +46,21 @@ import com.google.cloud.bigtable.data.v2.internal.session.VRpcDescriptor; import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager; import com.google.common.annotations.VisibleForTesting; -import com.google.common.collect.ImmutableList; import com.google.protobuf.Message; import io.grpc.CallOptions; import io.grpc.Context; 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; @@ -89,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 = @@ -120,7 +165,8 @@ static TableBase createAndStart( mutateRowsDescriptor, metrics, timer, - userCallbackExecutor); + userCallbackExecutor, + mutateRowsRateLimiter); } @VisibleForTesting @@ -136,6 +182,33 @@ static TableBase createAndStart( 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; @@ -145,6 +218,7 @@ static TableBase createAndStart( this.metrics = metrics; this.timer = timer; this.userCallbackExecutor = userCallbackExecutor; + this.mutateRowsRateLimiter = mutateRowsRateLimiter; } @Override @@ -219,13 +293,16 @@ public void mutateRows( SessionMutateRowsRequest req, VRpcListener listener, Deadline deadline) { - RetryingVRpc retry = - new RetryingVRpc<>(() -> sessionPool.newCall(mutateRowsDescriptor), timer); - ImmutableList mutations = + 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() - .flatMap(entry -> entry.getMutationsList().stream()) - .collect(toImmutableList()); - boolean idempotent = Util.isIdempotent(mutations); + .allMatch(entry -> Util.isIdempotent(entry.getMutationsList())); VRpcTracer tracer = metrics.newTableTracer(sessionPool.getInfo(), mutateRowsDescriptor, deadline); 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 beabeb1d69b4..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 @@ -23,6 +23,7 @@ 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 { @@ -58,8 +59,9 @@ public UnaryCallable decorateReadModifyWriteRow } @Override - public UnaryCallable decorateMutateRows( - UnaryCallable classic, UnaryCallSettings settings) { + 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 2fbfde7a1e6e..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 @@ -23,6 +23,7 @@ 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 { @@ -44,6 +45,7 @@ UnaryCallable decorateReadModifyWriteRow( RowAdapter rowAdapter, UnaryCallSettings settings); - UnaryCallable decorateMutateRows( - UnaryCallable classic, 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 5742f2906565..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 @@ -58,7 +58,9 @@ 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; @@ -112,7 +114,22 @@ public class ShimImpl implements Shim { private final MutateRowShim mutateRowShim; private final CheckAndMutateRowShim checkAndMutateRowShim; private final ReadModifyWriteRowShim readModifyWriteRowShim; - private final MutateRowsShim mutateRowsShim; + + @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, @@ -150,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( @@ -158,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( @@ -176,7 +195,7 @@ public static Shim create( ClientConfigurationManager configManager = new ClientConfigurationManager( - configChannelProvider.updateFeatureFlags(Client.BASE_FEATURE_FLAGS), + configChannelProvider.updateFeatureFlags(baseFeatureFlags), clientInfo, configChannelProvider, metrics.getDebugTagTracer(), @@ -191,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(); @@ -231,7 +250,6 @@ public ShimImpl( this.mutateRowShim = new MutateRowShim(sessionMap); this.checkAndMutateRowShim = new CheckAndMutateRowShim(sessionMap); this.readModifyWriteRowShim = new ReadModifyWriteRowShim(sessionMap); - this.mutateRowsShim = new MutateRowsShim(sessionMap); } /** @@ -441,12 +459,13 @@ public UnaryCallable decorateReadModifyWriteRow } @Override - public UnaryCallable decorateMutateRows( - UnaryCallable classic, UnaryCallSettings settings) { + public UnaryCallable decorateMutateRows( + UnaryCallable classic, + UnaryCallSettings settings) { return new DivertingUnaryCallable<>( configManager, classic, - mutateRowsShim, + 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 index ad843a02c584..b89a908de93c 100644 --- 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 @@ -15,21 +15,50 @@ */ 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 { +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 @@ -51,27 +80,147 @@ public boolean supports(BulkMutation request) { } @Override - public CompletableFuture call(BulkMutation request, Deadline deadline) { + public CompletableFuture call(BulkMutation request, Deadline deadline) { TargetId targetId = request.getTargetId(); - SessionMutateRowsRequest innerReq = request.toSessionProto(); + CompletableFuture responseFuture; if (targetId instanceof TableId) { - return sessionMap.apply( - (TableId) targetId, - SessionKey.Permission.WRITE, - t -> t.mutateRows(innerReq, deadline).thenApply(r -> null)); + 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; } - if (targetId instanceof AuthorizedViewId) { - return sessionMap.apply( - (AuthorizedViewId) targetId, - SessionKey.Permission.WRITE, - v -> v.mutateRows(innerReq, deadline).thenApply(r -> null)); + + 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; + } } - CompletableFuture f = new CompletableFuture<>(); - f.completeExceptionally( - new UnsupportedOperationException("Unsupported targetId type: " + targetId)); - return f; + 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/stub/EnhancedBigtableStub.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index 11f7cb18576d..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 @@ -203,15 +203,11 @@ public EnhancedBigtableStub( bulkMutateRowsCallable = createMutateRowsBaseCallable(); UnaryCallable bulkMutateRowsVoidCallable = new MutateRowsErrorConverterUnaryCallable(bulkMutateRowsCallable); - UnaryCallable maybePointWriteCallable = - new MaybePointWriteCallable( - bulkMutateRowsVoidCallable, - createPointWriteCallable(bulkMutateRowsVoidCallable), - requestContext); externalBulkMutateRowsCallable = - bigtableClientContext - .getSessionShim() - .decorateMutateRows(maybePointWriteCallable, perOpSettings.bulkMutateRowsSettings) + new MaybePointWriteCallable( + bulkMutateRowsVoidCallable, + createPointWriteCallable(bulkMutateRowsVoidCallable), + requestContext) .withDefaultCallContext( bigtableClientContext .getClientContext() @@ -802,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 1207a47da84e..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,8 +17,10 @@ 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; @@ -26,17 +28,23 @@ 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; @@ -234,6 +242,113 @@ public void testMutateRowsNotIdempotent() { 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; @@ -275,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 index 17ea5e6c7301..024485a7036c 100644 --- 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 @@ -16,23 +16,49 @@ 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.assertNull; 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; @@ -111,27 +137,224 @@ void supports_unsupportedTargetId_returnsFalse() { } @Test - void call_tableId_returnsNull() throws Exception { - SessionMutateRowsResponse response = SessionMutateRowsResponse.getDefaultInstance(); + 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")); + 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); - assertNull(shim.call(request, deadline).get()); + MutateRowsAttemptResult result = shim.call(request, deadline).get(); + assertThat(result.getFailedMutations()).isEmpty(); + assertThat(result.getIsRetryable()).isFalse(); } @Test - void call_authorizedViewId_returnsNull() throws Exception { - SessionMutateRowsResponse response = SessionMutateRowsResponse.getDefaultInstance(); + 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")); + BulkMutation request = + BulkMutation.create(AuthorizedViewId.of("table1", "view1")) + .add("k0", Mutation.create().deleteRow()); Deadline deadline = Deadline.after(1, TimeUnit.SECONDS); - assertNull(shim.call(request, deadline).get()); + 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 @@ -140,10 +363,17 @@ void call_unsupportedTargetId_futureFailsWithUnsupportedOperationException() { BulkMutation request = BulkMutation.create(unsupported); Deadline deadline = Deadline.after(1, TimeUnit.SECONDS); - CompletableFuture future = shim.call(request, deadline); + 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()); + } + } +} From a698bd2ffe5a30d119260ba4b243d7c64e913c8a Mon Sep 17 00:00:00 2001 From: Alex O'Neill Date: Wed, 28 Oct 2026 10:55:53 -0400 Subject: [PATCH 3/3] Revert "fix(bigtable): handle entry errors, partial retries, and RateLimitInfo in SessionMutateRows" This reverts commit 76335925c121a986154c2c3e2dfb8f33b998d871. --- .../v2/internal/api/AuthorizedViewAsync.java | 35 +- .../bigtable/data/v2/internal/api/Client.java | 29 +- .../data/v2/internal/api/ClientSettings.java | 6 +- .../data/v2/internal/api/TableAsync.java | 33 +- .../data/v2/internal/api/TableBase.java | 99 +-- .../data/v2/internal/compat/DisabledShim.java | 6 +- .../data/v2/internal/compat/Shim.java | 6 +- .../data/v2/internal/compat/ShimImpl.java | 35 +- .../compat/ops/DivertingUnaryCallable.java | 4 - .../internal/compat/ops/MutateRowsShim.java | 181 +---- .../data/v2/internal/csm/NoopMetrics.java | 7 - .../csm/tracers/CompositeVRpcTracer.java | 32 - .../csm/tracers/UserApiVRpcTracer.java | 18 - .../v2/internal/csm/tracers/VRpcTracer.java | 6 - .../internal/csm/tracers/VRpcTracerImpl.java | 14 - .../middleware/MutateRowsRetryingVRpc.java | 612 ---------------- .../internal/middleware/RateLimitingVRpc.java | 216 ------ .../SessionMutateRowsRateLimiter.java | 262 ------- .../data/v2/stub/EnhancedBigtableStub.java | 57 +- .../data/v2/internal/api/TableBaseTest.java | 119 +--- .../data/v2/internal/compat/ShimImplTest.java | 18 - .../ops/DivertingUnaryCallableTest.java | 50 +- .../compat/ops/MutateRowsShimTest.java | 250 +------ .../csm/tracers/CompositeVRpcTracerTest.java | 14 - .../internal/csm/tracers/VRpcTracerTest.java | 59 -- .../MutateRowsRetryingVRpcTest.java | 655 ------------------ .../SessionMutateRowsRateLimiterTest.java | 386 ----------- 27 files changed, 83 insertions(+), 3126 deletions(-) delete mode 100644 java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/MutateRowsRetryingVRpc.java delete mode 100644 java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RateLimitingVRpc.java delete mode 100644 java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/SessionMutateRowsRateLimiter.java delete mode 100644 java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/MutateRowsRetryingVRpcTest.java delete mode 100644 java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/SessionMutateRowsRateLimiterTest.java 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 72997953cc0e..0b3ef63e30b8 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 @@ -31,7 +31,6 @@ 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; @@ -41,7 +40,6 @@ import java.io.Closeable; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executor; -import javax.annotation.Nullable; public class AuthorizedViewAsync implements AutoCloseable, Closeable { @@ -60,36 +58,6 @@ 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() @@ -124,8 +92,7 @@ static AuthorizedViewAsync createAndStart( metrics, timer, backgroundExecutor, - userCallbackExecutor, - mutateRowsRateLimiter); + userCallbackExecutor); return new AuthorizedViewAsync(base); } 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 f984f1da5ac3..b6f72659c092 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,12 +33,10 @@ 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; @@ -58,7 +56,6 @@ 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()); @@ -109,7 +106,6 @@ 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 @@ -117,22 +113,9 @@ 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(computeBaseFeatureFlags(settings)); + settings.getChannelProvider().updateFeatureFlags(BASE_FEATURE_FLAGS); ClientInfo clientInfo = ClientInfo.builder() @@ -233,8 +216,6 @@ 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; @@ -281,8 +262,6 @@ 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; } @@ -389,8 +368,7 @@ public TableAsync openTableAsync(String tableId, Permission permission) { metrics.get(), sessionTimer, backgroundExecutor.get(), - userCallbackExecutor.get(), - mutateRowsRateLimiter); + userCallbackExecutor.get()); sessionPools.add(tableAsync.getSessionPool()); return tableAsync; } @@ -415,8 +393,7 @@ public AuthorizedViewAsync openAuthorizedViewAsync( metrics.get(), sessionTimer, backgroundExecutor.get(), - userCallbackExecutor.get(), - mutateRowsRateLimiter); + userCallbackExecutor.get()); 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 427805cce5f0..6b4448945eef 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,10 +32,8 @@ public abstract class ClientSettings { public abstract String getAppProfileId(); - public abstract boolean isServerInitiatedFlowControlEnabled(); - public static Builder builder() { - return new AutoValue_ClientSettings.Builder().setServerInitiatedFlowControlEnabled(false); + return new AutoValue_ClientSettings.Builder(); } @AutoValue.Builder @@ -49,8 +47,6 @@ 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/TableAsync.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableAsync.java index d44848f02ee0..086dd79bdf07 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 @@ -32,7 +32,6 @@ 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; @@ -42,7 +41,6 @@ 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; @@ -59,34 +57,6 @@ 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() @@ -120,8 +90,7 @@ public static TableAsync createAndStart( metrics, timer, backgroundExecutor, - userCallbackExecutor, - mutateRowsRateLimiter); + userCallbackExecutor); return new TableAsync(base); } 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 91395e38376a..f8d6ec880927 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 @@ -15,8 +15,11 @@ */ package com.google.cloud.bigtable.data.v2.internal.api; +import static com.google.common.collect.ImmutableList.toImmutableList; + import com.google.bigtable.v2.CloseSessionRequest; import com.google.bigtable.v2.FeatureFlags; +import com.google.bigtable.v2.Mutation; import com.google.bigtable.v2.SessionCheckAndMutateRowRequest; import com.google.bigtable.v2.SessionCheckAndMutateRowResponse; import com.google.bigtable.v2.SessionMutateRowRequest; @@ -33,12 +36,8 @@ 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,21 +45,19 @@ import com.google.cloud.bigtable.data.v2.internal.session.VRpcDescriptor; import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager; import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableList; import com.google.protobuf.Message; import io.grpc.CallOptions; import io.grpc.Context; 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; @@ -92,48 +89,6 @@ 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 = @@ -165,8 +120,7 @@ static TableBase createAndStart( mutateRowsDescriptor, metrics, timer, - userCallbackExecutor, - mutateRowsRateLimiter); + userCallbackExecutor); } @VisibleForTesting @@ -182,33 +136,6 @@ static TableBase createAndStart( 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; @@ -218,7 +145,6 @@ static TableBase createAndStart( this.metrics = metrics; this.timer = timer; this.userCallbackExecutor = userCallbackExecutor; - this.mutateRowsRateLimiter = mutateRowsRateLimiter; } @Override @@ -293,16 +219,13 @@ 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 = + RetryingVRpc retry = + new RetryingVRpc<>(() -> sessionPool.newCall(mutateRowsDescriptor), timer); + ImmutableList mutations = req.getEntriesList().stream() - .allMatch(entry -> Util.isIdempotent(entry.getMutationsList())); + .flatMap(entry -> entry.getMutationsList().stream()) + .collect(toImmutableList()); + boolean idempotent = Util.isIdempotent(mutations); VRpcTracer tracer = metrics.newTableTracer(sessionPool.getInfo(), mutateRowsDescriptor, deadline); 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 bbaed487c0ff..beabeb1d69b4 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 @@ -23,7 +23,6 @@ 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 { @@ -59,9 +58,8 @@ public UnaryCallable decorateReadModifyWriteRow } @Override - public UnaryCallable decorateMutateRows( - UnaryCallable classic, - UnaryCallSettings settings) { + 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 474369bd6085..2fbfde7a1e6e 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 @@ -23,7 +23,6 @@ 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 { @@ -45,7 +44,6 @@ UnaryCallable decorateReadModifyWriteRow( RowAdapter rowAdapter, UnaryCallSettings settings); - UnaryCallable decorateMutateRows( - UnaryCallable classic, - 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 91a880a5fab6..5742f2906565 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 @@ -58,9 +58,7 @@ 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; @@ -114,22 +112,7 @@ public class ShimImpl implements Shim { private final MutateRowShim mutateRowShim; 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; - } + private final MutateRowsShim mutateRowsShim; public static Shim create( ClientInfo clientInfo, @@ -167,8 +150,6 @@ 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( @@ -177,7 +158,7 @@ public static Shim create( metrics); if (!checkDirectAccessAvailable( - baseFeatureFlags, clientInfo, clientChannelProvider, metrics, bgExecutor)) { + Client.BASE_FEATURE_FLAGS, clientInfo, clientChannelProvider, metrics, bgExecutor)) { // Directpath failed, revert to basic channel provider for the client. clientChannelProvider = configureChannelProvider( @@ -195,7 +176,7 @@ public static Shim create( ClientConfigurationManager configManager = new ClientConfigurationManager( - configChannelProvider.updateFeatureFlags(baseFeatureFlags), + configChannelProvider.updateFeatureFlags(Client.BASE_FEATURE_FLAGS), clientInfo, configChannelProvider, metrics.getDebugTagTracer(), @@ -210,7 +191,7 @@ public static Shim create( return new DisabledShim(); } - FeatureFlags featureFlags = baseFeatureFlags; + FeatureFlags featureFlags = Client.BASE_FEATURE_FLAGS; if (configManager.areSessionsRequired()) { featureFlags = featureFlags.toBuilder().setSessionsRequired(true).build(); @@ -250,6 +231,7 @@ public ShimImpl( this.mutateRowShim = new MutateRowShim(sessionMap); this.checkAndMutateRowShim = new CheckAndMutateRowShim(sessionMap); this.readModifyWriteRowShim = new ReadModifyWriteRowShim(sessionMap); + this.mutateRowsShim = new MutateRowsShim(sessionMap); } /** @@ -459,13 +441,12 @@ public UnaryCallable decorateReadModifyWriteRow } @Override - public UnaryCallable decorateMutateRows( - UnaryCallable classic, - UnaryCallSettings settings) { + public UnaryCallable decorateMutateRows( + UnaryCallable classic, UnaryCallSettings settings) { return new DivertingUnaryCallable<>( configManager, classic, - new MutateRowsShim(sessionMap, settings.getRetryableCodes()), + mutateRowsShim, 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 90accab465c1..6500add3dee2 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,10 +144,6 @@ 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 index b89a908de93c..ad843a02c584 100644 --- 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 @@ -15,50 +15,21 @@ */ 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); +public class MutateRowsShim implements UnaryShim { 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 @@ -80,147 +51,27 @@ public boolean supports(BulkMutation request) { } @Override - public CompletableFuture call(BulkMutation request, Deadline deadline) { + 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 sessionMap.apply( + (TableId) targetId, + SessionKey.Permission.WRITE, + t -> t.mutateRows(innerReq, deadline).thenApply(r -> null)); } - - 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; - } + if (targetId instanceof AuthorizedViewId) { + return sessionMap.apply( + (AuthorizedViewId) targetId, + SessionKey.Permission.WRITE, + v -> v.mutateRows(innerReq, deadline).thenApply(r -> null)); } - return MutateRowsException.create( - rpcApiException, rpcApiException.getStatusCode(), failedMutations, anyEntryRetryable); + CompletableFuture f = new CompletableFuture<>(); + f.completeExceptionally( + new UnsupportedOperationException("Unsupported targetId type: " + targetId)); + return f; } } 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 b43aedab4683..4990987ea3bc 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,13 +109,6 @@ 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 e402ce75a03f..473c8c56811c 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,7 +18,6 @@ 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; @@ -107,37 +106,6 @@ 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 5fc44e056f57..e794036d5aa0 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,8 +21,6 @@ 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; @@ -71,22 +69,6 @@ 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 f1b92a661dce..924bdd4d6006 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,7 +18,6 @@ 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 { @@ -33,11 +32,6 @@ 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 72091e429c46..0f20ae1a0214 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,7 +26,6 @@ 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; @@ -134,19 +133,6 @@ 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 deleted file mode 100644 index f013faae0222..000000000000 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/MutateRowsRetryingVRpc.java +++ /dev/null @@ -1,612 +0,0 @@ -/* - * 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 deleted file mode 100644 index 8c6ad256bc96..000000000000 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RateLimitingVRpc.java +++ /dev/null @@ -1,216 +0,0 @@ -/* - * 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 deleted file mode 100644 index 872e6675c5da..000000000000 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/SessionMutateRowsRateLimiter.java +++ /dev/null @@ -1,262 +0,0 @@ -/* - * 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/stub/EnhancedBigtableStub.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index 5dee50d8c877..11f7cb18576d 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 @@ -203,11 +203,15 @@ public EnhancedBigtableStub( bulkMutateRowsCallable = createMutateRowsBaseCallable(); UnaryCallable bulkMutateRowsVoidCallable = new MutateRowsErrorConverterUnaryCallable(bulkMutateRowsCallable); - externalBulkMutateRowsCallable = + UnaryCallable maybePointWriteCallable = new MaybePointWriteCallable( - bulkMutateRowsVoidCallable, - createPointWriteCallable(bulkMutateRowsVoidCallable), - requestContext) + bulkMutateRowsVoidCallable, + createPointWriteCallable(bulkMutateRowsVoidCallable), + requestContext); + externalBulkMutateRowsCallable = + bigtableClientContext + .getSessionShim() + .decorateMutateRows(maybePointWriteCallable, perOpSettings.bulkMutateRowsSettings) .withDefaultCallContext( bigtableClientContext .getClientContext() @@ -798,42 +802,33 @@ 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(withCookie, requestContext); + new BulkMutateRowsUserFacingCallable( + flowControlCallable != null ? flowControlCallable : withCookie, requestContext); SpanName spanName = getSpanName("MutateRows"); UnaryCallable tracedBatcherUnaryCallable = new TracedBatcherUnaryCallable<>(userFacing); - ApiCallContext defaultCallContext = - clientContext - .getDefaultCallContext() - .withRetrySettings(perOpSettings.bulkMutateRowsSettings.getRetrySettings()); - - UnaryCallable classic = + UnaryCallable traced = new TracedUnaryCallable<>( - tracedBatcherUnaryCallable, clientContext.getTracerFactory(), spanName) - .withDefaultCallContext(defaultCallContext); + tracedBatcherUnaryCallable, clientContext.getTracerFactory(), spanName); - 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); + return traced.withDefaultCallContext( + clientContext + .getDefaultCallContext() + .withRetrySettings(perOpSettings.bulkMutateRowsSettings.getRetrySettings())); } /** 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 ca3b4f20af77..1207a47da84e 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,10 +17,8 @@ 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; @@ -28,23 +26,17 @@ 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; @@ -242,113 +234,6 @@ public void testMutateRowsNotIdempotent() { 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; @@ -390,12 +275,10 @@ public boolean hasSession() { static class FakeVRpc implements VRpc { private VRpcCallContext ctx; - private VRpcListener listener; @Override - public void start(ReqT req, VRpcCallContext ctx, VRpcListener listener) { + public void start(Object req, VRpcCallContext ctx, VRpcListener ignored) { 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 e95a07f4adc3..4dfab36d0694 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,11 +22,9 @@ 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; @@ -118,20 +116,4 @@ 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 e1de16abaabe..729f378b7a07 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,18 +16,13 @@ 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; @@ -39,12 +34,7 @@ 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; @@ -53,7 +43,6 @@ 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; @@ -195,42 +184,6 @@ 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() { @@ -299,12 +252,11 @@ 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 nextResponse; // never completes by default — test doesn't wait on it + return new CompletableFuture<>(); // never completes — 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 index 024485a7036c..17ea5e6c7301 100644 --- 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 @@ -16,49 +16,23 @@ 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.assertNull; 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; @@ -137,224 +111,27 @@ void supports_unsupportedTargetId_returnsFalse() { } @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(); + void call_tableId_returnsNull() throws Exception { + SessionMutateRowsResponse response = SessionMutateRowsResponse.getDefaultInstance(); 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")); + BulkMutation request = BulkMutation.create(TableId.of("table1")); Deadline deadline = Deadline.after(1, TimeUnit.SECONDS); - MutateRowsAttemptResult result = shim.call(request, deadline).get(); - assertThat(result.getFailedMutations()).isEmpty(); - assertThat(result.getIsRetryable()).isFalse(); + assertNull(shim.call(request, deadline).get()); } @Test - void call_authorizedViewId_allOk_returnsSuccess() throws Exception { - SessionMutateRowsResponse response = - SessionMutateRowsResponse.newBuilder().addEntries(entry(0, Code.OK_VALUE, "")).build(); + void call_authorizedViewId_returnsNull() throws Exception { + SessionMutateRowsResponse response = SessionMutateRowsResponse.getDefaultInstance(); 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")); + BulkMutation request = BulkMutation.create(AuthorizedViewId.of("table1", "view1")); 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); + assertNull(shim.call(request, deadline).get()); } @Test @@ -363,17 +140,10 @@ void call_unsupportedTargetId_futureFailsWithUnsupportedOperationException() { BulkMutation request = BulkMutation.create(unsupported); Deadline deadline = Deadline.after(1, TimeUnit.SECONDS); - CompletableFuture future = shim.call(request, deadline); + 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 c0f1d0d8e68b..12f041fa851c 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,18 +89,4 @@ 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 2120ae5ed5ca..54685b907700 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,13 +69,11 @@ 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; @@ -412,36 +410,6 @@ 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 = @@ -516,33 +484,6 @@ 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 deleted file mode 100644 index 654823457353..000000000000 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/MutateRowsRetryingVRpcTest.java +++ /dev/null @@ -1,655 +0,0 @@ -/* - * 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 deleted file mode 100644 index 49038390e08c..000000000000 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/SessionMutateRowsRateLimiterTest.java +++ /dev/null @@ -1,386 +0,0 @@ -/* - * 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()); - } - } -}