Skip to content

Commit

Permalink
fix: set wait timeout on watchdog (#1913)
Browse files Browse the repository at this point in the history
And add wait timeout on CDC settings.

Copy of https://togithub.com/googleapis/java-bigtable/pull/1898 with lower cdc timeouts
  • Loading branch information
igorbernstein2 committed Sep 13, 2023
1 parent 8843813 commit f8ba6fe
Show file tree
Hide file tree
Showing 3 changed files with 83 additions and 6 deletions.
Expand Up @@ -458,6 +458,7 @@ public Map<String, String> extract(ReadRowsRequest readRowsRequest) {
.setRetryableCodes(readRowsSettings.getRetryableCodes())
.setRetrySettings(readRowsSettings.getRetrySettings())
.setIdleTimeout(readRowsSettings.getIdleTimeout())
.setWaitTimeout(readRowsSettings.getWaitTimeout())
.build();

ServerStreamingCallable<ReadRowsRequest, RowT> watched =
Expand Down Expand Up @@ -906,6 +907,8 @@ public Map<String, String> extract(
settings.generateInitialChangeStreamPartitionsSettings().getRetrySettings())
.setIdleTimeout(
settings.generateInitialChangeStreamPartitionsSettings().getIdleTimeout())
.setWaitTimeout(
settings.generateInitialChangeStreamPartitionsSettings().getWaitTimeout())
.build();

ServerStreamingCallable<String, ByteStringRange> watched =
Expand Down Expand Up @@ -980,6 +983,7 @@ public Map<String, String> extract(
.setRetryableCodes(settings.readChangeStreamSettings().getRetryableCodes())
.setRetrySettings(settings.readChangeStreamSettings().getRetrySettings())
.setIdleTimeout(settings.readChangeStreamSettings().getIdleTimeout())
.setWaitTimeout(settings.readChangeStreamSettings().getWaitTimeout())
.build();

ServerStreamingCallable<ReadChangeStreamRequest, ChangeStreamRecordT> watched =
Expand Down
Expand Up @@ -714,13 +714,15 @@ private Builder() {
generateInitialChangeStreamPartitionsSettings
.setRetryableCodes(GENERATE_INITIAL_CHANGE_STREAM_PARTITIONS_RETRY_CODES)
.setRetrySettings(GENERATE_INITIAL_CHANGE_STREAM_PARTITIONS_RETRY_SETTINGS)
.setIdleTimeout(Duration.ofMinutes(5));
.setIdleTimeout(Duration.ofMinutes(5))
.setWaitTimeout(Duration.ofMinutes(1));

readChangeStreamSettings = ServerStreamingCallSettings.newBuilder();
readChangeStreamSettings
.setRetryableCodes(READ_CHANGE_STREAM_RETRY_CODES)
.setRetrySettings(READ_CHANGE_STREAM_RETRY_SETTINGS)
.setIdleTimeout(Duration.ofMinutes(5));
.setIdleTimeout(Duration.ofMinutes(5))
.setWaitTimeout(Duration.ofMinutes(1));

pingAndWarmSettings = UnaryCallSettings.newUnaryCallSettingsBuilder();
pingAndWarmSettings.setRetrySettings(
Expand Down
Expand Up @@ -30,14 +30,18 @@
import com.google.api.gax.grpc.GrpcCallContext;
import com.google.api.gax.grpc.GrpcTransportChannel;
import com.google.api.gax.rpc.FixedTransportChannelProvider;
import com.google.api.gax.rpc.InstantiatingWatchdogProvider;
import com.google.api.gax.rpc.ServerStreamingCallable;
import com.google.api.gax.rpc.WatchdogTimeoutException;
import com.google.auth.oauth2.ServiceAccountJwtAccessCredentials;
import com.google.bigtable.v2.BigtableGrpc;
import com.google.bigtable.v2.FeatureFlags;
import com.google.bigtable.v2.MutateRowsRequest;
import com.google.bigtable.v2.MutateRowsResponse;
import com.google.bigtable.v2.PingAndWarmRequest;
import com.google.bigtable.v2.PingAndWarmResponse;
import com.google.bigtable.v2.ReadChangeStreamRequest;
import com.google.bigtable.v2.ReadChangeStreamResponse;
import com.google.bigtable.v2.ReadRowsRequest;
import com.google.bigtable.v2.ReadRowsResponse;
import com.google.bigtable.v2.RowSet;
Expand All @@ -46,11 +50,8 @@
import com.google.cloud.bigtable.data.v2.BigtableDataSettings;
import com.google.cloud.bigtable.data.v2.FakeServiceBuilder;
import com.google.cloud.bigtable.data.v2.internal.RequestContext;
import com.google.cloud.bigtable.data.v2.models.BulkMutation;
import com.google.cloud.bigtable.data.v2.models.DefaultRowAdapter;
import com.google.cloud.bigtable.data.v2.models.Query;
import com.google.cloud.bigtable.data.v2.models.*;
import com.google.cloud.bigtable.data.v2.models.Row;
import com.google.cloud.bigtable.data.v2.models.RowMutationEntry;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Queues;
import com.google.common.io.BaseEncoding;
Expand Down Expand Up @@ -82,11 +83,13 @@
import java.security.NoSuchAlgorithmException;
import java.util.Base64;
import java.util.Collection;
import java.util.Iterator;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
Expand All @@ -101,6 +104,8 @@ public class EnhancedBigtableStubTest {
private static final String TABLE_NAME =
NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, "fake-table");
private static final String APP_PROFILE_ID = "app-profile-id";
private static final String WAIT_TIME_TABLE_ID = "test-wait-timeout";
private static final Duration WATCHDOG_CHECK_DURATION = Duration.ofMillis(100);

private Server server;
private MetadataInterceptor metadataInterceptor;
Expand Down Expand Up @@ -544,6 +549,46 @@ public void testBulkMutationFlowControlFeatureFlagIsNotSet() throws Exception {
assertThat(featureFlags.getMutateRowsRateLimit()).isFalse();
}

@Test
public void testWaitTimeoutIsSet() throws Exception {
EnhancedBigtableStubSettings.Builder settings = defaultSettings.toBuilder();
// Set a shorter wait timeout and make watchdog checks more frequently
settings.readRowsSettings().setWaitTimeout(WATCHDOG_CHECK_DURATION.dividedBy(2));
settings.setStreamWatchdogProvider(
InstantiatingWatchdogProvider.create().withCheckInterval(WATCHDOG_CHECK_DURATION));

EnhancedBigtableStub stub = EnhancedBigtableStub.create(settings.build());
Iterator<Row> iterator =
stub.readRowsCallable().call(Query.create(WAIT_TIME_TABLE_ID)).iterator();
try {
iterator.next();
Assert.fail("Should throw watchdog timeout exception");
} catch (WatchdogTimeoutException e) {
assertThat(e.getMessage()).contains("Canceled due to timeout waiting for next response");
}
}

@Test
public void testReadChangeStreamWaitTimeoutIsSet() throws Exception {
EnhancedBigtableStubSettings.Builder settings = defaultSettings.toBuilder();
// Set a shorter wait timeout and make watchdog checks more frequently
settings.readChangeStreamSettings().setWaitTimeout(WATCHDOG_CHECK_DURATION.dividedBy(2));
settings.setStreamWatchdogProvider(
InstantiatingWatchdogProvider.create().withCheckInterval(WATCHDOG_CHECK_DURATION));

EnhancedBigtableStub stub = EnhancedBigtableStub.create(settings.build());
Iterator<ChangeStreamRecord> iterator =
stub.readChangeStreamCallable()
.call(ReadChangeStreamQuery.create(WAIT_TIME_TABLE_ID))
.iterator();
try {
iterator.next();
Assert.fail("Should throw watchdog timeout exception");
} catch (WatchdogTimeoutException e) {
assertThat(e.getMessage()).contains("Canceled due to timeout waiting for next response");
}
}

private static class MetadataInterceptor implements ServerInterceptor {
final BlockingQueue<Metadata> headers = Queues.newLinkedBlockingDeque();

Expand Down Expand Up @@ -572,6 +617,8 @@ public <ReqT, RespT> Listener<ReqT> interceptCall(

private static class FakeDataService extends BigtableGrpc.BigtableImplBase {
final BlockingQueue<ReadRowsRequest> requests = Queues.newLinkedBlockingDeque();
final BlockingQueue<ReadChangeStreamRequest> readChangeReadStreamRequests =
Queues.newLinkedBlockingDeque();
final BlockingQueue<PingAndWarmRequest> pingRequests = Queues.newLinkedBlockingDeque();

@SuppressWarnings("unchecked")
Expand All @@ -593,6 +640,13 @@ public void mutateRows(
@Override
public void readRows(
ReadRowsRequest request, StreamObserver<ReadRowsResponse> responseObserver) {
if (request.getTableName().contains(WAIT_TIME_TABLE_ID)) {
try {
Thread.sleep(WATCHDOG_CHECK_DURATION.toMillis() * 2);
} catch (Exception e) {

}
}
requests.add(request);
// Dummy row for stream
responseObserver.onNext(
Expand All @@ -608,6 +662,23 @@ public void readRows(
responseObserver.onCompleted();
}

@Override
public void readChangeStream(
ReadChangeStreamRequest request,
StreamObserver<ReadChangeStreamResponse> responseObserver) {
if (request.getTableName().contains(WAIT_TIME_TABLE_ID)) {
try {
Thread.sleep(WATCHDOG_CHECK_DURATION.toMillis() * 2);
} catch (Exception e) {

}
}
readChangeReadStreamRequests.add(request);
// Dummy row for stream
responseObserver.onNext(ReadChangeStreamResponse.getDefaultInstance());
responseObserver.onCompleted();
}

@Override
public void pingAndWarm(
PingAndWarmRequest request, StreamObserver<PingAndWarmResponse> responseObserver) {
Expand Down

0 comments on commit f8ba6fe

Please sign in to comment.