Skip to content

Commit 588761b

Browse files
schiemonschiemon
committed
fix: fix RetryingRpcClientTest.doNotRetryWhenResponseIsCancelled
1 parent 7a7c0a6 commit 588761b

1 file changed

Lines changed: 108 additions & 11 deletions

File tree

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

Lines changed: 108 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import static org.assertj.core.api.Assertions.assertThat;
2020
import static org.assertj.core.api.Assertions.assertThatThrownBy;
2121
import static org.assertj.core.api.Assertions.catchThrowable;
22+
import static org.assertj.core.api.Fail.fail;
2223
import static org.awaitility.Awaitility.await;
2324
import static org.mockito.ArgumentMatchers.anyString;
2425
import static org.mockito.Mockito.doThrow;
@@ -28,8 +29,10 @@
2829
import static org.mockito.Mockito.verify;
2930
import static org.mockito.Mockito.when;
3031

32+
import java.time.Duration;
3133
import java.util.concurrent.BlockingQueue;
3234
import java.util.concurrent.CancellationException;
35+
import java.util.concurrent.CountDownLatch;
3336
import java.util.concurrent.Executors;
3437
import java.util.concurrent.LinkedTransferQueue;
3538
import java.util.concurrent.TimeUnit;
@@ -39,6 +42,8 @@
3942
import org.apache.thrift.TApplicationException;
4043
import org.junit.jupiter.api.Test;
4144
import org.junit.jupiter.api.extension.RegisterExtension;
45+
import org.junit.jupiter.params.ParameterizedTest;
46+
import org.junit.jupiter.params.provider.EnumSource;
4247

4348
import com.linecorp.armeria.client.ClientFactory;
4449
import com.linecorp.armeria.client.ClientRequestContext;
@@ -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;
@@ -62,10 +68,10 @@
6268
import testing.thrift.main.HelloService;
6369

