Skip to content

feat(bigtable): plumb per-operation timeouts through CbtTestProxy #14997

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
alexoneill wants to merge 1 commit into
base: main
Choose a base branch
from
+656 −17
Open

feat(bigtable): plumb per-operation timeouts through CbtTestProxy #14997

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions ...ain/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsAttemptCallable.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,8 +146,7 @@ public List<MutateRowsResponse> apply(Throwable throwable) {
this.callContext = Preconditions.checkNotNull(callContext, "callContext");
this.retryableCodes = Preconditions.checkNotNull(retryableCodes, "retryableCodes");
this.retryAlgorithm = retryAlgorithm;
// TODO: pass in the callContext so that the retry setting can be overridden per call
this.attemptSettings = retryAlgorithm.createFirstAttempt(null);
this.attemptSettings = retryAlgorithm.createFirstAttempt(callContext);

permanentFailures = Lists.newArrayList();
}
Expand Down
1 change: 1 addition & 0 deletions ...java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsAttemptCallableTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,7 @@ public void singleEntrySuccessTest() throws Exception {
assertThat(result.getIsRetryable()).isFalse();
// innerCallable received the request
assertThat(innerCallable.lastRequest).isEqualTo(request);
Mockito.verify(mockRetryAlgorithm).createFirstAttempt(callContext);
}

