diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsAttemptCallable.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsAttemptCallable.java index 4796ad48a094..8a324b4cd7a9 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsAttemptCallable.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsAttemptCallable.java @@ -146,8 +146,7 @@ public List 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(); } diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsAttemptCallableTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsAttemptCallableTest.java index 65f4f8a324f2..dee69dd439ec 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsAttemptCallableTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsAttemptCallableTest.java @@ -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 diff --git a/java-bigtable/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java b/java-bigtable/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java index 7a89c5d94140..c9f17bfa76a8 100644 --- a/java-bigtable/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java +++ b/java-bigtable/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java @@ -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; @@ -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; @@ -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; @@ -71,12 +75,18 @@ import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.logging.Logger; import java.util.regex.Matcher; import java.util.regex.Pattern; import java.util.stream.Collectors; +import javax.annotation.Nullable; /** Java implementation of the CBT test proxy. Used to test the Java CBT client. */ public class CbtTestProxy extends CloudBigtableV2TestProxyImplBase implements Closeable { @@ -88,12 +98,21 @@ public class CbtTestProxy extends CloudBigtableV2TestProxyImplBase implements Cl @AutoValue abstract static class CbtClient { static CbtClient create(BigtableDataSettings settings, BigtableDataClient dataClient) { - return new AutoValue_CbtTestProxy_CbtClient(settings, dataClient); + return create(settings, dataClient, Optional.empty()); + } + + static CbtClient create( + BigtableDataSettings settings, + BigtableDataClient dataClient, + Optional perOperationTimeout) { + return new AutoValue_CbtTestProxy_CbtClient(settings, dataClient, perOperationTimeout); } abstract BigtableDataSettings settings(); abstract BigtableDataClient dataClient(); + + abstract Optional perOperationTimeout(); } private static final Logger logger = Logger.getLogger(CbtTestProxy.class.getName()); @@ -107,45 +126,30 @@ public static CbtTestProxy create() { return new CbtTestProxy(); } - /** - * Helper method to override the timeout settings of data APIs. TODO(developer): per-attempt - * timeout may also be overridden, which will involve test harness update. - * - * @param settingsBuilder The Builder object of BigtableDataSettings. - * @param newTimeout The value that is used to set the timeout. - */ - private static BigtableDataSettings.Builder overrideTimeoutSetting( - Duration newTimeout, BigtableDataSettings.Builder settingsBuilder) { - - updateTimeout( - settingsBuilder.stubSettings().bulkMutateRowsSettings().retrySettings(), newTimeout); - updateTimeout(settingsBuilder.stubSettings().mutateRowSettings().retrySettings(), newTimeout); - updateTimeout(settingsBuilder.stubSettings().readRowSettings().retrySettings(), newTimeout); - updateTimeout(settingsBuilder.stubSettings().readRowsSettings().retrySettings(), newTimeout); - updateTimeout( - settingsBuilder.stubSettings().checkAndMutateRowSettings().retrySettings(), newTimeout); - updateTimeout( - settingsBuilder.stubSettings().readModifyWriteRowSettings().retrySettings(), newTimeout); - updateTimeout( - settingsBuilder.stubSettings().sampleRowKeysSettings().retrySettings(), newTimeout); - updateTimeout( - settingsBuilder.stubSettings().executeQuerySettings().retrySettings(), newTimeout); - updateTimeout( - settingsBuilder.stubSettings().prepareQuerySettings().retrySettings(), newTimeout); - - return settingsBuilder; + private static Optional 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 void updateTimeout(RetrySettings.Builder settings, Duration newTimeout) { - Duration rpcTimeout = settings.getInitialRpcTimeoutDuration(); - - // 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); + 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); } - - settings.setTotalTimeoutDuration(newTimeout); + Duration maxRpcTimeout = builder.getMaxRpcTimeoutDuration(); + if (maxRpcTimeout != null && maxRpcTimeout.compareTo(timeout) > 0) { + builder.setMaxRpcTimeoutDuration(timeout); + } + builder.setTotalTimeoutDuration(timeout); + return GrpcCallContext.createDefault().withRetrySettings(builder.build()); } /** Helper method to get a client object by its id. */ @@ -187,9 +191,9 @@ public synchronized void createClient( .setInstanceId(request.getInstanceId()) .setAppProfileId(request.getAppProfileId()); - if (request.hasPerOperationTimeout()) { - Duration newTimeout = Duration.ofMillis(Durations.toMillis(request.getPerOperationTimeout())); - settingsBuilder = overrideTimeoutSetting(newTimeout, settingsBuilder); + Optional perOperationTimeout = + resolveTimeout(null, request.hasPerOperationTimeout(), request.getPerOperationTimeout()); + if (perOperationTimeout.isPresent()) { logger.info( String.format( "Total timeout is set to %s for all the methods", @@ -230,7 +234,7 @@ public synchronized void createClient( } BigtableDataSettings settings = settingsBuilder.build(); BigtableDataClient client = BigtableDataClient.create(settings); - CbtClient cbtClient = CbtClient.create(settings, client); + CbtClient cbtClient = CbtClient.create(settings, client, perOperationTimeout); idClientMap.put(request.getClientId(), cbtClient); } catch (IOException e) { responseObserver.onError(Status.INTERNAL.withDescription(e.getMessage()).asException()); @@ -288,7 +292,18 @@ public void mutateRow( RowMutation mutation = RowMutation.fromProto(request.getRequest()); try { // This response is empty. - client.dataClient().mutateRow(mutation); + Optional 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); + } } catch (ApiException e) { responseObserver.onNext(MutateRowResult.newBuilder().setStatus(convertStatus(e)).build()); responseObserver.onCompleted(); @@ -318,7 +333,18 @@ public void bulkMutateRows( BulkMutation batch = BulkMutation.fromProto(request.getRequest()); try { - client.dataClient().bulkMutateRows(batch); + Optional timeout = + resolveTimeout( + client, request.hasPerOperationTimeout(), request.getPerOperationTimeout()); + if (timeout.isPresent()) { + ApiCallContext callContext = + createCallContextWithTimeout( + client.settings().getStubSettings().bulkMutateRowsSettings().getRetrySettings(), + timeout.get()); + client.dataClient().bulkMutationCallable().call(batch, callContext); + } else { + client.dataClient().bulkMutateRows(batch); + } } catch (MutateRowsException e) { MutateRowsResult.Builder resultBuilder = MutateRowsResult.newBuilder(); com.google.rpc.Status topLevelStatus = @@ -385,10 +411,25 @@ public void readRow(ReadRowRequest request, StreamObserver responseOb com.google.cloud.bigtable.data.v2.models.Row row; try { - row = - client - .dataClient() - .readRow(tableId, request.getRowKey(), FILTERS.fromProto(request.getFilter())); + Optional timeout = + resolveTimeout( + client, request.hasPerOperationTimeout(), request.getPerOperationTimeout()); + if (timeout.isPresent()) { + Query query = + Query.create(TableId.of(tableId)) + .rowKey(request.getRowKey()) + .filter(FILTERS.fromProto(request.getFilter())); + ApiCallContext callContext = + createCallContextWithTimeout( + client.settings().getStubSettings().readRowSettings().getRetrySettings(), + timeout.get()); + 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( @@ -438,7 +479,18 @@ public void readRows(ReadRowsRequest request, StreamObserver respons ServerStream rows; Query query = Query.fromProto(request.getRequest()); try { - rows = client.dataClient().readRows(query); + Optional timeout = + resolveTimeout( + client, request.hasPerOperationTimeout(), request.getPerOperationTimeout()); + if (timeout.isPresent()) { + ApiCallContext callContext = + createCallContextWithTimeout( + client.settings().getStubSettings().readRowsSettings().getRetrySettings(), + timeout.get()); + 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( @@ -560,7 +612,18 @@ public void sampleRowKeys( List keyOffsets; try { - keyOffsets = client.dataClient().sampleRowKeys(tableId); + Optional timeout = + resolveTimeout( + client, request.hasPerOperationTimeout(), request.getPerOperationTimeout()); + if (timeout.isPresent()) { + ApiCallContext callContext = + createCallContextWithTimeout( + client.settings().getStubSettings().sampleRowKeysSettings().getRetrySettings(), + timeout.get()); + 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(); @@ -599,7 +662,18 @@ public void checkAndMutateRow( ConditionalRowMutation mutation = ConditionalRowMutation.fromProto(request.getRequest()); Boolean matched; try { - matched = client.dataClient().checkAndMutateRow(mutation); + Optional timeout = + resolveTimeout( + client, request.hasPerOperationTimeout(), request.getPerOperationTimeout()); + if (timeout.isPresent()) { + ApiCallContext callContext = + createCallContextWithTimeout( + client.settings().getStubSettings().checkAndMutateRowSettings().getRetrySettings(), + timeout.get()); + matched = client.dataClient().checkAndMutateRowCallable().call(mutation, callContext); + } else { + matched = client.dataClient().checkAndMutateRow(mutation); + } } catch (ApiException e) { responseObserver.onNext( CheckAndMutateRowResult.newBuilder().setStatus(convertStatus(e)).build()); @@ -634,7 +708,18 @@ public void readModifyWriteRow( com.google.cloud.bigtable.data.v2.models.Row row; ReadModifyWriteRow mutation = ReadModifyWriteRow.fromProto(request.getRequest()); try { - row = client.dataClient().readModifyWriteRow(mutation); + Optional timeout = + resolveTimeout( + client, request.hasPerOperationTimeout(), request.getPerOperationTimeout()); + if (timeout.isPresent()) { + ApiCallContext callContext = + createCallContextWithTimeout( + client.settings().getStubSettings().readModifyWriteRowSettings().getRetrySettings(), + timeout.get()); + row = client.dataClient().readModifyWriteRowCallable().call(mutation, callContext); + } else { + row = client.dataClient().readModifyWriteRow(mutation); + } if (row != null) { RowResult.Builder resultBuilder = convertRowResult(row); responseObserver.onNext( @@ -684,7 +769,27 @@ public void executeQuery( return; } ResultSet resultSet = null; + Context.CancellableContext cancellableContext = null; + Context previousContext = null; try { + Optional timeout = + resolveTimeout( + client, request.hasPerOperationTimeout(), request.getPerOperationTimeout()); + if (timeout.isPresent()) { + Thread callerThread = Thread.currentThread(); + cancellableContext = + Context.current() + .withDeadlineAfter( + timeout.get().toNanos(), TimeUnit.NANOSECONDS, deadlineScheduler); + cancellableContext.addListener( + context -> { + if (context.cancellationCause() instanceof TimeoutException) { + callerThread.interrupt(); + } + }, + Runnable::run); + previousContext = cancellableContext.attach(); + } Map> paramTypes = new HashMap<>(); for (Map.Entry entry : request.getRequest().getParamsMap().entrySet()) { paramTypes.put(entry.getKey(), SqlType.fromProto(entry.getValue().getType())); @@ -699,6 +804,20 @@ public void executeQuery( responseObserver.onNext( new ResultSetSerializer(request.getProtoDescriptors()).toExecuteQueryResult(resultSet)); } catch (InterruptedException e) { + if (cancellableContext != null + && cancellableContext.cancellationCause() instanceof TimeoutException) { + Thread.interrupted(); + responseObserver.onNext( + ExecuteQueryResult.newBuilder() + .setStatus( + com.google.rpc.Status.newBuilder() + .setCode(Code.DEADLINE_EXCEEDED.getNumber()) + .setMessage(cancellableContext.cancellationCause().getMessage()) + .build()) + .build()); + responseObserver.onCompleted(); + return; + } responseObserver.onError(e); return; } catch (ExecutionException e) { @@ -714,6 +833,20 @@ public void executeQuery( responseObserver.onCompleted(); return; } catch (RuntimeException e) { + if (cancellableContext != null + && cancellableContext.cancellationCause() instanceof TimeoutException) { + Thread.interrupted(); + responseObserver.onNext( + ExecuteQueryResult.newBuilder() + .setStatus( + com.google.rpc.Status.newBuilder() + .setCode(Code.DEADLINE_EXCEEDED.getNumber()) + .setMessage(cancellableContext.cancellationCause().getMessage()) + .build()) + .build()); + responseObserver.onCompleted(); + return; + } // If client encounters problem, don't return any results. responseObserver.onNext( ExecuteQueryResult.newBuilder() @@ -729,6 +862,13 @@ public void executeQuery( if (resultSet != null) { resultSet.close(); } + if (previousContext != null) { + cancellableContext.detach(previousContext); + } + if (cancellableContext != null) { + cancellableContext.close(); + } + Thread.interrupted(); } responseObserver.onCompleted(); } @@ -741,6 +881,7 @@ public synchronized void close() { entry.getValue().dataClient().close(); it.remove(); } + deadlineScheduler.shutdownNow(); } private static String extractTableIdFromTableName(String fullTableName) @@ -824,6 +965,9 @@ private CredentialsProvider getCredentialsProvider(String accessToken) { @VisibleForTesting final ConcurrentHashMap idClientMap; + private final ScheduledExecutorService deadlineScheduler = + Executors.newSingleThreadScheduledExecutor(); + private static final Pattern tablePattern = Pattern.compile("projects/([^/]+)/instances/([^/]+)/tables/([^/]+)"); } diff --git a/java-bigtable/test-proxy/src/main/proto/test_proxy.proto b/java-bigtable/test-proxy/src/main/proto/test_proxy.proto index 34cf534425c2..f97861e1fdbe 100644 --- a/java-bigtable/test-proxy/src/main/proto/test_proxy.proto +++ b/java-bigtable/test-proxy/src/main/proto/test_proxy.proto @@ -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 @@ -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 @@ -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. @@ -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. @@ -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. @@ -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. @@ -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. @@ -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. @@ -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. diff --git a/java-bigtable/test-proxy/src/test/java/com/google/cloud/bigtable/testproxy/CbtTestProxyTest.java b/java-bigtable/test-proxy/src/test/java/com/google/cloud/bigtable/testproxy/CbtTestProxyTest.java index 25cc29e1d252..eba8cf9fc439 100644 --- a/java-bigtable/test-proxy/src/test/java/com/google/cloud/bigtable/testproxy/CbtTestProxyTest.java +++ b/java-bigtable/test-proxy/src/test/java/com/google/cloud/bigtable/testproxy/CbtTestProxyTest.java @@ -22,7 +22,21 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import com.google.bigtable.v2.BigtableGrpc; +import com.google.bigtable.v2.CheckAndMutateRowResponse; +import com.google.bigtable.v2.ColumnMetadata; +import com.google.bigtable.v2.ExecuteQueryResponse; +import com.google.bigtable.v2.MutateRowResponse; +import com.google.bigtable.v2.MutateRowsResponse; +import com.google.bigtable.v2.Mutation; +import com.google.bigtable.v2.PrepareQueryResponse; +import com.google.bigtable.v2.ProtoSchema; +import com.google.bigtable.v2.ReadModifyWriteRowResponse; import com.google.bigtable.v2.ReadModifyWriteRule; +import com.google.bigtable.v2.ReadRowsResponse; +import com.google.bigtable.v2.ResultSetMetadata; +import com.google.bigtable.v2.SampleRowKeysResponse; +import com.google.bigtable.v2.Type; import com.google.cloud.bigtable.data.v2.BigtableDataClient; import com.google.cloud.bigtable.data.v2.BigtableDataSettings; import com.google.cloud.bigtable.data.v2.models.Filters.Filter; @@ -32,8 +46,14 @@ import com.google.cloud.bigtable.testproxy.CbtTestProxy.CbtClient; import com.google.common.collect.ImmutableList; import com.google.protobuf.ByteString; +import com.google.protobuf.Timestamp; +import com.google.protobuf.util.Durations; import com.google.rpc.Code; +import io.grpc.Server; +import io.grpc.ServerBuilder; import io.grpc.stub.StreamObserver; +import java.time.Instant; +import org.junit.After; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; @@ -153,4 +173,561 @@ public void testReadModifyWriteRowNullEmitsEmptyRowResultWithOkStatus() { assertThat(result.getStatus().getCode()).isEqualTo(Code.OK.getNumber()); assertThat(result.hasRow()).isFalse(); } + + @After + public void tearDown() { + if (proxy != null) { + proxy.close(); + } + } + + private static class SlowBigtableService extends BigtableGrpc.BigtableImplBase { + private final long delayMillis; + + SlowBigtableService(long delayMillis) { + this.delayMillis = delayMillis; + } + + private boolean sleepDelay() { + try { + Thread.sleep(delayMillis); + return true; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return false; + } + } + + @Override + public void readRows( + com.google.bigtable.v2.ReadRowsRequest request, + StreamObserver responseObserver) { + if (!sleepDelay()) { + return; + } + responseObserver.onNext(ReadRowsResponse.getDefaultInstance()); + responseObserver.onCompleted(); + } + + @Override + public void mutateRow( + com.google.bigtable.v2.MutateRowRequest request, + StreamObserver responseObserver) { + if (!sleepDelay()) { + return; + } + responseObserver.onNext(MutateRowResponse.getDefaultInstance()); + responseObserver.onCompleted(); + } + + @Override + public void mutateRows( + com.google.bigtable.v2.MutateRowsRequest request, + StreamObserver responseObserver) { + if (!sleepDelay()) { + return; + } + MutateRowsResponse.Builder builder = MutateRowsResponse.newBuilder(); + for (int i = 0; i < request.getEntriesCount(); i++) { + builder + .addEntriesBuilder() + .setIndex(i) + .setStatus(com.google.rpc.Status.newBuilder().setCode(Code.OK.getNumber())); + } + responseObserver.onNext(builder.build()); + responseObserver.onCompleted(); + } + + @Override + public void sampleRowKeys( + com.google.bigtable.v2.SampleRowKeysRequest request, + StreamObserver responseObserver) { + if (!sleepDelay()) { + return; + } + responseObserver.onNext(SampleRowKeysResponse.getDefaultInstance()); + responseObserver.onCompleted(); + } + + @Override + public void checkAndMutateRow( + com.google.bigtable.v2.CheckAndMutateRowRequest request, + StreamObserver responseObserver) { + if (!sleepDelay()) { + return; + } + responseObserver.onNext( + CheckAndMutateRowResponse.newBuilder().setPredicateMatched(true).build()); + responseObserver.onCompleted(); + } + + @Override + public void readModifyWriteRow( + com.google.bigtable.v2.ReadModifyWriteRowRequest request, + StreamObserver responseObserver) { + if (!sleepDelay()) { + return; + } + responseObserver.onNext(ReadModifyWriteRowResponse.getDefaultInstance()); + responseObserver.onCompleted(); + } + + @Override + public void prepareQuery( + com.google.bigtable.v2.PrepareQueryRequest request, + StreamObserver responseObserver) { + responseObserver.onNext( + PrepareQueryResponse.newBuilder() + .setPreparedQuery(ByteString.copyFromUtf8("plan")) + .setValidUntil( + Timestamp.newBuilder() + .setSeconds(Instant.now().plusSeconds(3600).getEpochSecond()) + .build()) + .setMetadata( + ResultSetMetadata.newBuilder() + .setProtoSchema( + ProtoSchema.newBuilder() + .addColumns( + ColumnMetadata.newBuilder() + .setName("col") + .setType( + Type.newBuilder() + .setStringType(Type.String.getDefaultInstance()))))) + .build()); + responseObserver.onCompleted(); + } + + @Override + public void executeQuery( + com.google.bigtable.v2.ExecuteQueryRequest request, + StreamObserver responseObserver) { + if (!sleepDelay()) { + return; + } + responseObserver.onCompleted(); + } + } + + @Test + public void testPerOperationTimeoutSupersedesShorterClientTimeout() throws Exception { + Server server = + ServerBuilder.forPort(0).addService(new SlowBigtableService(200)).build().start(); + try { + String clientId = "short-timeout-client"; + @SuppressWarnings("unchecked") + StreamObserver createObserver = mock(StreamObserver.class); + proxy.createClient( + CreateClientRequest.newBuilder() + .setClientId(clientId) + .setDataTarget("localhost:" + server.getPort()) + .setProjectId("my-project") + .setInstanceId("my-instance") + .setPerOperationTimeout(Durations.fromMillis(50)) + .build(), + createObserver); + verify(createObserver).onCompleted(); + + // Without a per-request override, the 50ms client timeout expires during the 200ms delay. + @SuppressWarnings("unchecked") + StreamObserver defaultReadObserver = mock(StreamObserver.class); + proxy.readRow( + ReadRowRequest.newBuilder() + .setClientId(clientId) + .setTableName(TABLE_NAME) + .setRowKey("row-1") + .build(), + defaultReadObserver); + ArgumentCaptor defaultReadCaptor = ArgumentCaptor.forClass(RowResult.class); + verify(defaultReadObserver).onNext(defaultReadCaptor.capture()); + assertThat(defaultReadCaptor.getValue().getStatus().getCode()) + .isEqualTo(Code.DEADLINE_EXCEEDED.getNumber()); + + // A longer 2s per-operation timeout on ReadRowRequest supersedes the 50ms client timeout. + @SuppressWarnings("unchecked") + StreamObserver overrideReadObserver = mock(StreamObserver.class); + proxy.readRow( + ReadRowRequest.newBuilder() + .setClientId(clientId) + .setTableName(TABLE_NAME) + .setRowKey("row-1") + .setPerOperationTimeout(Durations.fromSeconds(2)) + .build(), + overrideReadObserver); + ArgumentCaptor overrideReadCaptor = ArgumentCaptor.forClass(RowResult.class); + verify(overrideReadObserver).onNext(overrideReadCaptor.capture()); + assertThat(overrideReadCaptor.getValue().getStatus().getCode()) + .isEqualTo(Code.OK.getNumber()); + + // A longer 2s per-operation timeout on MutateRowRequest supersedes the 50ms client timeout. + @SuppressWarnings("unchecked") + StreamObserver mutateObserver = mock(StreamObserver.class); + proxy.mutateRow( + MutateRowRequest.newBuilder() + .setClientId(clientId) + .setRequest( + com.google.bigtable.v2.MutateRowRequest.newBuilder() + .setTableName(TABLE_NAME) + .setRowKey(ByteString.copyFromUtf8("row-1")) + .addMutations( + Mutation.newBuilder() + .setSetCell( + Mutation.SetCell.newBuilder() + .setFamilyName("cf") + .setColumnQualifier(ByteString.copyFromUtf8("q")) + .setTimestampMicros(1000L) + .setValue(ByteString.copyFromUtf8("v"))))) + .setPerOperationTimeout(Durations.fromSeconds(2)) + .build(), + mutateObserver); + ArgumentCaptor mutateCaptor = ArgumentCaptor.forClass(MutateRowResult.class); + verify(mutateObserver).onNext(mutateCaptor.capture()); + assertThat(mutateCaptor.getValue().getStatus().getCode()).isEqualTo(Code.OK.getNumber()); + + // A longer 2s per-operation timeout on MutateRowsRequest supersedes the 50ms client timeout. + @SuppressWarnings("unchecked") + StreamObserver bulkObserver = mock(StreamObserver.class); + proxy.bulkMutateRows( + MutateRowsRequest.newBuilder() + .setClientId(clientId) + .setRequest( + com.google.bigtable.v2.MutateRowsRequest.newBuilder() + .setTableName(TABLE_NAME) + .addEntries( + com.google.bigtable.v2.MutateRowsRequest.Entry.newBuilder() + .setRowKey(ByteString.copyFromUtf8("row-1")) + .addMutations( + Mutation.newBuilder() + .setSetCell( + Mutation.SetCell.newBuilder() + .setFamilyName("cf") + .setColumnQualifier(ByteString.copyFromUtf8("q")) + .setTimestampMicros(1000L) + .setValue(ByteString.copyFromUtf8("v")))))) + .setPerOperationTimeout(Durations.fromSeconds(2)) + .build(), + bulkObserver); + ArgumentCaptor bulkCaptor = ArgumentCaptor.forClass(MutateRowsResult.class); + verify(bulkObserver).onNext(bulkCaptor.capture()); + assertThat(bulkCaptor.getValue().getStatus().getCode()).isEqualTo(Code.OK.getNumber()); + + // A longer 2s per-operation timeout on ReadRowsRequest supersedes the 50ms client timeout. + @SuppressWarnings("unchecked") + StreamObserver readRowsObserver = mock(StreamObserver.class); + proxy.readRows( + ReadRowsRequest.newBuilder() + .setClientId(clientId) + .setRequest( + com.google.bigtable.v2.ReadRowsRequest.newBuilder().setTableName(TABLE_NAME)) + .setPerOperationTimeout(Durations.fromSeconds(2)) + .build(), + readRowsObserver); + ArgumentCaptor readRowsCaptor = ArgumentCaptor.forClass(RowsResult.class); + verify(readRowsObserver).onNext(readRowsCaptor.capture()); + assertThat(readRowsCaptor.getValue().getStatus().getCode()).isEqualTo(Code.OK.getNumber()); + + // A longer 2s per-operation timeout on SampleRowKeysRequest supersedes the 50ms client + // timeout. + @SuppressWarnings("unchecked") + StreamObserver sampleObserver = mock(StreamObserver.class); + proxy.sampleRowKeys( + SampleRowKeysRequest.newBuilder() + .setClientId(clientId) + .setRequest( + com.google.bigtable.v2.SampleRowKeysRequest.newBuilder().setTableName(TABLE_NAME)) + .setPerOperationTimeout(Durations.fromSeconds(2)) + .build(), + sampleObserver); + ArgumentCaptor sampleCaptor = + ArgumentCaptor.forClass(SampleRowKeysResult.class); + verify(sampleObserver).onNext(sampleCaptor.capture()); + assertThat(sampleCaptor.getValue().getStatus().getCode()).isEqualTo(Code.OK.getNumber()); + + // A longer 2s per-operation timeout on CheckAndMutateRowRequest supersedes the 50ms client + // timeout. + @SuppressWarnings("unchecked") + StreamObserver camObserver = mock(StreamObserver.class); + proxy.checkAndMutateRow( + CheckAndMutateRowRequest.newBuilder() + .setClientId(clientId) + .setRequest( + com.google.bigtable.v2.CheckAndMutateRowRequest.newBuilder() + .setTableName(TABLE_NAME) + .setRowKey(ByteString.copyFromUtf8("row-1")) + .addTrueMutations( + Mutation.newBuilder() + .setSetCell( + Mutation.SetCell.newBuilder() + .setFamilyName("cf") + .setColumnQualifier(ByteString.copyFromUtf8("q")) + .setTimestampMicros(1000L) + .setValue(ByteString.copyFromUtf8("v"))))) + .setPerOperationTimeout(Durations.fromSeconds(2)) + .build(), + camObserver); + ArgumentCaptor camCaptor = + ArgumentCaptor.forClass(CheckAndMutateRowResult.class); + verify(camObserver).onNext(camCaptor.capture()); + assertThat(camCaptor.getValue().getStatus().getCode()).isEqualTo(Code.OK.getNumber()); + + // A longer 2s per-operation timeout on ReadModifyWriteRowRequest supersedes the 50ms client + // timeout. + @SuppressWarnings("unchecked") + StreamObserver rmwObserver = mock(StreamObserver.class); + proxy.readModifyWriteRow( + ReadModifyWriteRowRequest.newBuilder() + .setClientId(clientId) + .setRequest( + com.google.bigtable.v2.ReadModifyWriteRowRequest.newBuilder() + .setTableName(TABLE_NAME) + .setRowKey(ByteString.copyFromUtf8("row-1")) + .addRules( + ReadModifyWriteRule.newBuilder() + .setFamilyName("cf") + .setColumnQualifier(ByteString.copyFromUtf8("q")) + .setAppendValue(ByteString.copyFromUtf8("v")))) + .setPerOperationTimeout(Durations.fromSeconds(2)) + .build(), + rmwObserver); + ArgumentCaptor rmwCaptor = ArgumentCaptor.forClass(RowResult.class); + verify(rmwObserver).onNext(rmwCaptor.capture()); + assertThat(rmwCaptor.getValue().getStatus().getCode()).isEqualTo(Code.OK.getNumber()); + + // A longer 2s per-operation timeout on ExecuteQueryRequest supersedes the 50ms client + // timeout, whereas without a per-request override the 50ms client timeout expires. + @SuppressWarnings("unchecked") + StreamObserver defaultQueryObserver = mock(StreamObserver.class); + proxy.executeQuery( + ExecuteQueryRequest.newBuilder() + .setClientId(clientId) + .setRequest( + com.google.bigtable.v2.ExecuteQueryRequest.newBuilder() + .setInstanceName("projects/my-project/instances/my-instance") + .setQuery("SELECT * FROM my-table")) + .build(), + defaultQueryObserver); + ArgumentCaptor defaultQueryCaptor = + ArgumentCaptor.forClass(ExecuteQueryResult.class); + verify(defaultQueryObserver).onNext(defaultQueryCaptor.capture()); + assertThat(defaultQueryCaptor.getValue().getStatus().getCode()) + .isEqualTo(Code.DEADLINE_EXCEEDED.getNumber()); + + @SuppressWarnings("unchecked") + StreamObserver queryObserver = mock(StreamObserver.class); + proxy.executeQuery( + ExecuteQueryRequest.newBuilder() + .setClientId(clientId) + .setRequest( + com.google.bigtable.v2.ExecuteQueryRequest.newBuilder() + .setInstanceName("projects/my-project/instances/my-instance") + .setQuery("SELECT * FROM my-table")) + .setPerOperationTimeout(Durations.fromSeconds(2)) + .build(), + queryObserver); + ArgumentCaptor queryCaptor = + ArgumentCaptor.forClass(ExecuteQueryResult.class); + verify(queryObserver).onNext(queryCaptor.capture()); + assertThat(queryCaptor.getValue().getStatus().getCode()).isEqualTo(Code.OK.getNumber()); + } finally { + server.shutdownNow(); + } + } + + @Test + public void testPerOperationTimeoutSupersedesLongerClientTimeout() throws Exception { + Server server = + ServerBuilder.forPort(0).addService(new SlowBigtableService(200)).build().start(); + try { + String clientId = "long-timeout-client"; + @SuppressWarnings("unchecked") + StreamObserver createObserver = mock(StreamObserver.class); + proxy.createClient( + CreateClientRequest.newBuilder() + .setClientId(clientId) + .setDataTarget("localhost:" + server.getPort()) + .setProjectId("my-project") + .setInstanceId("my-instance") + .setPerOperationTimeout(Durations.fromSeconds(5)) + .build(), + createObserver); + verify(createObserver).onCompleted(); + + // A shorter 50ms per-operation timeout on ReadRowRequest supersedes the 5s client timeout. + @SuppressWarnings("unchecked") + StreamObserver readObserver = mock(StreamObserver.class); + proxy.readRow( + ReadRowRequest.newBuilder() + .setClientId(clientId) + .setTableName(TABLE_NAME) + .setRowKey("row-1") + .setPerOperationTimeout(Durations.fromMillis(50)) + .build(), + readObserver); + ArgumentCaptor readCaptor = ArgumentCaptor.forClass(RowResult.class); + verify(readObserver).onNext(readCaptor.capture()); + assertThat(readCaptor.getValue().getStatus().getCode()) + .isEqualTo(Code.DEADLINE_EXCEEDED.getNumber()); + + // A shorter 50ms per-operation timeout on MutateRowRequest supersedes the 5s client timeout. + @SuppressWarnings("unchecked") + StreamObserver mutateObserver = mock(StreamObserver.class); + proxy.mutateRow( + MutateRowRequest.newBuilder() + .setClientId(clientId) + .setRequest( + com.google.bigtable.v2.MutateRowRequest.newBuilder() + .setTableName(TABLE_NAME) + .setRowKey(ByteString.copyFromUtf8("row-1")) + .addMutations( + Mutation.newBuilder() + .setSetCell( + Mutation.SetCell.newBuilder() + .setFamilyName("cf") + .setColumnQualifier(ByteString.copyFromUtf8("q")) + .setTimestampMicros(1000L) + .setValue(ByteString.copyFromUtf8("v"))))) + .setPerOperationTimeout(Durations.fromMillis(50)) + .build(), + mutateObserver); + ArgumentCaptor mutateCaptor = ArgumentCaptor.forClass(MutateRowResult.class); + verify(mutateObserver).onNext(mutateCaptor.capture()); + assertThat(mutateCaptor.getValue().getStatus().getCode()) + .isEqualTo(Code.DEADLINE_EXCEEDED.getNumber()); + + // A shorter 50ms per-operation timeout on MutateRowsRequest supersedes the 5s client timeout. + @SuppressWarnings("unchecked") + StreamObserver bulkObserver = mock(StreamObserver.class); + proxy.bulkMutateRows( + MutateRowsRequest.newBuilder() + .setClientId(clientId) + .setRequest( + com.google.bigtable.v2.MutateRowsRequest.newBuilder() + .setTableName(TABLE_NAME) + .addEntries( + com.google.bigtable.v2.MutateRowsRequest.Entry.newBuilder() + .setRowKey(ByteString.copyFromUtf8("row-1")) + .addMutations( + Mutation.newBuilder() + .setSetCell( + Mutation.SetCell.newBuilder() + .setFamilyName("cf") + .setColumnQualifier(ByteString.copyFromUtf8("q")) + .setTimestampMicros(1000L) + .setValue(ByteString.copyFromUtf8("v")))))) + .setPerOperationTimeout(Durations.fromMillis(50)) + .build(), + bulkObserver); + ArgumentCaptor bulkCaptor = ArgumentCaptor.forClass(MutateRowsResult.class); + verify(bulkObserver).onNext(bulkCaptor.capture()); + assertThat(bulkCaptor.getValue().getStatus().getCode()) + .isEqualTo(Code.DEADLINE_EXCEEDED.getNumber()); + + // A shorter 50ms per-operation timeout on ReadRowsRequest supersedes the 5s client timeout. + @SuppressWarnings("unchecked") + StreamObserver readRowsObserver = mock(StreamObserver.class); + proxy.readRows( + ReadRowsRequest.newBuilder() + .setClientId(clientId) + .setRequest( + com.google.bigtable.v2.ReadRowsRequest.newBuilder().setTableName(TABLE_NAME)) + .setPerOperationTimeout(Durations.fromMillis(50)) + .build(), + readRowsObserver); + ArgumentCaptor readRowsCaptor = ArgumentCaptor.forClass(RowsResult.class); + verify(readRowsObserver).onNext(readRowsCaptor.capture()); + assertThat(readRowsCaptor.getValue().getStatus().getCode()) + .isEqualTo(Code.DEADLINE_EXCEEDED.getNumber()); + + // A shorter 50ms per-operation timeout on SampleRowKeysRequest supersedes the 5s client + // timeout. + @SuppressWarnings("unchecked") + StreamObserver sampleObserver = mock(StreamObserver.class); + proxy.sampleRowKeys( + SampleRowKeysRequest.newBuilder() + .setClientId(clientId) + .setRequest( + com.google.bigtable.v2.SampleRowKeysRequest.newBuilder().setTableName(TABLE_NAME)) + .setPerOperationTimeout(Durations.fromMillis(50)) + .build(), + sampleObserver); + ArgumentCaptor sampleCaptor = + ArgumentCaptor.forClass(SampleRowKeysResult.class); + verify(sampleObserver).onNext(sampleCaptor.capture()); + assertThat(sampleCaptor.getValue().getStatus().getCode()) + .isEqualTo(Code.DEADLINE_EXCEEDED.getNumber()); + + // A shorter 50ms per-operation timeout on CheckAndMutateRowRequest supersedes the 5s client + // timeout. + @SuppressWarnings("unchecked") + StreamObserver camObserver = mock(StreamObserver.class); + proxy.checkAndMutateRow( + CheckAndMutateRowRequest.newBuilder() + .setClientId(clientId) + .setRequest( + com.google.bigtable.v2.CheckAndMutateRowRequest.newBuilder() + .setTableName(TABLE_NAME) + .setRowKey(ByteString.copyFromUtf8("row-1")) + .addTrueMutations( + Mutation.newBuilder() + .setSetCell( + Mutation.SetCell.newBuilder() + .setFamilyName("cf") + .setColumnQualifier(ByteString.copyFromUtf8("q")) + .setTimestampMicros(1000L) + .setValue(ByteString.copyFromUtf8("v"))))) + .setPerOperationTimeout(Durations.fromMillis(50)) + .build(), + camObserver); + ArgumentCaptor camCaptor = + ArgumentCaptor.forClass(CheckAndMutateRowResult.class); + verify(camObserver).onNext(camCaptor.capture()); + assertThat(camCaptor.getValue().getStatus().getCode()) + .isEqualTo(Code.DEADLINE_EXCEEDED.getNumber()); + + // A shorter 50ms per-operation timeout on ReadModifyWriteRowRequest supersedes the 5s client + // timeout. + @SuppressWarnings("unchecked") + StreamObserver rmwObserver = mock(StreamObserver.class); + proxy.readModifyWriteRow( + ReadModifyWriteRowRequest.newBuilder() + .setClientId(clientId) + .setRequest( + com.google.bigtable.v2.ReadModifyWriteRowRequest.newBuilder() + .setTableName(TABLE_NAME) + .setRowKey(ByteString.copyFromUtf8("row-1")) + .addRules( + ReadModifyWriteRule.newBuilder() + .setFamilyName("cf") + .setColumnQualifier(ByteString.copyFromUtf8("q")) + .setAppendValue(ByteString.copyFromUtf8("v")))) + .setPerOperationTimeout(Durations.fromMillis(50)) + .build(), + rmwObserver); + ArgumentCaptor rmwCaptor = ArgumentCaptor.forClass(RowResult.class); + verify(rmwObserver).onNext(rmwCaptor.capture()); + assertThat(rmwCaptor.getValue().getStatus().getCode()) + .isEqualTo(Code.DEADLINE_EXCEEDED.getNumber()); + + // A shorter 50ms per-operation timeout on ExecuteQueryRequest supersedes the 5s client + // timeout. + @SuppressWarnings("unchecked") + StreamObserver queryObserver = mock(StreamObserver.class); + proxy.executeQuery( + ExecuteQueryRequest.newBuilder() + .setClientId(clientId) + .setRequest( + com.google.bigtable.v2.ExecuteQueryRequest.newBuilder() + .setInstanceName("projects/my-project/instances/my-instance") + .setQuery("SELECT * FROM my-table")) + .setPerOperationTimeout(Durations.fromMillis(50)) + .build(), + queryObserver); + ArgumentCaptor queryCaptor = + ArgumentCaptor.forClass(ExecuteQueryResult.class); + verify(queryObserver).onNext(queryCaptor.capture()); + assertThat(queryCaptor.getValue().getStatus().getCode()) + .isEqualTo(Code.DEADLINE_EXCEEDED.getNumber()); + } finally { + server.shutdownNow(); + } + } }