diff --git a/core/src/main/java/com/linecorp/armeria/internal/client/DefaultClientRequestContext.java b/core/src/main/java/com/linecorp/armeria/internal/client/DefaultClientRequestContext.java index 9a6490c8866..6af10aad221 100644 --- a/core/src/main/java/com/linecorp/armeria/internal/client/DefaultClientRequestContext.java +++ b/core/src/main/java/com/linecorp/armeria/internal/client/DefaultClientRequestContext.java @@ -538,6 +538,8 @@ private DefaultClientRequestContext(DefaultClientRequestContext ctx, log = RequestLog.builder(this); log.startRequest(); + // Cancel the original timeout and create a new scheduler for the derived context. + ctx.responseCancellationScheduler.cancelScheduled(); responseCancellationScheduler = CancellationScheduler.ofClient(TimeUnit.MILLISECONDS.toNanos(ctx.responseTimeoutMillis())); writeTimeoutMillis = ctx.writeTimeoutMillis(); diff --git a/core/src/main/java/com/linecorp/armeria/internal/common/CancellationScheduler.java b/core/src/main/java/com/linecorp/armeria/internal/common/CancellationScheduler.java index 18dc7276153..a5b91af0950 100644 --- a/core/src/main/java/com/linecorp/armeria/internal/common/CancellationScheduler.java +++ b/core/src/main/java/com/linecorp/armeria/internal/common/CancellationScheduler.java @@ -83,6 +83,11 @@ static CancellationScheduler noop() { */ boolean cancelScheduled(); + /** + * Returns true if a timeout task is scheduled. + */ + boolean isScheduled(); + void setTimeoutNanos(TimeoutMode mode, long timeoutNanos); default void finishNow() { diff --git a/core/src/main/java/com/linecorp/armeria/internal/common/DefaultCancellationScheduler.java b/core/src/main/java/com/linecorp/armeria/internal/common/DefaultCancellationScheduler.java index 7682ac42861..698aa7e523f 100644 --- a/core/src/main/java/com/linecorp/armeria/internal/common/DefaultCancellationScheduler.java +++ b/core/src/main/java/com/linecorp/armeria/internal/common/DefaultCancellationScheduler.java @@ -157,6 +157,11 @@ public boolean cancelScheduled() { } } + @Override + public boolean isScheduled() { + return scheduledFuture != null; + } + @Override public void setTimeoutNanos(TimeoutMode mode, long timeoutNanos) { lock.lock(); diff --git a/core/src/main/java/com/linecorp/armeria/internal/common/NoopCancellationScheduler.java b/core/src/main/java/com/linecorp/armeria/internal/common/NoopCancellationScheduler.java index 046d246278e..c6f6ac71b83 100644 --- a/core/src/main/java/com/linecorp/armeria/internal/common/NoopCancellationScheduler.java +++ b/core/src/main/java/com/linecorp/armeria/internal/common/NoopCancellationScheduler.java @@ -57,6 +57,11 @@ public boolean cancelScheduled() { return false; } + @Override + public boolean isScheduled() { + return false; + } + @Override public void setTimeoutNanos(TimeoutMode mode, long timeoutNanos) { } diff --git a/core/src/test/java/com/linecorp/armeria/client/retry/RetryTimeoutCancellationTest.java b/core/src/test/java/com/linecorp/armeria/client/retry/RetryTimeoutCancellationTest.java new file mode 100644 index 00000000000..5eb6734a5e0 --- /dev/null +++ b/core/src/test/java/com/linecorp/armeria/client/retry/RetryTimeoutCancellationTest.java @@ -0,0 +1,102 @@ +/* + * Copyright 2024 LINE Corporation + * + * LINE Corporation licenses this file to you under the Apache License, + * version 2.0 (the "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at: + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT + * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the + * License for the specific language governing permissions and limitations + * under the License. + */ + +package com.linecorp.armeria.client.retry; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Function; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +import com.linecorp.armeria.client.BlockingWebClient; +import com.linecorp.armeria.client.ClientRequestContext; +import com.linecorp.armeria.client.ClientRequestContextCaptor; +import com.linecorp.armeria.client.Clients; +import com.linecorp.armeria.client.HttpClient; +import com.linecorp.armeria.client.ResponseTimeoutMode; +import com.linecorp.armeria.client.WebClient; +import com.linecorp.armeria.common.AggregatedHttpResponse; +import com.linecorp.armeria.common.HttpResponse; +import com.linecorp.armeria.common.HttpStatus; +import com.linecorp.armeria.common.logging.RequestLogAccess; +import com.linecorp.armeria.internal.client.ClientRequestContextExtension; +import com.linecorp.armeria.internal.common.CancellationScheduler; +import com.linecorp.armeria.server.ServerBuilder; +import com.linecorp.armeria.testing.junit5.server.ServerExtension; + +class RetryTimeoutCancellationTest { + + private static AtomicInteger counter = new AtomicInteger(); + + @RegisterExtension + static ServerExtension server = new ServerExtension() { + @Override + protected void configure(ServerBuilder sb) { + sb.service("/foo", (ctx, req) -> { + return HttpResponse.of(req.aggregate().thenApply(unused -> { + if (counter.getAndIncrement() < 2) { + return HttpResponse.of(HttpStatus.INTERNAL_SERVER_ERROR); + } else { + return HttpResponse.of("hello"); + } + })); + }); + } + }; + + @BeforeEach + void setUp() { + counter.set(0); + } + + @Test + void shouldCancelTimoutScheduler() { + final Function retryingClient = + RetryingClient.builder(RetryRule.builder() + .onServerErrorStatus() + .thenBackoff(Backoff.fixed(100))) + .maxTotalAttempts(3) + .responseTimeoutMillisForEachAttempt(30_000) + .newDecorator(); + final BlockingWebClient client = WebClient.builder(server.httpUri()) + .decorator(retryingClient) + .responseTimeoutMode(ResponseTimeoutMode.FROM_START) + .responseTimeoutMillis(80_000) + .build() + .blocking(); + + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + final AggregatedHttpResponse res = client.post("/foo", "hello"); + final ClientRequestContext ctx = captor.get(); + assertThat(res.status()).isEqualTo(HttpStatus.OK); + assertThat(res.contentUtf8()).isEqualTo("hello"); + assertTimeoutNotScheduled(ctx); + for (RequestLogAccess child : ctx.log().children()) { + assertTimeoutNotScheduled((ClientRequestContext) child.whenComplete().join().context()); + } + } + } + + private static void assertTimeoutNotScheduled(ClientRequestContext ctx) { + final CancellationScheduler scheduler = ctx.as(ClientRequestContextExtension.class) + .responseCancellationScheduler(); + assertThat(scheduler.isScheduled()).isFalse(); + } +}