@Test
Expand Down
144 changes: 131 additions & 13 deletions java-bigtable/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,10 @@
import com.google.api.gax.core.CredentialsProvider;
import com.google.api.gax.core.FixedCredentialsProvider;
import com.google.api.gax.core.NoCredentialsProvider;
import com.google.api.gax.grpc.GrpcCallContext;
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
import com.google.api.gax.retrying.RetrySettings;
import com.google.api.gax.rpc.ApiCallContext;
import com.google.api.gax.rpc.ApiException;
import com.google.api.gax.rpc.ServerStream;
import com.google.auth.oauth2.AccessToken;
Expand All @@ -43,6 +45,7 @@
import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
import com.google.cloud.bigtable.data.v2.models.RowCell;
import com.google.cloud.bigtable.data.v2.models.RowMutation;
import com.google.cloud.bigtable.data.v2.models.TableId;
import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement;
import com.google.cloud.bigtable.data.v2.models.sql.ResultSet;
import com.google.cloud.bigtable.data.v2.models.sql.SqlType;
Expand All @@ -53,6 +56,7 @@
import com.google.protobuf.ByteString;
import com.google.protobuf.util.Durations;
import com.google.rpc.Code;
import io.grpc.Context;
import io.grpc.ManagedChannelBuilder;
import io.grpc.Status;
import io.grpc.StatusException;
Expand All @@ -73,6 +77,7 @@
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.logging.Logger;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
Expand Down Expand Up @@ -138,16 +143,37 @@ private static BigtableDataSettings.Builder overrideTimeoutSetting(

private static void updateTimeout(RetrySettings.Builder settings, Duration newTimeout) {
Duration rpcTimeout = settings.getInitialRpcTimeoutDuration();
Duration maxRpcTimeout = settings.getMaxRpcTimeoutDuration();
Duration oldTotalTimeout = settings.getTotalTimeoutDuration();

// TODO: this should happen in gax
// Clamp the rpcTimeout to the overall timeout
if (rpcTimeout != null && rpcTimeout.compareTo(newTimeout) > 0) {
settings.setInitialRpcTimeoutDuration(newTimeout).setMaxRpcTimeoutDuration(newTimeout);
// Clamp the rpcTimeout to the overall timeout, or raise it if it was previously clamped to the
// old total timeout.
if (rpcTimeout != null
&& (rpcTimeout.compareTo(newTimeout) > 0 || rpcTimeout.equals(oldTotalTimeout))) {
settings.setInitialRpcTimeoutDuration(newTimeout);
}
if (maxRpcTimeout != null
&& (maxRpcTimeout.compareTo(newTimeout) > 0 || maxRpcTimeout.equals(oldTotalTimeout))) {
settings.setMaxRpcTimeoutDuration(newTimeout);
}

settings.setTotalTimeoutDuration(newTimeout);
}

private static boolean hasPositiveTimeout(
boolean hasTimeout, com.google.protobuf.Duration timeout) {
return hasTimeout && (timeout.getSeconds() > 0 || timeout.getNanos() > 0);
}

private static ApiCallContext createCallContextWithTimeout(
RetrySettings baseRetrySettings, com.google.protobuf.Duration protoTimeout) {
Duration timeout = Duration.ofSeconds(protoTimeout.getSeconds(), protoTimeout.getNanos());
RetrySettings.Builder builder = baseRetrySettings.toBuilder();
updateTimeout(builder, timeout);
return GrpcCallContext.createDefault().withRetrySettings(builder.build());
}
Comment on lines 144 to +175

@mutianf mutianf Nov 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of mutating BigtableDataSettings in createClient (which overwrites the library defaults and requires the equals(oldTotalTimeout) heuristic to unclamp rpcTimeout / maxRpcTimeout), could we store Optional<Duration> perOperationTimeout on CbtClient and leave client.settings() at the library defaults (similar to the C++ test proxy)?

Then every operation can resolve the effective timeout (request override -> client override -> empty) and only clamp down from the untouched library defaults:

  @AutoValue
  abstract static class CbtClient {
    static CbtClient create(BigtableDataSettings settings, BigtableDataClient dataClient) {
      return create(settings, dataClient, Optional.empty());
    }

    static CbtClient create(
        BigtableDataSettings settings,
        BigtableDataClient dataClient,
        Optional<Duration> perOperationTimeout) {
      return new AutoValue_CbtTestProxy_CbtClient(settings, dataClient, perOperationTimeout);
    }

    abstract BigtableDataSettings settings();

    abstract BigtableDataClient dataClient();

    abstract Optional<Duration> perOperationTimeout();
  }

  private static Optional<Duration> resolveTimeout(
      @Nullable CbtClient client,
      boolean hasRequestTimeout,
      com.google.protobuf.Duration requestTimeout) {
    if (hasRequestTimeout && Durations.compare(requestTimeout, Durations.ZERO) > 0) {
      return Optional.of(
          Duration.ofSeconds(requestTimeout.getSeconds(), requestTimeout.getNanos()));
    }
    return client != null ? client.perOperationTimeout() : Optional.empty();
  }

  private static ApiCallContext createCallContextWithTimeout(
      RetrySettings defaultRetrySettings, Duration timeout) {
    RetrySettings.Builder builder = defaultRetrySettings.toBuilder();
    Duration initialRpcTimeout = builder.getInitialRpcTimeoutDuration();
    if (initialRpcTimeout != null && initialRpcTimeout.compareTo(timeout) > 0) {
      builder.setInitialRpcTimeoutDuration(timeout);
    }
    Duration maxRpcTimeout = builder.getMaxRpcTimeoutDuration();
    if (maxRpcTimeout != null && maxRpcTimeout.compareTo(timeout) > 0) {
      builder.setMaxRpcTimeoutDuration(timeout);
    }
    builder.setTotalTimeoutDuration(timeout);
    return GrpcCallContext.createDefault().withRetrySettings(builder.build());
  }

And in each operation handler:

    Optional<Duration> timeout =
        resolveTimeout(client, request.hasPerOperationTimeout(), request.getPerOperationTimeout());
    if (timeout.isPresent()) {
      ApiCallContext callContext =
          createCallContextWithTimeout(
              client.settings().getStubSettings().mutateRowSettings().getRetrySettings(),
              timeout.get());
      client.dataClient().mutateRowCallable().call(mutation, callContext);
    } else {
      client.dataClient().mutateRow(mutation);
    }


/** Helper method to get a client object by its id. */
private CbtClient getClient(String id) throws StatusException {
CbtClient client = idClientMap.get(id);
Expand Down Expand Up @@ -288,7 +314,16 @@ public void mutateRow(
RowMutation mutation = RowMutation.fromProto(request.getRequest());
try {
// This response is empty.
client.dataClient().mutateRow(mutation);
if (hasPositiveTimeout(
request.hasPerOperationTimeout(), request.getPerOperationTimeout())) {
ApiCallContext callContext =
createCallContextWithTimeout(
client.settings().getStubSettings().mutateRowSettings().getRetrySettings(),
request.getPerOperationTimeout());
client.dataClient().mutateRowCallable().call(mutation, callContext);
} else {
client.dataClient().mutateRow(mutation);
}
} catch (ApiException e) {
responseObserver.onNext(MutateRowResult.newBuilder().setStatus(convertStatus(e)).build());
responseObserver.onCompleted();
Expand Down Expand Up @@ -318,7 +353,16 @@ public void bulkMutateRows(

BulkMutation batch = BulkMutation.fromProto(request.getRequest());
try {
client.dataClient().bulkMutateRows(batch);
if (hasPositiveTimeout(
request.hasPerOperationTimeout(), request.getPerOperationTimeout())) {
ApiCallContext callContext =
createCallContextWithTimeout(
client.settings().getStubSettings().bulkMutateRowsSettings().getRetrySettings(),
request.getPerOperationTimeout());
client.dataClient().bulkMutationCallable().call(batch, callContext);
} else {
client.dataClient().bulkMutateRows(batch);
}
} catch (MutateRowsException e) {
MutateRowsResult.Builder resultBuilder = MutateRowsResult.newBuilder();
com.google.rpc.Status topLevelStatus =
Expand Down Expand Up @@ -385,10 +429,23 @@ public void readRow(ReadRowRequest request, StreamObserver<RowResult> responseOb

com.google.cloud.bigtable.data.v2.models.Row row;
try {
row =
client
.dataClient()
.readRow(tableId, request.getRowKey(), FILTERS.fromProto(request.getFilter()));
if (hasPositiveTimeout(
request.hasPerOperationTimeout(), request.getPerOperationTimeout())) {
Query query =
Query.create(TableId.of(tableId))
.rowKey(request.getRowKey())
.filter(FILTERS.fromProto(request.getFilter()));
ApiCallContext callContext =
createCallContextWithTimeout(
client.settings().getStubSettings().readRowSettings().getRetrySettings(),
request.getPerOperationTimeout());
row = client.dataClient().readRowCallable().call(query, callContext);
} else {
row =
client
.dataClient()
.readRow(tableId, request.getRowKey(), FILTERS.fromProto(request.getFilter()));
}
if (row != null) {
RowResult.Builder resultBuilder = convertRowResult(row);
responseObserver.onNext(
Expand Down Expand Up @@ -438,7 +495,16 @@ public void readRows(ReadRowsRequest request, StreamObserver<RowsResult> respons
ServerStream<com.google.cloud.bigtable.data.v2.models.Row> rows;
Query query = Query.fromProto(request.getRequest());
try {
rows = client.dataClient().readRows(query);
if (hasPositiveTimeout(
request.hasPerOperationTimeout(), request.getPerOperationTimeout())) {
ApiCallContext callContext =
createCallContextWithTimeout(
client.settings().getStubSettings().readRowsSettings().getRetrySettings(),
request.getPerOperationTimeout());
rows = client.dataClient().readRowsCallable().call(query, callContext);
} else {
rows = client.dataClient().readRows(query);
}
int cancelAfterRows = request.getCancelAfterRows();
RowsResult.Builder resultBuilder = convertRowsResult(rows, cancelAfterRows);
responseObserver.onNext(
Expand Down Expand Up @@ -560,7 +626,16 @@ public void sampleRowKeys(

List<KeyOffset> keyOffsets;
try {
keyOffsets = client.dataClient().sampleRowKeys(tableId);
if (hasPositiveTimeout(
request.hasPerOperationTimeout(), request.getPerOperationTimeout())) {
ApiCallContext callContext =
createCallContextWithTimeout(
client.settings().getStubSettings().sampleRowKeysSettings().getRetrySettings(),
request.getPerOperationTimeout());
keyOffsets = client.dataClient().sampleRowKeysCallable().call(tableId, callContext);
} else {
keyOffsets = client.dataClient().sampleRowKeys(tableId);
}
} catch (ApiException e) {
responseObserver.onNext(SampleRowKeysResult.newBuilder().setStatus(convertStatus(e)).build());
responseObserver.onCompleted();
Expand Down Expand Up @@ -599,7 +674,20 @@ public void checkAndMutateRow(
ConditionalRowMutation mutation = ConditionalRowMutation.fromProto(request.getRequest());
Boolean matched;
try {
matched = client.dataClient().checkAndMutateRow(mutation);
if (hasPositiveTimeout(
request.hasPerOperationTimeout(), request.getPerOperationTimeout())) {
ApiCallContext callContext =
createCallContextWithTimeout(
client
.settings()
.getStubSettings()
.checkAndMutateRowSettings()
.getRetrySettings(),
request.getPerOperationTimeout());
matched = client.dataClient().checkAndMutateRowCallable().call(mutation, callContext);
} else {
matched = client.dataClient().checkAndMutateRow(mutation);
}
} catch (ApiException e) {
responseObserver.onNext(
CheckAndMutateRowResult.newBuilder().setStatus(convertStatus(e)).build());
Expand Down Expand Up @@ -634,7 +722,20 @@ public void readModifyWriteRow(
com.google.cloud.bigtable.data.v2.models.Row row;
ReadModifyWriteRow mutation = ReadModifyWriteRow.fromProto(request.getRequest());
try {
row = client.dataClient().readModifyWriteRow(mutation);
if (hasPositiveTimeout(
request.hasPerOperationTimeout(), request.getPerOperationTimeout())) {
ApiCallContext callContext =
createCallContextWithTimeout(
client
.settings()
.getStubSettings()
.readModifyWriteRowSettings()
.getRetrySettings(),
request.getPerOperationTimeout());
row = client.dataClient().readModifyWriteRowCallable().call(mutation, callContext);
} else {
row = client.dataClient().readModifyWriteRow(mutation);
}
if (row != null) {
RowResult.Builder resultBuilder = convertRowResult(row);
responseObserver.onNext(
Expand Down Expand Up @@ -684,7 +785,20 @@ public void executeQuery(
return;
}
ResultSet resultSet = null;
Context.CancellableContext cancellableContext = null;
Context previousContext = null;
try {
if (hasPositiveTimeout(
request.hasPerOperationTimeout(), request.getPerOperationTimeout())) {
Duration timeout =
Duration.ofSeconds(
request.getPerOperationTimeout().getSeconds(),
request.getPerOperationTimeout().getNanos());
cancellableContext =
Context.current()
.withDeadlineAfter(timeout.toNanos(), TimeUnit.NANOSECONDS, null);

@mutianf mutianf Nov 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, use MoreExecutores.directExecutor() instead of null?

previousContext = cancellableContext.attach();
}
Map<String, SqlType<?>> paramTypes = new HashMap<>();
for (Map.Entry<String, Value> entry : request.getRequest().getParamsMap().entrySet()) {
paramTypes.put(entry.getKey(), SqlType.fromProto(entry.getValue().getType()));
Expand Down Expand Up @@ -729,6 +843,10 @@ public void executeQuery(
if (resultSet != null) {
resultSet.close();
}
if (cancellableContext != null) {
cancellableContext.detach(previousContext);
cancellableContext.close();
}
Comment on lines +846 to +849

@gemini-code-assist gemini-code-assist Bot Nov 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

If cancellableContext is initialized but previousContext remains null (for example, if an exception occurs or the thread is interrupted before attach() completes), calling cancellableContext.detach(previousContext) will throw a NullPointerException because Context.detach does not accept null. This NPE would mask the original exception.

To prevent this, guard the detach call by checking if previousContext is not null.

      if (previousContext != null) {
        cancellableContext.detach(previousContext);
      }
      if (cancellableContext != null) {
        cancellableContext.close();
      }

}
responseObserver.onCompleted();
}
Expand Down
46 changes: 44 additions & 2 deletions java-bigtable/test-proxy/src/main/proto/test_proxy.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,8 +81,10 @@ message CreateClientRequest {
string app_profile_id = 5;

// If provided, a custom timeout will be set for each API call conducted by
// the created client. Otherwise, the default timeout from the client library
// will be used. Note that the override applies to all the methods.
// the created client, unless overridden by per_operation_timeout on an
// individual operation request. Otherwise, the default timeout from the
// client library will be used. Note that the override applies to all the
// methods.
google.protobuf.Duration per_operation_timeout = 6;

// Optional config that dictates how the optional features should be enabled
Expand Down Expand Up @@ -137,6 +139,11 @@ message ReadRowRequest {

// The row filter to be applied to the target row.
google.bigtable.v2.RowFilter filter = 3;

// Optional per-operation timeout override. When set and positive, supersedes
// CreateClientRequest.per_operation_timeout and the client library default
// timeout for this call.
google.protobuf.Duration per_operation_timeout = 5;
}

// Response from test proxy service for ReadRowRequest or
Expand All @@ -160,6 +167,11 @@ message ReadRowsRequest {
// The streaming read can be canceled before all items are seen.
// Has no effect if non-positive.
int32 cancel_after_rows = 3;

// Optional per-operation timeout override. When set and positive, supersedes
// CreateClientRequest.per_operation_timeout and the client library default
// timeout for this call.
google.protobuf.Duration per_operation_timeout = 4;
}

// Response from test proxy service for ReadRowsRequest.
Expand All @@ -178,6 +190,11 @@ message MutateRowRequest {

// The raw request to the Bigtable server.
google.bigtable.v2.MutateRowRequest request = 2;

// Optional per-operation timeout override. When set and positive, supersedes
// CreateClientRequest.per_operation_timeout and the client library default
// timeout for this call.
google.protobuf.Duration per_operation_timeout = 3;
}

// Response from test proxy service for MutateRowRequest.
Expand All @@ -193,6 +210,11 @@ message MutateRowsRequest {

// The raw request to the Bigtable server.
google.bigtable.v2.MutateRowsRequest request = 2;

// Optional per-operation timeout override. When set and positive, supersedes
// CreateClientRequest.per_operation_timeout and the client library default
// timeout for this call.
google.protobuf.Duration per_operation_timeout = 3;
}

// Response from test proxy service for MutateRowsRequest.
Expand All @@ -212,6 +234,11 @@ message CheckAndMutateRowRequest {

// The raw request to the Bigtable server.
google.bigtable.v2.CheckAndMutateRowRequest request = 2;

// Optional per-operation timeout override. When set and positive, supersedes
// CreateClientRequest.per_operation_timeout and the client library default
// timeout for this call.
google.protobuf.Duration per_operation_timeout = 3;
}

// Response from test proxy service for CheckAndMutateRowRequest.
Expand All @@ -230,6 +257,11 @@ message SampleRowKeysRequest {

// The raw request to the Bigtable server.
google.bigtable.v2.SampleRowKeysRequest request = 2;

// Optional per-operation timeout override. When set and positive, supersedes
// CreateClientRequest.per_operation_timeout and the client library default
// timeout for this call.
google.protobuf.Duration per_operation_timeout = 3;
}

// Response from test proxy service for SampleRowKeysRequest.
Expand All @@ -248,6 +280,11 @@ message ReadModifyWriteRowRequest {

// The raw request to the Bigtable server.
google.bigtable.v2.ReadModifyWriteRowRequest request = 2;

// Optional per-operation timeout override. When set and positive, supersedes
// CreateClientRequest.per_operation_timeout and the client library default
// timeout for this call.
google.protobuf.Duration per_operation_timeout = 3;
}

// Request to test proxy service to execute a query.
Expand All @@ -260,6 +297,11 @@ message ExecuteQueryRequest {

// The file descriptor set for the query.
google.protobuf.FileDescriptorSet proto_descriptors = 3;

// Optional per-operation timeout override. When set and positive, supersedes
// CreateClientRequest.per_operation_timeout and the client library default
// timeout for this call.
google.protobuf.Duration per_operation_timeout = 4;
}

// Response from test proxy service for ExecuteQueryRequest.
Expand Down
Loading
Loading