diff --git a/core/src/main/java/com/linecorp/armeria/client/retry/AbortedAttemptException.java b/core/src/main/java/com/linecorp/armeria/client/retry/AbortedAttemptException.java new file mode 100644 index 00000000000..053c84300f7 --- /dev/null +++ b/core/src/main/java/com/linecorp/armeria/client/retry/AbortedAttemptException.java @@ -0,0 +1,44 @@ +/* + * Copyright 2025 LY Corporation + * + * LY 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 com.linecorp.armeria.common.Flags; + +/** + * A {@link RuntimeException} that is raised by a {@link RetriedRequest} to signal to the caller of + * {@link RetriedRequest#executeAttempt} that the attempt has been aborted because the request has been + * completed - either at call time or during the execution of the attempt. + */ +public final class AbortedAttemptException extends RuntimeException { + private static final long serialVersionUID = -1L; + private static final AbortedAttemptException INSTANCE = new AbortedAttemptException(true); + + /** + * Returns a {@link AbortedAttemptException} which may be a singleton or a new instance, depending on + * {@link Flags#verboseExceptionSampler()}'s decision. + */ + public static AbortedAttemptException get() { + return Flags.verboseExceptionSampler().isSampled( + AbortedAttemptException.class) ? + new AbortedAttemptException() : INSTANCE; + } + + private AbortedAttemptException() {} + + private AbortedAttemptException(@SuppressWarnings("unused") boolean dummy) { + super(null, null, false, false); + } +} diff --git a/core/src/main/java/com/linecorp/armeria/client/retry/AbstractRetryingClient.java b/core/src/main/java/com/linecorp/armeria/client/retry/AbstractRetryingClient.java index e8e5cc277e3..0e040313341 100644 --- a/core/src/main/java/com/linecorp/armeria/client/retry/AbstractRetryingClient.java +++ b/core/src/main/java/com/linecorp/armeria/client/retry/AbstractRetryingClient.java @@ -15,32 +15,21 @@ */ package com.linecorp.armeria.client.retry; -import static com.google.common.base.Preconditions.checkState; import static java.util.Objects.requireNonNull; -import java.util.concurrent.TimeUnit; -import java.util.function.Consumer; - -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; +import java.util.concurrent.CompletableFuture; import com.linecorp.armeria.client.Client; -import com.linecorp.armeria.client.ClientFactory; import com.linecorp.armeria.client.ClientRequestContext; -import com.linecorp.armeria.client.Endpoint; import com.linecorp.armeria.client.SimpleDecoratingClient; import com.linecorp.armeria.common.HttpHeaderNames; -import com.linecorp.armeria.common.HttpRequest; import com.linecorp.armeria.common.Request; import com.linecorp.armeria.common.Response; -import com.linecorp.armeria.common.RpcRequest; import com.linecorp.armeria.common.annotation.Nullable; -import com.linecorp.armeria.common.util.TimeoutMode; -import com.linecorp.armeria.internal.client.ClientUtil; +import com.linecorp.armeria.common.stream.AbortedStreamException; +import io.netty.channel.EventLoop; import io.netty.util.AsciiString; -import io.netty.util.AttributeKey; -import io.netty.util.concurrent.ScheduledFuture; /** * A {@link Client} decorator that handles failures of remote invocation and retries requests. @@ -48,21 +37,16 @@ * @param the {@link Request} type * @param the {@link Response} type */ -public abstract class AbstractRetryingClient +public abstract class AbstractRetryingClient + extends SimpleDecoratingClient { - - private static final Logger logger = LoggerFactory.getLogger(AbstractRetryingClient.class); - /** * The header which indicates the retry count of a {@link Request}. * The server might use this value to reject excessive retries, etc. */ public static final AsciiString ARMERIA_RETRY_COUNT = HttpHeaderNames.of("armeria-retry-count"); - private static final AttributeKey STATE = - AttributeKey.valueOf(AbstractRetryingClient.class, "STATE"); - - private final RetryConfigMapping mapping; + private final RetryConfigMapping retryMapping; @Nullable private final RetryConfig retryConfig; @@ -71,271 +55,255 @@ public abstract class AbstractRetryingClient delegate, RetryConfigMapping mapping, @Nullable RetryConfig retryConfig) { + Client delegate, RetryConfigMapping retryMapping, @Nullable RetryConfig retryConfig) { super(delegate); - this.mapping = requireNonNull(mapping, "mapping"); + this.retryMapping = requireNonNull(retryMapping, "retryMapping"); this.retryConfig = retryConfig; } + abstract RetryContext newRetryContext( + Client delegate, + ClientRequestContext ctx, + I req, + RetryConfig config); + @Override public final O execute(ClientRequestContext ctx, I req) throws Exception { - final RetryConfig config = mapping.get(ctx, req); - requireNonNull(config, "mapping.get() returned null"); + final RetryConfig config = + retryConfig != null ? retryConfig + : requireNonNull(retryMapping.get(ctx, req), + "retryMapping.get() returned null"); + final RetryContext rctx = newRetryContext(unwrap(), ctx, req, config); + + if (rctx.retryEventLoop().inEventLoop()) { + prepareAndRetry(rctx); + } else { + rctx.retryEventLoop().execute(() -> prepareAndRetry(rctx)); + } - final State state = new State(config, ctx.responseTimeoutMillis()); - ctx.setAttr(STATE, state); - return doExecute(ctx, req); + return rctx.res(); } - /** - * Returns the current {@link RetryConfigMapping} set for this client. - */ - protected final RetryConfigMapping mapping() { - return mapping; - } + private void prepareAndRetry(RetryContext rctx) { + assert rctx.retryEventLoop().inEventLoop(); - /** - * Invoked by {@link #execute(ClientRequestContext, Request)} - * after the deadline for response timeout is set. - */ - protected abstract O doExecute(ClientRequestContext ctx, I req) throws Exception; + try { + prepareRetry(rctx); + } catch (Throwable cause) { + handleUnexpectedException(rctx, cause); + return; + } - /** - * This should be called when retrying is finished. - */ - protected static void onRetryingComplete(ClientRequestContext ctx) { - ctx.logBuilder().endResponseWithLastChild(); + // Let us execute retry() through the scheduler so that it sees the first retry task. + assert rctx.scheduler().trySchedule(() -> retry(rctx, null), 0); } - /** - * Returns the {@link RetryRule}. - * - * @throws IllegalStateException if the {@link RetryRule} is not set - */ - protected final RetryRule retryRule() { - checkState(retryConfig != null, "No retryRule set. Are you using RetryConfigMapping?"); - final RetryRule retryRule = retryConfig.retryRule(); - checkState(retryRule != null, "retryRule is not set."); - return retryRule; - } + private void prepareRetry(RetryContext rctx) { + rctx.res().whenComplete().handle((result, cause) -> { + final Throwable abortCause; + if (cause != null) { + abortCause = cause; + } else { + abortCause = AbortedStreamException.get(); + } + if (rctx.retryEventLoop().inEventLoop()) { + rctx.req().abort(abortCause); + } else { + rctx.retryEventLoop().execute(() -> rctx.req().abort(abortCause)); + } + return null; + }); + + // The request could complete at any moment. + // In that case, let us make sure that we close the scheduler so we + // do not run any retry task unnecessarily. + // That said, running retry() even with a completed response is handled by + // rctx.request().executeAttempt() which throws an AbortedAttemptException which we handle gracefully. + rctx.req().whenComplete().handle((res, cause) -> { + assert rctx.retryEventLoop().inEventLoop(); + // Make sure we do not unnecessarily run any retry task. + rctx.scheduler().close(); + + if (cause != null) { + rctx.resFuture().completeExceptionally(cause); + } else { + rctx.resFuture().complete(res); + } - /** - * Fetches the {@link RetryConfig} that was mapped by the configured {@link RetryConfigMapping} for a given - * logical request. - */ - final RetryConfig mappedRetryConfig(ClientRequestContext ctx) { - @SuppressWarnings("unchecked") - final RetryConfig config = (RetryConfig) state(ctx).config; - return config; - } + return null; + }); - /** - * Returns the {@link RetryRuleWithContent}. - * - * @throws IllegalStateException if the {@link RetryRuleWithContent} is not set - */ - protected final RetryRuleWithContent retryRuleWithContent() { - checkState(retryConfig != null, "No retryRuleWithContent set. Are you using RetryConfigMapping?"); - final RetryRuleWithContent retryRuleWithContent = retryConfig.retryRuleWithContent(); - checkState(retryRuleWithContent != null, "retryRuleWithContent is not set."); - return retryRuleWithContent; + // If something goes wrong with the scheduler, e.g. when the ClientFactory closes and the scheduler is + // unable to schedule retry tasks, we want to make sure we gracefully abort the request with that + // exception. + rctx.scheduler().whenClosed().handle((unused, cause) -> { + assert rctx.retryEventLoop().inEventLoop(); + + if (cause == null) { + cause = new IllegalStateException( + "retry scheduler was closed before the request was completed"); + } + + rctx.req().abort(cause); + return null; + }); } - /** - * Schedules next retry. - */ - protected static void scheduleNextRetry(ClientRequestContext ctx, - Consumer actionOnException, - Runnable retryTask, long nextDelayMillis) { + // NOTE: + // - Does not throw. + // - Must run on the retryEventLoop. + // - The first call must be done from prepareAndRetry() above. + // - Subsequent calls must only be issued in the retry task scheduled by `rctx.scheduler()`. + // The corresponding `scheduler.schedule()` calls must all be done from ´tryScheduleRetryAfter()´ below. + private void retry( + RetryContext rctx, + @Nullable Backoff previousBackoff + ) { try { - if (nextDelayMillis == 0) { - ctx.eventLoop().execute(retryTask); - } else { - @SuppressWarnings("unchecked") - final ScheduledFuture scheduledFuture = (ScheduledFuture) ctx - .eventLoop().schedule(retryTask, nextDelayMillis, TimeUnit.MILLISECONDS); - scheduledFuture.addListener(future -> { - if (future.isCancelled()) { - // future is cancelled when the client factory is closed. - actionOnException.accept(new IllegalStateException( - ClientFactory.class.getSimpleName() + " has been closed.")); - } else if (future.cause() != null) { - // Other unexpected exceptions. - actionOnException.accept(future.cause()); + // First record the execution of the following attempt. This increases the attempt count + // and the attempt count for the backoff. Also see (A) for a correctness argument. + rctx.counter().consumeAttemptFrom(previousBackoff); + + rctx.req() + .executeAttempt(rctx.delegate()) + .handle((executionResult, cause) -> { + assert rctx.retryEventLoop().inEventLoop(); + + if (cause != null) { + rctx.req().abort( + new IllegalStateException("expected RetriedRequest to be completed")); + + if (cause instanceof AbortedAttemptException) { + // The attempt was aborted in the course of executing it. This can happen when + // we execute an attempt on a completed request or when the request is completed + // while RetriedRequest is waiting for the response of the attempt. + return null; + } + + return null; } - }); - } - } catch (Throwable t) { - actionOnException.accept(t); - } - } - /** - * Resets the {@link ClientRequestContext#responseTimeoutMillis()}. - * - * @return {@code true} if the response timeout is set, {@code false} if it can't be set due to the timeout - */ - @SuppressWarnings("MethodMayBeStatic") // Intentionally left non-static for better user experience. - protected final boolean setResponseTimeout(ClientRequestContext ctx) { - requireNonNull(ctx, "ctx"); - final long responseTimeoutMillis = state(ctx).responseTimeoutMillis(); - if (responseTimeoutMillis < 0) { - return false; - } else if (responseTimeoutMillis == 0) { - ctx.clearResponseTimeout(); - return true; - } else { - ctx.setResponseTimeoutMillis(TimeoutMode.SET_FROM_NOW, responseTimeoutMillis); - return true; - } - } + // An empty backoff means that the RetryRule to commit this attempt. + if (executionResult.decision().backoff() == null) { + rctx.req().commit(executionResult.attemptNumber()); + return null; + } - /** - * Returns the next delay which retry will be made after. The delay will be: - * - *

{@code Math.min(responseTimeoutMillis, Backoff.nextDelayMillis(int))} - * - * @return the number of milliseconds to wait for before attempting a retry. -1 if the - * {@code currentAttemptNo} exceeds the {@code maxAttempts} or the {@code nextDelay} is after - * the moment which timeout happens. - */ - protected final long getNextDelay(ClientRequestContext ctx, Backoff backoff) { - return getNextDelay(ctx, backoff, -1); - } + // Note that applying the pushback needs to be done before the call + // to tryScheduleRetryAfter`. + if (executionResult.minimumBackoffMillis() > 0) { + rctx.scheduler().applyMinimumBackoffMillisForNextRetry( + executionResult.minimumBackoffMillis()); + } - /** - * Returns the next delay which retry will be made after. The delay will be: - * - *

{@code Math.min(responseTimeoutMillis, Math.max(Backoff.nextDelayMillis(int), - * millisAfterFromServer))} - * - * @return the number of milliseconds to wait for before attempting a retry. -1 if the - * {@code currentAttemptNo} exceeds the {@code maxAttempts} or the {@code nextDelay} is after - * the moment which timeout happens. - */ - @SuppressWarnings("MethodMayBeStatic") // Intentionally left non-static for better user experience. - protected final long getNextDelay(ClientRequestContext ctx, Backoff backoff, long millisAfterFromServer) { - requireNonNull(ctx, "ctx"); - requireNonNull(backoff, "backoff"); - final State state = state(ctx); - final int currentAttemptNo = state.currentAttemptNoWith(backoff); - - if (currentAttemptNo < 0) { - logger.debug("Exceeded the default number of max attempt: {}", state.config.maxTotalAttempts()); - return -1; - } + final boolean isNextRetryScheduledOrExecuted = + tryScheduleRetryAfter(rctx, executionResult.decision().backoff()); + + if (!isNextRetryScheduledOrExecuted) { + // We are not going to retry again so we are the last attempt. Let us commit it even if + // the response was deemed to be unsatisfactory by the RetryRule (backoff != null). + rctx.req().commit(executionResult.attemptNumber()); + } else { + // The responsibility of aborting the RetriedRequest is now with the retry + // scheduled by `tryScheduleRetryAfter()`. Thus, we can safely abort this attempt now. + rctx.req().abort(executionResult.attemptNumber(), AbortedStreamException.get()); + } - long nextDelay = backoff.nextDelayMillis(currentAttemptNo); - if (nextDelay < 0) { - logger.debug("Exceeded the number of max attempts in the backoff: {}", backoff); - return -1; + return null; + }) + .exceptionally(cause -> { + handleUnexpectedException(rctx, cause); + return null; + }); + } catch (Throwable cause) { + handleUnexpectedException(rctx, cause); } + } - nextDelay = Math.max(nextDelay, millisAfterFromServer); - if (state.timeoutForWholeRetryEnabled() && nextDelay > state.actualResponseTimeoutMillis()) { - // The nextDelay will be after the moment which timeout will happen. So return just -1. - return -1; + // NOTE: Must run on the retryEventLoop. + private boolean tryScheduleRetryAfter(RetryContext rctx, Backoff nextBackoff) { + if (rctx.counter().hasReachedMaxAttempts()) { + return false; } - return nextDelay; - } + final long nextRetryDelayMillisFromBackoff = nextBackoff.nextDelayMillis( + // NOTE: `attemptsSoFarWithBackoff` is read-only and in particular it does not increase + // the attempt count for the backoff. + // +1 because nextDelayMillis counts the original request as one attempt for the backoff. + rctx.counter().attemptsSoFarWithBackoff(nextBackoff) + 1 + ); - /** - * Returns the total number of attempts of the current request represented by the specified - * {@link ClientRequestContext}. - */ - protected static int getTotalAttempts(ClientRequestContext ctx) { - final State state = ctx.attr(STATE); - if (state == null) { - return 0; + if (nextRetryDelayMillisFromBackoff < 0) { + // We exceeded the attempt limit for the backoff. + return false; } - return state.totalAttemptNo; - } - /** - * Creates a new derived {@link ClientRequestContext}, replacing the requests. - * If {@link ClientRequestContext#endpointGroup()} exists, a new {@link Endpoint} will be selected. - */ - protected static ClientRequestContext newDerivedContext(ClientRequestContext ctx, - @Nullable HttpRequest req, - @Nullable RpcRequest rpcReq, - boolean initialAttempt) { - return ClientUtil.newDerivedContext(ctx, req, rpcReq, initialAttempt); + // (A): We are under `maxAttempts` have also not exceeded the backoff attempt limit so let us + // try to schedule the next retry task. + // NOTE: The scheduler guarantees *if* we run this retry task, there will be no retry task run between + // this call and the execution of the retry task. This guarantees two things: + // 1. The attempt number and the attempt number for the backoff now and upon execution of the retry task + // are the same. + // 2. Based on 1., the delay we use to schedule the retry task is indeed the one we consume in the retry + // task via `counter.consumeAttemptFrom(nextBackoff)`. + return rctx.scheduler().trySchedule(/* retry task */ () -> retry(rctx, nextBackoff), + nextRetryDelayMillisFromBackoff); } - private static State state(ClientRequestContext ctx) { - final State state = ctx.attr(STATE); - assert state != null; - return state; + private void handleUnexpectedException(RetryContext rctx, Throwable cause) { + assert rctx.retryEventLoop().inEventLoop(); + // Aborting the request will trigger the cleanup logic in prepareRetry(). + rctx.req().abort(new IllegalStateException("Unexpected exception during retrying", cause)); } - private static final class State { - - private final RetryConfig config; - private final long deadlineNanos; - private final boolean isTimeoutEnabled; - - @Nullable - private Backoff lastBackoff; - private int currentAttemptNoWithLastBackoff; - private int totalAttemptNo; - - State(RetryConfig config, long responseTimeoutMillis) { - this.config = config; - - if (responseTimeoutMillis <= 0 || responseTimeoutMillis == Long.MAX_VALUE) { - deadlineNanos = 0; - isTimeoutEnabled = false; - } else { - deadlineNanos = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(responseTimeoutMillis); - isTimeoutEnabled = true; - } - totalAttemptNo = 1; + protected final class RetryContext { + final EventLoop retryEventLoop; + final RetriedRequest req; + final RetryScheduler scheduler; + final RetryCounter counter; + final Client delegate; + final O res; + final CompletableFuture resFuture; + + RetryContext(EventLoop retryEventLoop, RetriedRequest req, + RetryScheduler scheduler, RetryCounter counter, Client delegate, + O res, + CompletableFuture resFuture) { + this.retryEventLoop = retryEventLoop; + this.req = req; + this.scheduler = scheduler; + this.counter = counter; + this.delegate = delegate; + this.res = res; + this.resFuture = resFuture; } - /** - * Returns the smaller value between {@link RetryConfig#responseTimeoutMillisForEachAttempt()} and - * remaining {@link #responseTimeoutMillis}. - * - * @return 0 if the response timeout for both of each request and whole retry is disabled or - * -1 if the elapsed time from the first request has passed {@code responseTimeoutMillis} - */ - long responseTimeoutMillis() { - if (!timeoutForWholeRetryEnabled()) { - return config.responseTimeoutMillisForEachAttempt(); - } - - final long actualResponseTimeoutMillis = actualResponseTimeoutMillis(); + EventLoop retryEventLoop() { + return retryEventLoop; + } - // Consider 0 or less than 0 of actualResponseTimeoutMillis as timed out. - if (actualResponseTimeoutMillis <= 0) { - return -1; - } + RetriedRequest req() { + return req; + } - if (config.responseTimeoutMillisForEachAttempt() > 0) { - return Math.min(config.responseTimeoutMillisForEachAttempt(), actualResponseTimeoutMillis); - } + O res() { + return res; + } - return actualResponseTimeoutMillis; + CompletableFuture resFuture() { + return resFuture; } - boolean timeoutForWholeRetryEnabled() { - return isTimeoutEnabled; + RetryScheduler scheduler() { + return scheduler; } - long actualResponseTimeoutMillis() { - return TimeUnit.NANOSECONDS.toMillis(deadlineNanos - System.nanoTime()); + RetryCounter counter() { + return counter; } - int currentAttemptNoWith(Backoff backoff) { - if (totalAttemptNo++ >= config.maxTotalAttempts()) { - return -1; - } - if (lastBackoff != backoff) { - lastBackoff = backoff; - currentAttemptNoWithLastBackoff = 1; - } - return currentAttemptNoWithLastBackoff++; + Client delegate() { + return delegate; } } } diff --git a/core/src/main/java/com/linecorp/armeria/client/retry/DefaultRetryCounter.java b/core/src/main/java/com/linecorp/armeria/client/retry/DefaultRetryCounter.java new file mode 100644 index 00000000000..c716245d324 --- /dev/null +++ b/core/src/main/java/com/linecorp/armeria/client/retry/DefaultRetryCounter.java @@ -0,0 +1,85 @@ +/* + * Copyright 2025 LINE Corporation + * + * Licensed 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 com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkState; +import static java.util.Objects.requireNonNull; + +import com.google.common.base.MoreObjects; + +import com.linecorp.armeria.common.annotation.Nullable; + +final class DefaultRetryCounter implements RetryCounter { + private final int maxAttempts; + + private int numberAttemptsSoFar; + @Nullable + private Backoff lastBackoff; + private int numberAttemptsSoFarForLastBackoff; + + DefaultRetryCounter(int maxAttempts) { + checkArgument(maxAttempts > 0, "maxAttempts: %s (expected: > 0)", maxAttempts); + this.maxAttempts = maxAttempts; + numberAttemptsSoFar = 0; + lastBackoff = null; + numberAttemptsSoFarForLastBackoff = 0; + } + + @Override + public void consumeAttemptFrom(@Nullable Backoff backoff) { + checkState(!hasReachedMaxAttempts(), "Exceeded the maximum number of attempts: %s", maxAttempts); + + ++numberAttemptsSoFar; + + if (backoff != null) { + if (lastBackoff != backoff) { + lastBackoff = backoff; + numberAttemptsSoFarForLastBackoff = 0; + } + numberAttemptsSoFarForLastBackoff++; + } else { + assert lastBackoff == null; + } + } + + @Override + public int attemptsSoFarWithBackoff(Backoff backoff) { + requireNonNull(backoff, "backoff"); + if (lastBackoff != backoff) { + return 0; + } else { + return numberAttemptsSoFarForLastBackoff; + } + } + + @Override + public boolean hasReachedMaxAttempts() { + return numberAttemptsSoFar >= maxAttempts; + } + + @Override + public String toString() { + return MoreObjects + .toStringHelper(this) + .add("maxAttempts", maxAttempts) + .add("numberAttemptsSoFar", numberAttemptsSoFar) + .add("lastBackoff", lastBackoff) + .add("numberAttemptsSoFarForLastBackoff", numberAttemptsSoFarForLastBackoff) + .toString(); + } +} diff --git a/core/src/main/java/com/linecorp/armeria/client/retry/DefaultRetryScheduler.java b/core/src/main/java/com/linecorp/armeria/client/retry/DefaultRetryScheduler.java new file mode 100644 index 00000000000..616837c16e0 --- /dev/null +++ b/core/src/main/java/com/linecorp/armeria/client/retry/DefaultRetryScheduler.java @@ -0,0 +1,234 @@ +/* + * Copyright 2025 LY Corporation + * + * LY 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 com.google.common.base.Preconditions.checkState; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import com.google.common.math.LongMath; + +import com.linecorp.armeria.client.ClientFactory; +import com.linecorp.armeria.client.ResponseTimeoutException; +import com.linecorp.armeria.common.annotation.Nullable; +import com.linecorp.armeria.common.util.UnmodifiableFuture; + +import io.netty.channel.EventLoop; +import io.netty.util.concurrent.ScheduledFuture; + +final class DefaultRetryScheduler implements RetryScheduler { + private static final Logger logger = LoggerFactory.getLogger(DefaultRetryScheduler.class); + + private final EventLoop retryEventLoop; + private final long deadlineTimeNanos; + + private final CompletableFuture closedFuture; + private final RetryTaskWrapper retryTaskWrapper; + + @Nullable + private Runnable nextRetryTask; + @Nullable + private ScheduledFuture nextRetryTaskFuture; + // Long.MIN_VALUE if not set. + private long earliestRetryTimeNanos; + private boolean isClosed; + + private final class RetryTaskWrapper implements Runnable { + @Override + public void run() { + assert retryEventLoop.inEventLoop(); + + if (isClosed) { + // Very unexpected as we would cancel this future on the same event loop here in the scheduler + // upon close() or closeExceptionally(). + logger.debug("Tried to run a retry task after the scheduler was closed. Skipping this task."); + return; + } + + if (System.nanoTime() > deadlineTimeNanos) { + closeExceptionally(ResponseTimeoutException.get()); + return; + } + + assert nextRetryTask != null; + final Runnable retryTaskToRun = nextRetryTask; + nextRetryTask = null; + nextRetryTaskFuture = null; + earliestRetryTimeNanos = Long.MIN_VALUE; + + try { + retryTaskToRun.run(); + } catch (Throwable t) { + // Normally we are running retry() which does not throw an exception + // but let us be defensive here. + closeExceptionally(t); + } + } + } + + DefaultRetryScheduler(EventLoop retryEventLoop, long deadlineTimeNanos) { + this.retryEventLoop = retryEventLoop; + this.deadlineTimeNanos = deadlineTimeNanos; + + retryTaskWrapper = new RetryTaskWrapper(); + closedFuture = new CompletableFuture<>(); + + nextRetryTask = null; + nextRetryTaskFuture = null; + earliestRetryTimeNanos = Long.MIN_VALUE; + isClosed = false; + } + + @Override + public boolean trySchedule(Runnable retryTask, long delayMillis) { + checkInRetryEventLoop(); + + if (isClosed) { + return false; + } + + // We are a sequential scheduler there must not be a nextRetryTask already set. + checkState(!hasNextRetryTask(), "cannot schedule a retry task when another is scheduled"); + + assert nextRetryTask == null; + assert nextRetryTaskFuture == null; + + final long retryTimeNanos = Math.max( + LongMath.saturatedAdd( + System.nanoTime(), + TimeUnit.MILLISECONDS.toNanos(delayMillis) + ), + earliestRetryTimeNanos + ); + + if (retryTimeNanos >= deadlineTimeNanos) { + return false; + } + + try { + final long nextRetryDelayMillis = TimeUnit.NANOSECONDS.toMillis( + retryTimeNanos - System.nanoTime()); + + nextRetryTask = retryTask; + if (nextRetryDelayMillis <= 0) { + // Run immediately. + nextRetryTaskFuture = null; + retryTaskWrapper.run(); + } else { + nextRetryTaskFuture = retryEventLoop.schedule( + retryTaskWrapper, nextRetryDelayMillis, + TimeUnit.MILLISECONDS); + nextRetryTaskFuture.addListener(future -> { + if (isClosed) { + return; + } + + if (future.isCancelled()) { + // future is cancelled when the client factory is closed. + closeExceptionally(new IllegalStateException( + ClientFactory.class.getSimpleName() + " has been closed.")); + } else if (future.cause() != null) { + // Other unexpected exceptions. + closeExceptionally(future.cause()); + } + }); + } + + return true; + } catch (Throwable t) { + closeExceptionally(t); + return false; + } + } + + @Override + public void applyMinimumBackoffMillisForNextRetry(long minimumBackoffMillis) { + checkInRetryEventLoop(); + + if (isClosed) { + return; + } + + // We explicitly disallow that just to avoid having to implement cancelling and rescheduling + // logic. A scheduler implementing hedging would need to do support that, however. + checkState(!hasNextRetryTask(), + "cannot apply minimum backoff when a retry task is scheduled"); + + earliestRetryTimeNanos = + Math.min( + Math.max(earliestRetryTimeNanos, + LongMath.saturatedAdd(System.nanoTime(), + TimeUnit.MILLISECONDS.toNanos( + minimumBackoffMillis)) + ), + deadlineTimeNanos + ); + } + + private boolean hasNextRetryTask() { + checkInRetryEventLoop(); + // NOTE: nextRetryTask is null when scheduler is closed. + return nextRetryTask != null; + } + + @Override + public void close() { + checkInRetryEventLoop(); + + if (isClosed) { + return; + } + + isClosed = true; + clearRetryTaskIfExists(); + closedFuture.complete(null); + } + + private void closeExceptionally(Throwable cause) { + if (isClosed) { + return; + } + + isClosed = true; + clearRetryTaskIfExists(); + closedFuture.completeExceptionally(cause); + } + + @Override + public CompletableFuture whenClosed() { + checkInRetryEventLoop(); + + return UnmodifiableFuture.wrap(closedFuture); + } + + private void clearRetryTaskIfExists() { + if (nextRetryTaskFuture != null) { + nextRetryTaskFuture.cancel(false); + } + nextRetryTaskFuture = null; + earliestRetryTimeNanos = Long.MIN_VALUE; + nextRetryTask = null; + } + + private void checkInRetryEventLoop() { + checkState(retryEventLoop.inEventLoop(), "not in the retryEventLoop: %s but in thread %s", + retryEventLoop, Thread.currentThread().getName()); + } +} diff --git a/core/src/main/java/com/linecorp/armeria/client/retry/HttpRetryAttempt.java b/core/src/main/java/com/linecorp/armeria/client/retry/HttpRetryAttempt.java new file mode 100644 index 00000000000..e31466f09f9 --- /dev/null +++ b/core/src/main/java/com/linecorp/armeria/client/retry/HttpRetryAttempt.java @@ -0,0 +1,645 @@ +/* + * Copyright 2025 LY Corporation + * + * LY 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 com.google.common.base.Preconditions.checkState; +import static com.linecorp.armeria.client.retry.AbstractRetryingClient.ARMERIA_RETRY_COUNT; +import static com.linecorp.armeria.internal.client.ClientUtil.executeWithFallback; +import static com.linecorp.armeria.internal.client.ClientUtil.initContextAndExecuteWithFallback; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; +import java.util.function.Function; + +import com.google.common.base.MoreObjects; + +import com.linecorp.armeria.client.Client; +import com.linecorp.armeria.client.ClientRequestContext; +import com.linecorp.armeria.common.AggregatedHttpResponse; +import com.linecorp.armeria.common.ContextAwareEventLoop; +import com.linecorp.armeria.common.HttpRequest; +import com.linecorp.armeria.common.HttpResponse; +import com.linecorp.armeria.common.HttpResponseDuplicator; +import com.linecorp.armeria.common.RequestHeadersBuilder; +import com.linecorp.armeria.common.ResponseHeaders; +import com.linecorp.armeria.common.SplitHttpResponse; +import com.linecorp.armeria.common.annotation.Nullable; +import com.linecorp.armeria.common.logging.RequestLog; +import com.linecorp.armeria.common.logging.RequestLogBuilder; +import com.linecorp.armeria.common.logging.RequestLogProperty; +import com.linecorp.armeria.common.util.Exceptions; +import com.linecorp.armeria.common.util.UnmodifiableFuture; +import com.linecorp.armeria.internal.client.ClientPendingThrowableUtil; +import com.linecorp.armeria.internal.client.ClientRequestContextExtension; +import com.linecorp.armeria.internal.client.TruncatingHttpResponse; + +/** + * A single attempt for a {@link RetriedHttpRequest}. + * + *

+ * NOTE: All methods of {@link HttpRetryAttempt} must be invoked from the + * {@code retryEventLoop}. + *

+ */ +final class HttpRetryAttempt { + /** + * Result of {@code State.PREPARING_AND_DECIDING} which contains the response of the attempt + * and the response or cause to be delivered to the {@link RetryRule} for decision + * processing the response until the point where the {@link RetryRule} can make a decision + * (e.g., receiving headers, content, or trailers). + */ + private static final class PreparationResult { + /** The response of the attempt. */ + final HttpResponse res; + + /** + * Response that is derived from {@link PreparationResult#res} which is delivered to + * the {@link RetryRule} for decision. May be {@link PreparationResult#res} itself. + * It is {@code null} iff {@link PreparationResult#causeToDecide} is not {@code null}. + */ + @Nullable + final HttpResponse resToDecide; + + /** + * Exception that occurred while preparing the {@link PreparationResult#res}. It is delivered to + * the {@link RetryRule} for decision. + * It is not {@code null} iff {@link PreparationResult#resToDecide} is {@code null}. + */ + @Nullable + Throwable causeToDecide; + + private PreparationResult(HttpResponse res, @Nullable HttpResponse resToDecide, + @Nullable Throwable causeToDecide) { + // TODO: remove? + assert res != null; + assert (resToDecide == null) != (causeToDecide == null) + : resToDecide + ", " + causeToDecide; + + this.res = res; + this.resToDecide = resToDecide; + this.causeToDecide = causeToDecide; + } + + static PreparationResult ofSuccess(HttpResponse res, HttpResponse resForRule) { + return new PreparationResult(res, resForRule, null); + } + + static PreparationResult ofFailure(Throwable cause) { + return new PreparationResult(HttpResponse.ofFailure(cause), null, cause); + } + } + + /** + * The state of this attempt. + * The following state machine diagram illustrates the possible transitions: + * 
+     *          Start
+     *            |
+     *           IDLE
+     *            |
+     *            v
+     *        EXECUTING----------+
+     *            |              |
+     *            v              |
+     *         DECIDED-----------+
+     *          /   \            |
+     *         /     \           |
+     *        /       \          |
+     *       v         v         |
+     *   COMMITTED  ABORTED<-----+
+     *       \         /
+     *        +---+---+
+     *            |
+     *            v
+     *           End
+     *
+ */ + enum State { + /** + * The attempt was not executed yet via {@link #executeAndDecide(Client)}. + */ + IDLE, + /** + * The attempt was executed by {@link #executeAndDecide(Client)}. + * + *

+ * The response is being prepared for and given to the {@link RetryRule} to make a decision. + * In this context, 'preparing' means processing the response until the point where + * the {@link RetryRule} can make a decision (e.g., receiving headers, content, or trailers). + * {@link #res} is available from this state onwards. + *

+ */ + EXECUTING, + /** + * The response was prepared and the {@link RetryRule} made a decision. + * The attempt is ready to be committed via {@link #commit()} or aborted via {@link #abort(Throwable)}. + */ + DECIDED, + /** + * The attempt was committed via {@link #commit()}. This is a terminal state and the response + * is final. + */ + COMMITTED, + /** + * The attempt was aborted via {@link #abort(Throwable)}. The response is + * aborted via the same cause. In this state and only in this state {@link #cause} is not-{@code null}. + * This is a terminal state. + */ + ABORTED + } + + private final RetryConfig config; + private final ContextAwareEventLoop retryEventLoop; + private final int attemptNumber; + private final ClientRequestContext ctx; + private final HttpRequest req; + + private State state; + + /** + * The response of the attempt. It is available in {@link State#EXECUTING}, + * {@link State#DECIDED}, and {@link State#COMMITTED} states. + */ + @Nullable + private HttpResponse res; + + /** + * The cause of the attempt failure. It is not-{@code null} iff we are in {@link State#ABORTED} state. + */ + @Nullable + Throwable cause; + + HttpRetryAttempt( + RetryConfig config, + ContextAwareEventLoop retryEventLoop, + int attemptNumber, + ClientRequestContext ctx, + HttpRequest req + ) { + this.config = config; + this.retryEventLoop = retryEventLoop; + this.attemptNumber = attemptNumber; + this.ctx = ctx; + this.req = req; + + state = State.IDLE; + } + + /** + * Executes the attempt, prepares and gives the response or the cause to the {@link RetryRule}. + * This method must be called at most once and except for calls to {@link #abort(Throwable)} + * it must be called before any other method, in particular before {@link #commit()}. + * + * @param delegate the next {@link Client} in the decoration chain + * @return a future that will be completed with the {@link RetryDecision} or an exception if failed during + * execution, preparation, or decision. + * In particular, it fails with a {@link AbortedAttemptException} + * if the attempt was aborted by {@link #abort(Throwable)}. + * + * @see #commit() + */ + CompletableFuture executeAndDecide(Client delegate) { + checkState(retryEventLoop.inEventLoop()); + + if (state == State.ABORTED) { + assert cause != null; + return UnmodifiableFuture.exceptionallyCompletedFuture(cause); + } + + checkState(state == State.IDLE); + state = State.EXECUTING; + + res = execute(delegate); + + return prepareForDecision(res) + .thenCompose(preparationResult -> { + assert retryEventLoop.inEventLoop(); + // Let us avoid calling decide (which may be expensive if we already know that we were + // aborted. + if (state == State.ABORTED) { + assert cause != null; + assert preparationResult.res.isComplete(); + + if (preparationResult.resToDecide != null && + preparationResult.res != preparationResult.resToDecide) { + preparationResult.resToDecide.abort(cause); + } + + return UnmodifiableFuture.exceptionallyCompletedFuture(cause); + } + + assert state == State.EXECUTING : state; + + return decide(preparationResult.resToDecide, + preparationResult.causeToDecide) + .whenComplete((unusedDecision, unusedDecisionCause) -> { + assert retryEventLoop.inEventLoop(); + // Let us abort the response that was delivered to the rule as we do not need + // it anymore. + if (preparationResult.resToDecide != null && + // In case we delivered the attempt response + preparationResult.res != preparationResult.resToDecide) { + preparationResult.resToDecide.abort(); + } + }) + .thenCompose( + decision -> { + if (state == State.ABORTED) { + assert cause != null; + // Was aborted by {@link #abort()} while we were deciding. + assert preparationResult.res.isComplete(); + return UnmodifiableFuture.exceptionallyCompletedFuture(cause); + } else { + assert state == State.EXECUTING : state; + assert cause == null : cause; // sanity check + assert res != null; + res = preparationResult.res; + state = State.DECIDED; + return UnmodifiableFuture.completedFuture(decision); + } + } + ); + }) + .handle( + (decision, preparationOrDecisionCause) -> { + if (preparationOrDecisionCause == null) { + return (CompletableFuture) UnmodifiableFuture.completedFuture( + decision); + } + + if (state == State.ABORTED) { + assert cause != null; + if (cause != preparationOrDecisionCause) { + cause.addSuppressed(preparationOrDecisionCause); + } + return UnmodifiableFuture.exceptionallyCompletedFuture( + preparationOrDecisionCause); + } + assert state == State.EXECUTING : state; + + abort(preparationOrDecisionCause); + return UnmodifiableFuture.exceptionallyCompletedFuture( + preparationOrDecisionCause); + } + ) + .thenCompose(Function.identity()); + } + + private HttpResponse execute(Client delegate) { + final boolean isInitialAttempt = attemptNumber <= 1; + + if (!isInitialAttempt) { + ctx.mutateAdditionalRequestHeaders( + mutator -> mutator.add(ARMERIA_RETRY_COUNT, Integer.toString(attemptNumber - 1))); + } + + final RequestHeadersBuilder attemptHeadersBuilder = req.headers().toBuilder(); + attemptHeadersBuilder.setInt(ARMERIA_RETRY_COUNT, attemptNumber - 1); + + final ClientRequestContextExtension ctxExt = + ctx.as(ClientRequestContextExtension.class); + if (!isInitialAttempt && ctxExt != null && ctx.endpoint() == null) { + // clear the pending throwable to retry endpoint selection + ClientPendingThrowableUtil.removePendingThrowable(ctx); + // if the endpoint hasn't been selected, + // try to initialize the attempCtx with a new endpoint/event loop + return initContextAndExecuteWithFallback( + delegate, ctxExt, HttpResponse::of, + (context, cause) -> + HttpResponse.ofFailure(cause), req, false); + } else { + return executeWithFallback(delegate, ctx, + (context, cause) -> + HttpResponse.ofFailure(cause), req, false); + } + } + + /** + * Prepares the response for decision by the {@link RetryRule}. + * + *

+ * In the context of this class, preparation means processing the response until the point where the + * {@link RetryRule} can make a decision (e.g., receiving headers, content, or trailers). + * The returned future is completed with a {@link PreparationResult} which contains the response + * for the {@link RetryRule} ({@link PreparationResult#resToDecide}) to make a decision + * and the new response of the attempt (@link PreparationResult#res}). + *

+ * + *

+ * Note that {@link PreparationResult#res} might be different from the + * initial response stored in {@link #res} as we might need to transform it during preparation + * (e.g. by aggregation). + *

+ * + * @param res the response of the attempt + * @return a future that will be completed with the {@link PreparationResult} + */ + private CompletableFuture prepareForDecision(HttpResponse res) { + final boolean aggregateResponse = + !ctx.exchangeType().isResponseStreaming() || config.requiresResponseTrailers(); + return withCompletionOnRetryEventLoop( + aggregateResponse ? prepareAggregatedResponse(res) : prepareStreamingResponse(res)); + } + + /** + * Marks this attempt as committed and returns the response of the attempt. + * The attempt must be decided, i.e. {@link #executeAndDecide(Client)} must have been called and completed + * successfully before calling this method. After this call, the attempt cannot be aborted anymore. + * This method is idempotent once the first call returns successfully. + * + * @return the response of the attempt + * + * @see #executeAndDecide(Client) + */ + public HttpResponse commit() { + checkState(retryEventLoop.inEventLoop()); + + if (state == State.COMMITTED) { + assert res != null; + return res; + } + + checkState(state == State.DECIDED); + assert res != null; + assert cause == null : cause; + state = State.COMMITTED; + return res; + } + + /** + * Aborts this attempt with the specified {@code cause}. + * {@link #executeAndDecide(Client)} must be called before this method is called. + * After this call, the attempt cannot be committed anymore. + * This method is idempotent once the first call returns successfully. + * + * @param cause the cause of the abortion + */ + public void abort(Throwable cause) { + checkState(retryEventLoop.inEventLoop()); + + if (state == State.ABORTED) { + return; + } + + assert state == State.IDLE || + state == State.EXECUTING || + state == State.DECIDED : state; + assert this.cause == null : state; + state = State.ABORTED; + + final RequestLogBuilder logBuilder = ctx.logBuilder(); + // Set response content with null to make sure that the log is complete. + logBuilder.responseContent(null, null); + logBuilder.responseContentPreview(null); + + if (res != null) { + res.abort(cause); + } else { + // If we are in EXECUTING or PREPARING_AND_DECIDING, the execution/preparation must have failed. + } + } + + /** + * Returns the {@link ClientRequestContext} of this attempt. + * + * @return the {@link ClientRequestContext} of this attempt + */ + ClientRequestContext ctx() { + return ctx; + } + + /** + * Returns the {@link State} of this attempt. + * + * @return the {@link State} of this attempt + */ + State state() { + return state; + } + + private CompletableFuture prepareAggregatedResponse(HttpResponse res) { + return res.aggregate() + .thenCompose(aggRes -> { + // NOTE: Might not run on the retry event loop. + completeLogIfBytesNotTransferred(aggRes); + return ctx.log() + .whenAvailable(RequestLogProperty.RESPONSE_END_TIME) + .thenApply(unused -> + // NOTE: Might not run on the retry event loop. + PreparationResult.ofSuccess( + aggRes.toHttpResponse(), + aggRes.toHttpResponse() + ) + ); + }).handle((preparationResult, resCause) -> { + // NOTE: Might not run on the retry event loop. + if (resCause == null) { + return UnmodifiableFuture.completedFuture(preparationResult); + } + + resCause = Exceptions.peel(resCause); + + ctx.logBuilder().endRequest(resCause); + ctx.logBuilder().endResponse(resCause); + return UnmodifiableFuture.completedFuture(PreparationResult.ofFailure(resCause)); + }) + .thenCompose(Function.identity()); + } + + private CompletableFuture prepareStreamingResponse(HttpResponse res) { + final SplitHttpResponse splitRes = res.split(); + + return splitRes.headers() + .handle((resHeaders, headersCause) -> { + // NOTE: Might not run on the retry event loop. + final Throwable resCause; + if (headersCause == null) { + final RequestLog log = ctx.log().getIfAvailable( + RequestLogProperty.RESPONSE_CAUSE); + resCause = log != null ? log.responseCause() : null; + } else { + resCause = Exceptions.peel(headersCause); + } + + final HttpResponse unsplitRes = splitRes.unsplit(); + completeLogIfBytesNotTransferred(unsplitRes, resHeaders, resCause); + + return ctx.log() + // NOTE: .whenAvailable is guaranteed to never complete exceptionally. + .whenAvailable(RequestLogProperty.RESPONSE_HEADERS) + .thenApply(unused -> { + // NOTE: Might not run on the retry event loop. + if (resCause != null) { + unsplitRes.abort(resCause); + return PreparationResult.ofFailure(resCause); + } + + if (config.needsContentInRule()) { + try ( + HttpResponseDuplicator resDuplicator = + unsplitRes.toDuplicator( + ctx.eventLoop().withoutContext(), + ctx.maxResponseLength() + ) + ) { + return PreparationResult.ofSuccess( + resDuplicator.duplicate(), + new TruncatingHttpResponse( + resDuplicator.duplicate(), + config.maxContentLength() + ) + ); + } + } else { + return PreparationResult.ofSuccess( + unsplitRes, + unsplitRes + ); + } + }); + }) + .thenCompose(Function.identity()); + } + + private void completeLogIfBytesNotTransferred(AggregatedHttpResponse aggRes) { + if (!ctx.log().isAvailable(RequestLogProperty.REQUEST_FIRST_BYTES_TRANSFERRED_TIME)) { + final RequestLogBuilder attemptLogBuilder = ctx.logBuilder(); + attemptLogBuilder.endRequest(); + attemptLogBuilder.responseHeaders(aggRes.headers()); + if (!aggRes.trailers().isEmpty()) { + attemptLogBuilder.responseTrailers(aggRes.trailers()); + } + attemptLogBuilder.endResponse(); + } + } + + private void completeLogIfBytesNotTransferred( + HttpResponse res, + @Nullable ResponseHeaders headers, + @Nullable Throwable resCause) { + if (!ctx.log().isAvailable(RequestLogProperty.REQUEST_FIRST_BYTES_TRANSFERRED_TIME)) { + final RequestLogBuilder logBuilder = ctx.logBuilder(); + if (resCause != null) { + logBuilder.endRequest(resCause); + logBuilder.endResponse(resCause); + } else { + logBuilder.endRequest(); + if (headers != null) { + logBuilder.responseHeaders(headers); + } + res.whenComplete().handle((unused, cause) -> { + if (cause != null) { + logBuilder.endResponse(cause); + } else { + logBuilder.endResponse(); + } + return null; + }); + } + } + } + + /** + * Calls {@link RetryRule#shouldRetry(ClientRequestContext, Throwable)} to receive a {@link RetryDecision}. + * + * @param resForRule the response to be delivered to the {@link RetryRule} for decision. + * @param causeForRule the cause to be delivered to the {@link RetryRule} for decision. + * @return a future that will be completed with the {@link RetryDecision} or an exception if failed during + * the decision. The future will be completed on the retry event loop. + */ + private CompletableFuture decide(@Nullable HttpResponse resForRule, + @Nullable Throwable causeForRule) { + if (causeForRule != null) { + resForRule = null; + causeForRule = Exceptions.peel(causeForRule); + } else { + assert resForRule != null; + causeForRule = null; + } + + final CompletionStage<@Nullable RetryDecision> maybeDecision; + + if (config.needsContentInRule()) { + final RetryRuleWithContent retryRuleWithContent = + config.retryRuleWithContent(); + assert retryRuleWithContent != null; + try { + maybeDecision = retryRuleWithContent.shouldRetry(ctx, resForRule, causeForRule); + } catch (Throwable t) { + return UnmodifiableFuture.exceptionallyCompletedFuture(t); + } + } else { + final RetryRule retryRule = config.retryRule(); + assert retryRule != null; + try { + maybeDecision = retryRule.shouldRetry(ctx, causeForRule); + } catch (Throwable t) { + return UnmodifiableFuture.exceptionallyCompletedFuture(t); + } + } + + // Remember that RetryRule.shouldRetry could be client code so we do have any guarantees + // on which thread we are completing. Let us make sure we are completing on running on the + // retry event loop again. + return withCompletionOnRetryEventLoop( + maybeDecision.thenApply( + decision -> decision == null ? RetryDecision.noRetry() : decision + ) + ); + } + + /** + * Ensures the given {@link CompletionStage} completes on the retry event loop. + * If already on the retry event loop, completes directly; otherwise schedules completion. + */ + private CompletableFuture withCompletionOnRetryEventLoop(CompletionStage future) { + final CompletableFuture futureOnTheRetryEventLoop = new CompletableFuture<>(); + future.whenComplete((result, cause) -> { + if (retryEventLoop.inEventLoop()) { + if (cause != null) { + futureOnTheRetryEventLoop.completeExceptionally(cause); + } else { + futureOnTheRetryEventLoop.complete(result); + } + } else { + retryEventLoop.execute(() -> { + if (cause != null) { + futureOnTheRetryEventLoop.completeExceptionally(cause); + } else { + futureOnTheRetryEventLoop.complete(result); + } + }); + } + }); + return futureOnTheRetryEventLoop; + } + + @Override + public String toString() { + checkState(retryEventLoop.inEventLoop()); + + return MoreObjects + .toStringHelper(this) + .add("state", state) + .add("attemptNumber", attemptNumber) + .add("ctx", ctx) + .add("req", req) + .add("res", res) + .add("cause", cause) + .toString(); + } +} diff --git a/core/src/main/java/com/linecorp/armeria/client/retry/RetriedHttpRequest.java b/core/src/main/java/com/linecorp/armeria/client/retry/RetriedHttpRequest.java new file mode 100644 index 00000000000..7d02094acfa --- /dev/null +++ b/core/src/main/java/com/linecorp/armeria/client/retry/RetriedHttpRequest.java @@ -0,0 +1,345 @@ +/* + * Copyright 2025 LY Corporation + * + * LY 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 com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkState; +import static com.linecorp.armeria.client.retry.AbstractRetryingClient.ARMERIA_RETRY_COUNT; + +import java.util.ArrayList; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; + +import com.linecorp.armeria.client.Client; +import com.linecorp.armeria.client.ClientRequestContext; +import com.linecorp.armeria.client.ResponseTimeoutException; +import com.linecorp.armeria.common.AggregationOptions; +import com.linecorp.armeria.common.ContextAwareEventLoop; +import com.linecorp.armeria.common.HttpRequest; +import com.linecorp.armeria.common.HttpRequestDuplicator; +import com.linecorp.armeria.common.HttpResponse; +import com.linecorp.armeria.common.RequestHeadersBuilder; +import com.linecorp.armeria.common.annotation.Nullable; +import com.linecorp.armeria.common.util.UnmodifiableFuture; +import com.linecorp.armeria.internal.client.AggregatedHttpRequestDuplicator; +import com.linecorp.armeria.internal.client.ClientUtil; + +class RetriedHttpRequest implements RetriedRequest { + private enum State { + /** + * Initial state. {@link RetriedHttpRequest} is instantiated but not yet initialized. + */ + IDLE, + /** + * {@link RetriedHttpRequest} is being initialized in {@link #init()} with the first call + * to {@link #executeAttempt(Client)}. See the javadoc of {@link #init()} for details. + */ + INITIALIZING, + /** + * {@link RetriedHttpRequest} is initialized and is ready to execute/ is executing attempts. + */ + PENDING, + /** + * Terminal state. Either {@link #commit(int)} or {@link #abort(Throwable)} was called. + * {@link RetriedHttpRequest} will not execute any more attempts in that it completes every call + * to {@link #executeAttempt(Client)} exceptionally with an {@link AbortedAttemptException}. + * {@link #completedFuture} is completed. + */ + COMPLETED + } + + private final ContextAwareEventLoop retryEventLoop; + private final RetryConfig config; + private final ClientRequestContext ctx; + private final HttpRequest req; + + private final long deadlineTimeNanos; + private final boolean useRetryAfter; + private final CompletableFuture completedFuture; + + private State state; + + /** + * Future that completes when the initialization is done. It always completes successfully. + */ + private final CompletableFuture<@Nullable Void> initFuture; + + int previousAttemptNumber; + + /** + * Duplicator for the original request. It is set from {@code State.PENDING} onwards. + */ + @Nullable + private HttpRequestDuplicator reqDuplicator; + + final ArrayList attempts; + + RetriedHttpRequest( + ContextAwareEventLoop retryEventLoop, + RetryConfig config, + ClientRequestContext ctx, + HttpRequest req, + long deadlineTimeNanos, + boolean useRetryAfter + ) { + this.retryEventLoop = retryEventLoop; + this.ctx = ctx; + this.req = req; + this.config = config; + this.useRetryAfter = useRetryAfter; + this.deadlineTimeNanos = deadlineTimeNanos; + + completedFuture = new CompletableFuture<>(); + initFuture = new CompletableFuture<>(); + + state = State.IDLE; + previousAttemptNumber = 0; + // We can assume that most of the time we will succeed with the first attempt. + attempts = new ArrayList<>(1); + reqDuplicator = null; + } + + @Override + public CompletionStage executeAttempt( + Client delegate) { + checkState(ctx.eventLoop().inEventLoop()); + return init().thenCompose(unused -> executeAttemptAfterInit(delegate)); + } + + /** + * Initializes the {@link RetriedHttpRequest}. This method is idempotent. + * Initialization includes + * - subscribing to {@link #req} to initiate abortion in case its fails, + * - building and setting {@link #reqDuplicator}. + * + * @return a future that completes when the initialization is done. After the future completes + * the {@link RetriedHttpRequest} is either in {@code PENDING} or {@code COMPLETED} state. + * */ + private CompletableFuture init() { + assert ctx.eventLoop().inEventLoop(); + + if (state == State.INITIALIZING || state == State.PENDING || state == State.COMPLETED) { + return initFuture; + } + + assert state == State.IDLE : state; + state = State.INITIALIZING; + + // We are not putting this in the constructor as we might get aborted before + // we are properly initialized. + req.whenComplete().handle((unused, cause) -> { + if (cause != null) { + if (retryEventLoop.inEventLoop()) { + abort(cause); + } else { + retryEventLoop.execute(() -> abort(cause)); + } + } + return null; + }); + + // Guaranteed to be completed on the retryEventLoop. + final CompletableFuture reqDuplicatorFuture; + + if (ctx.exchangeType().isRequestStreaming()) { + reqDuplicatorFuture = UnmodifiableFuture.completedFuture( + req.toDuplicator(retryEventLoop.withoutContext(), 0)); + } else { + reqDuplicatorFuture = req.aggregate( + // TODO: Do we need to run this with the context? + AggregationOptions.usePooledObjects(ctx.alloc(), retryEventLoop)) + .thenApply(AggregatedHttpRequestDuplicator::new); + } + + reqDuplicatorFuture + .thenAccept( + reqDuplicator0 -> { + assert retryEventLoop.inEventLoop(); + assert reqDuplicator == null : reqDuplicator; + + if (state == State.COMPLETED) { + reqDuplicator0.close(); + } else { + assert state == State.INITIALIZING : state; + reqDuplicator = reqDuplicator0; + state = State.PENDING; + } + } + ) + .exceptionally(cause -> { + abort(cause); + return null; + }) + .thenRun(() -> initFuture.complete(null)); + + return initFuture; + } + + private CompletableFuture executeAttemptAfterInit( + Client delegate) { + assert retryEventLoop.inEventLoop(); + + if (state == State.COMPLETED) { + return UnmodifiableFuture.exceptionallyCompletedFuture(AbortedAttemptException.get()); + } + + // We need to be initialized/ the reqDuplicator must be present. + checkState(state == State.PENDING); + assert reqDuplicator != null; + checkState(previousAttemptNumber < config.maxTotalAttempts()); + + final int attemptNumber = ++previousAttemptNumber; + + assert attemptNumber == attempts.size() + 1 : attemptNumber + ", " + attempts.size(); + assert attemptNumber <= config.maxTotalAttempts() : attemptNumber + ", " + config.maxTotalAttempts(); + + if (!ClientUtil.checkAndSetResponseTimeout(ctx, deadlineTimeNanos, + config.responseTimeoutMillisForEachAttempt())) { + final ResponseTimeoutException timeoutException = ResponseTimeoutException.get(); + abort(timeoutException); + return UnmodifiableFuture.exceptionallyCompletedFuture(timeoutException); + } + + final HttpRetryAttempt attempt = newAttempt(attemptNumber); + attempts.add(attempt); + return attempt + .executeAndDecide(delegate) + .thenApply(decision -> { + assert retryEventLoop.inEventLoop(); + return new AttemptExecutionResult( + attemptNumber, + decision, + useRetryAfter ? + ClientUtil.retryAfterMillis(attempt.ctx().log()) + : -1L + ); + }) + .exceptionally(cause -> { + assert retryEventLoop.inEventLoop(); + abort(cause); + return null; + }); + } + + private HttpRetryAttempt newAttempt(int attemptNumber) { + assert reqDuplicator != null; + + final boolean isInitialAttempt = attemptNumber <= 1; + + final HttpRequest attemptReq; + final ClientRequestContext attemptCtx; + if (isInitialAttempt) { + attemptReq = reqDuplicator.duplicate(); + } else { + final RequestHeadersBuilder attemptHeadersBuilder = req.headers().toBuilder(); + attemptHeadersBuilder.setInt(ARMERIA_RETRY_COUNT, attemptNumber - 1); + attemptReq = reqDuplicator.duplicate(attemptHeadersBuilder.build()); + } + + attemptCtx = ClientUtil.newDerivedContext(ctx, attemptReq, ctx.rpcRequest(), isInitialAttempt); + + return new HttpRetryAttempt(config, retryEventLoop, attemptNumber, attemptCtx, attemptReq); + } + + @Override + public void commit(int attemptNumber) { + checkState(retryEventLoop.inEventLoop()); + + if (state == State.COMPLETED) { + return; + } + checkState(state == State.PENDING); + assert reqDuplicator != null; + assert !completedFuture.isDone(); + + checkArgument(attemptNumber >= 1); + checkArgument(attemptNumber <= attempts.size()); + + final HttpRetryAttempt attemptToCommit = attempts.get(attemptNumber - 1); + + checkState(attemptToCommit.state() == HttpRetryAttempt.State.DECIDED); + + state = State.COMPLETED; + + abortAllExcept(attemptToCommit); + + final HttpResponse res = attemptToCommit.commit(); + + reqDuplicator.close(); + ctx.logBuilder().endResponseWithChild(attemptToCommit.ctx().log()); + + completedFuture.complete(res); + } + + @Override + public void abort(Throwable cause) { + checkState(retryEventLoop.inEventLoop()); + + if (state == State.COMPLETED) { + return; + } + assert !completedFuture.isDone(); + + state = State.COMPLETED; + + abortAllExcept(null); + + if (reqDuplicator != null) { + reqDuplicator.close(); + } + + if (!ctx.log().isRequestComplete()) { + ctx.logBuilder().endRequest(cause); + } + ctx.logBuilder().endResponse(cause); + completedFuture.completeExceptionally(cause); + } + + @Override + public void abort(int attemptNumber, Throwable cause) { + checkState(retryEventLoop.inEventLoop()); + checkArgument(attemptNumber >= 1); + + final HttpRetryAttempt attemptToAbort; + if (state == State.COMPLETED) { + return; + } + + checkState(state == State.PENDING); + checkArgument(attemptNumber <= attempts.size()); + + attemptToAbort = attempts.get(attemptNumber - 1); + attemptToAbort.abort(cause); + } + + private void abortAllExcept(@Nullable HttpRetryAttempt winningAttempt) { + assert retryEventLoop.inEventLoop(); + assert state == State.COMPLETED : state; + + for (final HttpRetryAttempt attempt : attempts) { + if (attempt == winningAttempt) { + continue; + } + + attempt.abort(AbortedAttemptException.get()); + } + } + + @Override + public CompletableFuture whenComplete() { + return completedFuture; + } +} diff --git a/core/src/main/java/com/linecorp/armeria/client/retry/RetriedRequest.java b/core/src/main/java/com/linecorp/armeria/client/retry/RetriedRequest.java new file mode 100644 index 00000000000..967074076d1 --- /dev/null +++ b/core/src/main/java/com/linecorp/armeria/client/retry/RetriedRequest.java @@ -0,0 +1,157 @@ +/* + * Copyright 2025 LY Corporation + * + * LY 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 java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; + +import com.linecorp.armeria.client.Client; +import com.linecorp.armeria.common.Request; +import com.linecorp.armeria.common.Response; +import com.linecorp.armeria.common.logging.RequestLog; + +/** + * A retried request that manages multiple retry attempts. + * + *

+ * NOTE: All methods of {@link RetriedRequest} must be invoked from a single-threaded event loop ("retry event + * loop"). + * Implementations of {@link RetriedRequest} therefore do not need to be thread-safe. + *

+ * + * @param the {@link Request} type + * @param the {@link Response} type + */ +interface RetriedRequest { + /** + * The result of an attempt execution. + */ + final class AttemptExecutionResult { + private final int attemptNumber; + private final RetryDecision decision; + private final long minimumBackoffMillis; + + AttemptExecutionResult(int attemptNumber, RetryDecision decision, + long minimumBackoffMillis) { + this.attemptNumber = attemptNumber; + this.decision = decision; + this.minimumBackoffMillis = minimumBackoffMillis; + } + + /** + * Returns the attempt number. It can be used to {@link RetriedRequest#abort(int, Throwable)} or + * {@link RetriedRequest#commit(int)} the attempt. + * + * @return the attempt number. + */ + public int attemptNumber() { + return attemptNumber; + } + + /** + * Returns the decision made by {@link RetryRule} from the {@link RetryConfig} on the attempt response. + * The attempt response was prepared enough for the {@link RetryRule} to make a decision. This could + * include awaiting receiving headers, content, or trailers for example. + * + * @return the decision made by the {@link RetryRule} + */ + public RetryDecision decision() { + return decision; + } + + /** + * Returns the minimum backoff in milliseconds for the next retry attempt + * returned by the remote peer that received the attempt. + * If negative, the minimum backoff is either not available or not requested. + * + * @return the minimum backoff in milliseconds. If negative, no minimum backoff should be applied for + * the next retry attempt. + */ + public long minimumBackoffMillis() { + return minimumBackoffMillis; + } + } + + /** + * Tries to execute a new attempt on {@code delegate} based on the {@link RetryConfig}. + * + *

+ * It completes with a {@link AttemptExecutionResult} if the attempt was successfully executed and decided + * by the {@link RetryRule} from the {@link RetryConfig}. + *

+ * + *

+ * It completes exceptionally with a {@link AbortedAttemptException} if the attempt + * was aborted by the {@link RetriedRequest}. This can happen when the {@link RetriedRequest} completes + * concurrently while the attempt is being executed. + *

+ * + * @param delegate the next {@link Client} in the decoration chain + * @return a future completed with the result of the attempt execution. It completes with a + * {@link AbortedAttemptException} in case the attempt was aborted in a controlled way. + */ + CompletionStage executeAttempt(Client delegate); + + /** + * Commits the attempt with the specified {@code attemptNumber}. Committing an attempt includes: + *
    + *
  • completing the original request future with the attempt's response,
    • + *
  • aborting all other pending attempts, and
    • + *
  • marking the {@link RetriedRequest} as completed.
    • + *
+ * Does nothing if the {@link RetriedRequest} is already completed. + * + * @param attemptNumber the attempt number to commit + */ + void commit(int attemptNumber); + + /** + * Aborts a specific attempt that was previously started via {@link #executeAttempt(Client)}. + * + *

+ * The attempt must have been previously started via + * {@link #executeAttempt(Client)}. Aborting an attempt ensures that the attempt's + * response and {@link RequestLog} are completed but not set as the original request's response or log. + *

+ * + *

If the {@link RetriedRequest} is already completed, this method has no effect. + * If the specified attempt has not been started, an {@link IllegalStateException} is thrown.

+ * + * @param attemptNumber the identifier of the attempt to abort + * @throws IllegalStateException if the attempt has not been executed + */ + void abort(int attemptNumber, Throwable cause); + + /** + * Aborts all attempts and completes the {@link RetriedRequest} exceptionally with the specified + * {@code cause}. Marks the {@link RetriedRequest} as completed. Does nothing if the {@link RetriedRequest} + * was already completed before. + */ + void abort(Throwable cause); + + /** + * Returns a future that completes when this {@link RetriedRequest} is completed. + *
    + *
  • Completes successfully with the response of the attempt committed by {@link #commit(int)}
    • + *
  • Completes exceptionally with the cause from {@link #abort(Throwable)}
    • + *
+ * Note: {@link #abort(Throwable)} may be called internally when unexpected errors occur. + * + * @return a future that completes successfully with the committed attempt's response or + * exceptionally with the abort cause. It is guaranteed to be completed on the retry event loop. + */ + CompletableFuture whenComplete(); +} diff --git a/core/src/main/java/com/linecorp/armeria/client/retry/RetriedRpcRequest.java b/core/src/main/java/com/linecorp/armeria/client/retry/RetriedRpcRequest.java new file mode 100644 index 00000000000..303fc3878a6 --- /dev/null +++ b/core/src/main/java/com/linecorp/armeria/client/retry/RetriedRpcRequest.java @@ -0,0 +1,240 @@ +/* + * Copyright 2025 LY Corporation + * + * LY 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 com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkState; + +import java.util.ArrayList; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; + +import com.linecorp.armeria.client.Client; +import com.linecorp.armeria.client.ClientRequestContext; +import com.linecorp.armeria.client.ResponseTimeoutException; +import com.linecorp.armeria.common.ContextAwareEventLoop; +import com.linecorp.armeria.common.HttpRequest; +import com.linecorp.armeria.common.RpcRequest; +import com.linecorp.armeria.common.RpcResponse; +import com.linecorp.armeria.common.annotation.Nullable; +import com.linecorp.armeria.common.util.UnmodifiableFuture; +import com.linecorp.armeria.internal.client.ClientUtil; + +/** + * Manages retry attempts for an RPC request, coordinating execution and decision-making + * across multiple attempts until one succeeds or all retry options are exhausted. + * + *

All state mutations must occur on the {@code retryEventLoop} thread. + */ +class RetriedRpcRequest implements RetriedRequest { + private enum State { + /** + * Initial state. {@link RetriedRpcRequest} is instantiated but not yet initialized. + */ + IDLE, + /** + * {@link RetriedRpcRequest} is initialized and is ready to execute/ is executing attempts. + */ + PENDING, + /** + * Terminal state. Either {@link #commit(int)} or {@link #abort(Throwable)} was called. + * {@link RetriedRpcRequest} will not execute any more attempts in that it completes every call + * to {@link #executeAttempt(Client)} with an {@link AbortedAttemptException}. + * {@link #completedFuture} is completed. + */ + COMPLETED + } + + private final ContextAwareEventLoop retryEventLoop; + private final RetryConfig config; + private final ClientRequestContext ctx; + private final RpcRequest req; + private final CompletableFuture completedFuture; + + private final long deadlineTimeNanos; + + private State state; + + int previousAttemptNumber; + + final ArrayList attempts; + + RetriedRpcRequest( + ContextAwareEventLoop retryEventLoop, + RetryConfig config, + ClientRequestContext ctx, + RpcRequest req, + long deadlineTimeNanos + ) { + this.retryEventLoop = retryEventLoop; + this.ctx = ctx; + this.req = req; + this.config = config; + this.deadlineTimeNanos = deadlineTimeNanos; + completedFuture = new CompletableFuture<>(); + + state = State.IDLE; + previousAttemptNumber = 0; + // We can assume that most of the time we will succeed with the first attempt. + attempts = new ArrayList<>(1); + } + + @Override + public CompletionStage executeAttempt( + Client delegate) { + checkState(retryEventLoop.inEventLoop()); + + if (state == State.COMPLETED) { + return UnmodifiableFuture.exceptionallyCompletedFuture(AbortedAttemptException.get()); + } + + if (state == State.IDLE) { + state = State.PENDING; + } + + // We need to be initialized/ the reqDuplicator must be present. + checkState(state == State.PENDING); + checkState(previousAttemptNumber < config.maxTotalAttempts()); + + final int attemptNumber = ++previousAttemptNumber; + + assert attemptNumber == attempts.size() + 1 : attemptNumber + ", " + attempts.size(); + assert attemptNumber <= config.maxTotalAttempts() : attemptNumber + ", " + config.maxTotalAttempts(); + + if (!ClientUtil.checkAndSetResponseTimeout(ctx, deadlineTimeNanos, + config.responseTimeoutMillisForEachAttempt())) { + final ResponseTimeoutException timeoutException = ResponseTimeoutException.get(); + abort(timeoutException); + return UnmodifiableFuture.exceptionallyCompletedFuture(timeoutException); + } + + final RpcRetryAttempt attempt = newAttempt(attemptNumber); + attempts.add(attempt); + return attempt + .executeAndDecide(delegate) + .thenApply(decision -> { + assert retryEventLoop.inEventLoop(); + return new AttemptExecutionResult( + attemptNumber, + decision, + ClientUtil.retryAfterMillis(attempt.ctx().log()) + ); + }) + .exceptionally(cause -> { + assert retryEventLoop.inEventLoop(); + abort(cause); + return null; + }); + } + + private RpcRetryAttempt newAttempt(int attemptNumber) { + final boolean isInitialAttempt = attemptNumber <= 1; + + final ClientRequestContext attemptCtx = ClientUtil.newDerivedContext(ctx, null, req, + isInitialAttempt); + + return new RpcRetryAttempt(config, retryEventLoop, attemptNumber, attemptCtx, req); + } + + @Override + public void commit(int attemptNumber) { + checkState(retryEventLoop.inEventLoop()); + + if (state == State.COMPLETED) { + return; + } + checkState(state == State.PENDING); + + checkArgument(attemptNumber >= 1); + checkArgument(attemptNumber <= attempts.size()); + + final RpcRetryAttempt attemptToCommit = attempts.get(attemptNumber - 1); + + checkState(attemptToCommit != null); + checkState(attemptToCommit.state() == RpcRetryAttempt.State.DECIDED); + + state = State.COMPLETED; + + abortAllExcept(attemptToCommit); + + final RpcResponse attemptRes = attemptToCommit.commit(); + + ctx.logBuilder().endResponseWithChild(attemptToCommit.ctx().log()); + final HttpRequest attemptReq = attemptToCommit.ctx().request(); + if (attemptReq != null) { + ctx.updateRequest(attemptReq); + } + + completedFuture.complete(attemptRes); + } + + @Override + public void abort(Throwable cause) { + checkState(retryEventLoop.inEventLoop()); + + if (state == State.COMPLETED) { + return; + } + + assert !completedFuture.isDone(); + + state = State.COMPLETED; + + abortAllExcept(null); + + if (!ctx.log().isRequestComplete()) { + ctx.logBuilder().endRequest(cause); + } + ctx.logBuilder().endResponse(cause); + + completedFuture.completeExceptionally(cause); + } + + @Override + public void abort(int attemptNumber, Throwable cause) { + checkState(retryEventLoop.inEventLoop()); + checkArgument(attemptNumber >= 1); + + final RpcRetryAttempt attemptToAbort; + if (state == State.COMPLETED) { + return; + } + + checkState(state == State.PENDING); + checkArgument(attemptNumber <= attempts.size()); + + attemptToAbort = attempts.get(attemptNumber - 1); + attemptToAbort.abort(cause); + } + + private void abortAllExcept(@Nullable RpcRetryAttempt winningAttempt) { + assert retryEventLoop.inEventLoop(); + assert state == State.COMPLETED : state; + + for (final RpcRetryAttempt attempt : attempts) { + if (attempt == winningAttempt) { + continue; + } + + attempt.abort(AbortedAttemptException.get()); + } + } + + @Override + public CompletableFuture whenComplete() { + return completedFuture; + } +} diff --git a/core/src/main/java/com/linecorp/armeria/client/retry/RetryCounter.java b/core/src/main/java/com/linecorp/armeria/client/retry/RetryCounter.java new file mode 100644 index 00000000000..f3360267974 --- /dev/null +++ b/core/src/main/java/com/linecorp/armeria/client/retry/RetryCounter.java @@ -0,0 +1,55 @@ +/* + * Copyright 2025 LINE Corporation + * + * Licensed 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 com.linecorp.armeria.common.annotation.Nullable; + +/** + * A counter that keeps track of the number of attempts made so far for a {@link RetryingClient}. + * In particular, it keeps track of the total number of attempts but also + * the number of attempts made with each {@link Backoff}. + * + *

+ * Implementors do not need to be thread-safe. + *

+ */ +interface RetryCounter { + /** + * Records an attempt in that it increases the total number of attempts by one. If {@code backoff} is not + * {@code null}, the number of attempts for that {@code backoff} is increased as well. + * + * @param backoff the backoff used for the attempt, or {@code null} if no backoff was used. + */ + void consumeAttemptFrom(@Nullable Backoff backoff); + + /** + * Returns the number of attempts executed so far with {@code backoff}. + * + * @param backoff the backoff whose number of attempts is requested + * @return the number of attempts executed so far with {@code backoff} + */ + int attemptsSoFarWithBackoff(Backoff backoff); + + /** + * Returns {@code true} if the total number of attempts has reached the maximum number of attempts + * configured in the {@link RetryConfig}. + * + * @return {@code true} if the total number of attempts has reached the maximum number of attempts + * configured in the {@link RetryConfig} + */ + boolean hasReachedMaxAttempts(); +} diff --git a/core/src/main/java/com/linecorp/armeria/client/retry/RetryScheduler.java b/core/src/main/java/com/linecorp/armeria/client/retry/RetryScheduler.java new file mode 100644 index 00000000000..72d46ec1064 --- /dev/null +++ b/core/src/main/java/com/linecorp/armeria/client/retry/RetryScheduler.java @@ -0,0 +1,85 @@ +/* + * Copyright 2025 LY Corporation + * + * LY 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 java.util.concurrent.CompletableFuture; + +/** + * A scheduler that runs retry tasks sequentially on a designated event loop, + * while enforcing a fixed request deadline and minimum backoff between attempts. + * + *

+ * NOTE: All methods of {@link RetryScheduler} must be invoked from the + * {@code retryEventLoop}. Implementations of {@link RetryScheduler} therefore do not need to be thread-safe. + *

+ */ +interface RetryScheduler { + /** + * Tries to schedule the given {@code retryTask} to be run after the given {@code delayMillis} + * or as early as the minimum backoff delays applied for the next schedule allow. + * Implementations are advised to execute the {@code retryTask} directly if + * {@code delayMillis} is {@code 0}. + * + *

+ * The method returns {@code true} if the {@code retryTask} was successfully scheduled or + * {@code false} if not. Note that a return value of {@code false} does not necessarily indicate that + * something exceptional happened during scheduling. For example, {@code delayMillis} might be too long + * for the deadline in which case this method returns {@code false} + * but those events would not {@link #close()} the scheduler exceptionally. + *

+ * + * @param retryTask the task to run to perform a retry + * @param delayMillis the delay in milliseconds after which the {@code retryTask} should be run + * + * @return {@code true} if the {@code retryTask} was successfully scheduled or {@code false} if not. + */ + boolean trySchedule(Runnable retryTask, long delayMillis); + + /** + * Applies the given {@code minimumBackoffMillis} for the next retry scheduling. + * This means that the next retry task will not be scheduled before the given + * {@code minimumBackoffMillis} has elapsed since the call to this method. + * + *

+ * This method must not be called while executing another retry task. + *

+ * + * @param minimumBackoffMillis the minimum backoff in milliseconds to apply for the next retry scheduling + */ + void applyMinimumBackoffMillisForNextRetry(long minimumBackoffMillis); + + /** + * Closes the scheduler which is cancelling the next retry task if any and + * completing the future from {@link #whenClosed()}. + */ + void close(); + + /** + * Returns a future that is completed when the scheduler is closed. It completes + *
    + *
  • normally, when the scheduler is closed after a successful call to {@link #close()}, or
    • + *
  • exceptionally, when it is not possible to run a retry task once it was meant to be scheduled. + *
    • + *
+ * + *

+ * The future is guaranteed to be completed on the {@code retryEventLoop}. + *

+ * + * @return a future that is completed when the scheduler is closed + */ + CompletableFuture whenClosed(); +} diff --git a/core/src/main/java/com/linecorp/armeria/client/retry/RetryingClient.java b/core/src/main/java/com/linecorp/armeria/client/retry/RetryingClient.java index d914bfece68..47c2d8a41ed 100644 --- a/core/src/main/java/com/linecorp/armeria/client/retry/RetryingClient.java +++ b/core/src/main/java/com/linecorp/armeria/client/retry/RetryingClient.java @@ -17,54 +17,25 @@ package com.linecorp.armeria.client.retry; import static com.google.common.base.Preconditions.checkArgument; -import static com.linecorp.armeria.internal.client.ClientUtil.executeWithFallback; -import static com.linecorp.armeria.internal.client.ClientUtil.initContextAndExecuteWithFallback; -import java.time.Duration; -import java.util.Date; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.CompletionStage; import java.util.function.Function; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - +import com.linecorp.armeria.client.Client; import com.linecorp.armeria.client.ClientRequestContext; import com.linecorp.armeria.client.HttpClient; -import com.linecorp.armeria.client.ResponseTimeoutException; -import com.linecorp.armeria.common.AggregatedHttpResponse; -import com.linecorp.armeria.common.AggregationOptions; -import com.linecorp.armeria.common.HttpHeaderNames; +import com.linecorp.armeria.common.ContextAwareEventLoop; import com.linecorp.armeria.common.HttpRequest; -import com.linecorp.armeria.common.HttpRequestDuplicator; import com.linecorp.armeria.common.HttpResponse; -import com.linecorp.armeria.common.HttpResponseDuplicator; import com.linecorp.armeria.common.Request; -import com.linecorp.armeria.common.RequestHeadersBuilder; -import com.linecorp.armeria.common.ResponseHeaders; -import com.linecorp.armeria.common.SplitHttpResponse; import com.linecorp.armeria.common.annotation.Nullable; -import com.linecorp.armeria.common.logging.RequestLog; -import com.linecorp.armeria.common.logging.RequestLogAccess; -import com.linecorp.armeria.common.logging.RequestLogBuilder; -import com.linecorp.armeria.common.logging.RequestLogProperty; -import com.linecorp.armeria.common.stream.AbortedStreamException; -import com.linecorp.armeria.common.util.Exceptions; -import com.linecorp.armeria.internal.client.AggregatedHttpRequestDuplicator; -import com.linecorp.armeria.internal.client.ClientPendingThrowableUtil; -import com.linecorp.armeria.internal.client.ClientRequestContextExtension; -import com.linecorp.armeria.internal.client.TruncatingHttpResponse; - -import io.netty.handler.codec.DateFormatter; +import com.linecorp.armeria.internal.client.ClientUtil; /** * An {@link HttpClient} decorator that handles failures of an invocation and retries HTTP requests. */ public final class RetryingClient extends AbstractRetryingClient implements HttpClient { - - private static final Logger logger = LoggerFactory.getLogger(RetryingClient.class); - /** * Returns a new {@link RetryingClientBuilder} with the specified {@link RetryConfig}. * The {@link RetryConfig} object encapsulates {@link RetryRule} or {@link RetryRuleWithContent}, @@ -241,354 +212,28 @@ public static Function newDecorator(RetryRul } @Override - protected HttpResponse doExecute(ClientRequestContext ctx, HttpRequest req) throws Exception { - final CompletableFuture responseFuture = new CompletableFuture<>(); - final HttpResponse res = HttpResponse.of(responseFuture, ctx.eventLoop()); - if (ctx.exchangeType().isRequestStreaming()) { - final HttpRequestDuplicator reqDuplicator = req.toDuplicator(ctx.eventLoop().withoutContext(), 0); - doExecute0(ctx, reqDuplicator, req, res, responseFuture); - } else { - req.aggregate(AggregationOptions.usePooledObjects(ctx.alloc(), ctx.eventLoop())) - .handle((agg, cause) -> { - if (cause != null) { - handleException(ctx, null, responseFuture, cause, true); - } else { - final HttpRequestDuplicator reqDuplicator = new AggregatedHttpRequestDuplicator(agg); - doExecute0(ctx, reqDuplicator, req, res, responseFuture); - } - return null; - }); - } - return res; - } - - private void doExecute0(ClientRequestContext ctx, HttpRequestDuplicator rootReqDuplicator, - HttpRequest originalReq, HttpResponse returnedRes, - CompletableFuture future) { - final int totalAttempts = getTotalAttempts(ctx); - final boolean initialAttempt = totalAttempts <= 1; - // The request or response has been aborted by the client before it receives a response, - // so stop retrying. - if (originalReq.whenComplete().isCompletedExceptionally()) { - originalReq.whenComplete().handle((unused, cause) -> { - handleException(ctx, rootReqDuplicator, future, cause, initialAttempt); - return null; - }); - return; - } - if (returnedRes.isComplete()) { - returnedRes.whenComplete().handle((result, cause) -> { - final Throwable abortCause; - if (cause != null) { - abortCause = cause; - } else { - abortCause = AbortedStreamException.get(); - } - handleException(ctx, rootReqDuplicator, future, abortCause, initialAttempt); - return null; - }); - return; - } - - if (!setResponseTimeout(ctx)) { - handleException(ctx, rootReqDuplicator, future, ResponseTimeoutException.get(), initialAttempt); - return; - } - - final HttpRequest duplicateReq; - if (initialAttempt) { - duplicateReq = rootReqDuplicator.duplicate(); - } else { - final RequestHeadersBuilder newHeaders = originalReq.headers().toBuilder(); - newHeaders.setInt(ARMERIA_RETRY_COUNT, totalAttempts - 1); - duplicateReq = rootReqDuplicator.duplicate(newHeaders.build()); - } - - final ClientRequestContext derivedCtx; - try { - derivedCtx = newDerivedContext(ctx, duplicateReq, ctx.rpcRequest(), initialAttempt); - } catch (Throwable t) { - handleException(ctx, rootReqDuplicator, future, t, initialAttempt); - return; - } - - final HttpRequest ctxReq = derivedCtx.request(); - assert ctxReq != null; - final HttpResponse response; - final ClientRequestContextExtension ctxExtension = derivedCtx.as(ClientRequestContextExtension.class); - if (!initialAttempt && ctxExtension != null && derivedCtx.endpoint() == null) { - // clear the pending throwable to retry endpoint selection - ClientPendingThrowableUtil.removePendingThrowable(derivedCtx); - // if the endpoint hasn't been selected, try to initialize the ctx with a new endpoint/event loop - response = initContextAndExecuteWithFallback( - unwrap(), ctxExtension, HttpResponse::of, - (context, cause) -> HttpResponse.ofFailure(cause), ctxReq, false); - } else { - response = executeWithFallback(unwrap(), derivedCtx, - (context, cause) -> HttpResponse.ofFailure(cause), ctxReq, false); - } - - final RetryConfig config = mappedRetryConfig(ctx); - if (!ctx.exchangeType().isResponseStreaming() || config.requiresResponseTrailers()) { - response.aggregate().handle((aggregated, cause) -> { - if (cause != null) { - derivedCtx.logBuilder().endRequest(cause); - derivedCtx.logBuilder().endResponse(cause); - handleResponseWithoutContent(config, ctx, rootReqDuplicator, originalReq, returnedRes, - future, derivedCtx, HttpResponse.ofFailure(cause), cause); - } else { - completeLogIfBytesNotTransferred(aggregated, derivedCtx); - derivedCtx.log().whenAvailable(RequestLogProperty.RESPONSE_END_TIME).thenRun(() -> { - handleAggregatedResponse(config, ctx, rootReqDuplicator, originalReq, returnedRes, - future, derivedCtx, aggregated); - }); - } - return null; - }); - } else { - handleStreamingResponse(config, ctx, rootReqDuplicator, originalReq, returnedRes, - future, derivedCtx, response); - } - } - - // TODO(ikhoon): Add a request-scope class such as RetryRequestContext to avoid passing too many parameters. - private void handleResponseWithoutContent(RetryConfig config, ClientRequestContext ctx, - HttpRequestDuplicator rootReqDuplicator, HttpRequest originalReq, - HttpResponse returnedRes, CompletableFuture future, - ClientRequestContext derivedCtx, HttpResponse response, - @Nullable Throwable responseCause) { - if (responseCause != null) { - responseCause = Exceptions.peel(responseCause); - } - try { - final RetryRule retryRule = retryRule(config); - final CompletionStage f = retryRule.shouldRetry(derivedCtx, responseCause); - f.handle((decision, shouldRetryCause) -> { - warnIfExceptionIsRaised(retryRule, shouldRetryCause); - handleRetryDecision(decision, ctx, derivedCtx, rootReqDuplicator, - originalReq, returnedRes, future, response); - return null; - }); - } catch (Throwable cause) { - response.abort(); - handleException(ctx, rootReqDuplicator, future, cause, false); - } - } - - private void handleStreamingResponse(RetryConfig retryConfig, ClientRequestContext ctx, - HttpRequestDuplicator rootReqDuplicator, - HttpRequest originalReq, HttpResponse returnedRes, - CompletableFuture future, - ClientRequestContext derivedCtx, - HttpResponse response) { - final SplitHttpResponse splitResponse = response.split(); - splitResponse.headers().handle((headers, headersCause) -> { - final Throwable responseCause; - if (headersCause == null) { - final RequestLog log = derivedCtx.log().getIfAvailable(RequestLogProperty.RESPONSE_CAUSE); - responseCause = log != null ? log.responseCause() : null; - } else { - responseCause = Exceptions.peel(headersCause); - } - completeLogIfBytesNotTransferred(response, headers, derivedCtx, responseCause); - - derivedCtx.log().whenAvailable(RequestLogProperty.RESPONSE_HEADERS).thenRun(() -> { - if (retryConfig.needsContentInRule() && responseCause == null) { - final HttpResponse response0 = splitResponse.unsplit(); - final HttpResponseDuplicator duplicator = - response0.toDuplicator(derivedCtx.eventLoop().withoutContext(), - derivedCtx.maxResponseLength()); - try { - final TruncatingHttpResponse truncatingHttpResponse = - new TruncatingHttpResponse(duplicator.duplicate(), - retryConfig.maxContentLength()); - final HttpResponse duplicated = duplicator.duplicate(); - duplicator.close(); - - final RetryRuleWithContent ruleWithContent = - retryConfig.retryRuleWithContent(); - assert ruleWithContent != null; - ruleWithContent.shouldRetry(derivedCtx, truncatingHttpResponse, null) - .handle((decision, cause) -> { - warnIfExceptionIsRaised(ruleWithContent, cause); - truncatingHttpResponse.abort(); - handleRetryDecision(decision, ctx, derivedCtx, rootReqDuplicator, - originalReq, returnedRes, future, duplicated); - return null; - }); - } catch (Throwable cause) { - duplicator.abort(cause); - handleException(ctx, rootReqDuplicator, future, cause, false); - } - } else { - final HttpResponse response0; - if (responseCause != null) { - splitResponse.body().abort(responseCause); - response0 = HttpResponse.ofFailure(responseCause); - } else { - response0 = splitResponse.unsplit(); - } - handleResponseWithoutContent(retryConfig, ctx, rootReqDuplicator, originalReq, returnedRes, - future, derivedCtx, response0, responseCause); - } - }); - return null; - }); - } + RetryContext newRetryContext( + Client delegate, + ClientRequestContext ctx, + HttpRequest req, + RetryConfig config) { + final CompletableFuture resFuture = new CompletableFuture<>(); + final HttpResponse res = HttpResponse.of(resFuture, ctx.eventLoop()); - private void handleAggregatedResponse(RetryConfig retryConfig, ClientRequestContext ctx, - HttpRequestDuplicator rootReqDuplicator, - HttpRequest originalReq, HttpResponse returnedRes, - CompletableFuture future, - ClientRequestContext derivedCtx, - AggregatedHttpResponse aggregatedRes) { - if (retryConfig.needsContentInRule()) { - final RetryRuleWithContent ruleWithContent = retryConfig.retryRuleWithContent(); - assert ruleWithContent != null; - try { - ruleWithContent.shouldRetry(derivedCtx, aggregatedRes.toHttpResponse(), null) - .handle((decision, cause) -> { - warnIfExceptionIsRaised(ruleWithContent, cause); - handleRetryDecision( - decision, ctx, derivedCtx, rootReqDuplicator, originalReq, - returnedRes, future, aggregatedRes.toHttpResponse()); - return null; - }); - } catch (Throwable cause) { - handleException(ctx, rootReqDuplicator, future, cause, false); - } - return; - } - handleResponseWithoutContent(retryConfig, ctx, rootReqDuplicator, originalReq, returnedRes, - future, derivedCtx, aggregatedRes.toHttpResponse(), null); - } + final ContextAwareEventLoop retryEventLoop = ctx.eventLoop(); - private static void completeLogIfBytesNotTransferred(AggregatedHttpResponse response, - ClientRequestContext ctx) { - if (!ctx.log().isAvailable(RequestLogProperty.REQUEST_FIRST_BYTES_TRANSFERRED_TIME)) { - final RequestLogBuilder logBuilder = ctx.logBuilder(); - logBuilder.endRequest(); - logBuilder.responseHeaders(response.headers()); - if (!response.trailers().isEmpty()) { - logBuilder.responseTrailers(response.trailers()); - } - logBuilder.endResponse(); - } - } + final long deadlineTimeNanos = ClientUtil.deadlineTimeNanos(ctx); + final RetriedHttpRequest retriedReq = new RetriedHttpRequest( + retryEventLoop, config, ctx, req, deadlineTimeNanos, useRetryAfter + ); - private static void completeLogIfBytesNotTransferred( - HttpResponse response, @Nullable ResponseHeaders headers, ClientRequestContext ctx, - @Nullable Throwable responseCause) { - if (!ctx.log().isAvailable(RequestLogProperty.REQUEST_FIRST_BYTES_TRANSFERRED_TIME)) { - final RequestLogBuilder logBuilder = ctx.logBuilder(); - if (responseCause != null) { - logBuilder.endRequest(responseCause); - logBuilder.endResponse(responseCause); - } else { - logBuilder.endRequest(); - if (headers != null) { - logBuilder.responseHeaders(headers); - } - response.whenComplete().handle((unused, cause) -> { - if (cause != null) { - logBuilder.endResponse(cause); - } else { - logBuilder.endResponse(); - } - return null; - }); - } - } - } + final RetryScheduler scheduler = new DefaultRetryScheduler( + retryEventLoop, + deadlineTimeNanos + ); - private static void warnIfExceptionIsRaised(Object retryRule, @Nullable Throwable cause) { - if (cause != null) { - logger.warn("Unexpected exception is raised from {}.", retryRule, cause); - } - } - - private static void handleException(ClientRequestContext ctx, - @Nullable HttpRequestDuplicator rootReqDuplicator, - CompletableFuture future, Throwable cause, - boolean endRequestLog) { - future.completeExceptionally(cause); - if (rootReqDuplicator != null) { - rootReqDuplicator.abort(cause); - } - if (endRequestLog) { - ctx.logBuilder().endRequest(cause); - } - ctx.logBuilder().endResponse(cause); - } - - private void handleRetryDecision(@Nullable RetryDecision decision, ClientRequestContext ctx, - ClientRequestContext derivedCtx, HttpRequestDuplicator rootReqDuplicator, - HttpRequest originalReq, HttpResponse returnedRes, - CompletableFuture future, HttpResponse originalRes) { - final Backoff backoff = decision != null ? decision.backoff() : null; - if (backoff != null) { - final long millisAfter = useRetryAfter ? getRetryAfterMillis(derivedCtx) : -1; - final long nextDelay = getNextDelay(ctx, backoff, millisAfter); - if (nextDelay >= 0) { - abortResponse(originalRes, derivedCtx); - scheduleNextRetry( - ctx, cause -> handleException(ctx, rootReqDuplicator, future, cause, false), - () -> doExecute0(ctx, rootReqDuplicator, originalReq, returnedRes, future), - nextDelay); - return; - } - } - onRetryingComplete(ctx); - future.complete(originalRes); - rootReqDuplicator.close(); - } - - private static void abortResponse(HttpResponse originalRes, ClientRequestContext derivedCtx) { - // Set response content with null to make sure that the log is complete. - final RequestLogBuilder logBuilder = derivedCtx.logBuilder(); - logBuilder.responseContent(null, null); - logBuilder.responseContentPreview(null); - originalRes.abort(); - } - - private static long getRetryAfterMillis(ClientRequestContext ctx) { - final RequestLogAccess log = ctx.log(); - final String value; - final RequestLog requestLog = log.getIfAvailable(RequestLogProperty.RESPONSE_HEADERS); - value = requestLog != null ? requestLog.responseHeaders().get(HttpHeaderNames.RETRY_AFTER) : null; - - if (value != null) { - try { - return Duration.ofSeconds(Integer.parseInt(value)).toMillis(); - } catch (Exception ignored) { - // Not a second value. - } - - try { - @SuppressWarnings("UseOfObsoleteDateTimeApi") - final Date date = DateFormatter.parseHttpDate(value); - if (date != null) { - return date.getTime() - System.currentTimeMillis(); - } - } catch (Exception ignored) { - // `parseHttpDate()` can raise an exception rather than returning `null` - // when the given value has more than 64 characters. - } - - logger.debug("The retryAfter: {}, from the server is neither an HTTP date nor a second.", - value); - } - - return -1; - } + final RetryCounter counter = new DefaultRetryCounter(config.maxTotalAttempts()); - private static RetryRule retryRule(RetryConfig retryConfig) { - if (retryConfig.needsContentInRule()) { - return retryConfig.fromRetryRuleWithContent(); - } else { - final RetryRule rule = retryConfig.retryRule(); - assert rule != null; - return rule; - } + return new RetryContext(retryEventLoop, retriedReq, scheduler, counter, delegate, res, resFuture); } } diff --git a/core/src/main/java/com/linecorp/armeria/client/retry/RetryingRpcClient.java b/core/src/main/java/com/linecorp/armeria/client/retry/RetryingRpcClient.java index 32146db4c68..9e40a826b1e 100644 --- a/core/src/main/java/com/linecorp/armeria/client/retry/RetryingRpcClient.java +++ b/core/src/main/java/com/linecorp/armeria/client/retry/RetryingRpcClient.java @@ -15,29 +15,23 @@ */ package com.linecorp.armeria.client.retry; -import static com.linecorp.armeria.internal.client.ClientUtil.executeWithFallback; -import static com.linecorp.armeria.internal.client.ClientUtil.initContextAndExecuteWithFallback; - -import java.util.concurrent.CancellationException; import java.util.concurrent.CompletableFuture; import java.util.function.Function; +import com.linecorp.armeria.client.Client; import com.linecorp.armeria.client.ClientRequestContext; -import com.linecorp.armeria.client.ResponseTimeoutException; import com.linecorp.armeria.client.RpcClient; -import com.linecorp.armeria.client.endpoint.EndpointGroup; -import com.linecorp.armeria.common.HttpRequest; +import com.linecorp.armeria.common.ContextAwareEventLoop; import com.linecorp.armeria.common.Request; import com.linecorp.armeria.common.RpcRequest; import com.linecorp.armeria.common.RpcResponse; -import com.linecorp.armeria.internal.client.ClientPendingThrowableUtil; -import com.linecorp.armeria.internal.client.ClientRequestContextExtension; -import com.linecorp.armeria.internal.common.util.StringUtil; +import com.linecorp.armeria.internal.client.ClientUtil; /** * An {@link RpcClient} decorator that handles failures of an invocation and retries RPC requests. */ -public final class RetryingRpcClient extends AbstractRetryingClient +public final class RetryingRpcClient + extends AbstractRetryingClient implements RpcClient { /** @@ -141,99 +135,29 @@ public static RetryingRpcClientBuilder builder(RetryConfigMapping m } @Override - protected RpcResponse doExecute(ClientRequestContext ctx, RpcRequest req) throws Exception { - final CompletableFuture future = new CompletableFuture<>(); - final RpcResponse res = RpcResponse.from(future); - doExecute0(ctx, req, res, future); - return res; - } + RetryContext newRetryContext( + Client delegate, + ClientRequestContext ctx, + RpcRequest req, + RetryConfig config) { + final CompletableFuture resFuture = new CompletableFuture<>(); + final RpcResponse res = RpcResponse.from(resFuture); - private void doExecute0(ClientRequestContext ctx, RpcRequest req, - RpcResponse returnedRes, CompletableFuture future) { - final int totalAttempts = getTotalAttempts(ctx); - final boolean initialAttempt = totalAttempts <= 1; - if (returnedRes.isDone()) { - // The response has been cancelled by the client before it receives a response, so stop retrying. - handleException(ctx, future, new CancellationException( - "the response returned to the client has been cancelled"), initialAttempt); - return; - } - if (!setResponseTimeout(ctx)) { - handleException(ctx, future, ResponseTimeoutException.get(), initialAttempt); - return; - } - - final ClientRequestContext derivedCtx = newDerivedContext(ctx, null, req, initialAttempt); - - if (!initialAttempt) { - derivedCtx.mutateAdditionalRequestHeaders( - mutator -> mutator.add(ARMERIA_RETRY_COUNT, StringUtil.toString(totalAttempts - 1))); - } - - final RpcResponse res; - - final ClientRequestContextExtension ctxExtension = derivedCtx.as(ClientRequestContextExtension.class); - final EndpointGroup endpointGroup = derivedCtx.endpointGroup(); - if (!initialAttempt && ctxExtension != null && - endpointGroup != null && derivedCtx.endpoint() == null) { - // clear the pending throwable to retry endpoint selection - ClientPendingThrowableUtil.removePendingThrowable(derivedCtx); - // if the endpoint hasn't been selected, try to initialize the ctx with a new endpoint/event loop - res = initContextAndExecuteWithFallback(unwrap(), ctxExtension, RpcResponse::from, - (context, cause) -> RpcResponse.ofFailure(cause), - req, true); - } else { - res = executeWithFallback(unwrap(), derivedCtx, - (context, cause) -> RpcResponse.ofFailure(cause), - req, true); - } - - final RetryConfig retryConfig = mappedRetryConfig(ctx); - final RetryRuleWithContent retryRule = - retryConfig.needsContentInRule() ? - retryConfig.retryRuleWithContent() : retryConfig.fromRetryRule(); - res.handle((unused1, cause) -> { - try { - assert retryRule != null; - retryRule.shouldRetry(derivedCtx, res, cause).handle((decision, unused3) -> { - final Backoff backoff = decision != null ? decision.backoff() : null; - if (backoff != null) { - final long nextDelay = getNextDelay(derivedCtx, backoff); - if (nextDelay < 0) { - onRetryComplete(ctx, derivedCtx, res, future); - return null; - } - - scheduleNextRetry(ctx, cause0 -> handleException(ctx, future, cause0, false), - () -> doExecute0(ctx, req, returnedRes, future), nextDelay); - } else { - onRetryComplete(ctx, derivedCtx, res, future); - } - return null; - }); - } catch (Throwable t) { - handleException(ctx, future, t, false); - } - return null; - }); - } + final ContextAwareEventLoop retryEventLoop = ctx.eventLoop(); - private static void onRetryComplete(ClientRequestContext ctx, ClientRequestContext derivedCtx, - RpcResponse res, CompletableFuture future) { - onRetryingComplete(ctx); - final HttpRequest actualHttpReq = derivedCtx.request(); - if (actualHttpReq != null) { - ctx.updateRequest(actualHttpReq); - } - future.complete(res); - } + final long deadlineTimeNanos = ClientUtil.deadlineTimeNanos(ctx); + + final RetriedRpcRequest retriedReq = new RetriedRpcRequest( + retryEventLoop, config, ctx, req, deadlineTimeNanos + ); + + final RetryScheduler scheduler = new DefaultRetryScheduler( + retryEventLoop, + deadlineTimeNanos + ); + + final RetryCounter counter = new DefaultRetryCounter(config.maxTotalAttempts()); - private static void handleException(ClientRequestContext ctx, CompletableFuture future, - Throwable cause, boolean endRequestLog) { - future.completeExceptionally(cause); - if (endRequestLog) { - ctx.logBuilder().endRequest(cause); - } - ctx.logBuilder().endResponse(cause); + return new RetryContext(retryEventLoop, retriedReq, scheduler, counter, delegate, res, resFuture); } } diff --git a/core/src/main/java/com/linecorp/armeria/client/retry/RpcRetryAttempt.java b/core/src/main/java/com/linecorp/armeria/client/retry/RpcRetryAttempt.java new file mode 100644 index 00000000000..0d0d617e753 --- /dev/null +++ b/core/src/main/java/com/linecorp/armeria/client/retry/RpcRetryAttempt.java @@ -0,0 +1,407 @@ +/* + * Copyright 2025 LY Corporation + * + * LY 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 com.google.common.base.Preconditions.checkState; +import static com.linecorp.armeria.client.retry.AbstractRetryingClient.ARMERIA_RETRY_COUNT; +import static com.linecorp.armeria.internal.client.ClientUtil.executeWithFallback; +import static com.linecorp.armeria.internal.client.ClientUtil.initContextAndExecuteWithFallback; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; +import java.util.function.Function; + +import com.google.common.base.MoreObjects; + +import com.linecorp.armeria.client.Client; +import com.linecorp.armeria.client.ClientRequestContext; +import com.linecorp.armeria.common.ContextAwareEventLoop; +import com.linecorp.armeria.common.RpcRequest; +import com.linecorp.armeria.common.RpcResponse; +import com.linecorp.armeria.common.annotation.Nullable; +import com.linecorp.armeria.common.util.Exceptions; +import com.linecorp.armeria.common.util.UnmodifiableFuture; +import com.linecorp.armeria.internal.client.ClientPendingThrowableUtil; +import com.linecorp.armeria.internal.client.ClientRequestContextExtension; + +/** + * A single attempt for a {@link RetriedRpcRequest}. + * + *

+ * NOTE: All methods of {@link RpcRetryAttempt} must be invoked from the + * {@code retryEventLoop}. + *

+ */ +final class RpcRetryAttempt { + /** + * The state of this attempt. + * + *

This state machine controls the retry attempt lifecycle. All state transitions + * must occur on the {@code retryEventLoop} thread to ensure consistency. + * + *

The following state machine diagram illustrates the possible transitions: + * 

+     *          Start
+     *            |
+     *           IDLE
+     *            |
+     *            v
+     *        EXECUTING----------+
+     *            |              |
+     *            v              |
+     *         DECIDED-----------+
+     *          /   \            |
+     *         /     \           |
+     *        /       \          |
+     *       v         v         |
+     *   COMMITTED  ABORTED<-----+
+     *       \         /
+     *        +---+---+
+     *            |
+     *            v
+     *           End
+     *
+ */ + enum State { + /** + * The attempt has not been executed yet. + * Call {@link #executeAndDecide(Client)} to transition to {@link #EXECUTING}. + */ + IDLE, + + /** + * The attempt is executing and waiting for a response or exception. + * The {@link #res} field becomes available in this state. + * Can transition to {@link #DECIDED} on completion or {@link #ABORTED} if aborted. + */ + EXECUTING, + + /** + * The {@link RetryRule} has made a decision about whether to retry. + * The attempt is ready to be either committed via {@link #commit()} + * or aborted via {@link #abort(Throwable)}. + */ + DECIDED, + + /** + * The attempt was successfully committed via {@link #commit()}. + * This is a terminal state and the response is final. + */ + COMMITTED, + + /** + * The attempt was aborted via {@link #abort(Throwable)}. + * The {@link #cause} field is non-null only in this state. + * This is a terminal state. + */ + ABORTED + } + + private final RetryConfig config; + private final ContextAwareEventLoop retryEventLoop; + private final int attemptNumber; + private final ClientRequestContext ctx; + private final RpcRequest req; + + private State state; + + /** + * The response of the attempt. It is available in {@link State#EXECUTING}, + * {@link State#DECIDED}, and {@link State#COMMITTED} states. + */ + @Nullable + private RpcResponse res; + + /** + * The cause of the attempt failure. It is not-{@code null} iff we are in {@link State#ABORTED} state. + */ + @Nullable + Throwable cause; + + RpcRetryAttempt( + RetryConfig config, + ContextAwareEventLoop retryEventLoop, + int attemptNumber, + ClientRequestContext ctx, + RpcRequest req + ) { + this.config = config; + this.retryEventLoop = retryEventLoop; + this.attemptNumber = attemptNumber; + this.ctx = ctx; + this.req = req; + + state = State.IDLE; + } + + /** + * Executes the attempt, prepares and gives the response or the cause to the {@link RetryRule}. + * This method must be called at most once and except for calls to {@link #abort(Throwable)} + * it must be called before any other method, in particular before {@link #commit()}. + * + * @param delegate the next {@link Client} in the decoration chain + * @return a future that will be completed with the {@link RetryDecision} or an exception if failed during + * execution, preparation, or decision. + * In particular, it fails with a {@link AbortedAttemptException} + * if the attempt was aborted by {@link #abort(Throwable)}. + * + * @see #commit() + */ + CompletionStage executeAndDecide(Client delegate) { + checkState(retryEventLoop.inEventLoop()); + + if (state == State.ABORTED) { + assert cause != null; + return UnmodifiableFuture.exceptionallyCompletedFuture(cause); + } + + checkState(state == State.IDLE); + state = State.EXECUTING; + + res = execute(delegate); + + return res + .handle((unused, causeToDecide) -> { + assert retryEventLoop.inEventLoop(); + // Let us avoid calling decide (which may be expensive if we already know that we were + // aborted. + if (state == State.ABORTED) { + assert cause != null; + return UnmodifiableFuture.exceptionallyCompletedFuture(cause); + } + + assert state == State.EXECUTING : state; + + return decide(res, causeToDecide) + .thenCompose( + decision -> { + if (state == State.ABORTED) { + assert cause != null; + return UnmodifiableFuture.exceptionallyCompletedFuture( + cause); + } else { + assert state == State.EXECUTING : state; + assert cause == null : cause; // sanity check + assert res != null; + state = State.DECIDED; + return UnmodifiableFuture.completedFuture(decision); + } + } + ); + }) + .thenCompose(Function.identity()) + .handle( + (decision, preparationOrDecisionCause) -> { + if (preparationOrDecisionCause == null) { + return (CompletableFuture) UnmodifiableFuture.completedFuture( + decision); + } + + if (state == State.ABORTED) { + assert cause != null; + if (cause != preparationOrDecisionCause) { + cause.addSuppressed(preparationOrDecisionCause); + } + return UnmodifiableFuture.exceptionallyCompletedFuture( + preparationOrDecisionCause); + } + assert state == State.EXECUTING : state; + + abort(preparationOrDecisionCause); + return UnmodifiableFuture.exceptionallyCompletedFuture( + preparationOrDecisionCause); + } + ) + .thenCompose(Function.identity()); + } + + private RpcResponse execute(Client delegate) { + final boolean isInitialAttempt = attemptNumber <= 1; + + if (!isInitialAttempt) { + ctx.mutateAdditionalRequestHeaders( + mutator -> mutator.add(ARMERIA_RETRY_COUNT, Integer.toString(attemptNumber - 1))); + } + + final ClientRequestContextExtension ctxExt = + ctx.as(ClientRequestContextExtension.class); + if (!isInitialAttempt && ctxExt != null && ctx.endpointGroup() != null && ctx.endpoint() == null) { + // clear the pending throwable to retry endpoint selection + ClientPendingThrowableUtil.removePendingThrowable(ctx); + // if the endpoint hasn't been selected, + // try to initialize the attempCtx with a new endpoint/event loop + return initContextAndExecuteWithFallback( + delegate, ctxExt, RpcResponse::from, + (context, cause) -> + RpcResponse.ofFailure(cause), req, true); + } else { + return executeWithFallback(delegate, ctx, + (context, cause) -> + RpcResponse.ofFailure(cause), req, true); + } + } + + /** + * Marks this attempt as committed and returns the response of the attempt. + * The attempt must be decided, i.e. {@link #executeAndDecide(Client)} must have been called and completed + * successfully before calling this method. After this call, the attempt cannot be aborted anymore. + * This method is idempotent once the first call returns successfully. + * + * @return the response of the attempt + * + * @see #executeAndDecide(Client) + */ + public RpcResponse commit() { + checkState(retryEventLoop.inEventLoop()); + + if (state == State.COMMITTED) { + assert res != null; + return res; + } + + checkState(state == State.DECIDED); + assert res != null; + assert cause == null : cause; + state = State.COMMITTED; + return res; + } + + /** + * Aborts this attempt with the specified {@code cause}. + * {@link #executeAndDecide(Client)} must be called before this method is called. + * After this call, the attempt cannot be committed anymore. + * This method is idempotent once the first call returns successfully. + * + * @param cause the cause of the abortion + */ + public void abort(Throwable cause) { + checkState(retryEventLoop.inEventLoop()); + + if (state == State.ABORTED) { + return; + } + + assert state == State.IDLE || + state == State.EXECUTING || + state == State.DECIDED : state; + assert this.cause == null : this.cause; + state = State.ABORTED; + this.cause = cause; + ctx.cancel(); + } + + /** + * Returns the {@link ClientRequestContext} of this attempt. + * + * @return the {@link ClientRequestContext} of this attempt + */ + ClientRequestContext ctx() { + return ctx; + } + + /** + * Returns the {@link State} of this attempt. + * + * @return the {@link State} of this attempt + */ + State state() { + return state; + } + + /** + * Calls {@link RetryRule#shouldRetry(ClientRequestContext, Throwable)} to receive a {@link RetryDecision}. + * + * @param resForRule the response to be delivered to the {@link RetryRule} for decision. + * @param causeForRule the cause to be delivered to the {@link RetryRule} for decision. + * @return a future that will be completed with the {@link RetryDecision} or an exception if failed during + * the decision. The future will be completed on the retry event loop. + */ + private CompletableFuture decide(@Nullable RpcResponse resForRule, + @Nullable Throwable causeForRule) { + if (causeForRule != null) { + resForRule = null; + causeForRule = Exceptions.peel(causeForRule); + } else { + assert resForRule != null; + causeForRule = null; + } + + final RetryRuleWithContent retryRule = + config.needsContentInRule() ? config.retryRuleWithContent() + : config.fromRetryRule(); + assert retryRule != null; + + try { + final CompletionStage<@Nullable RetryDecision> maybeDecision = retryRule + .shouldRetry(ctx, + resForRule, + causeForRule); + + // Remember that RetryRule.shouldRetry could be client code so we do have any guarantees + // on which thread we are completing. Let us make sure we are completing on running on the + // retry event loop again. + return withCompletionOnRetryEventLoop( + maybeDecision.thenApply( + decision -> decision == null ? RetryDecision.noRetry() : decision + ) + ); + } catch (Throwable t) { + return UnmodifiableFuture.exceptionallyCompletedFuture(t); + } + } + + /** + * Ensures the given {@link CompletionStage} completes on the retry event loop. + * If already on the retry event loop, completes directly; otherwise schedules completion. + */ + private CompletableFuture withCompletionOnRetryEventLoop(CompletionStage future) { + final CompletableFuture futureOnTheRetryEventLoop = new CompletableFuture<>(); + future.whenComplete((result, cause) -> { + if (retryEventLoop.inEventLoop()) { + if (cause != null) { + futureOnTheRetryEventLoop.completeExceptionally(cause); + } else { + futureOnTheRetryEventLoop.complete(result); + } + } else { + retryEventLoop.execute(() -> { + if (cause != null) { + futureOnTheRetryEventLoop.completeExceptionally(cause); + } else { + futureOnTheRetryEventLoop.complete(result); + } + }); + } + }); + return futureOnTheRetryEventLoop; + } + + @Override + public String toString() { + checkState(retryEventLoop.inEventLoop()); + + return MoreObjects + .toStringHelper(this) + .add("state", state) + .add("attemptNumber", attemptNumber) + .add("ctx", ctx) + .add("req", req) + .add("res", res) + .add("cause", cause) + .toString(); + } +} diff --git a/core/src/main/java/com/linecorp/armeria/internal/client/ClientUtil.java b/core/src/main/java/com/linecorp/armeria/internal/client/ClientUtil.java index a4f65d598fc..e88e955329e 100644 --- a/core/src/main/java/com/linecorp/armeria/internal/client/ClientUtil.java +++ b/core/src/main/java/com/linecorp/armeria/internal/client/ClientUtil.java @@ -18,10 +18,18 @@ import static java.util.Objects.requireNonNull; import java.net.URI; +import java.time.Duration; +import java.util.Date; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; import java.util.function.BiFunction; import java.util.function.Function; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import com.google.common.math.LongMath; + import com.linecorp.armeria.client.Client; import com.linecorp.armeria.client.ClientRequestContext; import com.linecorp.armeria.client.Endpoint; @@ -30,6 +38,7 @@ import com.linecorp.armeria.client.UnprocessedRequestException; import com.linecorp.armeria.client.WebClient; import com.linecorp.armeria.client.endpoint.EndpointGroup; +import com.linecorp.armeria.common.HttpHeaderNames; import com.linecorp.armeria.common.HttpRequest; import com.linecorp.armeria.common.Request; import com.linecorp.armeria.common.RequestId; @@ -44,8 +53,12 @@ import com.linecorp.armeria.common.stream.StreamMessage; import com.linecorp.armeria.common.util.Exceptions; import com.linecorp.armeria.common.util.SafeCloseable; +import com.linecorp.armeria.common.util.TimeoutMode; + +import io.netty.handler.codec.DateFormatter; public final class ClientUtil { + private static final Logger logger = LoggerFactory.getLogger(ClientUtil.class); /** * An undefined {@link URI} to create {@link WebClient} without specifying {@link URI}. @@ -281,5 +294,111 @@ public static ClientRequestContext newDerivedContext(ClientRequestContext ctx, return derived; } + /** + * Tries to parse the {@code "Retry-After"} response header from the given {@link RequestLogAccess}. + * If the header is present and valid, returns the number of milliseconds after which a retry + * should be attempted. Otherwise, returns {@code -1}. + * + * @param log the {@link RequestLogAccess} which may contain the {@code "Retry-After"} header + * @return the number of milliseconds after which a retry should be attempted, or {@code -1} if + * the header is not present or invalid + */ + public static long retryAfterMillis(RequestLogAccess log) { + final String retryAfterValue; + final RequestLog requestLog = log.getIfAvailable(RequestLogProperty.RESPONSE_HEADERS); + retryAfterValue = requestLog != null ? + requestLog.responseHeaders().get(HttpHeaderNames.RETRY_AFTER) : null; + + if (retryAfterValue != null) { + try { + return Duration.ofSeconds(Integer.parseInt(retryAfterValue)).toMillis(); + } catch (Exception ignored) { + // Not a second value. + } + + try { + @SuppressWarnings("UseOfObsoleteDateTimeApi") + final Date retryAfterDate = DateFormatter.parseHttpDate(retryAfterValue); + if (retryAfterDate != null) { + return retryAfterDate.getTime() - System.currentTimeMillis(); + } + } catch (Exception ignored) { + // `parseHttpDate()` can raise an exception rather than returning `null` + // when the given value has more than 64 characters. + } + + logger.debug("The retryAfter: {}, from the server is neither an HTTP date nor a second.", + retryAfterValue); + } + + return -1; + } + + /** + * Returns the deadline in nanoseconds for the given {@link ClientRequestContext} based on its + * response timeout. If no response timeout is set, {@link Long#MAX_VALUE} will be returned. + * + * @param ctx the {@link ClientRequestContext} to get the response timeout from + * @return the deadline in nanoseconds, or {@link Long#MAX_VALUE} if no response timeout is set + */ + public static long deadlineTimeNanos(ClientRequestContext ctx) { + final long responseTimeoutMillis = ctx.responseTimeoutMillis(); + if (responseTimeoutMillis <= 0) { + return Long.MAX_VALUE; + } else { + return LongMath.saturatedAdd(System.nanoTime(), + TimeUnit.MILLISECONDS.toNanos(responseTimeoutMillis)); + } + } + + /** + * Sets the response timeout on a {@link ClientRequestContext} based on the minimum of a deadline + * and response timeout, or clears the timeout if neither is specified. + * + *

This method respects both absolute deadlines (in nanoseconds) and relative response timeouts + * (in milliseconds), using whichever is more restrictive. If the calculated timeout has already + * expired, the method returns {@code false} to indicate the request should be considered timed out. + * + * @param ctx the {@link ClientRequestContext} to configure + * @param deadlineTimeNanos absolute deadline in nanoseconds from {@link System#nanoTime()}, + * or {@link Long#MAX_VALUE} if no deadline is set + * @param responseTimeoutMillis relative timeout in milliseconds, or {@code 0} if no timeout is set + * @return {@code true} if the timeout was successfully set/cleared and the request is still valid, + * {@code false} if the request has already exceeded the calculated timeout + * + * @see ClientRequestContext#setResponseTimeoutMillis(TimeoutMode, long) + * @see ClientRequestContext#clearResponseTimeout() + */ + public static boolean checkAndSetResponseTimeout(ClientRequestContext ctx, long deadlineTimeNanos, + long responseTimeoutMillis) { + long remainingTimeUntilDeadlineMillis = Long.MAX_VALUE; + boolean hasTimeout = false; + + if (deadlineTimeNanos < Long.MAX_VALUE) { + hasTimeout = true; + remainingTimeUntilDeadlineMillis = TimeUnit.NANOSECONDS.toMillis( + LongMath.saturatedSubtract(deadlineTimeNanos, System.nanoTime()) + ); + } + + if (responseTimeoutMillis != 0) { + hasTimeout = true; + remainingTimeUntilDeadlineMillis = + Math.min(remainingTimeUntilDeadlineMillis, responseTimeoutMillis); + } + + if (hasTimeout) { + if (remainingTimeUntilDeadlineMillis > 0) { + ctx.setResponseTimeoutMillis(TimeoutMode.SET_FROM_NOW, remainingTimeUntilDeadlineMillis); + return true; + } else { + return false; + } + } else { + ctx.clearResponseTimeout(); + return true; + } + } + private ClientUtil() {} } diff --git a/core/src/test/java/com/linecorp/armeria/client/retry/RetrySchedulerTest.java b/core/src/test/java/com/linecorp/armeria/client/retry/RetrySchedulerTest.java new file mode 100644 index 00000000000..ddce2c3284f --- /dev/null +++ b/core/src/test/java/com/linecorp/armeria/client/retry/RetrySchedulerTest.java @@ -0,0 +1,874 @@ +/* + * Copyright 2025 LY Corporation + * + * LY 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 static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.assertj.core.api.Assertions.fail; +import static org.awaitility.Awaitility.await; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import java.util.stream.LongStream; +import java.util.stream.Stream; + +import org.jetbrains.annotations.NotNull; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +import com.google.common.collect.ImmutableList; + +import com.linecorp.armeria.client.ClientFactory; +import com.linecorp.armeria.client.ResponseTimeoutException; +import com.linecorp.armeria.common.annotation.Nullable; +import com.linecorp.armeria.internal.testing.AnticipatedException; +import com.linecorp.armeria.internal.testing.BlockingUtils; + +import io.netty.channel.DefaultEventLoop; +import io.netty.util.concurrent.ScheduledFuture; + +class RetrySchedulerTest { + private static final long RETRY_SCHEDULER_SCHEDULE_ADJUSTMENT_TOLERANCE_MILLIS = 5L; + // Number of millis we expect the retry event loop will take to actually execute a delayed task + // *after* its delay. + private static final long MAX_EXECUTION_DELAY_MILLIS = 100L; + + private final List executedRetryTasks = Collections.synchronizedList(new ArrayList<>()); + private ManagedRetryEventLoop retryEventLoop; + private DefaultRetryScheduler scheduler; + + @BeforeEach + void setUp() { + retryEventLoop = new ManagedRetryEventLoop(); + scheduler = new DefaultRetryScheduler( + retryEventLoop, + System.nanoTime() + TimeUnit.SECONDS.toNanos(10) + ); + } + + @AfterEach + void tearDown() throws ExecutionException, InterruptedException { + // If the test terminates the event loop it needs to a local copy. + runOnRetryEventLoop(() -> { + if (!scheduler.whenClosed().isDone()) { + scheduler.close(); + assertThat(scheduler.whenClosed()).isCompleted(); + } + }); + + retryEventLoop.shutdownGracefully(0, 3, TimeUnit.SECONDS).get(); + assertNoExceptionsOnRetryEventLoop(retryEventLoop); + + ManagedRetryTask.nextRetryTaskNumber.set(0); + executedRetryTasks.clear(); + } + + @Test + void schedule_noDelay_executeImmediately() { + runOnRetryEventLoop(() -> { + assertThat(scheduler.trySchedule(nextRetryTask(), 0)).isTrue(); + assertThat(scheduler.whenClosed()).isNotDone(); + + assertNumRetryTaskExecutions(1); + assertNoScheduleCalls(); + }); + } + + @Test + void schedule_withDelay_executeAfterDelay() { + runOnRetryEventLoop(() -> { + assertThat(scheduler.trySchedule(nextRetryTask(), 200L)).isTrue(); + assertThat(scheduler.whenClosed()).isNotDone(); + }); + + await().untilAsserted( + () -> { + assertNumRetryTaskExecutions(1); + assertEventLoopScheduleCalls(200L); + } + ); + } + + // We only support sequential scheduling at the moment. + @Test + void schedule_whileScheduling_throwsIllegalStateException() throws InterruptedException { + final AtomicLong schedulingTimeNanos = new AtomicLong(); + + runOnRetryEventLoop(() -> { + schedulingTimeNanos.set(System.nanoTime()); + assertThat(scheduler.trySchedule(nextRetryTask(), 1_000L)).isTrue(); + assertThat(scheduler.whenClosed()).isNotDone(); + }); + + assertNumRetryTaskExecutions(0); + assertEventLoopScheduleCalls(1_000L); + + Thread.sleep(200); + + assertNumRetryTaskExecutions(0); + assertEventLoopScheduleCalls(1_000L); + + runOnRetryEventLoop(() -> { + assertThatThrownBy(() -> scheduler.trySchedule(nextRetryTask(), 0)) + .isInstanceOf(IllegalStateException.class); + assertThat(scheduler.whenClosed()).isNotDone(); + assertThatThrownBy(() -> scheduler.trySchedule(nextRetryTask(), 100)) + .isInstanceOf(IllegalStateException.class); + assertThat(scheduler.whenClosed()).isNotDone(); + assertThatThrownBy(() -> scheduler.trySchedule(nextRetryTask(), 1_000)) + .isInstanceOf(IllegalStateException.class); + assertThat(scheduler.whenClosed()).isNotDone(); + }); + + assertNumRetryTaskExecutions(0); + assertEventLoopScheduleCalls(1_000L); + + await().untilAsserted(() -> { + assertNumRetryTaskExecutions(1); + assertEventLoopScheduleCalls(1_000L); + }); + } + + private static Stream schedule_multiple_executeMultipleAfterDelay_args() { + return Stream.of( + Arguments.of( + new RetryPlan( + "All increasing delays", + Arrays.asList( + new ScheduleCall(100L), + new ScheduleCall(200L), + new ScheduleCall(300L) + ), + Arrays.asList(100L, 200L, 300L) + ) + ), + Arguments.of( + new RetryPlan( + "All decreasing delays", + Arrays.asList( + new ScheduleCall(300L), + new ScheduleCall(200L), + new ScheduleCall(100L) + ), + Arrays.asList(300L, 200L, 100L) + ) + ), + Arguments.of( + new RetryPlan( + "No order in delays, some direct", + Arrays.asList( + new ScheduleCall(200L), + new ScheduleCall(-10L), + new ScheduleCall(0L), + new ScheduleCall(100L) + ), + Arrays.asList(200L, /* direct, */ /* direct, */ 100L + ) + ) + ), + Arguments.of( + new RetryPlan( + "Many direct", + IntStream.range(0, 128) + .mapToObj(i -> new ScheduleCall(0L)) + .collect(Collectors.toList()), + Collections.emptyList() // All direct + ) + ), + Arguments.of( + new RetryPlan( + "Mixed delays and minimum backoffs", + Arrays.asList( + new ScheduleCall(0L), + new ScheduleCall(100L, Collections.singletonList(50L)), + new ScheduleCall(-42L), + new ScheduleCall(50L, Collections.singletonList(100L)), + new ScheduleCall(0L) + ), + Arrays.asList(/* direct, */ 100L, /* direct, */ 100L/*, direct */) + ) + ), + Arguments.of( + new RetryPlan( + "Multiple minimum backoffs", + Arrays.asList( + new ScheduleCall(0L, + Arrays.asList(Long.MIN_VALUE, -50L, 100L) + ), + new ScheduleCall(100L, + Arrays.asList(50L, Long.MIN_VALUE) + ), + new ScheduleCall(-10L, Collections.singletonList(200L)) + ), + Arrays.asList( + 100L, + 100L, + 200L + ) + ) + ) + ); + } + + @ParameterizedTest + @MethodSource("schedule_multiple_executeMultipleAfterDelay_args") + void schedule_multiple_executeMultipleAfterDelay(RetryPlan retryPlan) { + final int numberOfAttempts = retryPlan.retryTaskDelaysMillis.size(); + assert numberOfAttempts >= 1; + + final CountDownLatch retryDone = execRetryPlan(retryPlan); + + await().untilAsserted( + () -> { + assertNumRetryTaskExecutions(numberOfAttempts); + assertEventLoopScheduleCalls(retryPlan.expectedScheduleDelayMillis); + }); + + assertThat(retryDone.getCount()).isZero(); + } + + @Test + void schedule_withDelayAndMinimumBackoff_executeAfterDelay() { + final long minimumBackoffMillis = 200L; + final long delayMillis = minimumBackoffMillis + 200L; + + runOnRetryEventLoop(() -> { + scheduler.applyMinimumBackoffMillisForNextRetry(Long.MIN_VALUE); + scheduler.applyMinimumBackoffMillisForNextRetry(-1); + scheduler.applyMinimumBackoffMillisForNextRetry(0); + scheduler.applyMinimumBackoffMillisForNextRetry(minimumBackoffMillis - 200); + // Should override all the previous calls. + scheduler.applyMinimumBackoffMillisForNextRetry(minimumBackoffMillis); + assertThat(scheduler.trySchedule(nextRetryTask(), delayMillis)).isTrue(); + assertThat(scheduler.whenClosed()).isNotDone(); + }); + + assertNumRetryTaskExecutions(0); + assertEventLoopScheduleCalls(delayMillis); + + await().untilAsserted( + () -> { + assertNumRetryTaskExecutions(1); + assertEventLoopScheduleCalls(delayMillis); + } + ); + } + + @Test + void schedule_withDelayAndMinimumBackoff_executeAfterMinimumBackoff() { + final long delayMillis = 200L; + final long minimumBackoffMillis = delayMillis + 200L; + + runOnRetryEventLoop(() -> { + scheduler.applyMinimumBackoffMillisForNextRetry(minimumBackoffMillis); + assertThat(scheduler.trySchedule(nextRetryTask(), delayMillis)).isTrue(); + assertThat(scheduler.whenClosed()).isNotDone(); + }); + + await().untilAsserted( + () -> { + assertNumRetryTaskExecutions(1); + assertEventLoopScheduleCalls(minimumBackoffMillis); + } + ); + } + + @Test + void schedule_beyondDeadline_returnFalse() { + runOnRetryEventLoop(() -> { + assertThat(scheduler.trySchedule( + nextRetryTask(), + 10_000 + + RETRY_SCHEDULER_SCHEDULE_ADJUSTMENT_TOLERANCE_MILLIS + 1 + ) + ).isFalse(); + assertThat(scheduler.whenClosed()).isNotDone(); + }); + + assertNumRetryTaskExecutions(0); + assertNoScheduleCalls(); + } + + @Test + void schedule_exceptionInRetryTask_closeExceptionally() { + final AtomicReference> whenClosedRef = new AtomicReference<>(); + + runOnRetryEventLoop(() -> { + assertThat(scheduler.trySchedule( + nextRetryTask(() -> { + throw new AnticipatedException("bad task"); + }), + 200L)).isTrue(); + assertThat(scheduler.whenClosed()).isNotDone(); + whenClosedRef.set(scheduler.whenClosed()); + }); + + await().untilAsserted(() -> { + try { + assertThat(whenClosedRef.get()).isCompletedExceptionally(); + whenClosedRef.get().get(); + fail(); + } catch (Throwable e) { + assertThat(e.getCause()).isInstanceOf(AnticipatedException.class) + .hasMessage("bad task"); + } + }); + + assertNumRetryTaskExecutions(1); + assertEventLoopScheduleCalls(200L); + } + + @Test + void schedule_closeRetryEventLoop_closeExceptionally() throws Exception { + final ManagedRetryEventLoop localRetryEventLoop = new ManagedRetryEventLoop(); + final DefaultRetryScheduler localScheduler = new DefaultRetryScheduler( + localRetryEventLoop, + System.nanoTime() + TimeUnit.SECONDS.toNanos(10) + ); + + final AtomicReference> whenClosedRef = new AtomicReference<>(); + + runOnRetryEventLoop(localRetryEventLoop, () -> { + assertThat(localScheduler.trySchedule(nextRetryTask(), 1_000)).isTrue(); + assertThat(localScheduler.whenClosed()).isNotDone(); + whenClosedRef.set(localScheduler.whenClosed()); + }); + + // Close the event loop before the task is executed. + localRetryEventLoop.shutdownGracefully(0, 500, TimeUnit.MILLISECONDS).get(); + + assertThat(whenClosedRef.get()).isCompletedExceptionally(); + try { + whenClosedRef.get().get(); + fail(); + } catch (Throwable e) { + assertThat(e.getCause()).isInstanceOf(IllegalStateException.class) + .hasMessageContaining(ClientFactory.class.getSimpleName()) + .hasMessageContaining("has been closed"); + } + + assertNumRetryTaskExecutions(0); + assertEventLoopScheduleCalls(localRetryEventLoop, Collections.singletonList(1_000L)); + + Thread.sleep(1_000 + MAX_EXECUTION_DELAY_MILLIS); + + assertNumRetryTaskExecutions(0); + assertEventLoopScheduleCalls(localRetryEventLoop, Collections.singletonList(1_000L)); + + assertNoExceptionsOnRetryEventLoop(localRetryEventLoop); + } + + @Test + void schedule_clogRetryEventLoop_closeExceptionally() throws Exception { + final DefaultRetryScheduler localScheduler = new DefaultRetryScheduler( + retryEventLoop, + System.nanoTime() + TimeUnit.SECONDS.toNanos(1) + ); + + runOnRetryEventLoop(() -> { + assertThat(localScheduler.trySchedule(nextRetryTask(), 500L)).isTrue(); + assertThat(localScheduler.whenClosed()).isNotDone(); + }); + + assertNumRetryTaskExecutions(0); + assertEventLoopScheduleCalls(500L); + + runOnRetryEventLoop(() -> { + BlockingUtils.blockingRun(() -> { + Thread.sleep(1_500L); + }); + }); + + await().untilAsserted(() -> { + assertNumRetryTaskExecutions(0); + assertEventLoopScheduleCalls(500L); + }); + + await().untilAsserted(() -> { + runOnRetryEventLoop(() -> { + try { + localScheduler.whenClosed().getNow(null); + fail(); + } catch (Throwable e) { + assertThat(e.getCause()).isInstanceOf(ResponseTimeoutException.class); + } + }); + }); + + Thread.sleep(500 + MAX_EXECUTION_DELAY_MILLIS); + + assertNumRetryTaskExecutions(0); + assertEventLoopScheduleCalls(500L); + } + + @Test + void schedule_retryEventLoopRejects_closeExceptionally() throws Exception { + retryEventLoop.setRejectSchedule(true); + + final AtomicReference> whenClosedRef = new AtomicReference<>(); + + runOnRetryEventLoop(() -> { + assertThat(scheduler.trySchedule(nextRetryTask(), 500L)).isFalse(); + assertThat(scheduler.whenClosed()).isDone(); + whenClosedRef.set(scheduler.whenClosed()); + }); + + try { + whenClosedRef.get().get(); + fail(); + } catch (Throwable e) { + assertThat(e.getCause()).isInstanceOf(RejectedExecutionException.class); + } + } + + @Test + void schedule_scheduledFutureCompletesExceptionally_closeExceptionally() { + final AtomicReference> whenClosedRef = new AtomicReference<>(); + + runOnRetryEventLoop(() -> { + assertThat(scheduler.trySchedule(nextRetryTask(), 1_000)).isTrue(); + assertThat(scheduler.whenClosed()).isNotDone(); + whenClosedRef.set(scheduler.whenClosed()); + }); + + assertNumRetryTaskExecutions(0); + assertEventLoopScheduleCalls(1_000L); + + runOnRetryEventLoop(() -> { + retryEventLoop.setExceptionWhenTaskRun(new AnticipatedException()); + }); + + await().untilAsserted(() -> { + try { + assertThat(whenClosedRef.get()).isCompletedExceptionally(); + whenClosedRef.get().get(); + fail(); + } catch (Throwable e) { + assertThat(e.getCause()).isInstanceOf(AnticipatedException.class); + } + }); + + assertNumRetryTaskExecutions(0); + assertEventLoopScheduleCalls(1_000L); + } + + @Test + void applyMinimumBackoff_whileScheduling_throwIllegalStateException() throws InterruptedException { + runOnRetryEventLoop(() -> { + assertThat(scheduler.trySchedule(nextRetryTask(), 1_000)).isTrue(); + assertThat(scheduler.whenClosed()).isNotDone(); + }); + + assertNumRetryTaskExecutions(0); + assertEventLoopScheduleCalls(1_000L); + Thread.sleep(200); + + runOnRetryEventLoop(() -> { + assertThatThrownBy(() -> scheduler.applyMinimumBackoffMillisForNextRetry(42)) + .isInstanceOf(IllegalStateException.class); + assertThat(scheduler.whenClosed()).isNotDone(); + }); + + assertNumRetryTaskExecutions(0); + assertEventLoopScheduleCalls(1_000L); + + await().untilAsserted( + () -> { + assertNumRetryTaskExecutions(1); + assertEventLoopScheduleCalls(1_000L); + } + ); + + runOnRetryEventLoop(() -> assertThat(scheduler.whenClosed()).isNotDone()); + } + + @Test + void applyMinimumBackoff_thatExceedsDeadline_rejectEverySchedule() { + runOnRetryEventLoop(() -> { + scheduler.applyMinimumBackoffMillisForNextRetry( + 10_000 + + RETRY_SCHEDULER_SCHEDULE_ADJUSTMENT_TOLERANCE_MILLIS + 1); + assertThat(scheduler.whenClosed()).isNotDone(); + assertThat(scheduler.trySchedule(nextRetryTask(), 0)).isFalse(); + assertThat(scheduler.trySchedule(nextRetryTask(), 100)).isFalse(); + }); + + assertNoScheduleCalls(); + assertNumRetryTaskExecutions(0); + + runOnRetryEventLoop(() -> { + assertThat(scheduler.whenClosed()).isNotDone(); + }); + } + + @Test + void close_afterNothing_returnTrue() { + runOnRetryEventLoop(() -> { + scheduler.close(); + assertThat(scheduler.whenClosed()).isCompleted(); + }); + + assertNoScheduleCalls(); + assertNumRetryTaskExecutions(0); + } + + @Test + void close_thenSchedule_returnFalse() { + runOnRetryEventLoop(() -> { + scheduler.close(); + assertThat(scheduler.whenClosed()).isCompleted(); + assertThat(scheduler.trySchedule(nextRetryTask(), 200)).isFalse(); + }); + + assertNoScheduleCalls(); + assertNumRetryTaskExecutions(0); + } + + @Test + void close_whileSchedule_cancelRetryTask() throws InterruptedException { + runOnRetryEventLoop(() -> { + assertThat(scheduler.trySchedule(nextRetryTask(), 1000L)).isTrue(); + assertThat(scheduler.whenClosed()).isNotDone(); + }); + + assertNumRetryTaskExecutions(0); + assertEventLoopScheduleCalls(1000L); + + Thread.sleep(200); + + assertEventLoopScheduleCalls(1000L); + assertNumRetryTaskExecutions(0); + + runOnRetryEventLoop(() -> { + scheduler.close(); + assertThat(scheduler.whenClosed()).isCompleted(); + }); + + assertEventLoopScheduleCalls(1000L); + assertNumRetryTaskExecutions(0); + + Thread.sleep(800 + MAX_EXECUTION_DELAY_MILLIS); + + assertEventLoopScheduleCalls(1000L); + assertNumRetryTaskExecutions(0); + } + + @Test + void close_thenDoSomething_doesNotScheduleAndThrow() throws InterruptedException { + runOnRetryEventLoop(() -> { + scheduler.close(); + assertThat(scheduler.whenClosed()).isCompleted(); + }); + + runOnRetryEventLoop(() -> { + assertThat(scheduler.trySchedule(nextRetryTask(), 0)).isFalse(); + assertThat(scheduler.trySchedule(nextRetryTask(), 200)).isFalse(); + }); + + Thread.sleep(200 + MAX_EXECUTION_DELAY_MILLIS); + + runOnRetryEventLoop(() -> { + scheduler.applyMinimumBackoffMillisForNextRetry(Long.MIN_VALUE); + scheduler.close(); + scheduler.applyMinimumBackoffMillisForNextRetry(-1); + scheduler.applyMinimumBackoffMillisForNextRetry(0); + scheduler.applyMinimumBackoffMillisForNextRetry(100); + scheduler.close(); + }); + + assertNoScheduleCalls(); + assertNumRetryTaskExecutions(0); + } + + @Test + void all_outsideRetryEventLoop_throwsIllegalStateException() { + assertThatThrownBy(() -> scheduler.trySchedule(nextRetryTask(), 100)).isInstanceOf( + IllegalStateException.class); + + assertNoScheduleCalls(); + assertNumRetryTaskExecutions(0); + runOnRetryEventLoop(() -> { + assertThat(scheduler.whenClosed()).isNotDone(); + }); + + assertThatThrownBy(() -> scheduler.applyMinimumBackoffMillisForNextRetry(200)) + .isInstanceOf(IllegalStateException.class); + + assertNoScheduleCalls(); + assertNumRetryTaskExecutions(0); + runOnRetryEventLoop(() -> { + assertThat(scheduler.whenClosed()).isNotDone(); + }); + + assertThatThrownBy(scheduler::close).isInstanceOf(IllegalStateException.class); + + assertNoScheduleCalls(); + assertNumRetryTaskExecutions(0); + runOnRetryEventLoop(() -> { + assertThat(scheduler.whenClosed()).isNotDone(); + }); + } + + private ManagedRetryTask nextRetryTask() { + return nextRetryTask(() -> { + // Default does nothing. + }); + } + + private ManagedRetryTask nextRetryTask(Runnable innerTask) { + return new ManagedRetryTask(task -> { + assertThat(retryEventLoop.inEventLoop()).isTrue(); + executedRetryTasks.add(task); + innerTask.run(); + }); + } + + private void runOnRetryEventLoop(Runnable runnable) { + runOnRetryEventLoop(retryEventLoop, runnable); + } + + private void runOnRetryEventLoop(ManagedRetryEventLoop retryEventLoop, Runnable runnable) { + final CountDownLatch latch = new CountDownLatch(1); + retryEventLoop.execute(() -> { + try { + runnable.run(); + } finally { + latch.countDown(); + } + }); + try { + assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue(); + } catch (InterruptedException e) { + fail(e); + } + + assertNoExceptionsOnRetryEventLoop(retryEventLoop); + } + + private static final class RetryPlan { + final String name; + final List retryTaskDelaysMillis; + // Does not include direct invocations as they don't call the retry event loop. + final List expectedScheduleDelayMillis; + + RetryPlan(String name, List retryTaskDelaysMillis, + List expectedScheduleDelayMillis) { + this.name = name; + this.retryTaskDelaysMillis = ImmutableList.copyOf(retryTaskDelaysMillis); + this.expectedScheduleDelayMillis = ImmutableList.copyOf(expectedScheduleDelayMillis); + } + + @Override + public String toString() { + return name; + } + } + + private static final class ScheduleCall { + final long delayMillis; + final List minimumBackoffMillis; + + ScheduleCall(long delayMillis, List minimumBackoffMillis) { + this.delayMillis = delayMillis; + this.minimumBackoffMillis = ImmutableList.copyOf(minimumBackoffMillis); + } + + ScheduleCall(long delayMillis) { + this.delayMillis = delayMillis; + minimumBackoffMillis = ImmutableList.of(); + } + } + + @NotNull + private CountDownLatch execRetryPlan(RetryPlan retryPlan) { + final int numberOfAttempts = retryPlan.retryTaskDelaysMillis.size(); + final List tasks = Collections.synchronizedList(new ArrayList<>()); + + final Consumer scheduleTask = retryNumber -> { + assert 0 <= retryNumber && retryNumber < numberOfAttempts; + assertThat(scheduler.whenClosed()).isNotDone(); + + final ScheduleCall scheduleCall = retryPlan.retryTaskDelaysMillis.get(retryNumber); + + for (long minimumBackoffMillis : scheduleCall.minimumBackoffMillis) { + scheduler.applyMinimumBackoffMillisForNextRetry(minimumBackoffMillis); + } + + // In the end we expect to execute that retry task after scheduledDelayMillis. + assertThat(scheduler.trySchedule( + tasks.get(retryNumber), + scheduleCall.delayMillis) + ) + .as("scheduling retry %d", retryNumber) + .isTrue(); + }; + + final CountDownLatch retryDone = new CountDownLatch(1); + + for (int retryNumber = 0; retryNumber < numberOfAttempts - 1; retryNumber++) { + final int nextRetryNumber = retryNumber + 1; + tasks.add(nextRetryTask(() -> scheduleTask.accept(nextRetryNumber))); + } + tasks.add(nextRetryTask(retryDone::countDown)); + assert tasks.size() == numberOfAttempts; + + // Execute the first task. + runOnRetryEventLoop(() -> scheduleTask.accept(0)); + + return retryDone; + } + + private void assertNoScheduleCalls() { + assertEventLoopScheduleCalls(Collections.emptyList()); + } + + private void assertEventLoopScheduleCalls(long... expectedScheduleDelaysMillisOnEventLoop) { + assert expectedScheduleDelaysMillisOnEventLoop != null; + assertEventLoopScheduleCalls(LongStream.of(expectedScheduleDelaysMillisOnEventLoop) + .boxed() + .collect(Collectors.toList())); + } + + private void assertEventLoopScheduleCalls(List expectedScheduleDelaysMillisOnEventLoop) { + assertEventLoopScheduleCalls(retryEventLoop, expectedScheduleDelaysMillisOnEventLoop); + } + + private void assertEventLoopScheduleCalls(ManagedRetryEventLoop retryEventLoop, + List expectedScheduleDelaysMillisOnEventLoop) { + assertThat(retryEventLoop.scheduleDelaysMillis()) + .as("Expected number of schedule() calls") + .hasSize(expectedScheduleDelaysMillisOnEventLoop.size()); + + for (int i = 0; i < expectedScheduleDelaysMillisOnEventLoop.size(); i++) { + final long expectedDelayMillis = expectedScheduleDelaysMillisOnEventLoop.get(i); + + assertThat(retryEventLoop.scheduleDelaysMillis().get(i)) + .isBetween( + expectedDelayMillis - RETRY_SCHEDULER_SCHEDULE_ADJUSTMENT_TOLERANCE_MILLIS, + expectedDelayMillis); + } + } + + private void assertNumRetryTaskExecutions(int numRetryTasks) { + assertThat(executedRetryTasks).hasSize(numRetryTasks); + + for (int i = 0; i < numRetryTasks; i++) { + final ManagedRetryTask task = executedRetryTasks.get(i); + assertThat(task.nextRetryTaskNumber()).isEqualTo(i); + } + } + + private void assertNoExceptionsOnRetryEventLoop(ManagedRetryEventLoop retryEventLoop) { + assertThat(retryEventLoop.exceptionsCaughtOnEventLoop()) + .as("No exceptions thrown on retry event loop") + .isEmpty(); + } + + private static final class ManagedRetryTask implements Runnable { + static final AtomicInteger nextRetryTaskNumber = new AtomicInteger(0); + + private final Consumer runnable; + private final int retryTaskNumber; + private final AtomicBoolean executed; + + ManagedRetryTask(Consumer runnable) { + this.runnable = runnable; + retryTaskNumber = nextRetryTaskNumber.getAndIncrement(); + executed = new AtomicBoolean(false); + } + + @Override + public void run() { + assertThat(executed.get()).isFalse(); + executed.set(true); + runnable.accept(this); + } + + int nextRetryTaskNumber() { + return retryTaskNumber; + } + } + + private static final class ManagedRetryEventLoop extends DefaultEventLoop { + private final List exceptionsCaughtOnEventLoop = Collections.synchronizedList( + new ArrayList<>()); + private final List scheduleDelaysMillis = Collections.synchronizedList( + new ArrayList<>()); + private final AtomicBoolean rejectSchedule = new AtomicBoolean(false); + private final AtomicReference<@Nullable RuntimeException> throwWhenTaskRun = new AtomicReference<>( + null); + + @Override + protected void run() { + try { + super.run(); + } catch (Throwable t) { + exceptionsCaughtOnEventLoop.add(t); + throw t; + } + } + + void setRejectSchedule(boolean rejectSchedule) { + this.rejectSchedule.set(rejectSchedule); + } + + void setExceptionWhenTaskRun(@Nullable RuntimeException t) { + throwWhenTaskRun.set(t); + } + + @Override + public ScheduledFuture schedule(Runnable command, long delay, TimeUnit unit) { + scheduleDelaysMillis.add(TimeUnit.MILLISECONDS.convert(delay, unit)); + if (rejectSchedule.get()) { + reject(); + } + + return super.schedule(() -> { + final RuntimeException t = throwWhenTaskRun.get(); + if (t != null) { + throw t; + } + command.run(); + }, delay, unit); + } + + List exceptionsCaughtOnEventLoop() { + return exceptionsCaughtOnEventLoop; + } + + List scheduleDelaysMillis() { + return scheduleDelaysMillis; + } + } +} diff --git a/core/src/test/java/com/linecorp/armeria/client/retry/RetryingClientLoadBalancingTest.java b/core/src/test/java/com/linecorp/armeria/client/retry/RetryingClientLoadBalancingTest.java index a2b52679da6..62d6a8fe43e 100644 --- a/core/src/test/java/com/linecorp/armeria/client/retry/RetryingClientLoadBalancingTest.java +++ b/core/src/test/java/com/linecorp/armeria/client/retry/RetryingClientLoadBalancingTest.java @@ -16,6 +16,7 @@ package com.linecorp.armeria.client.retry; import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.linecorp.armeria.client.retry.AbstractRetryingClient.ARMERIA_RETRY_COUNT; import static org.assertj.core.api.Assertions.assertThat; import java.net.InetSocketAddress; @@ -104,8 +105,10 @@ void test(TestMode mode) { status = null; } + final boolean isRetry = ctx.request().headers().contains(ARMERIA_RETRY_COUNT); + // Retry only once on failure. - if (!HttpStatus.OK.equals(status) && AbstractRetryingClient.getTotalAttempts(ctx) <= 1) { + if (!HttpStatus.OK.equals(status) && !isRetry) { return UnmodifiableFuture.completedFuture(RetryDecision.retry(Backoff.withoutDelay())); } else { return UnmodifiableFuture.completedFuture(RetryDecision.noRetry()); @@ -127,9 +130,9 @@ void test(TestMode mode) { case FAILURE: final List expectedPortsWhenRetried = ImmutableList.builder() - .addAll(expectedPorts) - .addAll(expectedPorts) - .build(); + .addAll(expectedPorts) + .addAll(expectedPorts) + .build(); assertThat(accessedPorts).isEqualTo(expectedPortsWhenRetried); break; } diff --git a/core/src/test/java/com/linecorp/armeria/client/retry/RetryingClientTest.java b/core/src/test/java/com/linecorp/armeria/client/retry/RetryingClientTest.java index c68446ad7e3..210dfec74de 100644 --- a/core/src/test/java/com/linecorp/armeria/client/retry/RetryingClientTest.java +++ b/core/src/test/java/com/linecorp/armeria/client/retry/RetryingClientTest.java @@ -21,6 +21,7 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.assertj.core.api.Assertions.catchThrowable; +import static org.awaitility.Awaitility.await; import java.time.Duration; import java.util.Arrays; @@ -55,6 +56,8 @@ import com.linecorp.armeria.client.ClientFactory; 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.ResponseTimeoutException; import com.linecorp.armeria.client.UnprocessedRequestException; @@ -526,15 +529,15 @@ void honorRetryMapping() { void evaluatesMappingOnce() { final AtomicInteger evaluations = new AtomicInteger(0); final RetryConfigMapping mapping = - (ctx, req) -> { - evaluations.incrementAndGet(); - return RetryConfig - .builder0(RetryRule.builder() - .onStatus(HttpStatus.valueOf(500)) - .thenBackoff()) - .maxTotalAttempts(2) - .build(); - }; + (ctx, req) -> { + evaluations.incrementAndGet(); + return RetryConfig + .builder0(RetryRule.builder() + .onStatus(HttpStatus.valueOf(500)) + .thenBackoff()) + .maxTotalAttempts(2) + .build(); + }; final WebClient client = client(mapping); @@ -712,11 +715,24 @@ void doNotRetryWhenResponseIsAborted() throws Exception { @Test void doNotRetryWhenSubscriberIsCancelled() throws Exception { final WebClient client = client(retryAlways); - client.get("/subscriber-cancel").subscribe( + final AtomicInteger subscriberCancelServiceCallCounterWhenCancelled = new AtomicInteger(-1); + + final HttpResponse res; + final ClientRequestContext ctx; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + res = client.get("/subscriber-cancel"); + ctx = captor.get(); + } + + res.subscribe( new Subscriber() { @Override public void onSubscribe(Subscription s) { - s.cancel(); // Cancel as soon as getting the subscription. + ctx.eventLoop().schedule(() -> { + subscriberCancelServiceCallCounterWhenCancelled.set( + subscriberCancelServiceCallCounter.get()); + s.cancel(); + }, 1000, TimeUnit.MILLISECONDS); } @Override @@ -729,8 +745,13 @@ public void onError(Throwable t) {} public void onComplete() {} }); - TimeUnit.SECONDS.sleep(1L); // Sleep to check if there's a retry. - assertThat(subscriberCancelServiceCallCounter.get()).isEqualTo(1); + await().untilAsserted(() -> { + // We should have made at least one service call during the second. + assertThat(subscriberCancelServiceCallCounter.get()).isGreaterThanOrEqualTo(1); + // After we cancelled we do not expect further service calls. + assertThat(subscriberCancelServiceCallCounterWhenCancelled.get()).isEqualTo( + subscriberCancelServiceCallCounter.get()); + }); } @Test diff --git a/thrift/thrift0.13/src/test/java/com/linecorp/armeria/it/client/retry/RetryingRpcClientTest.java b/thrift/thrift0.13/src/test/java/com/linecorp/armeria/it/client/retry/RetryingRpcClientTest.java index 99ec07c9752..f10044b1c14 100644 --- a/thrift/thrift0.13/src/test/java/com/linecorp/armeria/it/client/retry/RetryingRpcClientTest.java +++ b/thrift/thrift0.13/src/test/java/com/linecorp/armeria/it/client/retry/RetryingRpcClientTest.java @@ -28,8 +28,10 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import java.time.Duration; import java.util.concurrent.BlockingQueue; import java.util.concurrent.CancellationException; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executors; import java.util.concurrent.LinkedTransferQueue; import java.util.concurrent.TimeUnit; @@ -39,6 +41,8 @@ import org.apache.thrift.TApplicationException; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; import com.linecorp.armeria.client.ClientFactory; import com.linecorp.armeria.client.ClientRequestContext; @@ -53,7 +57,9 @@ import com.linecorp.armeria.common.HttpRequest; import com.linecorp.armeria.common.RpcResponse; import com.linecorp.armeria.common.logging.RequestLog; +import com.linecorp.armeria.common.util.TimeoutMode; import com.linecorp.armeria.common.util.UnmodifiableFuture; +import com.linecorp.armeria.internal.testing.BlockingUtils; import com.linecorp.armeria.server.ServerBuilder; import com.linecorp.armeria.server.thrift.THttpService; import com.linecorp.armeria.testing.junit5.server.ServerExtension; @@ -62,10 +68,10 @@ import testing.thrift.main.HelloService; class RetryingRpcClientTest { - + private static final Backoff fixedBackoff = Backoff.fixed(500); private static final RetryRuleWithContent retryAlways = (ctx, response, cause) -> - UnmodifiableFuture.completedFuture(RetryDecision.retry(Backoff.fixed(500))); + UnmodifiableFuture.completedFuture(RetryDecision.retry(fixedBackoff)); private static final RetryRuleWithContent retryOnException = RetryRuleWithContent.onException(Backoff.withoutDelay()); @@ -326,34 +332,123 @@ void shouldGetExceptionWhenFactoryIsClosed() throws Exception { "(?i).*(factory has been closed|not accepting a task).*")); } - @Test - void doNotRetryWhenResponseIsCancelled() throws Exception { + enum DoNotRetryWhenResponseIsCancelledTestParams { + // Cancel delays for a backoff of 50 milliseconds (quickBackoffMillis). + CANCEL_FIRST_REQUEST_NO_DELAY(true, 0), + CANCEL_FIRST_REQUEST_WITH_DELAY(true, 500), + CANCEL_AFTER_FIRST_REQUEST_NO_DELAY(false, 0), + CANCEL_AFTER_FIRST_REQUEST_WITH_DELAY(false, 500); + + static final int BACKOFF_MILLIS = 50; + final boolean ensureCancelBeforeFirstRequest; + final long cancelDelayMillis; + + DoNotRetryWhenResponseIsCancelledTestParams(boolean ensureCancelBeforeFirstRequest, + long cancelDelayMillis) { + this.ensureCancelBeforeFirstRequest = ensureCancelBeforeFirstRequest; + this.cancelDelayMillis = cancelDelayMillis; + } + } + + @ParameterizedTest + @EnumSource(DoNotRetryWhenResponseIsCancelledTestParams.class) + void doNotRetryWhenResponseIsCancelled(DoNotRetryWhenResponseIsCancelledTestParams param) throws Exception { serviceRetryCount.set(0); + + final RetryRuleWithContent quickRetryAlways = + RetryRuleWithContent.builder() + .onException() + .thenBackoff(Backoff.fixed( + DoNotRetryWhenResponseIsCancelledTestParams.BACKOFF_MILLIS)); + + final int maxExpectedAttempts = + (int) (param.cancelDelayMillis / DoNotRetryWhenResponseIsCancelledTestParams.BACKOFF_MILLIS) + + 5; + final AtomicInteger serviceRetryCountWhenCancelled = new AtomicInteger(); + final CountDownLatch canRetry = new CountDownLatch(1); try (ClientFactory factory = ClientFactory.builder().build()) { final AtomicReference context = new AtomicReference<>(); final HelloService.Iface client = ThriftClients.builder(server.httpUri()) .path("/thrift") .factory(factory) - .rpcDecorator(RetryingRpcClient.builder(retryAlways).newDecorator()) + .rpcDecorator(RetryingRpcClient.builder(quickRetryAlways) + // We want to cancel the request before + // we quit because of reaching max attempts. + .maxTotalAttempts(maxExpectedAttempts) + .newDecorator()) + .rpcDecorator((delegate, ctx, req) -> { + // Clog the retry event loop so we do not retry until canRetry.countDown() + // is called. + // If you see failure of this test, and you altered AbstractRetryingClient, + // make sure you are executing (prepare)Retry() on the retry event loop and + // that the retry event loop is ctx.eventLoop(). + ctx.eventLoop().execute(() -> { + BlockingUtils.blockingRun(() -> { + canRetry.await(); + }); + }); + + return delegate.execute(ctx, req); + }) .rpcDecorator((delegate, ctx, req) -> { - context.set(ctx); final RpcResponse res = delegate.execute(ctx, req); - res.cancel(true); + + if (param.ensureCancelBeforeFirstRequest) { + Thread.sleep(param.cancelDelayMillis); + assertThat(res.isDone()).isFalse(); + res.cancel(true); + serviceRetryCountWhenCancelled.set(serviceRetryCount.get()); + canRetry.countDown(); + } else { + canRetry.countDown(); + Thread.sleep(param.cancelDelayMillis); + assertThat(res.isDone()).isFalse(); + res.cancel(true); + serviceRetryCountWhenCancelled.set(serviceRetryCount.get()); + } + return res; }) + .rpcDecorator((delegate, ctx, req) -> { + context.set(ctx); + ctx.setResponseTimeout( + TimeoutMode.EXTEND, + Duration.ofMillis(param.cancelDelayMillis + 1000) + ); + + return delegate.execute(ctx, req); + }) .build(HelloService.Iface.class); when(serviceHandler.hello(anyString())).thenThrow(new IllegalArgumentException()); assertThatThrownBy(() -> client.hello("hello")).isInstanceOf(CancellationException.class); await().untilAsserted(() -> { - verify(serviceHandler, only()).hello("hello"); + assertThat(serviceRetryCountWhenCancelled.get()).isIn(serviceRetryCount.get(), + serviceRetryCount.get() - 1); + verify(serviceHandler, times(serviceRetryCount.get())).hello("hello"); }); + + final RequestLog log = context.get().log().whenComplete().join(); + if (param.ensureCancelBeforeFirstRequest) { + assertThat(serviceRetryCount.get()).isZero(); + assertThat(log.requestCause()).isExactlyInstanceOf(CancellationException.class); + assertThat(log.responseCause()).isExactlyInstanceOf(CancellationException.class); + } else { + // We still could cancel the before the first request so we do not have a guarantee for + // requestCause() to be null. + assertThat(log.responseCause()).isExactlyInstanceOf(CancellationException.class); + } + // Sleep 1 second more to check if there was another retry. TimeUnit.SECONDS.sleep(1); - verify(serviceHandler, only()).hello("hello"); - assertThat(serviceRetryCount).hasValue(1); + if (param.ensureCancelBeforeFirstRequest) { + assertThat(serviceRetryCount.get()).isZero(); + } + assertThat(serviceRetryCountWhenCancelled.get()).isIn(serviceRetryCount.get(), + serviceRetryCount.get() - 1); + verify(serviceHandler, times(serviceRetryCount.get())).hello("hello"); } } }