Skip to content

Commit

Permalink
fix: time based flakiness in execute query deadline test (#2358)
Browse files Browse the repository at this point in the history
Change-Id: I93c1c03a0c41c92dbe65b5ec4888e7df526ad457
  • Loading branch information
jackdingilian authored Sep 26, 2024
1 parent 6bc9820 commit b474173
Showing 1 changed file with 17 additions and 9 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import static org.junit.Assert.assertThrows;

import com.google.api.gax.retrying.RetrySettings;
import com.google.api.gax.rpc.DeadlineExceededException;
import com.google.api.gax.rpc.UnavailableException;
import com.google.bigtable.v2.BigtableGrpc;
import com.google.bigtable.v2.ExecuteQueryRequest;
Expand All @@ -36,7 +37,6 @@
import com.google.cloud.bigtable.data.v2.models.sql.Statement;
import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub;
import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi.ServerStreamingStashCallable;
import com.google.common.collect.Range;
import io.grpc.Context;
import io.grpc.Deadline;
import io.grpc.Server;
Expand Down Expand Up @@ -141,12 +141,13 @@ public void testExecuteQueryRequestsSetDefaultDeadline() {
Iterator<SqlRow> iterator = stream.rows().iterator();
// We don't care about this but are reusing the fake service that tests retries
assertThrows(UnavailableException.class, iterator::next).getCause();
// We have 30s default, we assume less than 1s has been burned when the fake service sets it
assertThat(fakeService.deadlineMillisRemaining).isIn(Range.closed(29000L, 30100L));
// We have 30s default, we give it a wide range to avoid flakiness, this is mostly just checking
// that some default is set
assertThat(fakeService.deadlineMillisRemaining).isLessThan(30001L);
}

@Test
public void testExecuteQueryRequestsRespectOverridenDeadline() throws IOException {
public void testExecuteQueryRequestsRespectDeadline() throws IOException {
BigtableDataSettings.Builder overrideSettings =
BigtableDataSettings.newBuilderForEmulator(server.getPort())
.setProjectId("fake-project")
Expand All @@ -156,18 +157,16 @@ public void testExecuteQueryRequestsRespectOverridenDeadline() throws IOExceptio
.executeQuerySettings()
.setRetrySettings(
RetrySettings.newBuilder()
.setInitialRpcTimeout(Duration.ofMinutes(5))
.setMaxRpcTimeout(Duration.ofMinutes(5))
.setInitialRpcTimeout(Duration.ofMillis(10))
.setMaxRpcTimeout(Duration.ofMillis(10))
.build());
EnhancedBigtableStub overrideDeadline =
EnhancedBigtableStub.create(overrideSettings.build().getStubSettings());
SqlServerStream streamOverride =
overrideDeadline.executeQueryCallable().call(Statement.of("SELECT * FROM table"));
Iterator<SqlRow> overrideIterator = streamOverride.rows().iterator();
// We don't care about this but are reusing the fake service that tests retries
assertThrows(UnavailableException.class, overrideIterator::next).getCause();
// We have 30s default, we assume less than 1s has been burned when the fake service sets it
assertThat(fakeService.deadlineMillisRemaining).isIn(Range.closed(299000L, 300100L));
assertThrows(DeadlineExceededException.class, overrideIterator::next).getCause();
}

private static class FakeService extends BigtableGrpc.BigtableImplBase {
Expand All @@ -181,6 +180,15 @@ public void executeQuery(
Deadline deadline = Context.current().getDeadline();
if (deadline != null) {
deadlineMillisRemaining = deadline.timeRemaining(TimeUnit.MILLISECONDS);
} else {
// set to max long when deadline isn't set
deadlineMillisRemaining = Long.MAX_VALUE;
}
// Sleep for 100ms to trigger deadline exceeded for tests with a shorter deadline
try {
Thread.sleep(100);
} catch (InterruptedException e) {
throw new RuntimeException(e);
}
attempts++;
responseObserver.onNext(metadata(columnMetadata("test", stringType())));
Expand Down

0 comments on commit b474173

Please sign in to comment.