Skip to content

Commit 5316498

Browse files
schiemonschiemon
committed
fix: doNotRetryWhenResponseIsCancelled
1 parent e2a8648 commit 5316498

File tree

2 files changed

+119
-11
lines changed

2 files changed

+119
-11
lines changed

‎core/src/main/java/com/linecorp/armeria/client/retry/RetryingRpcClient.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -152,7 +152,7 @@ private void doExecute0(ClientRequestContext ctx, RpcRequest req,
152152
RpcResponse returnedRes, CompletableFuture<RpcResponse> future) {
153153
final int totalAttempts = getTotalAttempts(ctx);
154154
final boolean initialAttempt = totalAttempts <= 1;
155-
if (returnedRes.isDone()) {
155+
if (ctx.isCancelled() || returnedRes.isDone()) {
156156
// The response has been cancelled by the client before it receives a response, so stop retrying.
157157
handleException(ctx, future, new CancellationException(
158158
"the response returned to the client has been cancelled"), initialAttempt);

‎thrift/thrift0.13/src/test/java/com/linecorp/armeria/it/client/retry/RetryingRpcClientTest.java‎

Lines changed: 118 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,10 @@
2828
import static org.mockito.Mockito.verify;
2929
import static org.mockito.Mockito.when;
3030

31+
import java.time.Duration;
3132
import java.util.concurrent.BlockingQueue;
3233
import java.util.concurrent.CancellationException;
34+
import java.util.concurrent.CompletableFuture;
3335
import java.util.concurrent.Executors;
3436
import java.util.concurrent.LinkedTransferQueue;
3537
import java.util.concurrent.TimeUnit;
@@ -39,9 +41,12 @@
3941
import org.apache.thrift.TApplicationException;
4042
import org.junit.jupiter.api.Test;
4143
import org.junit.jupiter.api.extension.RegisterExtension;
44+
import org.junit.jupiter.params.ParameterizedTest;
45+
import org.junit.jupiter.params.provider.EnumSource;
4246

4347
import com.linecorp.armeria.client.ClientFactory;
4448
import com.linecorp.armeria.client.ClientRequestContext;
49+
import com.linecorp.armeria.client.ResponseCancellationException;
4550
import com.linecorp.armeria.client.UnprocessedRequestException;
4651
import com.linecorp.armeria.client.retry.Backoff;
4752
import com.linecorp.armeria.client.retry.RetryConfig;
@@ -53,6 +58,7 @@
5358
import com.linecorp.armeria.common.HttpRequest;
5459
import com.linecorp.armeria.common.RpcResponse;
5560
import com.linecorp.armeria.common.logging.RequestLog;
61+
import com.linecorp.armeria.common.util.TimeoutMode;
5662
import com.linecorp.armeria.common.util.UnmodifiableFuture;
5763
import com.linecorp.armeria.server.ServerBuilder;
5864
import com.linecorp.armeria.server.thrift.THttpService;
@@ -326,34 +332,136 @@ void shouldGetExceptionWhenFactoryIsClosed() throws Exception {
326332
"(?i).*(factory has been closed|not accepting a task).*"));
327333
}
328334

329-
@Test
330-
void doNotRetryWhenResponseIsCancelled() throws Exception {
335+
enum DoNotRetryWhenResponseIsCancelledTestParams {
336+
// Cancel delays for a backoff of 50 milliseconds (quickBackoffMillis).
337+
CANCEL_CTX_FIRST_REQUEST_NO_DELAY(true, true, 0),
338+
CANCEL_CTX_FIRST_REQUEST_WITH_DELAY(true, true, 500),
339+
CANCEL_CTX_AFTER_FIRST_REQUEST_NO_DELAY(false, true, 0),
340+
CANCEL_CTX_AFTER_FIRST_REQUEST_WITH_DELAY(false, false, 500),
341+
CANCEL_RES_AFTER_FIRST_REQUEST_NO_DELAY(false, true, 0),
342+
CANCEL_RES_AFTER_FIRST_REQUEST_WITH_DELAY(false, false, 500);
343+
344+
static final int BACKOFF_MILLIS = 50;
345+
final boolean ensureCancelBeforeFirstRequest;
346+
final boolean cancelViaCtx;
347+
final long cancelDelayMillis;
348+
349+
DoNotRetryWhenResponseIsCancelledTestParams(boolean ensureCancelBeforeFirstRequest,
350+
boolean cancelViaCtx,
351+
long cancelDelayMillis) {
352+
this.ensureCancelBeforeFirstRequest = ensureCancelBeforeFirstRequest;
353+
if (this.ensureCancelBeforeFirstRequest) {
354+
this.cancelViaCtx = true;
355+
} else {
356+
this.cancelViaCtx = cancelViaCtx;
357+
}
358+
359+
this.cancelDelayMillis = cancelDelayMillis;
360+
}
361+
}
362+
363+
@ParameterizedTest
364+
@EnumSource(DoNotRetryWhenResponseIsCancelledTestParams.class)
365+
void doNotRetryWhenResponseIsCancelled(DoNotRetryWhenResponseIsCancelledTestParams param) throws Exception {
331366
serviceRetryCount.set(0);
367+
368+
final RetryRuleWithContent<RpcResponse> quickRetryAlways =
369+
RetryRuleWithContent.<RpcResponse>builder()
370+
.onException()
371+
.thenBackoff(Backoff.fixed(
372+
DoNotRetryWhenResponseIsCancelledTestParams.BACKOFF_MILLIS));
373+
374+
final int maxExpectedAttempts =
375+
(int) (param.cancelDelayMillis / DoNotRetryWhenResponseIsCancelledTestParams.BACKOFF_MILLIS) +
376+
5;
377+
final AtomicInteger serviceRetryCountWhenCancelled = new AtomicInteger();
332378
try (ClientFactory factory = ClientFactory.builder().build()) {
333379
final AtomicReference<ClientRequestContext> context = new AtomicReference<>();
334380
final HelloService.Iface client =
335381
ThriftClients.builder(server.httpUri())
336382
.path("/thrift")
337383
.factory(factory)
338-
.rpcDecorator(RetryingRpcClient.builder(retryAlways).newDecorator())
384+
.rpcDecorator(RetryingRpcClient.builder(quickRetryAlways)
385+
// We want to cancel the request before
386+
// we quit because of reaching max attempts.
387+
.maxTotalAttempts(maxExpectedAttempts)
388+
.newDecorator())
339389
.rpcDecorator((delegate, ctx, req) -> {
340-
context.set(ctx);
341-
final RpcResponse res = delegate.execute(ctx, req);
342-
res.cancel(true);
390+
final CompletableFuture<RpcResponse> resFuture = new CompletableFuture<>();
391+
final RpcResponse res = RpcResponse.from(resFuture);
392+
393+
if (param.ensureCancelBeforeFirstRequest) {
394+
Thread.sleep(param.cancelDelayMillis);
395+
396+
assertThat(ctx.isCancelled()).isFalse();
397+
assertThat(res.isDone()).isFalse();
398+
assert param.cancelViaCtx;
399+
ctx.cancel();
400+
serviceRetryCountWhenCancelled.set(serviceRetryCount.get());
401+
402+
resFuture.complete(delegate.execute(ctx, req));
403+
} else {
404+
final RpcResponse nextRes = delegate.execute(ctx, req);
405+
resFuture.complete(nextRes);
406+
407+
Thread.sleep(param.cancelDelayMillis);
408+
409+
assertThat(ctx.isCancelled()).isFalse();
410+
assertThat(nextRes.isDone()).isFalse();
411+
412+
if (param.cancelViaCtx) {
413+
ctx.cancel();
414+
} else {
415+
nextRes.cancel(true);
416+
}
417+
418+
serviceRetryCountWhenCancelled.set(serviceRetryCount.get());
419+
}
420+
343421
return res;
344422
})
423+
.rpcDecorator((delegate, ctx, req) -> {
424+
context.set(ctx);
425+
// Make sure we do not get cancelled while delaying the cancel.
426+
ctx.setResponseTimeout(
427+
TimeoutMode.EXTEND,
428+
Duration.ofMillis(param.cancelDelayMillis + 5000)
429+
);
430+
431+
return delegate.execute(ctx, req);
432+
})
345433
.build(HelloService.Iface.class);
346434
when(serviceHandler.hello(anyString())).thenThrow(new IllegalArgumentException());
347435

348-
assertThatThrownBy(() -> client.hello("hello")).isInstanceOf(CancellationException.class);
436+
assertThatThrownBy(() -> client.hello("hello"))
437+
.isOfAnyClassIn(CancellationException.class, ResponseCancellationException.class);
349438

350439
await().untilAsserted(() -> {
351-
verify(serviceHandler, only()).hello("hello");
440+
assertThat(serviceRetryCountWhenCancelled.get()).isIn(serviceRetryCount.get(),
441+
serviceRetryCount.get() - 1);
442+
verify(serviceHandler, times(serviceRetryCount.get())).hello("hello");
352443
});
444+
445+
final RequestLog log = context.get().log().whenComplete().join();
446+
if (param.ensureCancelBeforeFirstRequest) {
447+
assertThat(serviceRetryCount.get()).isZero();
448+
assertThat(log.requestCause()).isExactlyInstanceOf(ResponseCancellationException.class);
449+
assertThat(log.responseCause()).isExactlyInstanceOf(ResponseCancellationException.class);
450+
} else {
451+
// We still could cancel the before the first request so we do not have a guarantee for
452+
// requestCause() to be null.
453+
assertThat(log.responseCause()).isOfAnyClassIn(ResponseCancellationException.class,
454+
CancellationException.class);
455+
}
456+
353457
// Sleep 1 second more to check if there was another retry.
354458
TimeUnit.SECONDS.sleep(1);
355-
verify(serviceHandler, only()).hello("hello");
356-
assertThat(serviceRetryCount).hasValue(1);
459+
if (param.ensureCancelBeforeFirstRequest) {
460+
assertThat(serviceRetryCount.get()).isZero();
461+
}
462+
assertThat(serviceRetryCountWhenCancelled.get()).isIn(serviceRetryCount.get(),
463+
serviceRetryCount.get() - 1);
464+
verify(serviceHandler, times(serviceRetryCount.get())).hello("hello");
357465
}
358466
}
359467
}

0 commit comments

Comments
 (0)