Skip to content

Commit 9b3c601

Browse files
igorbernstein2igorbernstein2
authored
feat: promote stream wait timeouts to deadlines for point reads (#848)
Special case point reads to use grpc's deadlines instead of relying on the watchdog
1 parent 66a9c9e commit 9b3c601

File tree

3 files changed

+276
-2
lines changed

3 files changed

+276
-2
lines changed

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java

+7-2
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@
7575
import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsBatchingDescriptor;
7676
import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsRetryingCallable;
7777
import com.google.cloud.bigtable.data.v2.stub.readrows.FilterMarkerRowsCallable;
78+
import com.google.cloud.bigtable.data.v2.stub.readrows.PointReadTimeoutCallable;
7879
import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsBatchingDescriptor;
7980
import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsConvertExceptionCallable;
8081
import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsResumptionStrategy;
@@ -336,7 +337,7 @@ public <RowT> UnaryCallable<Query, RowT> createReadRowCallable(RowAdapter<RowT>
336337
private <ReqT, RowT> ServerStreamingCallable<ReadRowsRequest, RowT> createReadRowsBaseCallable(
337338
ServerStreamingCallSettings<ReqT, Row> readRowsSettings, RowAdapter<RowT> rowAdapter) {
338339

339-
ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> base =
340+
final ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> base =
340341
GrpcRawCallableFactory.createServerStreamingCallable(
341342
GrpcCallSettings.<ReadRowsRequest, ReadRowsResponse>newBuilder()
342343
.setMethodDescriptor(BigtableGrpc.getReadRowsMethod())
@@ -352,11 +353,15 @@ public Map<String, String> extract(ReadRowsRequest readRowsRequest) {
352353
.build(),
353354
readRowsSettings.getRetryableCodes());
354355

356+
// Promote streamWaitTimeout to deadline for point reads
357+
ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> withPointTimeouts =
358+
new PointReadTimeoutCallable<>(base);
359+
355360
// Sometimes ReadRows connections are disconnected via an RST frame. This error is transient and
356361
// should be treated similar to UNAVAILABLE. However, this exception has an INTERNAL error code
357362
// which by default is not retryable. Convert the exception so it can be retried in the client.
358363
ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> convertException =
359-
new ReadRowsConvertExceptionCallable<>(base);
364+
new ReadRowsConvertExceptionCallable<>(withPointTimeouts);
360365

361366
ServerStreamingCallable<ReadRowsRequest, RowT> merging =
362367
new RowMergingCallable<>(convertException, rowAdapter);

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/PointReadTimeoutCallable.java

+86
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
/*
2+
* Copyright 2021 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.data.v2.stub.readrows;
17+
18+
import com.google.api.core.InternalApi;
19+
import com.google.api.gax.rpc.ApiCallContext;
20+
import com.google.api.gax.rpc.ResponseObserver;
21+
import com.google.api.gax.rpc.ServerStreamingCallable;
22+
import com.google.bigtable.v2.ReadRowsRequest;
23+
import javax.annotation.Nullable;
24+
import org.threeten.bp.Duration;
25+
26+
/**
27+
* Specialization of ReadRows streams for point reads.
28+
*
29+
* <p>Under normal circumstances, the ReadRows RPC can't make any assumptions about deadlines. In
30+
* general case the end user can be issuing a full table scan. However, when dealing with point
31+
* reads, the client can make assumptions and promote the per row timeout to be a per attempt
32+
* timeout.
33+
*
34+
* <p>This callable will check if the request is a point read and promote the timeout to be a
35+
* deadline.
36+
*/
37+
@InternalApi
38+
public class PointReadTimeoutCallable<RespT>
39+
extends ServerStreamingCallable<ReadRowsRequest, RespT> {
40+
private final ServerStreamingCallable<ReadRowsRequest, RespT> inner;
41+
42+
public PointReadTimeoutCallable(ServerStreamingCallable<ReadRowsRequest, RespT> inner) {
43+
this.inner = inner;
44+
}
45+
46+
@Override
47+
public void call(ReadRowsRequest request, ResponseObserver<RespT> observer, ApiCallContext ctx) {
48+
if (isPointRead(request)) {
49+
Duration effectiveTimeout = getEffectivePointReadTimeout(ctx);
50+
if (effectiveTimeout != null) {
51+
ctx = ctx.withTimeout(effectiveTimeout);
52+
}
53+
}
54+
inner.call(request, observer, ctx);
55+
}
56+
57+
private boolean isPointRead(ReadRowsRequest request) {
58+
if (request.getRowsLimit() == 1) {
59+
return true;
60+
}
61+
if (!request.getRows().getRowRangesList().isEmpty()) {
62+
return false;
63+
}
64+
return request.getRows().getRowKeysCount() == 1;
65+
}
66+
67+
/**
68+
* Extracts the effective timeout for a point read.
69+
*
70+
* <p>The effective time is the minimum of a streamWaitTimeout and a user set attempt timeout.
71+
*/
72+
@Nullable
73+
private Duration getEffectivePointReadTimeout(ApiCallContext ctx) {
74+
Duration streamWaitTimeout = ctx.getStreamWaitTimeout();
75+
Duration attemptTimeout = ctx.getTimeout();
76+
77+
if (streamWaitTimeout == null) {
78+
return attemptTimeout;
79+
}
80+
81+
if (attemptTimeout == null) {
82+
return streamWaitTimeout;
83+
}
84+
return (attemptTimeout.compareTo(streamWaitTimeout) <= 0) ? attemptTimeout : streamWaitTimeout;
85+
}
86+
}

‎google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/PointReadTimeoutCallableTest.java

+183
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
/*
2+
* Copyright 2021 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.data.v2.stub.readrows;
17+
18+
import static com.google.common.truth.Truth.assertThat;
19+
20+
import com.google.api.gax.grpc.GrpcCallContext;
21+
import com.google.api.gax.rpc.ApiCallContext;
22+
import com.google.api.gax.rpc.ResponseObserver;
23+
import com.google.api.gax.rpc.ServerStreamingCallable;
24+
import com.google.bigtable.v2.ReadRowsRequest;
25+
import com.google.bigtable.v2.RowRange;
26+
import com.google.bigtable.v2.RowSet;
27+
import com.google.common.collect.ImmutableList;
28+
import com.google.protobuf.ByteString;
29+
import java.util.Arrays;
30+
import java.util.List;
31+
import org.junit.Before;
32+
import org.junit.Rule;
33+
import org.junit.Test;
34+
import org.junit.runner.RunWith;
35+
import org.junit.runners.JUnit4;
36+
import org.mockito.ArgumentCaptor;
37+
import org.mockito.Mock;
38+
import org.mockito.Mockito;
39+
import org.mockito.junit.MockitoJUnit;
40+
import org.mockito.junit.MockitoRule;
41+
import org.threeten.bp.Duration;
42+
43+
@RunWith(JUnit4.class)
44+
public class PointReadTimeoutCallableTest {
45+
@Rule public final MockitoRule moo = MockitoJUnit.rule();
46+
47+
@Mock private ServerStreamingCallable<ReadRowsRequest, Object> inner;
48+
private ArgumentCaptor<ApiCallContext> ctxCaptor;
49+
@Mock private ResponseObserver<Object> responseObserver;
50+
51+
@Before
52+
public void setUp() throws Exception {
53+
ctxCaptor = ArgumentCaptor.forClass(ApiCallContext.class);
54+
55+
Mockito.doNothing()
56+
.when(inner)
57+
.call(
58+
Mockito.isA(ReadRowsRequest.class),
59+
Mockito.any(ResponseObserver.class),
60+
ctxCaptor.capture());
61+
}
62+
63+
@Test
64+
public void promotesStreamWaitTimeout() {
65+
Duration duration = Duration.ofMillis(100);
66+
PointReadTimeoutCallable<Object> callable = new PointReadTimeoutCallable<>(inner);
67+
68+
for (ReadRowsRequest req : createPointReadRequests()) {
69+
callable.call(
70+
req, responseObserver, GrpcCallContext.createDefault().withStreamWaitTimeout(duration));
71+
72+
assertThat(ctxCaptor.getValue().getTimeout()).isEqualTo(duration);
73+
}
74+
}
75+
76+
@Test
77+
public void promotesStreamWaitTimeoutForRowLimit() {
78+
Duration duration = Duration.ofMillis(100);
79+
PointReadTimeoutCallable<Object> callable = new PointReadTimeoutCallable<>(inner);
80+
81+
for (ReadRowsRequest req : createPointReadRequests()) {
82+
callable.call(
83+
createRowsLimitRequest(),
84+
responseObserver,
85+
GrpcCallContext.createDefault().withStreamWaitTimeout(duration));
86+
87+
assertThat(ctxCaptor.getValue().getTimeout()).isEqualTo(duration);
88+
}
89+
}
90+
91+
@Test
92+
public void respectsExistingTimeout() {
93+
Duration duration = Duration.ofMillis(100);
94+
PointReadTimeoutCallable<Object> callable = new PointReadTimeoutCallable<>(inner);
95+
96+
List<ReadRowsRequest> requests =
97+
ImmutableList.<ReadRowsRequest>builder()
98+
.addAll(createPointReadRequests())
99+
.add(ReadRowsRequest.getDefaultInstance())
100+
.build();
101+
102+
for (ReadRowsRequest req : requests) {
103+
callable.call(req, responseObserver, GrpcCallContext.createDefault().withTimeout(duration));
104+
assertThat(ctxCaptor.getValue().getTimeout()).isEqualTo(duration);
105+
}
106+
}
107+
108+
@Test
109+
public void usesMinimum1() {
110+
Duration attemptTimeout = Duration.ofMillis(100);
111+
Duration streamTimeout = Duration.ofMillis(200);
112+
PointReadTimeoutCallable<Object> callable = new PointReadTimeoutCallable<>(inner);
113+
114+
for (ReadRowsRequest req : createPointReadRequests()) {
115+
GrpcCallContext ctx =
116+
GrpcCallContext.createDefault()
117+
.withTimeout(attemptTimeout)
118+
.withStreamWaitTimeout(streamTimeout);
119+
callable.call(req, responseObserver, ctx);
120+
121+
assertThat(ctxCaptor.getValue().getTimeout()).isEqualTo(attemptTimeout);
122+
}
123+
}
124+
125+
@Test
126+
public void usesMinimum2() {
127+
Duration attemptTimeout = Duration.ofMillis(200);
128+
Duration streamTimeout = Duration.ofMillis(100);
129+
PointReadTimeoutCallable<Object> callable = new PointReadTimeoutCallable<>(inner);
130+
131+
for (ReadRowsRequest req : createPointReadRequests()) {
132+
GrpcCallContext ctx =
133+
GrpcCallContext.createDefault()
134+
.withTimeout(attemptTimeout)
135+
.withStreamWaitTimeout(streamTimeout);
136+
137+
callable.call(req, responseObserver, ctx);
138+
139+
assertThat(ctxCaptor.getValue().getTimeout()).isEqualTo(streamTimeout);
140+
}
141+
}
142+
143+
@Test
144+
public void nonPointReadsAreUntouched() {
145+
Duration streamTimeout = Duration.ofMillis(100);
146+
PointReadTimeoutCallable<Object> callable = new PointReadTimeoutCallable<>(inner);
147+
148+
List<ReadRowsRequest> requests =
149+
Arrays.<ReadRowsRequest>asList(
150+
ReadRowsRequest.getDefaultInstance(),
151+
ReadRowsRequest.newBuilder()
152+
.setRows(
153+
RowSet.newBuilder()
154+
.addRowKeys(ByteString.copyFromUtf8("a"))
155+
.addRowKeys(ByteString.copyFromUtf8("ab")))
156+
.build(),
157+
ReadRowsRequest.newBuilder()
158+
.setRows(RowSet.newBuilder().addRowRanges(RowRange.getDefaultInstance()))
159+
.build());
160+
161+
for (ReadRowsRequest req : requests) {
162+
callable.call(
163+
req,
164+
responseObserver,
165+
GrpcCallContext.createDefault().withStreamWaitTimeout(streamTimeout));
166+
assertThat(ctxCaptor.getValue().getTimeout()).isNull();
167+
}
168+
}
169+
170+
private List<ReadRowsRequest> createPointReadRequests() {
171+
return Arrays.asList(createRowsLimitRequest(), createRowKeyRequest());
172+
}
173+
174+
private ReadRowsRequest createRowsLimitRequest() {
175+
return ReadRowsRequest.newBuilder().setRowsLimit(1).build();
176+
}
177+
178+
private ReadRowsRequest createRowKeyRequest() {
179+
return ReadRowsRequest.newBuilder()
180+
.setRows(RowSet.newBuilder().addRowKeys(ByteString.copyFromUtf8("key")))
181+
.build();
182+
}
183+
}

0 commit comments

Comments
 (0)