Thanks to visit codestin.com
Credit goes to github.com

Skip to content
This repository was archived by the owner on Sep 26, 2023. It is now read-only.

Commit 51c40ab

Browse files
noahdietznoahdietz
authored
feat: wrap non-retryable RPCs in retry machinery (#1328)
* feat: wrap non-retryable RPCs in retry machinery * address feedback * add contextual retry settings to timeout tests * add tracing test for non-retryable callable * add server streaming callable tests
1 parent 751ccf3 commit 51c40ab

File tree

4 files changed

+395
-35
lines changed

4 files changed

+395
-35
lines changed

‎gax-grpc/src/test/java/com/google/api/gax/grpc/TimeoutTest.java‎

Lines changed: 110 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -74,9 +74,12 @@ public class TimeoutTest {
7474
private static final int DEADLINE_IN_MINUTES = 10;
7575
private static final int DEADLINE_IN_SECONDS = 20;
7676
private static final ImmutableSet<StatusCode.Code> emptyRetryCodes = ImmutableSet.of();
77+
private static final ImmutableSet<StatusCode.Code> retryUnknownCode =
78+
ImmutableSet.of(StatusCode.Code.UNKNOWN);
7779
private static final Duration totalTimeout = Duration.ofDays(DEADLINE_IN_DAYS);
7880
private static final Duration maxRpcTimeout = Duration.ofMinutes(DEADLINE_IN_MINUTES);
7981
private static final Duration initialRpcTimeout = Duration.ofSeconds(DEADLINE_IN_SECONDS);
82+
private static final GrpcCallContext defaultCallContext = GrpcCallContext.createDefault();
8083

8184
@Rule public MockitoRule mockitoRule = MockitoJUnit.rule().strictness(Strictness.STRICT_STUBS);
8285
@Mock private Marshaller<String> stringMarshaller;
@@ -97,7 +100,8 @@ public void testNonRetryUnarySettings() {
97100
.setRpcTimeoutMultiplier(1.0)
98101
.setMaxRpcTimeout(maxRpcTimeout)
99102
.build();
100-
CallOptions callOptionsUsed = setupUnaryCallable(retrySettings);
103+
CallOptions callOptionsUsed =
104+
setupUnaryCallable(retrySettings, emptyRetryCodes, defaultCallContext);
101105

102106
// Verify that the gRPC channel used the CallOptions with our custom timeout of ~2 Days.
103107
assertThat(callOptionsUsed.getDeadline()).isNotNull();
@@ -108,6 +112,46 @@ public void testNonRetryUnarySettings() {
108112
assertThat(callOptionsUsed.getAuthority()).isEqualTo(CALL_OPTIONS_AUTHORITY);
109113
}
110114

115+
@Test
116+
public void testNonRetryUnarySettingsContextWithRetry() {
117+
RetrySettings retrySettings =
118+
RetrySettings.newBuilder()
119+
.setTotalTimeout(totalTimeout)
120+
.setInitialRetryDelay(Duration.ZERO)
121+
.setRetryDelayMultiplier(1.0)
122+
.setMaxRetryDelay(Duration.ZERO)
123+
.setMaxAttempts(1)
124+
.setJittered(true)
125+
.setInitialRpcTimeout(initialRpcTimeout)
126+
.setRpcTimeoutMultiplier(1.0)
127+
.setMaxRpcTimeout(maxRpcTimeout)
128+
.build();
129+
Duration newTimeout = Duration.ofSeconds(5);
130+
RetrySettings contextRetrySettings =
131+
retrySettings
132+
.toBuilder()
133+
.setInitialRpcTimeout(newTimeout)
134+
.setMaxRpcTimeout(newTimeout)
135+
.setMaxAttempts(3)
136+
.build();
137+
GrpcCallContext retryingContext =
138+
defaultCallContext
139+
.withRetrySettings(contextRetrySettings)
140+
.withRetryableCodes(retryUnknownCode);
141+
CallOptions callOptionsUsed =
142+
setupUnaryCallable(retrySettings, emptyRetryCodes, retryingContext);
143+
144+
// Verify that the gRPC channel used the CallOptions the initial timeout of ~5 seconds.
145+
// This indicates that the context retry settings were used on a callable that was instantiated
146+
// with non-retryable settings.
147+
assertThat(callOptionsUsed.getDeadline()).isNotNull();
148+
assertThat(callOptionsUsed.getDeadline())
149+
.isGreaterThan(Deadline.after(newTimeout.toSecondsPart() - 1, TimeUnit.SECONDS));
150+
assertThat(callOptionsUsed.getDeadline())
151+
.isLessThan(Deadline.after(newTimeout.toSecondsPart(), TimeUnit.SECONDS));
152+
assertThat(callOptionsUsed.getAuthority()).isEqualTo(CALL_OPTIONS_AUTHORITY);
153+
}
154+
111155
@Test
112156
public void testNonRetryUnarySettingsWithoutInitialRpcTimeout() {
113157
RetrySettings retrySettings =
@@ -121,7 +165,8 @@ public void testNonRetryUnarySettingsWithoutInitialRpcTimeout() {
121165
.setRpcTimeoutMultiplier(1.0)
122166
.setMaxRpcTimeout(maxRpcTimeout)
123167
.build();
124-
CallOptions callOptionsUsed = setupUnaryCallable(retrySettings);
168+
CallOptions callOptionsUsed =
169+
setupUnaryCallable(retrySettings, emptyRetryCodes, defaultCallContext);
125170

126171
// Verify that the gRPC channel used the CallOptions with our custom timeout of ~2 Days.
127172
assertThat(callOptionsUsed.getDeadline()).isNotNull();
@@ -145,7 +190,8 @@ public void testNonRetryUnarySettingsWithoutIndividualRpcTimeout() {
145190
.setRpcTimeoutMultiplier(1.0)
146191
.setRpcTimeoutMultiplier(1.0)
147192
.build();
148-
CallOptions callOptionsUsed = setupUnaryCallable(retrySettings);
193+
CallOptions callOptionsUsed =
194+
setupUnaryCallable(retrySettings, emptyRetryCodes, defaultCallContext);
149195

150196
// Verify that the gRPC channel used the CallOptions with our custom timeout of ~2 Days.
151197
assertThat(callOptionsUsed.getDeadline()).isNotNull();
@@ -170,7 +216,8 @@ public void testNonRetryServerStreamingSettings() {
170216
.setRpcTimeoutMultiplier(1.0)
171217
.setMaxRpcTimeout(maxRpcTimeout)
172218
.build();
173-
CallOptions callOptionsUsed = setupServerStreamingCallable(retrySettings);
219+
CallOptions callOptionsUsed =
220+
setupServerStreamingCallable(retrySettings, emptyRetryCodes, defaultCallContext);
174221

175222
// Verify that the gRPC channel used the CallOptions with our custom timeout of ~2 Days.
176223
assertThat(callOptionsUsed.getDeadline()).isNotNull();
@@ -181,6 +228,41 @@ public void testNonRetryServerStreamingSettings() {
181228
assertThat(callOptionsUsed.getAuthority()).isEqualTo(CALL_OPTIONS_AUTHORITY);
182229
}
183230

231+
@Test
232+
public void testNonRetryServerStreamingSettingsContextWithRetry() {
233+
RetrySettings retrySettings =
234+
RetrySettings.newBuilder()
235+
.setTotalTimeout(totalTimeout)
236+
.setInitialRetryDelay(Duration.ZERO)
237+
.setRetryDelayMultiplier(1.0)
238+
.setMaxRetryDelay(Duration.ZERO)
239+
.setMaxAttempts(1)
240+
.setJittered(true)
241+
.setInitialRpcTimeout(initialRpcTimeout)
242+
.setRpcTimeoutMultiplier(1.0)
243+
.setMaxRpcTimeout(maxRpcTimeout)
244+
.build();
245+
Duration newTimeout = Duration.ofSeconds(5);
246+
RetrySettings contextRetrySettings =
247+
retrySettings.toBuilder().setTotalTimeout(newTimeout).setMaxAttempts(3).build();
248+
GrpcCallContext retryingContext =
249+
defaultCallContext
250+
.withRetrySettings(contextRetrySettings)
251+
.withRetryableCodes(retryUnknownCode);
252+
CallOptions callOptionsUsed =
253+
setupServerStreamingCallable(retrySettings, emptyRetryCodes, retryingContext);
254+
255+
// Verify that the gRPC channel used the CallOptions the total timeout of ~5 seconds.
256+
// This indicates that the context retry settings were used on a callable that was instantiated
257+
// with non-retryable settings.
258+
assertThat(callOptionsUsed.getDeadline()).isNotNull();
259+
assertThat(callOptionsUsed.getDeadline())
260+
.isGreaterThan(Deadline.after(newTimeout.toSecondsPart() - 1, TimeUnit.SECONDS));
261+
assertThat(callOptionsUsed.getDeadline())
262+
.isLessThan(Deadline.after(newTimeout.toSecondsPart(), TimeUnit.SECONDS));
263+
assertThat(callOptionsUsed.getAuthority()).isEqualTo(CALL_OPTIONS_AUTHORITY);
264+
}
265+
184266
@Test
185267
public void testNonRetryServerStreamingSettingsWithoutInitialRpcTimeout() {
186268
RetrySettings retrySettings =
@@ -194,7 +276,8 @@ public void testNonRetryServerStreamingSettingsWithoutInitialRpcTimeout() {
194276
.setRpcTimeoutMultiplier(1.0)
195277
.setMaxRpcTimeout(maxRpcTimeout)
196278
.build();
197-
CallOptions callOptionsUsed = setupServerStreamingCallable(retrySettings);
279+
CallOptions callOptionsUsed =
280+
setupServerStreamingCallable(retrySettings, emptyRetryCodes, defaultCallContext);
198281

199282
// Verify that the gRPC channel used the CallOptions with our custom timeout of ~2 Days.
200283
assertThat(callOptionsUsed.getDeadline()).isNotNull();
@@ -218,7 +301,8 @@ public void testNonRetryServerStreamingSettingsWithoutIndividualRpcTimeout() {
218301
.setRpcTimeoutMultiplier(1.0)
219302
.setRpcTimeoutMultiplier(1.0)
220303
.build();
221-
CallOptions callOptionsUsed = setupServerStreamingCallable(retrySettings);
304+
CallOptions callOptionsUsed =
305+
setupServerStreamingCallable(retrySettings, emptyRetryCodes, defaultCallContext);
222306

223307
// Verify that the gRPC channel used the CallOptions with our custom timeout of ~2 Days.
224308
assertThat(callOptionsUsed.getDeadline()).isNotNull();
@@ -229,7 +313,10 @@ public void testNonRetryServerStreamingSettingsWithoutIndividualRpcTimeout() {
229313
assertThat(callOptionsUsed.getAuthority()).isEqualTo(CALL_OPTIONS_AUTHORITY);
230314
}
231315

232-
private CallOptions setupUnaryCallable(RetrySettings retrySettings) {
316+
private CallOptions setupUnaryCallable(
317+
RetrySettings retrySettings,
318+
ImmutableSet<StatusCode.Code> retryableCodes,
319+
GrpcCallContext callContext) {
233320
MethodDescriptor<String, String> methodDescriptor =
234321
MethodDescriptor.<String, String>newBuilder()
235322
.setSchemaDescriptor("yaml")
@@ -248,8 +335,8 @@ private CallOptions setupUnaryCallable(RetrySettings retrySettings) {
248335
// Clobber the "authority" property with an identifier that allows us to trace
249336
// the use of this CallOptions variable.
250337
CallOptions spyCallOptions = CallOptions.DEFAULT.withAuthority("RETRYING_TEST");
251-
GrpcCallContext grpcCallContext =
252-
GrpcCallContext.createDefault().withChannel(managedChannel).withCallOptions(spyCallOptions);
338+
GrpcCallContext context =
339+
callContext.withChannel(managedChannel).withCallOptions(spyCallOptions);
253340

254341
ArgumentCaptor<CallOptions> callOptionsArgumentCaptor =
255342
ArgumentCaptor.forClass(CallOptions.class);
@@ -266,16 +353,16 @@ private CallOptions setupUnaryCallable(RetrySettings retrySettings) {
266353
.setMethodDescriptor(methodDescriptor)
267354
.setParamsExtractor(paramsExtractor)
268355
.build();
269-
UnaryCallSettings<String, String> nonRetriedCallSettings =
356+
UnaryCallSettings<String, String> unaryCallSettings =
270357
UnaryCallSettings.<String, String>newUnaryCallSettingsBuilder()
271358
.setRetrySettings(retrySettings)
272-
.setRetryableCodes(emptyRetryCodes)
359+
.setRetryableCodes(retryableCodes)
273360
.build();
274361
UnaryCallable<String, String> callable =
275362
GrpcCallableFactory.createUnaryCallable(
276363
grpcCallSettings,
277-
nonRetriedCallSettings,
278-
ClientContext.newBuilder().setDefaultCallContext(grpcCallContext).build());
364+
unaryCallSettings,
365+
ClientContext.newBuilder().setDefaultCallContext(context).build());
279366

280367
try {
281368
ApiFuture<String> future = callable.futureCall("Is your refrigerator running?");
@@ -287,7 +374,10 @@ private CallOptions setupUnaryCallable(RetrySettings retrySettings) {
287374
return callOptionsArgumentCaptor.getValue();
288375
}
289376

290-
private CallOptions setupServerStreamingCallable(RetrySettings retrySettings) {
377+
private CallOptions setupServerStreamingCallable(
378+
RetrySettings retrySettings,
379+
ImmutableSet<StatusCode.Code> retryableCodes,
380+
GrpcCallContext callContext) {
291381
MethodDescriptor<String, String> methodDescriptor =
292382
MethodDescriptor.<String, String>newBuilder()
293383
.setSchemaDescriptor("yaml")
@@ -306,8 +396,8 @@ private CallOptions setupServerStreamingCallable(RetrySettings retrySettings) {
306396
// Clobber the "authority" property with an identifier that allows us to trace
307397
// the use of this CallOptions variable.
308398
CallOptions spyCallOptions = CallOptions.DEFAULT.withAuthority("RETRYING_TEST");
309-
GrpcCallContext grpcCallContext =
310-
GrpcCallContext.createDefault().withChannel(managedChannel).withCallOptions(spyCallOptions);
399+
GrpcCallContext context =
400+
callContext.withChannel(managedChannel).withCallOptions(spyCallOptions);
311401

312402
ArgumentCaptor<CallOptions> callOptionsArgumentCaptor =
313403
ArgumentCaptor.forClass(CallOptions.class);
@@ -324,16 +414,16 @@ private CallOptions setupServerStreamingCallable(RetrySettings retrySettings) {
324414
.setMethodDescriptor(methodDescriptor)
325415
.setParamsExtractor(paramsExtractor)
326416
.build();
327-
ServerStreamingCallSettings<String, String> nonRetriedCallSettings =
417+
ServerStreamingCallSettings<String, String> serverStreamingCallSettings =
328418
ServerStreamingCallSettings.<String, String>newBuilder()
329419
.setRetrySettings(retrySettings)
330-
.setRetryableCodes(emptyRetryCodes)
420+
.setRetryableCodes(retryableCodes)
331421
.build();
332422
ServerStreamingCallable<String, String> callable =
333423
GrpcCallableFactory.createServerStreamingCallable(
334424
grpcCallSettings,
335-
nonRetriedCallSettings,
336-
ClientContext.newBuilder().setDefaultCallContext(grpcCallContext).build());
425+
serverStreamingCallSettings,
426+
ClientContext.newBuilder().setDefaultCallContext(context).build());
337427

338428
try {
339429
ServerStream<String> stream = callable.call("Is your refrigerator running?");

‎gax/src/main/java/com/google/api/gax/rpc/Callables.java‎

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -56,19 +56,21 @@ public static <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> retrying(
5656
UnaryCallSettings<?, ?> callSettings,
5757
ClientContext clientContext) {
5858

59-
if (areRetriesDisabled(callSettings.getRetryableCodes(), callSettings.getRetrySettings())) {
59+
UnaryCallSettings<?, ?> settings = callSettings;
60+
61+
if (areRetriesDisabled(settings.getRetryableCodes(), settings.getRetrySettings())) {
6062
// When retries are disabled, the total timeout can be treated as the rpc timeout.
61-
return innerCallable.withDefaultCallContext(
62-
clientContext
63-
.getDefaultCallContext()
64-
.withTimeout(callSettings.getRetrySettings().getTotalTimeout()));
63+
settings =
64+
settings
65+
.toBuilder()
66+
.setSimpleTimeoutNoRetries(settings.getRetrySettings().getTotalTimeout())
67+
.build();
6568
}
6669

6770
RetryAlgorithm<ResponseT> retryAlgorithm =
6871
new RetryAlgorithm<>(
6972
new ApiResultRetryAlgorithm<ResponseT>(),
70-
new ExponentialRetryAlgorithm(
71-
callSettings.getRetrySettings(), clientContext.getClock()));
73+
new ExponentialRetryAlgorithm(settings.getRetrySettings(), clientContext.getClock()));
7274
ScheduledRetryingExecutor<ResponseT> retryingExecutor =
7375
new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor());
7476
return new RetryingCallable<>(
@@ -81,25 +83,26 @@ public static <RequestT, ResponseT> ServerStreamingCallable<RequestT, ResponseT>
8183
ServerStreamingCallSettings<RequestT, ResponseT> callSettings,
8284
ClientContext clientContext) {
8385

84-
if (areRetriesDisabled(callSettings.getRetryableCodes(), callSettings.getRetrySettings())) {
86+
ServerStreamingCallSettings<RequestT, ResponseT> settings = callSettings;
87+
if (areRetriesDisabled(settings.getRetryableCodes(), settings.getRetrySettings())) {
8588
// When retries are disabled, the total timeout can be treated as the rpc timeout.
86-
return innerCallable.withDefaultCallContext(
87-
clientContext
88-
.getDefaultCallContext()
89-
.withTimeout(callSettings.getRetrySettings().getTotalTimeout()));
89+
settings =
90+
settings
91+
.toBuilder()
92+
.setSimpleTimeoutNoRetries(settings.getRetrySettings().getTotalTimeout())
93+
.build();
9094
}
9195

9296
StreamingRetryAlgorithm<Void> retryAlgorithm =
9397
new StreamingRetryAlgorithm<>(
9498
new ApiResultRetryAlgorithm<Void>(),
95-
new ExponentialRetryAlgorithm(
96-
callSettings.getRetrySettings(), clientContext.getClock()));
99+
new ExponentialRetryAlgorithm(settings.getRetrySettings(), clientContext.getClock()));
97100

98101
ScheduledRetryingExecutor<Void> retryingExecutor =
99102
new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor());
100103

101104
return new RetryingServerStreamingCallable<>(
102-
innerCallable, retryingExecutor, callSettings.getResumptionStrategy());
105+
innerCallable, retryingExecutor, settings.getResumptionStrategy());
103106
}
104107

105108
@BetaApi("The surface for streaming is not stable yet and may change in the future.")

0 commit comments

Comments
 (0)