6470
class RetryingRpcClientTest {
65-
71+
private static final Backoff fixedBackoff = Backoff.fixed(500);
6672
private static final RetryRuleWithContent<RpcResponse> retryAlways =
6773
(ctx, response, cause) ->
68-
UnmodifiableFuture.completedFuture(RetryDecision.retry(Backoff.fixed(500)));
74+
UnmodifiableFuture.completedFuture(RetryDecision.retry(fixedBackoff));
6975

7076
private static final RetryRuleWithContent<RpcResponse> retryOnException =
7177
RetryRuleWithContent.onException(Backoff.withoutDelay());
@@ -326,34 +332,125 @@ 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_FIRST_REQUEST_NO_DELAY(true, 0),
338+
CANCEL_FIRST_REQUEST_WITH_DELAY(true, 500),
339+
CANCEL_AFTER_FIRST_REQUEST_NO_DELAY(false, 0),
340+
CANCEL_AFTER_FIRST_REQUEST_WITH_DELAY(false, 500);
341+
342+
static final int BACKOFF_MILLIS = 50;
343+
final boolean ensureCancelBeforeFirstRequest;
344+
final long cancelDelayMillis;
345+
346+
DoNotRetryWhenResponseIsCancelledTestParams(boolean ensureCancelBeforeFirstRequest,
347+
long cancelDelayMillis) {
348+
this.ensureCancelBeforeFirstRequest = ensureCancelBeforeFirstRequest;
349+
this.cancelDelayMillis = cancelDelayMillis;
350+
}
351+
}
352+
353+
@ParameterizedTest
354+
@EnumSource(DoNotRetryWhenResponseIsCancelledTestParams.class)
355+
void doNotRetryWhenResponseIsCancelled(DoNotRetryWhenResponseIsCancelledTestParams param) throws Exception {
331356
serviceRetryCount.set(0);
357+
358+
final RetryRuleWithContent<RpcResponse> quickRetryAlways =
359+
RetryRuleWithContent.<RpcResponse>builder()
360+
.onException()
361+
.thenBackoff(Backoff.fixed(
362+
DoNotRetryWhenResponseIsCancelledTestParams.BACKOFF_MILLIS));
363+
364+
final int maxExpectedAttempts =
365+
(int) (param.cancelDelayMillis / DoNotRetryWhenResponseIsCancelledTestParams.BACKOFF_MILLIS) +
366+
5;
367+
final AtomicInteger serviceRetryCountWhenCancelled = new AtomicInteger();
368+
final CountDownLatch canRetry = new CountDownLatch(1);
332369
try (ClientFactory factory = ClientFactory.builder().build()) {
333370
final AtomicReference<ClientRequestContext> context = new AtomicReference<>();
334371
final HelloService.Iface client =
335372
ThriftClients.builder(server.httpUri())
336373
.path("/thrift")
337374
.factory(factory)
338-
.rpcDecorator(RetryingRpcClient.builder(retryAlways).newDecorator())
375+
.rpcDecorator(RetryingRpcClient.builder(quickRetryAlways)
376+
// We want to cancel the request before
377+
// we quit because of reaching max attempts.
378+
.maxTotalAttempts(maxExpectedAttempts)
379+
.newDecorator())
380+
.rpcDecorator((delegate, ctx, req) -> {
381+
// Clog the retry event loop so we do not retry until canRetry.countDown()
382+
// is called.
383+
// If you see failure of this test, and you altered AbstractRetryingClient,
384+
// make sure you are executing (prepare)Retry() on the retry event loop and
385+
// that the retry event loop is ctx.eventLoop().
386+
ctx.eventLoop().execute(() -> {
387+
try {
388+
canRetry.await();
389+
} catch (InterruptedException e) {
390+
fail(e);
391+
}
392+
});
393+
394+
return delegate.execute(ctx, req);
395+
})
339396
.rpcDecorator((delegate, ctx, req) -> {
340-
context.set(ctx);
341397
final RpcResponse res = delegate.execute(ctx, req);
342-
res.cancel(true);
398+
399+
if (param.ensureCancelBeforeFirstRequest) {
400+
Thread.sleep(param.cancelDelayMillis);
401+
assertThat(res.isDone()).isFalse();
402+
res.cancel(true);
403+
serviceRetryCountWhenCancelled.set(serviceRetryCount.get());
404+
canRetry.countDown();
405+
} else {
406+
canRetry.countDown();
407+
Thread.sleep(param.cancelDelayMillis);
408+
assertThat(res.isDone()).isFalse();
409+
res.cancel(true);
410+
serviceRetryCountWhenCancelled.set(serviceRetryCount.get());
411+
}
412+
343413
return res;
344414
})
415+
.rpcDecorator((delegate, ctx, req) -> {
416+
context.set(ctx);
417+
ctx.setResponseTimeout(
418+
TimeoutMode.EXTEND,
419+
Duration.ofMillis(param.cancelDelayMillis + 1000)
420+
);
421+
422+
return delegate.execute(ctx, req);
423+
})
345424
.build(HelloService.Iface.class);
346425
when(serviceHandler.hello(anyString())).thenThrow(new IllegalArgumentException());
347426

348427
assertThatThrownBy(() -> client.hello("hello")).isInstanceOf(CancellationException.class);
349428

350429
await().untilAsserted(() -> {
351-
verify(serviceHandler, only()).hello("hello");
430+
assertThat(serviceRetryCountWhenCancelled.get()).isIn(serviceRetryCount.get(),
431+
serviceRetryCount.get() + 1);
432+
verify(serviceHandler, times(serviceRetryCount.get())).hello("hello");
352433
});
434+
435+
final RequestLog log = context.get().log().whenComplete().join();
436+
if (param.ensureCancelBeforeFirstRequest) {
437+
assertThat(serviceRetryCount.get()).isZero();
438+
assertThat(log.requestCause()).isExactlyInstanceOf(CancellationException.class);
439+
assertThat(log.responseCause()).isExactlyInstanceOf(CancellationException.class);
440+
} else {
441+
// We still could cancel the before the first request so we do not have a guarantee for
442+
// requestCause() to be null.
443+
assertThat(log.responseCause()).isExactlyInstanceOf(CancellationException.class);
444+
}
445+
353446
// Sleep 1 second more to check if there was another retry.
354-
TimeUnit.SECONDS.sleep(1);
355-
verify(serviceHandler, only()).hello("hello");
356-
assertThat(serviceRetryCount).hasValue(1);
447+
TimeUnit.MILLISECONDS.sleep(100);
448+
if (param.ensureCancelBeforeFirstRequest) {
449+
assertThat(serviceRetryCount.get()).isZero();
450+
}
451+
assertThat(serviceRetryCountWhenCancelled.get()).isIn(serviceRetryCount.get(),
452+
serviceRetryCount.get() + 1);
453+
verify(serviceHandler, times(serviceRetryCount.get())).hello("hello");
357454
}
358455
}
359456
}

0 commit comments

Comments
 (0)