Skip to content
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.

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

fix: set wait timeout on watchdog #1913

Merged
merged 8 commits into from Sep 13, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
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));
igorbernstein2 marked this conversation as resolved.
Show resolved Hide resolved
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