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..87e9f771ca0 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 @@ -18,17 +18,26 @@ import static com.google.common.base.Preconditions.checkState; import static java.util.Objects.requireNonNull; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.locks.ReentrantLock; import java.util.function.Consumer; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableList; + 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.client.retry.RetrySchedulingException.Type; import com.linecorp.armeria.common.HttpHeaderNames; import com.linecorp.armeria.common.HttpRequest; import com.linecorp.armeria.common.Request; @@ -36,11 +45,12 @@ import com.linecorp.armeria.common.RpcRequest; import com.linecorp.armeria.common.annotation.Nullable; import com.linecorp.armeria.common.util.TimeoutMode; +import com.linecorp.armeria.common.util.UnmodifiableFuture; import com.linecorp.armeria.internal.client.ClientUtil; +import com.linecorp.armeria.internal.common.util.ReentrantShortLock; 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. @@ -59,7 +69,7 @@ public abstract class AbstractRetryingClient STATE = + private static final AttributeKey> STATE = AttributeKey.valueOf(AbstractRetryingClient.class, "STATE"); private final RetryConfigMapping mapping; @@ -82,7 +92,19 @@ public final O execute(ClientRequestContext ctx, I req) throws Exception { final RetryConfig config = mapping.get(ctx, req); requireNonNull(config, "mapping.get() returned null"); - final State state = new State(config, ctx.responseTimeoutMillis()); + final State state; + final ReentrantShortLock retryLock = new ReentrantShortLock(); + if (ctx.responseTimeoutMillis() <= 0 || ctx.responseTimeoutMillis() == Long.MAX_VALUE) { + final RetryScheduler scheduler = new RetryScheduler(retryLock, ctx.eventLoop()); + state = new State<>(retryLock, config, scheduler); + } else { + final long responseTimeoutTimeNanos = + System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(ctx.responseTimeoutMillis()); + final RetryScheduler scheduler = new RetryScheduler(retryLock, ctx.eventLoop(), + responseTimeoutTimeNanos); + state = new State<>(retryLock, config, scheduler, responseTimeoutTimeNanos); + } + ctx.setAttr(STATE, state); return doExecute(ctx, req); } @@ -101,10 +123,20 @@ protected final RetryConfigMapping mapping() { protected abstract O doExecute(ClientRequestContext ctx, I req) throws Exception; /** - * This should be called when retrying is finished. + * todo(szymon): [doc]. + */ + protected static void completeRetryingIfNoPendingAttempts(ClientRequestContext ctx) { + final State state = state(ctx); + state.completeIfNoPendingAttempts(); + } + + /** + * todo(szymon): [doc]. */ - protected static void onRetryingComplete(ClientRequestContext ctx) { - ctx.logBuilder().endResponseWithLastChild(); + protected static void completeRetryingExceptionally(ClientRequestContext ctx, + Throwable cause) { + final State state = state(ctx); + state.completeExceptionally(cause); } /** @@ -142,102 +174,216 @@ protected final RetryRuleWithContent retryRuleWithContent() { } /** - * Schedules next retry. + * todo(szymon): [doc]. */ - protected static void scheduleNextRetry(ClientRequestContext ctx, - Consumer actionOnException, - Runnable retryTask, long nextDelayMillis) { - 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()); - } - }); - } - } catch (Throwable t) { - actionOnException.accept(t); + protected void startRetryAttempt(ClientRequestContext ctx, + ClientRequestContext attemptCtx, + O attemptRes, + AttemptWinHandler onWinHandler, + AttemptAbortHandler + onAbortHandler + ) { + requireNonNull(ctx, "ctx"); + requireNonNull(attemptCtx, "attemptCtx"); + requireNonNull(onWinHandler, "onWinHandler"); + requireNonNull(onAbortHandler, "onAbortHandler"); + + final State state = state(ctx); + final ReentrantLock retryLock = state.getLock(); + + retryLock.lock(); + + if (isRetryingComplete(ctx)) { + // This really should not happen. However, let us call the abort handler to free + // potentially expensive resources here. + retryLock.unlock(); + onAbortHandler.accept(attemptCtx, + attemptRes, + new IllegalStateException("Retrying is already complete.")); + return; } + + state.startAttempt(attemptCtx, attemptRes, onWinHandler, onAbortHandler); + retryLock.unlock(); } /** - * 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 + * todo(szymon): [doc]. */ - @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; + protected void completeRetryAttempt(ClientRequestContext ctx, + ClientRequestContext attemptCtx, + O attemptRes, boolean isWinning) { + if (isRetryingComplete(ctx)) { + // The complete handler of this attempt was already executed (provided that `attemptCtx` + // was registered with `startRetryAttempt` as by the contract of `completeRetryAttempt`. + return; } + + state(ctx).completeAttempt(attemptCtx, attemptRes, isWinning); } /** - * 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. + * todo(szymon): [doc]. */ - protected final long getNextDelay(ClientRequestContext ctx, Backoff backoff) { - return getNextDelay(ctx, backoff, -1); + protected static boolean isRetryingComplete(ClientRequestContext ctx) { + requireNonNull(ctx, "ctx"); + + final State state = state(ctx); + + return state.isRetryingComplete(); } /** - * 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. + * todo(szymon): [doc]. */ - @SuppressWarnings("MethodMayBeStatic") // Intentionally left non-static for better user experience. - protected final long getNextDelay(ClientRequestContext ctx, Backoff backoff, long millisAfterFromServer) { + protected void scheduleNextRetry(ClientRequestContext ctx, + Consumer retryTask, + Backoff backoff, + Consumer actionOnException) { + scheduleNextRetry(ctx, retryTask, backoff, -1, actionOnException); + } + + /** + * todo(szymon): [doc]. + */ + protected static void scheduleNextRetry(ClientRequestContext ctx, + Consumer retryTask, + Backoff backoff, + long retryDelayFromServerMillis, + Consumer actionOnException) { requireNonNull(ctx, "ctx"); + requireNonNull(retryTask, "retryTask"); requireNonNull(backoff, "backoff"); - final State state = state(ctx); - final int currentAttemptNo = state.currentAttemptNoWith(backoff); + requireNonNull(actionOnException, "actionOnException"); + + if (isRetryingComplete(ctx)) { + actionOnException.accept(new RetrySchedulingException(Type.RETRYING_ALREADY_COMPLETED)); + return; + } + + final State state = state(ctx); + final ReentrantLock retryLock = state.getLock(); + + retryLock.lock(); + + final long nowTimeNanos = System.nanoTime(); + + long earliestRetryTimeNanos = retryDelayFromServerMillis >= 0 ? + (nowTimeNanos + TimeUnit.MILLISECONDS.toNanos( + retryDelayFromServerMillis)) + : Long.MIN_VALUE; + + if (state.timeoutForWholeRetryEnabled() && + earliestRetryTimeNanos > state.responseTimeoutTimeNanos()) { + retryLock.unlock(); + actionOnException.accept( + new RetrySchedulingException( + Type.DELAY_FROM_SERVER_EXCEEDS_RESPONSE_TIMEOUT)); + return; + } - if (currentAttemptNo < 0) { - logger.debug("Exceeded the default number of max attempt: {}", state.config.maxTotalAttempts()); - return -1; + // Even when we cannot schedule the retry task, we want to respect the + // minimum retry delay from the server. + final RetryScheduler scheduler = state.scheduler(); + earliestRetryTimeNanos = scheduler.addEarliestNextRetryTimeNanos(earliestRetryTimeNanos); + + final int attemptNoWithBackoff = state.numAttemptsSoFarWithBackoff(backoff); + if (attemptNoWithBackoff < 0) { + retryLock.unlock(); + scheduler.rescheduleCurrentRetryTaskIfTooEarly(); + actionOnException.accept( + new RetrySchedulingException(RetrySchedulingException.Type.NO_MORE_ATTEMPTS_IN_RETRY)); + return; } - long nextDelay = backoff.nextDelayMillis(currentAttemptNo); - if (nextDelay < 0) { - logger.debug("Exceeded the number of max attempts in the backoff: {}", backoff); - return -1; + final long retryDelayMillis = backoff.nextDelayMillis(attemptNoWithBackoff); + if (retryDelayMillis < 0) { + retryLock.unlock(); + scheduler.rescheduleCurrentRetryTaskIfTooEarly(); + actionOnException.accept( + new RetrySchedulingException( + RetrySchedulingException.Type.NO_MORE_ATTEMPTS_IN_BACKOFF)); + return; } - 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; + final long retryTimeNanos = Math.max(nowTimeNanos + TimeUnit.MILLISECONDS.toNanos(retryDelayMillis), + earliestRetryTimeNanos); + if (state.timeoutForWholeRetryEnabled() && retryTimeNanos > state.responseTimeoutTimeNanos()) { + retryLock.unlock(); + scheduler.rescheduleCurrentRetryTaskIfTooEarly(); + // This cannot be the minimum delay from the server, because we already checked it above. + actionOnException.accept( + new RetrySchedulingException( + Type.DELAY_FROM_BACKOFF_EXCEEDS_RESPONSE_TIMEOUT)); + return; } - return nextDelay; + final AtomicBoolean exceptionDuringScheduling = new AtomicBoolean(true); + state.startRetryTask(); + scheduler.schedule(retryLock0 -> { + assert retryLock == retryLock0; + assert retryLock.isHeldByCurrentThread(); + assert retryLock.getHoldCount() == 1; + + if (isRetryingComplete(ctx)) { + state.completeRetryTask(); + retryLock.unlock(); + actionOnException.accept(new RetrySchedulingException(Type.RETRYING_ALREADY_COMPLETED)); + return; + } + + final int thisAttemptNo = state.acquireAttemptNoWithCurrentBackoff(backoff); + assert thisAttemptNo >= 1; + + retryLock.unlock(); + + try { + assert !retryLock.isHeldByCurrentThread(); + retryTask.accept(thisAttemptNo); + } finally { + state.completeRetryTask(); + } + + completeRetryingIfNoPendingAttempts(ctx); + }, retryTimeNanos, earliestRetryTimeNanos, cause -> { + state.completeRetryTask(); + + if (exceptionDuringScheduling.get()) { + assert state(ctx).getLock().isHeldByCurrentThread(); + retryLock.unlock(); + } + + actionOnException.accept(cause); + + if (exceptionDuringScheduling.get()) { + retryLock.lock(); + } + }); + + exceptionDuringScheduling.set(false); + retryLock.unlock(); + } + + // todo(szymon): [doc] + + /** + * 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).responseTimeoutMillisForAttempt(); + if (responseTimeoutMillis < 0) { + return false; + } else if (responseTimeoutMillis == 0) { + ctx.clearResponseTimeout(); + return true; + } else { + ctx.setResponseTimeoutMillis(TimeoutMode.SET_FROM_NOW, responseTimeoutMillis); + return true; + } } /** @@ -245,7 +391,7 @@ protected final long getNextDelay(ClientRequestContext ctx, Backoff backoff, lon * {@link ClientRequestContext}. */ protected static int getTotalAttempts(ClientRequestContext ctx) { - final State state = ctx.attr(STATE); + final State state = ctx.attr(STATE); if (state == null) { return 0; } @@ -263,79 +409,443 @@ protected static ClientRequestContext newDerivedContext(ClientRequestContext ctx return ClientUtil.newDerivedContext(ctx, req, rpcReq, initialAttempt); } - private static State state(ClientRequestContext ctx) { - final State state = ctx.attr(STATE); + @FunctionalInterface + protected interface AttemptWinHandler { + /** + * todo(szymon): [doc]. + */ + void accept(ClientRequestContext attemptCtx, O attemptResOnComplete); + } + + @FunctionalInterface + protected interface AttemptAbortHandler { + /** + * todo(szymon): [doc]. + */ + CompletableFuture accept(ClientRequestContext attemptCtx, O attemptRes, + @Nullable Throwable abortCause); + } + + /** + * todo(szymon): [doc]. + */ + @VisibleForTesting + @SuppressWarnings("unchecked") + public static State state(ClientRequestContext ctx) { + final State state = (State) ctx.attr(STATE); assert state != null; return state; } - private static final class State { + /** + * todo(szymon): [doc]. + */ + @VisibleForTesting + public static final class State { + /** + * todo(szymon): [doc]. + */ + @VisibleForTesting + public static class Attempt { + private final ClientRequestContext attemptCtx; + private final O attemptRes; + private @Nullable O attemptResOnComplete; + + private final AttemptWinHandler onWinHandler; + private final AttemptAbortHandler + onAbortHandler; + + private boolean hasCalledHandler; + + Attempt(ClientRequestContext attemptCtx, O attemptRes, + AttemptWinHandler onWinHandler, + AttemptAbortHandler onAbortHandler) { + this.attemptCtx = attemptCtx; + this.attemptRes = attemptRes; + this.onWinHandler = onWinHandler; + this.onAbortHandler = onAbortHandler; + + hasCalledHandler = false; + attemptResOnComplete = null; + } + ClientRequestContext ctx() { + return attemptCtx; + } + + /** + * todo(szymon): [doc]. + */ + public O attemptRes() { + return attemptRes; + } + + void completeWithRes(O res) { + assert attemptResOnComplete == null; + attemptResOnComplete = res; + } + + CompletableFuture callHandler(boolean isWinning, @Nullable Throwable cause) { + if (hasCalledHandler) { + return UnmodifiableFuture.completedFuture(null); + } + + hasCalledHandler = true; + + if (isWinning) { + assert attemptResOnComplete != null; + onWinHandler.accept(attemptCtx, attemptResOnComplete); + return UnmodifiableFuture.completedFuture(null); + } else { + return onAbortHandler.accept(attemptCtx, attemptRes, cause); + } + } + } + + private final ReentrantLock lock; + private boolean isRetryingComplete; private final RetryConfig config; private final long deadlineNanos; private final boolean isTimeoutEnabled; + private final Map> startedAttempts = new HashMap<>(); + + private final RetryScheduler retryScheduler; + + private int numPendingAttempts; + + // An attempt is considered scheduled between two points: + // - right before its retry task is scheduled + // - right after the retry task finishes execution + // + // As a consequence, the retry task must call startAttempt() synchronously. + // If it does not, we may assume there are no active or scheduled attempts + // once the task ends, and stop retrying too early. + private int numScheduledAttempts; + + private @Nullable Attempt lastAttempt; + @Nullable private Backoff lastBackoff; private int currentAttemptNoWithLastBackoff; + // Starting with 1 private int totalAttemptNo; - State(RetryConfig config, long responseTimeoutMillis) { + State(ReentrantLock retryingLock, RetryConfig config, RetryScheduler retryScheduler) { + lock = retryingLock; + isRetryingComplete = false; this.config = config; + this.retryScheduler = retryScheduler; + totalAttemptNo = 1; + numPendingAttempts = 0; + numScheduledAttempts = 0; - if (responseTimeoutMillis <= 0 || responseTimeoutMillis == Long.MAX_VALUE) { - deadlineNanos = 0; - isTimeoutEnabled = false; - } else { - deadlineNanos = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(responseTimeoutMillis); - isTimeoutEnabled = true; - } + deadlineNanos = 0; + isTimeoutEnabled = false; + } + + State(ReentrantLock retryingLock, RetryConfig config, RetryScheduler retryScheduler, + long responseTimeoutTimeNanos) { + lock = retryingLock; + isRetryingComplete = false; + this.config = config; + this.retryScheduler = retryScheduler; totalAttemptNo = 1; + numPendingAttempts = 0; + numScheduledAttempts = 0; + + deadlineNanos = responseTimeoutTimeNanos; + isTimeoutEnabled = true; + } + + RetryScheduler scheduler() { + return retryScheduler; + } + + ReentrantLock getLock() { + return lock; } /** * Returns the smaller value between {@link RetryConfig#responseTimeoutMillisForEachAttempt()} and - * remaining {@link #responseTimeoutMillis}. + * remaining {@link #responseTimeoutMillisForAttempt}. This method is thread-safe. * * @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 responseTimeoutMillisForAttempt() { + lock.lock(); + + try { + if (!timeoutForWholeRetryEnabled()) { + return config.responseTimeoutMillisForEachAttempt(); + } + + final long actualResponseTimeoutMillis = responseTimeoutMillis(); + + // Consider 0 or less than 0 of actualResponseTimeoutMillis as timed out. + if (actualResponseTimeoutMillis <= 0) { + return -1; + } + + if (config.responseTimeoutMillisForEachAttempt() > 0) { + return Math.min(config.responseTimeoutMillisForEachAttempt(), actualResponseTimeoutMillis); + } + + return actualResponseTimeoutMillis; + } finally { + lock.unlock(); + } + } + + boolean timeoutForWholeRetryEnabled() { + return isTimeoutEnabled; + } + long responseTimeoutMillis() { - if (!timeoutForWholeRetryEnabled()) { - return config.responseTimeoutMillisForEachAttempt(); + assert isTimeoutEnabled; + return Math.max(TimeUnit.NANOSECONDS.toMillis(deadlineNanos - System.nanoTime()), -1); + } + + long responseTimeoutTimeNanos() { + assert isTimeoutEnabled; + return deadlineNanos; + } + + void startAttempt(ClientRequestContext attemptCtx, O attemptRes, + AttemptWinHandler onWinHandler, + AttemptAbortHandler + onAbortHandler) { + lock.lock(); + try { + checkState(!isRetryingComplete()); + checkState(!startedAttempts.containsKey(attemptCtx), + "Attempt %s already is already pending.", attemptCtx); + final Attempt attempt = new Attempt<>(attemptCtx, attemptRes, onWinHandler, + onAbortHandler); + startedAttempts.put(attempt.ctx(), attempt); + numPendingAttempts++; + } finally { + lock.unlock(); } + } - final long actualResponseTimeoutMillis = actualResponseTimeoutMillis(); + void completeAttempt(ClientRequestContext attemptCtx, O attemptRes, boolean isWinning) { + lock.lock(); + try { + if (isRetryingComplete()) { + lock.unlock(); + return; + } + + checkState(startedAttempts.containsKey(attemptCtx), + "Attempt %s is not registered as pending.", attemptCtx); + + lastAttempt = startedAttempts.get(attemptCtx); + lastAttempt.completeWithRes(attemptRes); + numPendingAttempts--; + } catch (Exception e) { + lock.unlock(); + throw e; + } - // Consider 0 or less than 0 of actualResponseTimeoutMillis as timed out. - if (actualResponseTimeoutMillis <= 0) { - return -1; + if (isWinning) { + // If this is the winning attempt, we can complete the retrying. + // transfers ownership of current lock + complete0(); + } else { + // transfers ownership of current lock + completeIfNoPendingAttempts0(); } + } - if (config.responseTimeoutMillisForEachAttempt() > 0) { - return Math.min(config.responseTimeoutMillisForEachAttempt(), actualResponseTimeoutMillis); + void startRetryTask() { + // This can get (temporarily) above max total attempts as there is a moment between scheduling + // a new retry task and cancelling the previous one. This is because startRetryTask() needs to be + // increment before we call the RetryScheduler.schedule() method. + lock.lock(); + numScheduledAttempts++; + lock.unlock(); + } + + void completeRetryTask() { + lock.lock(); + try { + checkState(numScheduledAttempts > 0); + numScheduledAttempts--; + } finally { + lock.unlock(); } + } + + int numAttemptsSoFarWithBackoff(Backoff backoff) { + lock.lock(); + try { + if (totalAttemptNo >= config.maxTotalAttempts()) { + return -1; + } - return actualResponseTimeoutMillis; + if (lastBackoff != backoff) { + return 1; + } + + return currentAttemptNoWithLastBackoff; + } finally { + lock.unlock(); + } } - boolean timeoutForWholeRetryEnabled() { - return isTimeoutEnabled; + int acquireAttemptNoWithCurrentBackoff(Backoff backoff) { + lock.lock(); + try { + checkState(!isRetryingComplete()); + checkState((totalAttemptNo + 1) <= config.maxTotalAttempts(), + "Exceeded the maximum number of attempts: %s", config.maxTotalAttempts()); + + totalAttemptNo++; + + if (lastBackoff != backoff) { + lastBackoff = backoff; + currentAttemptNoWithLastBackoff = 1; + } + + currentAttemptNoWithLastBackoff++; + + return totalAttemptNo; + } finally { + lock.unlock(); + } } - long actualResponseTimeoutMillis() { - return TimeUnit.NANOSECONDS.toMillis(deadlineNanos - System.nanoTime()); + int numPendingAttempts() { + lock.lock(); + try { + return numPendingAttempts + (numScheduledAttempts > 0 ? 1 : 0); + } finally { + lock.unlock(); + } + } + + void completeIfNoPendingAttempts() { + lock.lock(); + + if (isRetryingComplete()) { + lock.unlock(); + return; + } + + // transfers ownership of current lock + completeIfNoPendingAttempts0(); } - int currentAttemptNoWith(Backoff backoff) { - if (totalAttemptNo++ >= config.maxTotalAttempts()) { - return -1; + // transfers ownership of current lock + private void completeIfNoPendingAttempts0() { + assert lock.isHeldByCurrentThread(); + assert !isRetryingComplete(); + + if (numPendingAttempts() > 0) { + lock.unlock(); + return; } - if (lastBackoff != backoff) { - lastBackoff = backoff; - currentAttemptNoWithLastBackoff = 1; + + // transfer ownership of current lock + complete0(); + } + + // Takes ownership of the current retry lock acquisition + private void complete0() { + assert lock.isHeldByCurrentThread(); + assert !isRetryingComplete(); + + final boolean hasLastAttempt = lastAttempt != null; + + if (!hasLastAttempt) { + // Takes ownership of the current retry lock acquisition + completeExceptionally0( + new IllegalStateException("Completed retrying without any attempts.")); + } else { + isRetryingComplete = true; + scheduler().shutdown(); + lock.unlock(); + // We do not want to call the handlers with the lock. + // Save to reference lastAttempt as this cannot change anymore. + notifyAllAttemptHandlers(lastAttempt, null); + } + } + + void completeExceptionally(Throwable cause) { + lock.lock(); + + if (isRetryingComplete()) { + lock.unlock(); + return; + } + + // Takes ownership of the current retry lock acquisition + completeExceptionally0(cause); + } + + private void notifyAllAttemptHandlers(@Nullable Attempt winningAttempt, @Nullable Throwable cause) { + assert isRetryingComplete(); + + final CompletableFuture allOtherAttemptsAbortedFuture = CompletableFuture.allOf( + startedAttempts.values().stream() + .filter(attempt -> attempt != winningAttempt) + .map(attempt -> + attempt.callHandler(false, cause) + ) + .toArray(CompletableFuture[]::new) + ); + + allOtherAttemptsAbortedFuture.handle( + (unused, throwable) -> { + if (throwable != null) { + logger.warn("Failed to notify all aborted attempt handlers.", throwable); + } + + if (winningAttempt != null) { + winningAttempt.callHandler(true, cause); + } + return null; + }); + } + + // Takes ownership of the current retry lock acquisition + private void completeExceptionally0(Throwable cause) { + assert lock.isHeldByCurrentThread(); + assert !isRetryingComplete(); + + isRetryingComplete = true; + scheduler().shutdown(); + + // We do not want to call the handlers with the lock. + lock.unlock(); + + notifyAllAttemptHandlers(null, cause); + } + + /** + * todo(szymon): [doc]. + */ + public boolean isRetryingComplete() { + try { + lock.lock(); + return isRetryingComplete; + } finally { + lock.unlock(); + } + } + + /** + * todo(szymon): [doc]. + */ + public List> startedAttempts() { + lock.lock(); + try { + return ImmutableList.copyOf(startedAttempts.values()); + } finally { + lock.unlock(); } - return currentAttemptNoWithLastBackoff++; } } } diff --git a/core/src/main/java/com/linecorp/armeria/client/retry/RetryConfig.java b/core/src/main/java/com/linecorp/armeria/client/retry/RetryConfig.java index 31f7892815d..315e35fd660 100644 --- a/core/src/main/java/com/linecorp/armeria/client/retry/RetryConfig.java +++ b/core/src/main/java/com/linecorp/armeria/client/retry/RetryConfig.java @@ -77,6 +77,7 @@ static RetryConfigBuilder builder0( private final int maxTotalAttempts; private final long responseTimeoutMillisForEachAttempt; + private final long hedgingDelayMillis; private final int maxContentLength; @Nullable @@ -90,7 +91,16 @@ static RetryConfigBuilder builder0( RetryConfig(RetryRule retryRule, int maxTotalAttempts, long responseTimeoutMillisForEachAttempt) { this(requireNonNull(retryRule, "retryRule"), null, - maxTotalAttempts, responseTimeoutMillisForEachAttempt, 0); + maxTotalAttempts, responseTimeoutMillisForEachAttempt, + 0, -1); + checkArguments(maxTotalAttempts, responseTimeoutMillisForEachAttempt); + } + + RetryConfig(RetryRule retryRule, int maxTotalAttempts, long responseTimeoutMillisForEachAttempt, + long hedgingDelayMillis) { + this(requireNonNull(retryRule, "retryRule"), null, + maxTotalAttempts, responseTimeoutMillisForEachAttempt, + 0, hedgingDelayMillis); checkArguments(maxTotalAttempts, responseTimeoutMillisForEachAttempt); } @@ -100,7 +110,19 @@ static RetryConfigBuilder builder0( int maxTotalAttempts, long responseTimeoutMillisForEachAttempt) { this(null, requireNonNull(retryRuleWithContent, "retryRuleWithContent"), - maxTotalAttempts, responseTimeoutMillisForEachAttempt, maxContentLength); + maxTotalAttempts, responseTimeoutMillisForEachAttempt, + maxContentLength, -1); + } + + RetryConfig( + RetryRuleWithContent retryRuleWithContent, + int maxContentLength, + int maxTotalAttempts, + long responseTimeoutMillisForEachAttempt, + long hedgingDelayMillis) { + this(null, requireNonNull(retryRuleWithContent, "retryRuleWithContent"), + maxTotalAttempts, responseTimeoutMillisForEachAttempt, + maxContentLength, hedgingDelayMillis); } private RetryConfig( @@ -108,13 +130,16 @@ private RetryConfig( @Nullable RetryRuleWithContent retryRuleWithContent, int maxTotalAttempts, long responseTimeoutMillisForEachAttempt, - int maxContentLength) { + int maxContentLength, + long hedgingDelayMillis + ) { checkArguments(maxTotalAttempts, responseTimeoutMillisForEachAttempt); this.retryRule = retryRule; this.retryRuleWithContent = retryRuleWithContent; this.maxTotalAttempts = maxTotalAttempts; this.responseTimeoutMillisForEachAttempt = responseTimeoutMillisForEachAttempt; this.maxContentLength = maxContentLength; + this.hedgingDelayMillis = hedgingDelayMillis; if (retryRuleWithContent == null) { fromRetryRuleWithContent = null; } else { @@ -145,9 +170,15 @@ public RetryConfigBuilder toBuilder() { assert retryRule != null; builder = builder0(retryRule); } - return builder + builder .maxTotalAttempts(maxTotalAttempts) .responseTimeoutMillisForEachAttempt(responseTimeoutMillisForEachAttempt); + + if (hedgingDelayMillis >= 0) { + builder.hedgingDelayMillis(hedgingDelayMillis); + } + + return builder; } /** @@ -166,6 +197,13 @@ public long responseTimeoutMillisForEachAttempt() { return responseTimeoutMillisForEachAttempt; } + /** + * todo(szymon): [doc]. + */ + public long hedgingDelayMillis() { + return hedgingDelayMillis; + } + /** * Returns the {@link RetryRule} which was specified with {@link RetryConfig#builder(RetryRule)}. */ diff --git a/core/src/main/java/com/linecorp/armeria/client/retry/RetryConfigBuilder.java b/core/src/main/java/com/linecorp/armeria/client/retry/RetryConfigBuilder.java index cbda418e44f..1208fea8def 100644 --- a/core/src/main/java/com/linecorp/armeria/client/retry/RetryConfigBuilder.java +++ b/core/src/main/java/com/linecorp/armeria/client/retry/RetryConfigBuilder.java @@ -36,6 +36,7 @@ public final class RetryConfigBuilder { private int maxTotalAttempts = Flags.defaultMaxTotalAttempts(); private long responseTimeoutMillisForEachAttempt = Flags.defaultResponseTimeoutMillis(); + private long hedgingDelayMillis = -1; private int maxContentLength; @Nullable @@ -84,6 +85,33 @@ public RetryConfigBuilder maxTotalAttempts(int maxTotalAttempts) { return this; } + /** + * todo(szymon): [doc]. + */ + public RetryConfigBuilder hedgingDelay(Duration hedgingDelay) { + final long millis = + requireNonNull(hedgingDelay, "hedgingDelay") + .toMillis(); + checkArgument( + millis >= 0, + "hedgingDelay.toMillis(): %s (expected: >= 0)", + millis); + hedgingDelayMillis = millis; + return this; + } + + /** + * todo(szymon): [doc]. + */ + public RetryConfigBuilder hedgingDelayMillis(long hedgingDelayMillis) { + checkArgument( + hedgingDelayMillis >= 0, + "hedgingDelayMillis: %s (expected: >= 0)", + hedgingDelayMillis); + this.hedgingDelayMillis = hedgingDelayMillis; + return this; + } + /** * Sets the specified {@code responseTimeoutMillisForEachAttempt}. */ @@ -116,14 +144,22 @@ public RetryConfigBuilder responseTimeoutForEachAttempt(Duration responseTime */ public RetryConfig build() { if (retryRule != null) { - return new RetryConfig<>(retryRule, maxTotalAttempts, responseTimeoutMillisForEachAttempt); + if (hedgingDelayMillis >= 0) { + return new RetryConfig<>(retryRule, maxTotalAttempts, responseTimeoutMillisForEachAttempt, + hedgingDelayMillis); + } else { + return new RetryConfig<>(retryRule, maxTotalAttempts, responseTimeoutMillisForEachAttempt); + } } assert retryRuleWithContent != null; - return new RetryConfig<>( - retryRuleWithContent, - maxContentLength, - maxTotalAttempts, - responseTimeoutMillisForEachAttempt); + + if (hedgingDelayMillis >= 0) { + return new RetryConfig<>(retryRuleWithContent, maxContentLength, maxTotalAttempts, + responseTimeoutMillisForEachAttempt, hedgingDelayMillis); + } else { + return new RetryConfig<>(retryRuleWithContent, maxContentLength, maxTotalAttempts, + responseTimeoutMillisForEachAttempt); + } } @Override @@ -139,6 +175,7 @@ ToStringHelper toStringHelper() { .add("retryRuleWithContent", retryRuleWithContent) .add("maxTotalAttempts", maxTotalAttempts) .add("responseTimeoutMillisForEachAttempt", responseTimeoutMillisForEachAttempt) - .add("maxContentLength", maxContentLength); + .add("maxContentLength", maxContentLength) + .add("hedgingDelayMillis", hedgingDelayMillis); } } 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..412e428b575 --- /dev/null +++ b/core/src/main/java/com/linecorp/armeria/client/retry/RetryScheduler.java @@ -0,0 +1,394 @@ +/* + * 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 java.util.Objects.requireNonNull; + +import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Consumer; + +import com.linecorp.armeria.client.retry.RetrySchedulingException.Type; +import com.linecorp.armeria.common.annotation.Nullable; + +import io.netty.channel.EventLoop; +import io.netty.util.concurrent.ScheduledFuture; + +class RetryScheduler { + private static class RetryTaskHandle { + enum CancellationReason { + OVERTAKEN, + RESCHEDULED, + SHUTDOWN + } + + private final Consumer retryTask; + private final ScheduledFuture scheduledFuture; + private @Nullable CancellationReason cancellationReason; + private final long retryTaskId; + + private final Consumer<@Nullable ? super Throwable> exceptionHandler; + + private final long retryTimeNanos; + + RetryTaskHandle(long retryTaskId, ScheduledFuture scheduledFuture, + Consumer retryTask, + long retryTimeNanos, + Consumer<@Nullable ? super Throwable> exceptionHandler) { + this.retryTaskId = retryTaskId; + this.retryTask = retryTask; + this.scheduledFuture = scheduledFuture; + this.exceptionHandler = exceptionHandler; + this.retryTimeNanos = retryTimeNanos; + cancellationReason = null; + } + + Consumer getRetryTaskRunnable() { + return retryTask; + } + + long id() { + return retryTaskId; + } + + long retryTimeNanos() { + return retryTimeNanos; + } + + void setCancellationReason(CancellationReason cancellationReason) { + checkState(this.cancellationReason == null, "cancellationReason"); + this.cancellationReason = cancellationReason; + } + + @Nullable + CancellationReason getCancellationReason() { + return cancellationReason; + } + + Consumer getExceptionHandler() { + return exceptionHandler; + } + + ScheduledFuture getFuture() { + return scheduledFuture; + } + } + + // todo: should we make this a flag? + // Number of nanoseconds that we allow the retry task to be scheduled earlier than + // the earliestNextRetryTimeNanos. + // This should avoid unnecessary rescheduling. + private static final long RESCHEDULING_OVERTAKING_TOLERANCE_NANOS = TimeUnit.MICROSECONDS.toNanos(500); + + // for locking + private final ReentrantLock retryLock; + private final EventLoop eventLoop; + private final long latestNextRetryTimeNanos; + + private long earliestNextRetryTimeNanos; + private long retryTaskId; + + // The retry task that is about to be executed next. + // It is possible that the delay of this task is shorter than earliestNextRetryTimeNanos, + // because of calls to addEarliestNextRetryTimeNanos with a pending call to schedule or + // rescheduleCurrentRetryTaskIfTooEarly. + private @Nullable RetryTaskHandle currentRetryTask; + + RetryScheduler(ReentrantLock retryLock, EventLoop eventLoop) { + this(retryLock, eventLoop, Long.MAX_VALUE); + } + + RetryScheduler(ReentrantLock retryLock, EventLoop eventLoop, long latestNextRetryTimeNanos) { + this.retryLock = requireNonNull(retryLock, "retryLock"); + this.eventLoop = requireNonNull(eventLoop, "eventLoop"); + currentRetryTask = null; + earliestNextRetryTimeNanos = Long.MIN_VALUE; + this.latestNextRetryTimeNanos = latestNextRetryTimeNanos; + retryTaskId = 0; + } + + public void schedule(Consumer retryTask, long nextRetryTimeNanos, + long earliestNextRetryTimeNanos, + Consumer exceptionHandler) { + requireNonNull(retryTask, "retryTask"); + checkArgument(nextRetryTimeNanos >= earliestNextRetryTimeNanos, + "nextRetryTimeNanos: %s (expected: >= %s)", + nextRetryTimeNanos, earliestNextRetryTimeNanos); + requireNonNull(exceptionHandler, "exceptionHandler"); + + retryLock.lock(); + + addEarliestNextRetryTimeNanos(earliestNextRetryTimeNanos); + nextRetryTimeNanos = Math.max(nextRetryTimeNanos, this.earliestNextRetryTimeNanos); + + // Math.max because we could have added an earliestNextRetryTimeNanos that is later than the + // currently scheduled retry task (even not by us in this method). + if (currentRetryTask == null || nextRetryTimeNanos < Math.max(this.earliestNextRetryTimeNanos, + currentRetryTask.retryTimeNanos())) { + // Takes ownership of the current retry lock acquisition + scheduleNextRetryTask(retryTask, nextRetryTimeNanos, exceptionHandler, false); + } else { + // Make sure the current retry task is not scheduled too early. + // Takes ownership of the current retry lock acquisition + rescheduleCurrentRetryTaskIfTooEarly0(); + + exceptionHandler.accept(new RetrySchedulingException( + RetrySchedulingException.Type.RETRY_TASK_OVERTAKEN)); + } + } + + private boolean cancelCurrentRetryTask(RetryTaskHandle.CancellationReason cancellationReason) { + retryLock.lock(); + + try { + if (currentRetryTask != null) { + final ScheduledFuture retryTaskFuture = currentRetryTask.getFuture(); + + currentRetryTask.setCancellationReason(cancellationReason); + // This should not invoke the user-defined exception handler with our lock as + // we marked the task accordingly. + if (!retryTaskFuture.cancel(false)) { + return false; + } else { + clearCurrentRetryTask(); + return true; + } + } + + return true; + } finally { + retryLock.unlock(); + } + } + + private void handleRetryTaskCompletion(RetryTaskHandle retryTaskHandle) { + retryLock.lock(); + + try { + final ScheduledFuture retryTaskFuture = retryTaskHandle.getFuture(); + assert retryTaskFuture.isDone(); + + if (currentRetryTask == retryTaskHandle) { + clearCurrentRetryTask(); + } + + if (retryTaskFuture.isCancelled()) { + final RetryTaskHandle.CancellationReason cancellationReason = + retryTaskHandle.getCancellationReason(); + + if (cancellationReason == RetryTaskHandle.CancellationReason.RESCHEDULED) { + return; + } + + if (cancellationReason == RetryTaskHandle.CancellationReason.OVERTAKEN) { + retryTaskHandle.getExceptionHandler().accept( + new RetrySchedulingException(Type.RETRY_TASK_OVERTAKEN) + ); + return; + } + + if (cancellationReason == RetryTaskHandle.CancellationReason.SHUTDOWN) { + retryTaskHandle.getExceptionHandler().accept( + new RetrySchedulingException(Type.RETRYING_ALREADY_COMPLETED) + ); + return; + } + + assert cancellationReason == null; + + // The retry task was cancelled by the user, not by the scheduler. + retryTaskHandle.getExceptionHandler().accept( + new RetrySchedulingException(Type.RETRY_TASK_CANCELLED)); + return; + } + + if (!retryTaskFuture.isSuccess()) { + retryTaskHandle.getExceptionHandler().accept(retryTaskFuture.cause()); + } + } finally { + retryLock.unlock(); + } + } + + private void clearCurrentRetryTask() { + retryLock.lock(); + try { + currentRetryTask = null; + earliestNextRetryTimeNanos = Long.MIN_VALUE; + } finally { + retryLock.unlock(); + } + } + + // Takes ownership of the current retry lock acquisition + private void scheduleNextRetryTask(Consumer retryRunnable, long retryTimeNanos, + Consumer exceptionHandler, + boolean isReschedule) { + + assert retryLock.isHeldByCurrentThread(); + assert earliestNextRetryTimeNanos <= retryTimeNanos; + + if (!cancelCurrentRetryTask( + isReschedule ? + RetryTaskHandle.CancellationReason.RESCHEDULED + : RetryTaskHandle.CancellationReason.OVERTAKEN + )) { + retryLock.unlock(); + exceptionHandler.accept(new IllegalStateException("Current retry task could not be cancelled.")); + return; + } + + assert currentRetryTask == null; + + try { + final long nowNanos = System.nanoTime(); + final long delayNanos; + if (retryTimeNanos <= nowNanos) { + delayNanos = 0; + } else { + delayNanos = retryTimeNanos - nowNanos; + } + + final long thisRetryId = retryTaskId++; + + final Runnable wrappedRetryRunnable = () -> { + retryLock.lock(); + assert retryLock.getHoldCount() == 1; + + // Let be defensive here. + if (currentRetryTask == null) { + clearCurrentRetryTask(); + retryLock.unlock(); + exceptionHandler.accept(new RetrySchedulingException(Type.RETRY_TASK_CANCELLED)); + return; + } + + if (currentRetryTask.id() != thisRetryId) { + retryLock.unlock(); + exceptionHandler.accept(new RetrySchedulingException(Type.RETRY_TASK_CANCELLED)); + return; + } + + final long taskRunTimeNanos = System.nanoTime(); + + // Math.max to be robust against overflows. + if (Math.max(taskRunTimeNanos, taskRunTimeNanos + RESCHEDULING_OVERTAKING_TOLERANCE_NANOS) < + earliestNextRetryTimeNanos) { + + final Consumer currentRetryRunnable = + currentRetryTask.getRetryTaskRunnable(); + final Consumer currentExceptionHandler = + currentRetryTask.getExceptionHandler(); + + clearCurrentRetryTask(); + // Do not invoke rescheduleCurrentRetryTaskIfTooEarly here as it will not be able + // to cancel us. + // Takes ownership of the current retry lock acquisition + scheduleNextRetryTask(currentRetryRunnable, + earliestNextRetryTimeNanos, + currentExceptionHandler, false); + return; + } + + clearCurrentRetryTask(); + + assert retryLock.isHeldByCurrentThread(); + assert retryLock.getHoldCount() == 1; + + // Run this with the lock. Expected to release the retry lock after consuming an attempt. + retryRunnable.accept(retryLock); + }; + + //noinspection unchecked + final ScheduledFuture nextRetryTaskFuture = + (ScheduledFuture) eventLoop.schedule(wrappedRetryRunnable, delayNanos, + TimeUnit.NANOSECONDS); + + // We are passing in the original to avoid multiple wrapping in case of the retry task being + // rescheduled multiple times. + final RetryTaskHandle nextRetryTask = new RetryTaskHandle(thisRetryId, nextRetryTaskFuture, + retryRunnable, + retryTimeNanos, + exceptionHandler); + + nextRetryTaskFuture.addListener(f -> handleRetryTaskCompletion(nextRetryTask)); + currentRetryTask = nextRetryTask; + retryLock.unlock(); + } catch (Throwable t) { + retryLock.unlock(); + exceptionHandler.accept(t); + } + } + + public long addEarliestNextRetryTimeNanos(long earliestNextRetryTimeNanos) { + retryLock.lock(); + + try { + checkState(earliestNextRetryTimeNanos <= latestNextRetryTimeNanos); + this.earliestNextRetryTimeNanos = Math.max(this.earliestNextRetryTimeNanos, + earliestNextRetryTimeNanos); + + return this.earliestNextRetryTimeNanos; + } finally { + retryLock.unlock(); + } + } + + // Takes ownership of the current retry lock acquisition (if any) + public void rescheduleCurrentRetryTaskIfTooEarly() { + retryLock.lock(); + // Takes ownership of the current retry lock acquisition + rescheduleCurrentRetryTaskIfTooEarly0(); + } + + // Takes ownership of the current retry lock acquisition + private void rescheduleCurrentRetryTaskIfTooEarly0() { + assert retryLock.isHeldByCurrentThread(); + + if (currentRetryTask == null) { + retryLock.unlock(); + return; + } + + if ( + Math.max(currentRetryTask.retryTimeNanos(), + currentRetryTask.retryTimeNanos() + RESCHEDULING_OVERTAKING_TOLERANCE_NANOS) < + earliestNextRetryTimeNanos + ) { + // The current retry task is going to be executed before earliestNextRetryTimeNanos so + // we need to reschedule it. + + // Takes ownership of the current retry lock acquisition + scheduleNextRetryTask(currentRetryTask.getRetryTaskRunnable(), + earliestNextRetryTimeNanos, + currentRetryTask.getExceptionHandler(), true); + } else { + retryLock.unlock(); + } + } + + public boolean shutdown() { + retryLock.lock(); + try { + return cancelCurrentRetryTask(RetryTaskHandle.CancellationReason.SHUTDOWN); + } finally { + retryLock.unlock(); + } + } +} diff --git a/core/src/main/java/com/linecorp/armeria/client/retry/RetrySchedulingException.java b/core/src/main/java/com/linecorp/armeria/client/retry/RetrySchedulingException.java new file mode 100644 index 00000000000..8dc0f0c6be8 --- /dev/null +++ b/core/src/main/java/com/linecorp/armeria/client/retry/RetrySchedulingException.java @@ -0,0 +1,54 @@ +/* + * 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; + +class RetrySchedulingException extends RuntimeException { + private static final long serialVersionUID = 1L; + private final Type type; + + enum Type { + RETRYING_ALREADY_COMPLETED("Retrying completed already"), + NO_MORE_ATTEMPTS_IN_RETRY("No more attempts with respect to `RetryConfig.maxAttempts`"), + NO_MORE_ATTEMPTS_IN_BACKOFF("No more attempts available in backoff"), + DELAY_FROM_BACKOFF_EXCEEDS_RESPONSE_TIMEOUT("Delay from backoff exceeds response timeout"), + DELAY_FROM_SERVER_EXCEEDS_RESPONSE_TIMEOUT("Delay from server exceeds response timeout"), + RETRY_TASK_OVERTAKEN( + "Retry task cannot run because of another retry task for the same attempt being" + + " scheduled earlier"), + RETRY_TASK_CANCELLED("A retry task was cancelled unexpectedly"); + + private final String message; + + Type(String message) { + this.message = message; + } + + public String getMessage() { + return message; + } + } + + RetrySchedulingException(Type type) { + super(type.getMessage()); + this.type = type; + } + + Type getType() { + return type; + } +} + 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..7f0e7fb4255 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 @@ -29,6 +29,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import com.linecorp.armeria.client.ClientFactory; import com.linecorp.armeria.client.ClientRequestContext; import com.linecorp.armeria.client.HttpClient; import com.linecorp.armeria.client.ResponseTimeoutException; @@ -246,32 +247,38 @@ protected HttpResponse doExecute(ClientRequestContext ctx, HttpRequest req) thro 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); + doExecute0(new RetryingContext(mappedRetryConfig(ctx), ctx, reqDuplicator, req, res, + responseFuture), 1); } else { req.aggregate(AggregationOptions.usePooledObjects(ctx.alloc(), ctx.eventLoop())) .handle((agg, cause) -> { if (cause != null) { - handleException(ctx, null, responseFuture, cause, true); + completeRetryingExceptionally(ctx, null, responseFuture, cause, true); } else { final HttpRequestDuplicator reqDuplicator = new AggregatedHttpRequestDuplicator(agg); - doExecute0(ctx, reqDuplicator, req, res, responseFuture); + doExecute0(new RetryingContext(mappedRetryConfig(ctx), ctx, reqDuplicator, req, res, + responseFuture), 1); } 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, + private void doExecute0(RetryingContext retryingContext, int attemptNo) { + final RetryConfig config = retryingContext.config(); + final ClientRequestContext ctx = retryingContext.ctx(); + final HttpRequestDuplicator rootReqDuplicator = retryingContext.reqDuplicator(); + final HttpRequest originalReq = retryingContext.req(); + final HttpResponse returnedRes = retryingContext.res(); + + final boolean initialAttempt = attemptNo <= 1; + // The request or attemptRes has been aborted by the client before it receives a attemptRes, // so stop retrying. if (originalReq.whenComplete().isCompletedExceptionally()) { originalReq.whenComplete().handle((unused, cause) -> { - handleException(ctx, rootReqDuplicator, future, cause, initialAttempt); + completeRetryingExceptionally(retryingContext, cause, initialAttempt); return null; }); return; @@ -284,212 +291,284 @@ private void doExecute0(ClientRequestContext ctx, HttpRequestDuplicator rootReqD } else { abortCause = AbortedStreamException.get(); } - handleException(ctx, rootReqDuplicator, future, abortCause, initialAttempt); + completeRetryingExceptionally(retryingContext, abortCause, initialAttempt); return null; }); return; } if (!setResponseTimeout(ctx)) { - handleException(ctx, rootReqDuplicator, future, ResponseTimeoutException.get(), initialAttempt); + completeRetryingExceptionally(retryingContext, ResponseTimeoutException.get(), initialAttempt); return; } - final HttpRequest duplicateReq; + final HttpRequest attemptReq; if (initialAttempt) { - duplicateReq = rootReqDuplicator.duplicate(); + attemptReq = rootReqDuplicator.duplicate(); } else { final RequestHeadersBuilder newHeaders = originalReq.headers().toBuilder(); - newHeaders.setInt(ARMERIA_RETRY_COUNT, totalAttempts - 1); - duplicateReq = rootReqDuplicator.duplicate(newHeaders.build()); + newHeaders.setInt(ARMERIA_RETRY_COUNT, attemptNo - 1); + attemptReq = rootReqDuplicator.duplicate(newHeaders.build()); } - final ClientRequestContext derivedCtx; + final ClientRequestContext attemptCtx; try { - derivedCtx = newDerivedContext(ctx, duplicateReq, ctx.rpcRequest(), initialAttempt); + attemptCtx = newDerivedContext(ctx, attemptReq, ctx.rpcRequest(), initialAttempt); } catch (Throwable t) { - handleException(ctx, rootReqDuplicator, future, t, initialAttempt); + completeRetryingExceptionally(retryingContext, 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) { + final HttpRequest attemptCtxReq = attemptCtx.request(); + assert attemptCtxReq != null; + final HttpResponse attemptRes; + final ClientRequestContextExtension ctxExtension = attemptCtx.as(ClientRequestContextExtension.class); + if (!initialAttempt && ctxExtension != null && attemptCtx.endpoint() == null) { // clear the pending throwable to retry endpoint selection - ClientPendingThrowableUtil.removePendingThrowable(derivedCtx); + ClientPendingThrowableUtil.removePendingThrowable(attemptCtx); // if the endpoint hasn't been selected, try to initialize the ctx with a new endpoint/event loop - response = initContextAndExecuteWithFallback( + attemptRes = initContextAndExecuteWithFallback( unwrap(), ctxExtension, HttpResponse::of, - (context, cause) -> HttpResponse.ofFailure(cause), ctxReq, false); + (context, cause) -> HttpResponse.ofFailure(cause), attemptCtxReq, false); } else { - response = executeWithFallback(unwrap(), derivedCtx, - (context, cause) -> HttpResponse.ofFailure(cause), ctxReq, false); + attemptRes = executeWithFallback(unwrap(), attemptCtx, + (context, cause) -> HttpResponse.ofFailure(cause), attemptCtxReq, + false); } - final RetryConfig config = mappedRetryConfig(ctx); + startRetryAttempt(ctx, attemptCtx, attemptRes, (winningAttemptCtx, + winningAttemptRes) -> { + assert attemptCtx == winningAttemptCtx; + + winningAttemptCtx.eventLoop().execute(() -> { + retryingContext.reqDuplicator().close(); + retryingContext + .ctx() + .logBuilder() + .endResponseWithChild(winningAttemptCtx.log()); + retryingContext.resFuture().complete(winningAttemptRes); + }); + }, + (abortingAttemptCtx, + abortingAttemptRes, + cause + ) -> { + assert attemptCtx == abortingAttemptCtx; + assert attemptRes == abortingAttemptRes; + + final CompletableFuture abortCompletedFuture = new CompletableFuture<>(); + + attemptCtx.eventLoop().execute(() -> { + final RequestLogBuilder logBuilder = abortingAttemptCtx.logBuilder(); + logBuilder.responseContent(null, null); + logBuilder.responseContentPreview(null); + + if (cause != null) { + attemptCtx.cancel(cause); + abortingAttemptRes.abort(cause); + } else { + attemptCtx.cancel(); + abortingAttemptRes.abort(); + } + + abortingAttemptRes.whenComplete().handle((unused, unusedCause) -> { + abortCompletedFuture.complete(null); + return null; + }); + }); + + return abortCompletedFuture; + } + ); + if (!ctx.exchangeType().isResponseStreaming() || config.requiresResponseTrailers()) { - response.aggregate().handle((aggregated, cause) -> { + attemptRes.aggregate().handle((attemptAggResponse, cause) -> { if (cause != null) { - derivedCtx.logBuilder().endRequest(cause); - derivedCtx.logBuilder().endResponse(cause); - handleResponseWithoutContent(config, ctx, rootReqDuplicator, originalReq, returnedRes, - future, derivedCtx, HttpResponse.ofFailure(cause), cause); + attemptCtx.logBuilder().endRequest(cause); + attemptCtx.logBuilder().endResponse(cause); + + handleResponseWithoutContent(retryingContext, attemptCtx, + 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); + completeAttemptLogIfBytesNotTransferred(attemptCtx, attemptAggResponse); + + attemptCtx.log().whenAvailable(RequestLogProperty.RESPONSE_END_TIME).thenRun(() -> { + handleAggregatedResponse(retryingContext, attemptCtx, attemptAggResponse); }); } return null; }); } else { - handleStreamingResponse(config, ctx, rootReqDuplicator, originalReq, returnedRes, - future, derivedCtx, response); + handleStreamingResponse(retryingContext, attemptCtx, attemptRes); + } + + @Nullable + final long hedgingDelayMillis = config.hedgingDelayMillis(); + + if (hedgingDelayMillis >= 0) { + final Backoff hedgingDelayBackoff = Backoff.fixed(hedgingDelayMillis); + scheduleNextRetry(ctx, hedgingAttemptNo -> doExecute0(retryingContext, hedgingAttemptNo), + hedgingDelayBackoff, + cause -> handleExceptionAfterScheduling(retryingContext, cause)); + } + } + + private static void handleExceptionAfterScheduling( + RetryingContext retryingContext, Throwable cause) { + if (cause instanceof RetrySchedulingException) { + switch (((RetrySchedulingException) cause).getType()) { + case RETRYING_ALREADY_COMPLETED: + return; + case NO_MORE_ATTEMPTS_IN_RETRY: + case NO_MORE_ATTEMPTS_IN_BACKOFF: + case DELAY_FROM_BACKOFF_EXCEEDS_RESPONSE_TIMEOUT: + case DELAY_FROM_SERVER_EXCEEDS_RESPONSE_TIMEOUT: + case RETRY_TASK_OVERTAKEN: + completeRetryingIfNoPendingAttempts(retryingContext.ctx()); + return; + case RETRY_TASK_CANCELLED: + cause = new IllegalStateException( + ClientFactory.class.getSimpleName() + " has been closed.", cause); + break; + } } + completeRetryingExceptionally(retryingContext, cause, false); } - // 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); + private void handleResponseWithoutContent(RetryingContext retryingContext, + ClientRequestContext attemptCtx, HttpResponse attemptRes, + @Nullable Throwable attemptResCause) { + if (attemptResCause != null) { + attemptResCause = Exceptions.peel(attemptResCause); } + try { - final RetryRule retryRule = retryRule(config); - final CompletionStage f = retryRule.shouldRetry(derivedCtx, responseCause); + final RetryRule retryRule = retryRule(retryingContext.config()); + final CompletionStage f = retryRule.shouldRetry(attemptCtx, attemptResCause); f.handle((decision, shouldRetryCause) -> { warnIfExceptionIsRaised(retryRule, shouldRetryCause); - handleRetryDecision(decision, ctx, derivedCtx, rootReqDuplicator, - originalReq, returnedRes, future, response); + handleRetryDecision(retryingContext, decision, attemptCtx, attemptRes); return null; }); } catch (Throwable cause) { - response.abort(); - handleException(ctx, rootReqDuplicator, future, cause, false); + completeRetryingExceptionally(retryingContext, 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) -> { + private void handleStreamingResponse(RetryingContext retryingContext, + ClientRequestContext attemptCtx, + HttpResponse attemptRes) { + final SplitHttpResponse attemptSplitRes = attemptRes.split(); + attemptSplitRes.headers().handle((headers, headersCause) -> { final Throwable responseCause; if (headersCause == null) { - final RequestLog log = derivedCtx.log().getIfAvailable(RequestLogProperty.RESPONSE_CAUSE); + final RequestLog log = attemptCtx.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()); + + completeAttemptLogIfBytesNotTransferred(attemptCtx, attemptRes, headers, responseCause); + + attemptCtx.log().whenAvailable(RequestLogProperty.RESPONSE_HEADERS).thenRun(() -> { + if (retryingContext.config().needsContentInRule() && responseCause == null) { + final HttpResponse attemptUnsplitRes = attemptSplitRes.unsplit(); + final HttpResponseDuplicator attemptResDuplicator = + attemptUnsplitRes.toDuplicator(attemptCtx.eventLoop().withoutContext(), + attemptCtx.maxResponseLength()); try { - final TruncatingHttpResponse truncatingHttpResponse = - new TruncatingHttpResponse(duplicator.duplicate(), - retryConfig.maxContentLength()); - final HttpResponse duplicated = duplicator.duplicate(); - duplicator.close(); + final TruncatingHttpResponse attemptTruncatedRes = + new TruncatingHttpResponse(attemptResDuplicator.duplicate(), + retryingContext.config().maxContentLength()); + final HttpResponse attemptDuplicatedRes = attemptResDuplicator.duplicate(); + attemptResDuplicator.close(); final RetryRuleWithContent ruleWithContent = - retryConfig.retryRuleWithContent(); + retryingContext.config().retryRuleWithContent(); assert ruleWithContent != null; - ruleWithContent.shouldRetry(derivedCtx, truncatingHttpResponse, null) + ruleWithContent.shouldRetry(attemptCtx, attemptTruncatedRes, null) .handle((decision, cause) -> { warnIfExceptionIsRaised(ruleWithContent, cause); - truncatingHttpResponse.abort(); - handleRetryDecision(decision, ctx, derivedCtx, rootReqDuplicator, - originalReq, returnedRes, future, duplicated); + attemptTruncatedRes.abort(); + + handleRetryDecision(retryingContext, decision, attemptCtx, + attemptDuplicatedRes); return null; }); } catch (Throwable cause) { - duplicator.abort(cause); - handleException(ctx, rootReqDuplicator, future, cause, false); + attemptResDuplicator.abort(cause); + completeRetryingExceptionally(retryingContext, cause, false); } } else { - final HttpResponse response0; + final HttpResponse attemptUnsplitRes; if (responseCause != null) { - splitResponse.body().abort(responseCause); - response0 = HttpResponse.ofFailure(responseCause); + attemptSplitRes.body().abort(responseCause); + attemptUnsplitRes = HttpResponse.ofFailure(responseCause); } else { - response0 = splitResponse.unsplit(); + attemptUnsplitRes = attemptSplitRes.unsplit(); } - handleResponseWithoutContent(retryConfig, ctx, rootReqDuplicator, originalReq, returnedRes, - future, derivedCtx, response0, responseCause); + handleResponseWithoutContent(retryingContext, attemptCtx, attemptUnsplitRes, responseCause); } }); return null; }); } - 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(); + private void handleAggregatedResponse(RetryingContext retryingContext, + ClientRequestContext attemptCtx, + AggregatedHttpResponse attemptAggRes) { + if (retryingContext.config().needsContentInRule()) { + final RetryRuleWithContent ruleWithContent = + retryingContext.config().retryRuleWithContent(); assert ruleWithContent != null; try { - ruleWithContent.shouldRetry(derivedCtx, aggregatedRes.toHttpResponse(), null) + ruleWithContent.shouldRetry(attemptCtx, attemptAggRes.toHttpResponse(), null) .handle((decision, cause) -> { warnIfExceptionIsRaised(ruleWithContent, cause); - handleRetryDecision( - decision, ctx, derivedCtx, rootReqDuplicator, originalReq, - returnedRes, future, aggregatedRes.toHttpResponse()); + + handleRetryDecision(retryingContext, + decision, attemptCtx, attemptAggRes.toHttpResponse()); return null; }); } catch (Throwable cause) { - handleException(ctx, rootReqDuplicator, future, cause, false); + completeRetryingExceptionally(retryingContext, cause, false); } return; } - handleResponseWithoutContent(retryConfig, ctx, rootReqDuplicator, originalReq, returnedRes, - future, derivedCtx, aggregatedRes.toHttpResponse(), null); + + handleResponseWithoutContent(retryingContext, attemptCtx, attemptAggRes.toHttpResponse(), + null); } - private static void completeLogIfBytesNotTransferred(AggregatedHttpResponse response, - ClientRequestContext ctx) { - if (!ctx.log().isAvailable(RequestLogProperty.REQUEST_FIRST_BYTES_TRANSFERRED_TIME)) { - final RequestLogBuilder logBuilder = ctx.logBuilder(); + private static void completeAttemptLogIfBytesNotTransferred(ClientRequestContext attemptCtx, + AggregatedHttpResponse attemptAggRes) { + if (!attemptCtx.log().isAvailable(RequestLogProperty.REQUEST_FIRST_BYTES_TRANSFERRED_TIME)) { + final RequestLogBuilder logBuilder = attemptCtx.logBuilder(); logBuilder.endRequest(); - logBuilder.responseHeaders(response.headers()); - if (!response.trailers().isEmpty()) { - logBuilder.responseTrailers(response.trailers()); + logBuilder.responseHeaders(attemptAggRes.headers()); + if (!attemptAggRes.trailers().isEmpty()) { + logBuilder.responseTrailers(attemptAggRes.trailers()); } logBuilder.endResponse(); } } - 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); + private static void completeAttemptLogIfBytesNotTransferred( + ClientRequestContext attemptCtx, HttpResponse attemptRes, + @Nullable ResponseHeaders attemptResHeaders, + @Nullable Throwable attemptResCause) { + if (!attemptCtx.log().isAvailable(RequestLogProperty.REQUEST_FIRST_BYTES_TRANSFERRED_TIME)) { + final RequestLogBuilder logBuilder = attemptCtx.logBuilder(); + if (attemptResCause != null) { + logBuilder.endRequest(attemptResCause); + logBuilder.endResponse(attemptResCause); } else { logBuilder.endRequest(); - if (headers != null) { - logBuilder.responseHeaders(headers); + if (attemptResHeaders != null) { + logBuilder.responseHeaders(attemptResHeaders); } - response.whenComplete().handle((unused, cause) -> { + attemptRes.whenComplete().handle((unused, cause) -> { if (cause != null) { logBuilder.endResponse(cause); } else { @@ -507,48 +586,54 @@ private static void warnIfExceptionIsRaised(Object retryRule, @Nullable Throwabl } } - private static void handleException(ClientRequestContext ctx, - @Nullable HttpRequestDuplicator rootReqDuplicator, - CompletableFuture future, Throwable cause, - boolean endRequestLog) { - future.completeExceptionally(cause); + private static void completeRetryingExceptionally(RetryingContext retryingContext, Throwable cause, + boolean endRequestLog) { + completeRetryingExceptionally( + retryingContext.ctx(), retryingContext.reqDuplicator(), retryingContext.resFuture(), + cause, endRequestLog); + } + + private static void completeRetryingExceptionally(ClientRequestContext ctx, + @Nullable HttpRequestDuplicator rootReqDuplicator, + CompletableFuture returnedResFuture, + Throwable cause, + boolean endRequestLog) { + if (isRetryingComplete(ctx)) { + return; + } + if (rootReqDuplicator != null) { rootReqDuplicator.abort(cause); } + if (endRequestLog) { ctx.logBuilder().endRequest(cause); } + ctx.logBuilder().endResponse(cause); + returnedResFuture.completeExceptionally(cause); + + completeRetryingExceptionally(ctx, cause); } - private void handleRetryDecision(@Nullable RetryDecision decision, ClientRequestContext ctx, - ClientRequestContext derivedCtx, HttpRequestDuplicator rootReqDuplicator, - HttpRequest originalReq, HttpResponse returnedRes, - CompletableFuture future, HttpResponse originalRes) { + private void handleRetryDecision(RetryingContext retryingContext, @Nullable RetryDecision decision, + ClientRequestContext attemptCtx, HttpResponse attemptRes) { + if (isRetryingComplete(retryingContext.ctx())) { + return; + } + 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; - } + final long millisAfter = useRetryAfter ? getRetryAfterMillis(attemptCtx) : -1; + scheduleNextRetry(retryingContext.ctx(), + attemptNo -> doExecute0(retryingContext, attemptNo), + backoff, + millisAfter, + cause -> handleExceptionAfterScheduling(retryingContext, cause)); } - 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(); + completeRetryAttempt(retryingContext.ctx(), attemptCtx, attemptRes, backoff == null); } private static long getRetryAfterMillis(ClientRequestContext ctx) { @@ -591,4 +676,49 @@ private static RetryRule retryRule(RetryConfig retryConfig) { return rule; } } + + private static class RetryingContext { + private final ClientRequestContext ctx; + private final HttpRequestDuplicator reqDuplicator; + private final HttpRequest req; + private final HttpResponse res; + private final CompletableFuture resFuture; + private final RetryConfig retryConfig; + + RetryingContext(RetryConfig retryConfig, ClientRequestContext ctx, + HttpRequestDuplicator reqDuplicator, + HttpRequest req, HttpResponse res, + CompletableFuture resFuture) { + this.retryConfig = retryConfig; + this.ctx = ctx; + this.reqDuplicator = reqDuplicator; + this.req = req; + this.res = res; + this.resFuture = resFuture; + } + + RetryConfig config() { + return retryConfig; + } + + ClientRequestContext ctx() { + return ctx; + } + + HttpRequestDuplicator reqDuplicator() { + return reqDuplicator; + } + + HttpRequest req() { + return req; + } + + HttpResponse res() { + return res; + } + + CompletableFuture resFuture() { + return 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..8ce4622f8ed 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 @@ -22,6 +22,7 @@ import java.util.concurrent.CompletableFuture; import java.util.function.Function; +import com.linecorp.armeria.client.ClientFactory; import com.linecorp.armeria.client.ClientRequestContext; import com.linecorp.armeria.client.ResponseTimeoutException; import com.linecorp.armeria.client.RpcClient; @@ -142,98 +143,159 @@ 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); + final CompletableFuture returnedResFuture = new CompletableFuture<>(); + final RpcResponse res = RpcResponse.from(returnedResFuture); + doExecute0(ctx, req, res, returnedResFuture, 1); return res; } private void doExecute0(ClientRequestContext ctx, RpcRequest req, - RpcResponse returnedRes, CompletableFuture future) { - final int totalAttempts = getTotalAttempts(ctx); - final boolean initialAttempt = totalAttempts <= 1; + RpcResponse returnedRes, CompletableFuture returnedResFuture, + int attemptNo) { + final boolean initialAttempt = attemptNo <= 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( + completeRetryingExceptionally(ctx, returnedResFuture, new CancellationException( "the response returned to the client has been cancelled"), initialAttempt); return; } if (!setResponseTimeout(ctx)) { - handleException(ctx, future, ResponseTimeoutException.get(), initialAttempt); + completeRetryingExceptionally(ctx, returnedResFuture, ResponseTimeoutException.get(), + initialAttempt); return; } - final ClientRequestContext derivedCtx = newDerivedContext(ctx, null, req, initialAttempt); + final ClientRequestContext attemptCtx = newDerivedContext(ctx, null, req, initialAttempt); if (!initialAttempt) { - derivedCtx.mutateAdditionalRequestHeaders( - mutator -> mutator.add(ARMERIA_RETRY_COUNT, StringUtil.toString(totalAttempts - 1))); + attemptCtx.mutateAdditionalRequestHeaders( + mutator -> mutator.add(ARMERIA_RETRY_COUNT, StringUtil.toString(attemptNo - 1))); } - final RpcResponse res; + final RpcResponse attemptRes; - final ClientRequestContextExtension ctxExtension = derivedCtx.as(ClientRequestContextExtension.class); - final EndpointGroup endpointGroup = derivedCtx.endpointGroup(); + final ClientRequestContextExtension ctxExtension = attemptCtx.as(ClientRequestContextExtension.class); + final EndpointGroup endpointGroup = attemptCtx.endpointGroup(); if (!initialAttempt && ctxExtension != null && - endpointGroup != null && derivedCtx.endpoint() == null) { + endpointGroup != null && attemptCtx.endpoint() == null) { // clear the pending throwable to retry endpoint selection - ClientPendingThrowableUtil.removePendingThrowable(derivedCtx); + ClientPendingThrowableUtil.removePendingThrowable(attemptCtx); // 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); + attemptRes = initContextAndExecuteWithFallback(unwrap(), ctxExtension, RpcResponse::from, + (context, cause) -> RpcResponse.ofFailure(cause), + req, true); } else { - res = executeWithFallback(unwrap(), derivedCtx, - (context, cause) -> RpcResponse.ofFailure(cause), - req, true); + attemptRes = executeWithFallback(unwrap(), attemptCtx, + (context, cause) -> RpcResponse.ofFailure(cause), + req, true); } + startRetryAttempt(ctx, attemptCtx, attemptRes, (winningAttemptCtx, winningAttemptRes) -> { + winningAttemptCtx.eventLoop().execute(() -> { + ctx.logBuilder().endResponseWithChild(winningAttemptCtx.log()); + final HttpRequest actualHttpReq = winningAttemptCtx.request(); + if (actualHttpReq != null) { + ctx.updateRequest(actualHttpReq); + } + + returnedResFuture.complete(winningAttemptRes); + }); + }, (abortingAttemptCtx, + abortingAttemptRes, + cause) -> { + final CompletableFuture abortComplete = new CompletableFuture<>(); + + abortingAttemptCtx.eventLoop().execute(() -> { + if (cause != null) { + abortingAttemptCtx.cancel(cause); + } else { + abortingAttemptCtx.cancel(); + } + + abortingAttemptRes.toCompletableFuture().handle((unused, unusedCause) -> { + abortComplete.complete(null); + return null; + }); + }); + + return abortComplete; + }); + final RetryConfig retryConfig = mappedRetryConfig(ctx); final RetryRuleWithContent retryRule = retryConfig.needsContentInRule() ? retryConfig.retryRuleWithContent() : retryConfig.fromRetryRule(); - res.handle((unused1, cause) -> { + attemptRes.handle((unused1, cause) -> { try { assert retryRule != null; - retryRule.shouldRetry(derivedCtx, res, cause).handle((decision, unused3) -> { + retryRule.shouldRetry(attemptCtx, attemptRes, 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); + scheduleNextRetry(ctx, + nextAttemptNo -> doExecute0(ctx, req, returnedRes, + returnedResFuture, nextAttemptNo), + backoff, + cause0 -> handleExceptionAfterScheduling(ctx, returnedResFuture, + cause0)); } + + completeRetryAttempt(ctx, attemptCtx, attemptRes, backoff == null); return null; }); } catch (Throwable t) { - handleException(ctx, future, t, false); + completeRetryingExceptionally(ctx, returnedResFuture, t, false); } return null; }); + + final long hedgingDelayMillis = retryConfig.hedgingDelayMillis(); + + if (hedgingDelayMillis >= 0) { + final Backoff hedgingBackoff = Backoff.fixed(hedgingDelayMillis); + scheduleNextRetry(ctx, hedgingAttemptNo -> + doExecute0(ctx, req, returnedRes, returnedResFuture, + hedgingAttemptNo), + hedgingBackoff, + cause -> handleExceptionAfterScheduling(ctx, returnedResFuture, cause)); + } } - 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); + private static void handleExceptionAfterScheduling( + ClientRequestContext ctx, CompletableFuture returnedResFuture, Throwable cause) { + if (cause instanceof RetrySchedulingException) { + switch (((RetrySchedulingException) cause).getType()) { + case RETRYING_ALREADY_COMPLETED: + return; + case NO_MORE_ATTEMPTS_IN_RETRY: + case NO_MORE_ATTEMPTS_IN_BACKOFF: + case DELAY_FROM_BACKOFF_EXCEEDS_RESPONSE_TIMEOUT: + case DELAY_FROM_SERVER_EXCEEDS_RESPONSE_TIMEOUT: + case RETRY_TASK_OVERTAKEN: + completeRetryingIfNoPendingAttempts(ctx); + return; + case RETRY_TASK_CANCELLED: + cause = new IllegalStateException( + ClientFactory.class.getSimpleName() + " has been closed.", cause); + break; + } } - future.complete(res); + completeRetryingExceptionally(ctx, returnedResFuture, cause, false); } - private static void handleException(ClientRequestContext ctx, CompletableFuture future, - Throwable cause, boolean endRequestLog) { - future.completeExceptionally(cause); + private static void completeRetryingExceptionally(ClientRequestContext ctx, + CompletableFuture returnedResFuture, + Throwable cause, boolean endRequestLog) { + if (isRetryingComplete(ctx)) { + return; + } + + returnedResFuture.completeExceptionally(cause); if (endRequestLog) { ctx.logBuilder().endRequest(cause); } ctx.logBuilder().endResponse(cause); + + completeRetryingExceptionally(ctx, cause); } } diff --git a/core/src/main/java/com/linecorp/armeria/common/logging/DefaultRequestLog.java b/core/src/main/java/com/linecorp/armeria/common/logging/DefaultRequestLog.java index 250eed2fae3..b9a8920a4a4 100644 --- a/core/src/main/java/com/linecorp/armeria/common/logging/DefaultRequestLog.java +++ b/core/src/main/java/com/linecorp/armeria/common/logging/DefaultRequestLog.java @@ -595,6 +595,15 @@ private void propagateRequestSideLog(RequestLogAccess child) { }); } + @Override + public void endResponseWithChild(RequestLogAccess child) { + checkState(!hasLastChild, "last child is already added"); + checkState(children != null && !children.isEmpty(), "at least one child should be already added"); + checkState(children.stream().anyMatch(c -> c == child), "child is not a child of this log: %s", child); + hasLastChild = true; + propagateResponseSideLog(child.partial()); + } + @Override public void endResponseWithLastChild() { checkState(!hasLastChild, "last child is already added"); @@ -604,58 +613,58 @@ public void endResponseWithLastChild() { propagateResponseSideLog(lastChild.partial()); } - private void propagateResponseSideLog(RequestLog lastChild) { - if (lastChild.isAvailable(RequestLogProperty.RESPONSE_CAUSE)) { + private void propagateResponseSideLog(RequestLog child) { + if (child.isAvailable(RequestLogProperty.RESPONSE_CAUSE)) { // Update responseCause first if available because callbacks of the other properties may need it // to retry or open circuit breakers. - final Throwable responseCause = lastChild.responseCause(); + final Throwable responseCause = child.responseCause(); if (responseCause != null) { responseCause(responseCause); } } - // Update the available properties without adding a callback if the lastChild already has them. - if (lastChild.isAvailable(RequestLogProperty.RESPONSE_START_TIME)) { - startResponse(lastChild.responseStartTimeNanos(), lastChild.responseStartTimeMicros(), true); + // Update the available properties without adding a callback if the child already has them. + if (child.isAvailable(RequestLogProperty.RESPONSE_START_TIME)) { + startResponse(child.responseStartTimeNanos(), child.responseStartTimeMicros(), true); } else { - lastChild.whenAvailable(RequestLogProperty.RESPONSE_START_TIME) - .thenAccept(log -> startResponse(log.responseStartTimeNanos(), - log.responseStartTimeMicros(), true)); + child.whenAvailable(RequestLogProperty.RESPONSE_START_TIME) + .thenAccept(log -> startResponse(log.responseStartTimeNanos(), + log.responseStartTimeMicros(), true)); } - if (lastChild.isAvailable(RequestLogProperty.RESPONSE_FIRST_BYTES_TRANSFERRED_TIME)) { - final Long timeNanos = lastChild.responseFirstBytesTransferredTimeNanos(); + if (child.isAvailable(RequestLogProperty.RESPONSE_FIRST_BYTES_TRANSFERRED_TIME)) { + final Long timeNanos = child.responseFirstBytesTransferredTimeNanos(); if (timeNanos != null) { responseFirstBytesTransferred(timeNanos); } } else { - lastChild.whenAvailable(RequestLogProperty.RESPONSE_FIRST_BYTES_TRANSFERRED_TIME) - .thenAccept(log -> { - final Long timeNanos = log.responseFirstBytesTransferredTimeNanos(); - if (timeNanos != null) { - responseFirstBytesTransferred(timeNanos); - } - }); + child.whenAvailable(RequestLogProperty.RESPONSE_FIRST_BYTES_TRANSFERRED_TIME) + .thenAccept(log -> { + final Long timeNanos = log.responseFirstBytesTransferredTimeNanos(); + if (timeNanos != null) { + responseFirstBytesTransferred(timeNanos); + } + }); } - if (lastChild.isAvailable(RequestLogProperty.RESPONSE_HEADERS)) { - responseHeaders(lastChild.responseHeaders()); + if (child.isAvailable(RequestLogProperty.RESPONSE_HEADERS)) { + responseHeaders(child.responseHeaders()); } else { - lastChild.whenAvailable(RequestLogProperty.RESPONSE_HEADERS) - .thenAccept(log -> responseHeaders(log.responseHeaders())); + child.whenAvailable(RequestLogProperty.RESPONSE_HEADERS) + .thenAccept(log -> responseHeaders(log.responseHeaders())); } - if (lastChild.isAvailable(RequestLogProperty.RESPONSE_TRAILERS)) { - responseTrailers(lastChild.responseTrailers()); + if (child.isAvailable(RequestLogProperty.RESPONSE_TRAILERS)) { + responseTrailers(child.responseTrailers()); } else { - lastChild.whenAvailable(RequestLogProperty.RESPONSE_TRAILERS) - .thenAccept(log -> responseTrailers(log.responseTrailers())); + child.whenAvailable(RequestLogProperty.RESPONSE_TRAILERS) + .thenAccept(log -> responseTrailers(log.responseTrailers())); } - if (lastChild.isComplete()) { - propagateResponseEndData(lastChild); + if (child.isComplete()) { + propagateResponseEndData(child); } else { - lastChild.whenComplete().thenAccept(this::propagateResponseEndData); + child.whenComplete().thenAccept(this::propagateResponseEndData); } } @@ -1068,8 +1077,8 @@ private void endRequest0(@Nullable Throwable requestCause, long requestEndTimeNa final int deferredFlags; if (requestCause != null) { // Will auto-fill request content and its preview if request has failed. - deferredFlags = this.deferredFlags & (~(RequestLogProperty.REQUEST_CONTENT.flag() | - RequestLogProperty.REQUEST_CONTENT_PREVIEW.flag())); + deferredFlags = this.deferredFlags & ~(RequestLogProperty.REQUEST_CONTENT.flag() | + RequestLogProperty.REQUEST_CONTENT_PREVIEW.flag()); } else { deferredFlags = this.deferredFlags; } @@ -1419,8 +1428,8 @@ private void endResponse0(@Nullable Throwable responseCause, long responseEndTim final int deferredFlags; if (responseCause != null) { // Will auto-fill response content and its preview if response has failed. - deferredFlags = this.deferredFlags & (~(RequestLogProperty.RESPONSE_CONTENT.flag() | - RequestLogProperty.RESPONSE_CONTENT_PREVIEW.flag())); + deferredFlags = this.deferredFlags & ~(RequestLogProperty.RESPONSE_CONTENT.flag() | + RequestLogProperty.RESPONSE_CONTENT_PREVIEW.flag()); } else { deferredFlags = this.deferredFlags; } diff --git a/core/src/main/java/com/linecorp/armeria/common/logging/RequestLogBuilder.java b/core/src/main/java/com/linecorp/armeria/common/logging/RequestLogBuilder.java index 61915cfaaa6..2156e9da01d 100644 --- a/core/src/main/java/com/linecorp/armeria/common/logging/RequestLogBuilder.java +++ b/core/src/main/java/com/linecorp/armeria/common/logging/RequestLogBuilder.java @@ -437,6 +437,12 @@ void session(@Nullable Channel channel, SessionProtocol sessionProtocol, @Nullab */ void addChild(RequestLogAccess child); + /** + * Fills the response-side logs from the specified child. Note that already collected properties + * in the child log will be propagated immediately. + */ + void endResponseWithChild(RequestLogAccess child); + /** * Fills the response-side logs from the last added child. Note that already collected properties * in the child log will be propagated immediately. 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..c86e72ba0dc --- /dev/null +++ b/core/src/test/java/com/linecorp/armeria/client/retry/RetrySchedulerTest.java @@ -0,0 +1,676 @@ +/* + * 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.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Consumer; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +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 org.mockito.ArgumentCaptor; + +import com.google.common.collect.ImmutableList; + +import com.linecorp.armeria.client.retry.RetrySchedulingException.Type; +import com.linecorp.armeria.internal.common.util.ReentrantShortLock; +import com.linecorp.armeria.internal.testing.AnticipatedException; + +import io.netty.channel.DefaultEventLoop; +import io.netty.channel.EventLoop; + +class RetrySchedulerTest { + private static final long SCHEDULING_TOLERANCE_NANOS = TimeUnit.MILLISECONDS.toNanos(200); + private static final long SCHEDULING_TOLERANCE_MILLIS = TimeUnit.NANOSECONDS.toMillis( + SCHEDULING_TOLERANCE_NANOS); + + private ReentrantLock retryLock; + private Consumer dummyRetryTask; + private Consumer dummyThrowingRetryTask; + + private EventLoop eventLoop; + + private RetryScheduler scheduler; + + private static class SpyableRetryTask implements Consumer { + private final Consumer delegate; + + SpyableRetryTask(Consumer delegate) { + this.delegate = delegate; + } + + @Override + public void accept(ReentrantLock retryLock) { + delegate.accept(retryLock); + } + } + + @BeforeEach + void setUp() { + eventLoop = spy(new DefaultEventLoop()); + retryLock = new ReentrantShortLock(); + dummyRetryTask = new SpyableRetryTask(retryLock -> { + assertThat(retryLock).isEqualTo(this.retryLock); + assertThat(retryLock.isHeldByCurrentThread()).isTrue(); + assertThat(retryLock.getHoldCount()).isOne(); + retryLock.unlock(); + }); + + dummyThrowingRetryTask = new SpyableRetryTask(retryLock -> { + assertThat(retryLock).isEqualTo(this.retryLock); + assertThat(retryLock.isHeldByCurrentThread()).isTrue(); + assertThat(retryLock.getHoldCount()).isOne(); + retryLock.unlock(); + throw new AnticipatedException(); + }); + scheduler = new RetryScheduler(retryLock, eventLoop); + } + + @AfterEach + void tearDown() throws Exception { + assertThat(retryLock.isLocked()).isFalse(); + assertThat(scheduler.shutdown()).isTrue(); + eventLoop.shutdownGracefully(); + } + + @ParameterizedTest + @MethodSource("scheduleTaskParameters") + void testScheduleTask(int taskDelayMs, int minDelayMs, int prevDelayMs, int expectedDelayMs) + throws Exception { + // Convert all delays to nanos + final long now = System.nanoTime(); + final long taskScheduledTimeNanos = now + TimeUnit.MILLISECONDS.toNanos(taskDelayMs); + final long newEarliestNextRetryTimeNanos = + minDelayMs >= 0 ? now + TimeUnit.MILLISECONDS.toNanos(minDelayMs) : Long.MIN_VALUE; + final long previousEarliestNextRetryTimeNanos = + prevDelayMs >= 0 ? now + TimeUnit.MILLISECONDS.toNanos(prevDelayMs) : Long.MIN_VALUE; + + final long expectedScheduledTimeNanos = now + TimeUnit.MILLISECONDS.toNanos(expectedDelayMs); + + // Set the previous delay if it exists + if (prevDelayMs >= 0) { + scheduler.addEarliestNextRetryTimeNanos(previousEarliestNextRetryTimeNanos); + } + + // Schedule the task + final Consumer task = spy(dummyRetryTask); + final Consumer exceptionHandler = mock(Consumer.class); + + scheduler.schedule(task, taskScheduledTimeNanos, newEarliestNextRetryTimeNanos, exceptionHandler); + + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(now, expectedScheduledTimeNanos) + ); + + Thread.sleep(expectedDelayMs + SCHEDULING_TOLERANCE_MILLIS); + + verify(task, times(1)).accept(retryLock); + verifyNoMoreInteractions(exceptionHandler); + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(now, expectedScheduledTimeNanos) + ); + } + + private static Stream scheduleTaskParameters() { + return Stream.of( + // taskDelay, minDelay, prevDelay, expectedDelay + + // No previous delay, no minimum delay + Arguments.of(0, -1, -1, 0), + Arguments.of(200, -1, -1, 200), + + Arguments.of(100, 100, -1, 100), + Arguments.of(150, 100, -1, 150), + + // With previous delay, no minimum delay + Arguments.of(50, -1, 100, 100), + Arguments.of(100, -1, 100, 100), + Arguments.of(150, -1, 100, 150), + + // With both previous and minimum delay + Arguments.of(100, 100, 200, 200), + Arguments.of(150, 100, 200, 200), + Arguments.of(200, 100, 200, 200), + Arguments.of(250, 100, 200, 250) + ); + } + + @Test + void testScheduleTaskInThePast() throws Exception { + final Consumer task = spy(dummyRetryTask); + final long taskSchedulingTime = System.nanoTime(); + final long expectedTaskRunTime = taskSchedulingTime; + final Consumer exceptionHandler = mock(Consumer.class); + + scheduler.schedule(task, expectedTaskRunTime - 1_000_000, Long.MIN_VALUE, exceptionHandler); + + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(taskSchedulingTime, expectedTaskRunTime) + ); + + Thread.sleep(SCHEDULING_TOLERANCE_MILLIS); + + verify(task, times(1)).accept(retryLock); + verifyNoMoreInteractions(exceptionHandler); + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(taskSchedulingTime, expectedTaskRunTime) + ); + + final Consumer task2 = spy(dummyRetryTask); + final long task2SchedulingTime = System.nanoTime(); + final long expectedTask2RunTime = task2SchedulingTime; + final Consumer task2ExceptionHandler = mock(Consumer.class); + + scheduler.schedule(task2, Long.MIN_VALUE, Long.MIN_VALUE, task2ExceptionHandler); + + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(taskSchedulingTime, expectedTaskRunTime), + EventLoopScheduleCall.of(task2SchedulingTime, expectedTask2RunTime) + ); + + Thread.sleep(SCHEDULING_TOLERANCE_MILLIS); + + verify(task2, times(1)).accept(retryLock); + verifyNoMoreInteractions(task2ExceptionHandler); + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(taskSchedulingTime, expectedTaskRunTime), + EventLoopScheduleCall.of(task2SchedulingTime, expectedTask2RunTime) + ); + } + + @Test + void testFailingScheduleWithInvalidArguments() { + final Consumer task = spy(dummyRetryTask); + final Consumer exceptionHandler = mock(Consumer.class); + + assertThatThrownBy(() -> scheduler.schedule(null, 200, 200, + exceptionHandler)).isInstanceOf( + NullPointerException.class).hasMessageContaining("retryTask"); + + assertThatThrownBy(() -> scheduler.schedule(task, 200, 200, null)).isInstanceOf( + NullPointerException.class).hasMessageContaining("exceptionHandler"); + + assertThatThrownBy(() -> scheduler.schedule(task, Long.MAX_VALUE - 1, Long.MAX_VALUE, + exceptionHandler)).isInstanceOf( + IllegalArgumentException.class).hasMessageContaining("nextRetryTimeNanos"); + + verifyNoMoreInteractions(task, exceptionHandler); + verifyEventLoopScheduleCalls(); + } + + @Test + void testScheduleTaskWithException() throws Exception { + final Consumer task1 = spy(dummyThrowingRetryTask); + + final long task1SchedulingTime = System.nanoTime(); + final long expectedTask1RunTime = task1SchedulingTime + TimeUnit.MILLISECONDS.toNanos(100); + final Consumer task1ExceptionHandler = mock(Consumer.class); + scheduler.schedule(task1, expectedTask1RunTime, Long.MIN_VALUE, task1ExceptionHandler); + + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(task1SchedulingTime, expectedTask1RunTime) + ); + + Thread.sleep(100 + SCHEDULING_TOLERANCE_MILLIS); + + verify(task1, times(1)).accept(retryLock); + final ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Throwable.class); + verify(task1ExceptionHandler, times(1)).accept(exceptionCaptor.capture()); + assertThat(exceptionCaptor.getValue()).isInstanceOf(AnticipatedException.class); + } + + @Test + void testEarlierRetryTaskOvertakesLaterOne() throws Exception { + final Consumer task1 = spy(dummyRetryTask); + + final long task1SchedulingTime = System.nanoTime(); + final long expectedTask1RunTime = task1SchedulingTime + TimeUnit.MILLISECONDS.toNanos(200); + final Consumer task1ExceptionHandler = mock(Consumer.class); + scheduler.schedule(task1, expectedTask1RunTime, Long.MIN_VALUE, task1ExceptionHandler); + + final Consumer task2 = spy(dummyRetryTask); + final long task2SchedulingTime = System.nanoTime(); + final long expectedTask2RunTime = task2SchedulingTime + TimeUnit.MILLISECONDS.toNanos(100); + final Consumer task2ExceptionHandler = mock(Consumer.class); + scheduler.schedule(task2, expectedTask2RunTime, Long.MIN_VALUE, task2ExceptionHandler); + + Thread.sleep(200 + SCHEDULING_TOLERANCE_MILLIS); + + verify(task1, times(0)).accept(retryLock); + verify(task2, times(1)).accept(retryLock); + + verifyExceptionHandlerCatchedSchedulingException(task1ExceptionHandler, + RetrySchedulingException.Type.RETRY_TASK_OVERTAKEN); + + verifyNoMoreInteractions(task2ExceptionHandler); + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(task1SchedulingTime, expectedTask1RunTime), + EventLoopScheduleCall.of(task2SchedulingTime, expectedTask2RunTime) + ); + } + + @Test + void testMultipleRetryTasksBeingOvertaken() throws Exception { + // Rationale of this test: + // - Schedule 10 tasks with decreasing run times (1000ms, 900ms, ..., 100ms). + // - Expect that the first 9 tasks are not executed as they are overtaken by the next task as + // next task is scheduled earlier. + // - Only the last task (100ms) should be executed. + + final List> tasks = new ArrayList<>(); + final List schedulingTimes = new ArrayList<>(); + final List expectedRunTimes = new ArrayList<>(); + final List> exceptionHandlers = new ArrayList<>(); + + for (int taskNo = 0; taskNo < 10; taskNo++) { + final Consumer task = spy(dummyRetryTask); + tasks.add(task); + final Consumer exceptionHandler = mock(Consumer.class); + exceptionHandlers.add(exceptionHandler); + + final long schedulingTime = System.nanoTime(); + schedulingTimes.add(schedulingTime); + final long expectedRunTime = schedulingTime + TimeUnit.MILLISECONDS.toNanos(1000 - taskNo * 100); + expectedRunTimes.add(expectedRunTime); + + scheduler.schedule(task, expectedRunTime, Long.MIN_VALUE, exceptionHandler); + } + + // Wait for the tasks to be scheduled + Thread.sleep(1000 + SCHEDULING_TOLERANCE_MILLIS); + + for (int taskNo = 0; taskNo < 9; taskNo++) { + final Consumer task = tasks.get(taskNo); + verify(task, times(0)).accept(retryLock); + verifyExceptionHandlerCatchedSchedulingException( + exceptionHandlers.get(taskNo), + RetrySchedulingException.Type.RETRY_TASK_OVERTAKEN + ); + } + + // Verify that the last task was executed + verify(tasks.get(9), times(1)).accept(retryLock); + verifyNoMoreInteractions(exceptionHandlers.get(9)); + + verifyEventLoopScheduleCalls( + ImmutableList.copyOf( + tasks.stream() + .map(task -> EventLoopScheduleCall.of(schedulingTimes.get(tasks.indexOf(task)), + expectedRunTimes.get(tasks.indexOf(task)))) + .collect(Collectors.toList()) + ) + ); + } + + @Test + void testLaterRetryTaskDoesNotOvertakeEarlierOne() throws Exception { + final Consumer task1 = spy(dummyRetryTask); + final long task1SchedulingTime = System.nanoTime(); + final long expectedTask1RunTime = task1SchedulingTime + TimeUnit.MILLISECONDS.toNanos(200); + final Consumer task1ExceptionHandler = mock(Consumer.class); + scheduler.schedule(task1, expectedTask1RunTime, Long.MIN_VALUE, task1ExceptionHandler); + + // Schedule a new task with a later time than the current one + final Consumer task2 = spy(dummyRetryTask); + final long task2SchedulingTime = System.nanoTime(); + final long expectedTask2RunTime = task2SchedulingTime + TimeUnit.MILLISECONDS.toNanos(300); + final Consumer task2ExceptionHandler = mock(Consumer.class); + scheduler.schedule(task2, expectedTask2RunTime, Long.MIN_VALUE, task2ExceptionHandler); + + Thread.sleep(300 + SCHEDULING_TOLERANCE_MILLIS); + + // Verify that the first task was executed + verify(task1, times(1)).accept(retryLock); + verifyNoMoreInteractions(task1ExceptionHandler); + + // Verify that the second task was not executed + verify(task2, times(0)).accept(retryLock); + verifyExceptionHandlerCatchedSchedulingException( + task2ExceptionHandler, RetrySchedulingException.Type.RETRY_TASK_OVERTAKEN); + + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(task1SchedulingTime, expectedTask1RunTime) + ); + } + + @Test + void testRescheduleTaskWhenEarliestNextRetryTimeUpdated() throws Exception { + // Set the earliest next retry time to 200ms from now + final long now = System.nanoTime(); + final long earliestTime = now + TimeUnit.MILLISECONDS.toNanos(200); + scheduler.addEarliestNextRetryTimeNanos(earliestTime); + + // Schedule a task to run after 300ms + final long task1SchedulingTime = System.nanoTime(); + final long taskTime = task1SchedulingTime + TimeUnit.MILLISECONDS.toNanos(300); + final Consumer task1 = spy(dummyRetryTask); + final Consumer exceptionHandler = mock(Consumer.class); + + scheduler.schedule(task1, taskTime, Long.MIN_VALUE, exceptionHandler); + + // Move the earliest next retry time to 100ms from now + final long earliestTimeUpdateTime = System.nanoTime(); + final long newEarliestTimeNanos = now + TimeUnit.MILLISECONDS.toNanos(400); + scheduler.addEarliestNextRetryTimeNanos(newEarliestTimeNanos); + scheduler.rescheduleCurrentRetryTaskIfTooEarly(); + + Thread.sleep(400 + SCHEDULING_TOLERANCE_MILLIS); + + verify(task1, times(1)).accept(retryLock); + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(task1SchedulingTime, taskTime), + EventLoopScheduleCall.of(earliestTimeUpdateTime, newEarliestTimeNanos) + ); + verifyNoMoreInteractions(exceptionHandler); + } + + @Test + void testRescheduleWithNoTasks() throws InterruptedException { + scheduler.rescheduleCurrentRetryTaskIfTooEarly(); + verifyEventLoopScheduleCalls(); + + scheduler.addEarliestNextRetryTimeNanos(System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(-100)); + scheduler.addEarliestNextRetryTimeNanos(System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(0)); + scheduler.addEarliestNextRetryTimeNanos(System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(100)); + + verifyEventLoopScheduleCalls(); + + scheduler.rescheduleCurrentRetryTaskIfTooEarly(); + + verifyEventLoopScheduleCalls(); + + final long taskScheduledTime = System.nanoTime(); + final long expectedTaskTime = taskScheduledTime + TimeUnit.MILLISECONDS.toNanos(100); + + scheduler.schedule( + retryLock0 -> { + assertThat(retryLock0).isEqualTo(retryLock); + assertThat(retryLock0.isHeldByCurrentThread()).isTrue(); + assertThat(retryLock0.getHoldCount()).isOne(); + retryLock0.unlock(); + // note that we inherit the earliest next retry time from the calls above + }, taskScheduledTime + TimeUnit.MILLISECONDS.toNanos(50), Long.MIN_VALUE, t -> { + // do nothing + }); + + Thread.sleep(100 + SCHEDULING_TOLERANCE_MILLIS); + + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(taskScheduledTime, expectedTaskTime) + ); + + scheduler.addEarliestNextRetryTimeNanos(System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(1)); + scheduler.addEarliestNextRetryTimeNanos(System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(100)); + + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(taskScheduledTime, expectedTaskTime) + ); + + scheduler.rescheduleCurrentRetryTaskIfTooEarly(); + + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(taskScheduledTime, expectedTaskTime) + ); + } + + @Test + void testIdempotentReschedule() throws Exception { + // Schedule a task to run after 200ms + final Consumer task = spy(dummyRetryTask); + final long taskScheduledTime = System.nanoTime(); + final long taskTime = taskScheduledTime + TimeUnit.MILLISECONDS.toNanos(100); + final Consumer exceptionHandler = mock(Consumer.class); + + scheduler.schedule(task, taskTime, Long.MIN_VALUE, exceptionHandler); + + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(taskScheduledTime, taskTime) + ); + + final long earliestTime1 = taskScheduledTime + TimeUnit.MILLISECONDS.toNanos(200); + + scheduler.addEarliestNextRetryTimeNanos(earliestTime1 - 100); + scheduler.addEarliestNextRetryTimeNanos(earliestTime1); + + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(taskScheduledTime, taskTime) + ); + + final long rescheduleTime1 = System.nanoTime(); + scheduler.rescheduleCurrentRetryTaskIfTooEarly(); + + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(taskScheduledTime, taskTime), + EventLoopScheduleCall.of(rescheduleTime1, earliestTime1) + ); + + scheduler.rescheduleCurrentRetryTaskIfTooEarly(); + scheduler.rescheduleCurrentRetryTaskIfTooEarly(); + + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(taskScheduledTime, taskTime), + EventLoopScheduleCall.of(rescheduleTime1, earliestTime1) + ); + + final long earliestTime2 = taskScheduledTime + TimeUnit.MILLISECONDS.toNanos(300); + + scheduler.addEarliestNextRetryTimeNanos(earliestTime2 - 100); + scheduler.addEarliestNextRetryTimeNanos(earliestTime2 - 200); + scheduler.addEarliestNextRetryTimeNanos(earliestTime2); + + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(taskScheduledTime, taskTime), + EventLoopScheduleCall.of(rescheduleTime1, earliestTime1) + ); + + Thread.sleep(20); + + final long rescheduleTime2 = System.nanoTime(); + scheduler.rescheduleCurrentRetryTaskIfTooEarly(); + + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(taskScheduledTime, taskTime), + EventLoopScheduleCall.of(rescheduleTime1, earliestTime1), + EventLoopScheduleCall.of(rescheduleTime2, earliestTime2) + ); + + scheduler.rescheduleCurrentRetryTaskIfTooEarly(); + scheduler.rescheduleCurrentRetryTaskIfTooEarly(); + + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(taskScheduledTime, taskTime), + EventLoopScheduleCall.of(rescheduleTime1, earliestTime1), + EventLoopScheduleCall.of(rescheduleTime2, earliestTime2) + ); + + Thread.sleep(300 + SCHEDULING_TOLERANCE_MILLIS); + + verify(task, times(1)).accept(retryLock); + verifyNoMoreInteractions(exceptionHandler); + + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(taskScheduledTime, taskTime), + EventLoopScheduleCall.of(rescheduleTime1, earliestTime1), + EventLoopScheduleCall.of(rescheduleTime2, earliestTime2) + ); + } + + @Test + void testSchedulerShutdownWithoutTask() throws Exception { + // Shutdown the scheduler without scheduling any tasks + assertThat(scheduler.shutdown()).isTrue(); + } + + @Test + void testSchedulerShutdownCancelsTask() throws Exception { + // Schedule a task that should run after 200ms + final Consumer task = spy(dummyRetryTask); + final long taskSchedulingTime = System.nanoTime(); + final long expectedTaskRunTime = taskSchedulingTime + TimeUnit.MILLISECONDS.toNanos(200); + final Consumer exceptionHandler = mock(Consumer.class); + scheduler.schedule(task, expectedTaskRunTime, Long.MIN_VALUE, exceptionHandler); + + // Shutdown the scheduler + assertThat(scheduler.shutdown()).isTrue(); + + Thread.sleep(SCHEDULING_TOLERANCE_MILLIS); + + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(taskSchedulingTime, expectedTaskRunTime) + ); + + Thread.sleep(400); + + verify(task, times(0)).accept(any()); + verifyExceptionHandlerCatchedSchedulingException( + exceptionHandler, Type.RETRYING_ALREADY_COMPLETED); + } + + @Test + void testEventLoopShutdownCancelsTask() throws Exception { + // Schedule a task that should run after 200ms + final Consumer task = spy(dummyRetryTask); + final long taskSchedulingTime = System.nanoTime(); + final long expectedTaskRunTime = taskSchedulingTime + TimeUnit.MILLISECONDS.toNanos(200); + final Consumer exceptionHandler = mock(Consumer.class); + scheduler.schedule(task, expectedTaskRunTime, Long.MIN_VALUE, exceptionHandler); + + // Shutdown the event loop + eventLoop.shutdownGracefully().sync(); + + Thread.sleep(SCHEDULING_TOLERANCE_MILLIS); + + verifyEventLoopScheduleCalls( + EventLoopScheduleCall.of(taskSchedulingTime, expectedTaskRunTime) + ); + + Thread.sleep(400); + + verify(task, times(0)).accept(any()); + verifyExceptionHandlerCatchedSchedulingException( + exceptionHandler, Type.RETRY_TASK_CANCELLED); + } + + @Test + void testScheduledOnShutdownEventLoop() throws InterruptedException { + eventLoop.shutdownGracefully().sync(); + + final Consumer task = spy(dummyRetryTask); + final long taskSchedulingTime = System.nanoTime(); + final long expectedTaskRunTime = taskSchedulingTime + TimeUnit.MILLISECONDS.toNanos(200); + final Consumer exceptionHandler = mock(Consumer.class); + scheduler.schedule(task, expectedTaskRunTime, Long.MIN_VALUE, exceptionHandler); + + final ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Throwable.class); + verify(exceptionHandler, times(1)).accept(exceptionCaptor.capture()); + final Throwable capturedException = exceptionCaptor.getValue(); + + verify(task, times(0)).accept(retryLock); + assertThat(capturedException).isInstanceOf(RejectedExecutionException.class); + assertThat(capturedException.getMessage()).contains("event executor terminated"); + } + + // todo(szymon): add test that verifies that when a task is about to run it gets rescheduled when the + // earliest next retry time was set to something later in the meantime. + + private static final class EventLoopScheduleCall { + private final long delayNanos; + + private EventLoopScheduleCall(long scheduledTimeNanos) { + this(System.nanoTime(), scheduledTimeNanos); + } + + private EventLoopScheduleCall(long schedulingTimeNanos, long scheduledTimeNanos) { + assert schedulingTimeNanos <= scheduledTimeNanos : "Scheduling time must be before scheduled time"; + delayNanos = scheduledTimeNanos - schedulingTimeNanos; + } + + public static EventLoopScheduleCall of(long scheduledTimeNanos) { + return new EventLoopScheduleCall(scheduledTimeNanos); + } + + public static EventLoopScheduleCall of(long schedulingTimeNanos, + long scheduledTimeNanos) { + return new EventLoopScheduleCall(schedulingTimeNanos, scheduledTimeNanos); + } + + public long delayNanos() { + return delayNanos; + } + } + + private static void verifyExceptionHandlerCatchedSchedulingException( + Consumer exceptionHandler, RetrySchedulingException.Type expectedType) { + final ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Throwable.class); + verify(exceptionHandler, times(1)).accept(exceptionCaptor.capture()); + + final Throwable capturedException = exceptionCaptor.getValue(); + assertThat(capturedException).isInstanceOf(RetrySchedulingException.class); + assertThat(((RetrySchedulingException) capturedException).getType()).isEqualTo(expectedType); + } + + private void verifyEventLoopScheduleCalls(EventLoopScheduleCall... expectedSchedules) { + verifyEventLoopScheduleCalls(ImmutableList.copyOf(expectedSchedules)); + } + + private void verifyEventLoopScheduleCalls(List expectedSchedules) { + final int expectedNumCalls = expectedSchedules.size(); + + final ArgumentCaptor scheduleDelayArgumentCaptor = ArgumentCaptor.forClass(Long.class); + final ArgumentCaptor scheduleTimeUnitArgumentCaptor = ArgumentCaptor.forClass(TimeUnit.class); + + verify(eventLoop, times(expectedNumCalls)).schedule(any(Runnable.class), + scheduleDelayArgumentCaptor.capture(), + scheduleTimeUnitArgumentCaptor.capture()); + + final List actualDelays = scheduleDelayArgumentCaptor.getAllValues(); + final List actualTimeUnits = scheduleTimeUnitArgumentCaptor.getAllValues(); + + assertThat(actualDelays).hasSize(expectedNumCalls); + assertThat(actualTimeUnits).hasSize(expectedNumCalls); + + // for simplicity, we assume all time units are NANOSECONDS + assertThat(actualTimeUnits).allMatch(unit -> unit == TimeUnit.NANOSECONDS); + + for (int i = 0; i < expectedSchedules.size(); i++) { + final EventLoopScheduleCall expected = expectedSchedules.get(i); + final long actualDelayNanos = actualDelays.get(i); + + assertThat(actualDelayNanos) + .isBetween(expected.delayNanos() - SCHEDULING_TOLERANCE_NANOS, + expected.delayNanos() + SCHEDULING_TOLERANCE_NANOS); + } + } +} 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..bd2c268f221 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,8 @@ 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.assertj.core.api.Assertions.fail; +import static org.awaitility.Awaitility.await; import java.time.Duration; import java.util.Arrays; @@ -31,6 +33,7 @@ import java.util.concurrent.CompletionStage; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; @@ -55,11 +58,15 @@ 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; import com.linecorp.armeria.client.WebClient; import com.linecorp.armeria.client.logging.LoggingClient; +import com.linecorp.armeria.client.retry.AbstractRetryingClient.State; +import com.linecorp.armeria.client.retry.AbstractRetryingClient.State.Attempt; import com.linecorp.armeria.common.AggregatedHttpResponse; import com.linecorp.armeria.common.HttpData; import com.linecorp.armeria.common.HttpHeaderNames; @@ -102,6 +109,9 @@ static void afterAll() { clientFactory.closeAsync(); } + @Nullable + ClientRequestContext ctx; + private final AtomicInteger responseAbortServiceCallCounter = new AtomicInteger(); private final AtomicInteger requestAbortServiceCallCounter = new AtomicInteger(); @@ -328,16 +338,25 @@ void retryWhenContentMatched() { .factory(clientFactory) .decorator(retryingDecorator) .build(); - - final AggregatedHttpResponse res = client.get("/retry-content").aggregate().join(); + final AggregatedHttpResponse res; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + res = client.get("/retry-content").aggregate().join(); + ctx = captor.get(); + } assertThat(res.contentUtf8()).isEqualTo("Succeeded after retry"); + awaitValidClientRequestContext(ctx, 3); } @Test void retryWhenStatusMatched() { final WebClient client = client(RetryRule.builder().onServerErrorStatus().onException().thenBackoff()); - final AggregatedHttpResponse res = client.get("/503-then-success").aggregate().join(); + final AggregatedHttpResponse res; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + res = client.get("/503-then-success").aggregate().join(); + ctx = captor.get(); + } assertThat(res.contentUtf8()).isEqualTo("Succeeded after retry"); + awaitValidClientRequestContext(ctx, 2); } @Test @@ -346,8 +365,13 @@ void retryWhenStatusMatchedWithContent() { .onServerErrorStatus() .onException() .thenBackoff(), 10000, 0, 100); - final AggregatedHttpResponse res = client.get("/503-then-success").aggregate().join(); + final AggregatedHttpResponse res; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + res = client.get("/503-then-success").aggregate().join(); + ctx = captor.get(); + } assertThat(res.contentUtf8()).isEqualTo("Succeeded after retry"); + awaitValidClientRequestContext(ctx, 2); } @Test @@ -358,8 +382,13 @@ void retryWhenTrailerMatched() { return trailers.getInt("grpc-status", -1) != 0; }) .thenBackoff()); - final AggregatedHttpResponse res = client.get("/trailers-then-success").aggregate().join(); + final AggregatedHttpResponse res; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + res = client.get("/trailers-then-success").aggregate().join(); + ctx = captor.get(); + } assertThat(res.contentUtf8()).isEqualTo("Succeeded after retry"); + awaitValidClientRequestContext(ctx, 2); } @Test @@ -368,16 +397,26 @@ void retryWhenTotalDurationIsHigh() { client(RetryRule.builder() .onTotalDuration((unused, duration) -> duration.toNanos() > 100) .thenBackoff()); - final AggregatedHttpResponse res = client.get("/1sleep-then-success").aggregate().join(); + final AggregatedHttpResponse res; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + res = client.get("/1sleep-then-success").aggregate().join(); + ctx = captor.get(); + } assertThat(res.contentUtf8()).isEqualTo("Succeeded after retry"); + awaitValidClientRequestContext(ctx); } @Test void disableResponseTimeout() { final WebClient client = client(RetryRule.failsafe(), 0, 0, 100); - final AggregatedHttpResponse res = client.get("/503-then-success").aggregate().join(); + final AggregatedHttpResponse res; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + res = client.get("/503-then-success").aggregate().join(); + ctx = captor.get(); + } assertThat(res.contentUtf8()).isEqualTo("Succeeded after retry"); // response timeout did not happen. + awaitValidClientRequestContext(ctx, 2); } @Test @@ -385,10 +424,16 @@ void respectRetryAfter() { final WebClient client = client(RetryRule.failsafe()); final Stopwatch sw = Stopwatch.createStarted(); - final AggregatedHttpResponse res = client.get("/retry-after-1-second").aggregate().join(); + final AggregatedHttpResponse res; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + res = client.get("/retry-after-1-second").aggregate().join(); + ctx = captor.get(); + } + assertThat(res.contentUtf8()).isEqualTo("Succeeded after retry"); assertThat(sw.elapsed(TimeUnit.MILLISECONDS)).isGreaterThanOrEqualTo( (long) (TimeUnit.SECONDS.toMillis(1) * 0.9)); + awaitValidClientRequestContext(ctx, 2); } @Test @@ -396,12 +441,17 @@ void respectRetryAfterWithHttpDate() { final WebClient client = client(RetryRule.failsafe()); final Stopwatch sw = Stopwatch.createStarted(); - final AggregatedHttpResponse res = client.get("/retry-after-with-http-date").aggregate().join(); - assertThat(res.contentUtf8()).isEqualTo("Succeeded after retry"); + final AggregatedHttpResponse res; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + res = client.get("/retry-after-with-http-date").aggregate().join(); + ctx = captor.get(); + } + assertThat(res.contentUtf8()).isEqualTo("Succeeded after retry"); // Since ZonedDateTime doesn't express exact time, // just check out whether it is retried after delayed some time. assertThat(sw.elapsed(TimeUnit.MILLISECONDS)).isGreaterThanOrEqualTo(1000); + awaitValidClientRequestContext(ctx, 2); } @Test @@ -410,27 +460,41 @@ void propagateLastResponseWhenNextRetryIsAfterTimeout() { .onServerErrorStatus() .onException() .thenBackoff(Backoff.fixed(10000000))); - final AggregatedHttpResponse res = client.get("/service-unavailable").aggregate().join(); + final AggregatedHttpResponse res; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + res = client.get("/service-unavailable").aggregate().join(); + ctx = captor.get(); + } assertThat(res.status()).isSameAs(HttpStatus.SERVICE_UNAVAILABLE); + awaitValidClientRequestContext(ctx, 1); } @Test void propagateLastResponseWhenExceedMaxAttempts() { final WebClient client = client( RetryRule.builder().onServerErrorStatus().onException().thenBackoff(Backoff.fixed(1)), 0, 0, 3); - final AggregatedHttpResponse res = client.get("/service-unavailable").aggregate().join(); + final AggregatedHttpResponse res; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + res = client.get("/service-unavailable").aggregate().join(); + ctx = captor.get(); + } assertThat(res.status()).isSameAs(HttpStatus.SERVICE_UNAVAILABLE); + awaitValidClientRequestContext(ctx, 3); // equal to max attempts } @Test void retryAfterOneYear() { final WebClient client = client(RetryRule.failsafe()); - // The response will be the last response whose headers contains HttpHeaderNames.RETRY_AFTER // because next retry is after timeout - final ResponseHeaders headers = client.get("retry-after-one-year").aggregate().join().headers(); + final ResponseHeaders headers; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + headers = client.get("retry-after-one-year").aggregate().join().headers(); + ctx = captor.get(); + } assertThat(headers.status()).isSameAs(HttpStatus.SERVICE_UNAVAILABLE); assertThat(headers.get(HttpHeaderNames.RETRY_AFTER)).isNotNull(); + awaitValidClientRequestContext(ctx, 1); } @Test @@ -445,8 +509,15 @@ void retryOnResponseTimeout() { }; final WebClient client = client(strategy, 0, 500, 100); - final AggregatedHttpResponse res = client.get("/1sleep-then-success").aggregate().join(); + + final AggregatedHttpResponse res; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + res = client.get("/1sleep-then-success").aggregate().join(); + ctx = captor.get(); + } + assertThat(res.contentUtf8()).isEqualTo("Succeeded after retry"); + awaitValidClientRequestContext(ctx, 2); } @Test @@ -467,10 +538,15 @@ void retryWithContentOnResponseTimeout() { .onException(ResponseTimeoutException.class) .thenBackoff(backoff))); final WebClient client = client(strategy, 0, 500, 100); - final AggregatedHttpResponse res = client.get("/1sleep-then-success").aggregate().join(); + final AggregatedHttpResponse res; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + res = client.get("/1sleep-then-success").aggregate().join(); + ctx = captor.get(); + } assertThat(res.contentUtf8()).isEqualTo("Succeeded after retry"); // Make sure that all customized RetryRuleWithContents are called. assertThat(queue).containsExactly(1, 2, 3); + awaitValidClientRequestContext(ctx, 2); } @Test @@ -507,49 +583,72 @@ void honorRetryMapping() { final WebClient client = client(mapping); Stopwatch stopwatch = Stopwatch.createStarted(); - assertThat(client.get("/500-always").aggregate().join().status()) - .isEqualTo(HttpStatus.valueOf(500)); + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + assertThat(client.get("/500-always").aggregate().join().status()) + .isEqualTo(HttpStatus.valueOf(500)); + ctx = captor.get(); + } assertThat(stopwatch.elapsed()).isBetween(Duration.ofSeconds(2), Duration.ofSeconds(6)); + awaitValidClientRequestContext(ctx, 2); + waitForEveryAttemptResponse(ctx); stopwatch = Stopwatch.createStarted(); - assertThat(client.get("/501-always").aggregate().join().status()) - .isEqualTo(HttpStatus.valueOf(501)); + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + assertThat(client.get("/501-always").aggregate().join().status()) + .isEqualTo(HttpStatus.valueOf(501)); + ctx = captor.get(); + } assertThat(stopwatch.elapsed()).isBetween(Duration.ofSeconds(14), Duration.ofSeconds(28)); + awaitValidClientRequestContext(ctx, 8); + waitForEveryAttemptResponse(ctx); stopwatch = Stopwatch.createStarted(); - assertThat(client.get("/502-always").aggregate().join().status()) - .isEqualTo(HttpStatus.valueOf(502)); + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + assertThat(client.get("/502-always").aggregate().join().status()) + .isEqualTo(HttpStatus.valueOf(502)); + ctx = captor.get(); + } assertThat(stopwatch.elapsed()).isBetween(Duration.ofSeconds(0), Duration.ofSeconds(2)); + awaitValidClientRequestContext(ctx, 1); } @Test 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); - assertThat(client.get("/500-then-success").aggregate().join().status()) - .isEqualTo(HttpStatus.valueOf(200)); + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + assertThat(client.get("/500-then-success").aggregate().join().status()) + .isEqualTo(HttpStatus.valueOf(200)); + ctx = captor.get(); + } // 1 logical request; 2 retries assertThat(evaluations.get()).isEqualTo(1); + awaitValidClientRequestContext(ctx, 2); + waitForEveryAttemptResponse(ctx); reqCount.set(0); - assertThat(client.get("/500-then-success").aggregate().join().status()) - .isEqualTo(HttpStatus.valueOf(200)); + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + assertThat(client.get("/500-then-success").aggregate().join().status()) + .isEqualTo(HttpStatus.valueOf(200)); + ctx = captor.get(); + } - // 2 logical requests; 4 retries + // 2 logical requests; 2 retries assertThat(evaluations.get()).isEqualTo(2); + awaitValidClientRequestContext(ctx, 2); } @Test @@ -582,10 +681,14 @@ void retryWithContentOnUnprocessedException() { .decorator(retryingDecorator) .build(); final Stopwatch stopwatch = Stopwatch.createStarted(); - assertThatThrownBy(() -> client.get("/unprocessed-exception").aggregate().join()) - .isInstanceOf(CompletionException.class) - .hasCauseInstanceOf(UnprocessedRequestException.class); + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + assertThatThrownBy(() -> client.get("/unprocessed-exception").aggregate().join()) + .isInstanceOf(CompletionException.class) + .hasCauseInstanceOf(UnprocessedRequestException.class); + ctx = captor.get(); + } assertThat(stopwatch.elapsed()).isBetween(Duration.ofSeconds(7), Duration.ofSeconds(20)); + awaitValidClientRequestContext(ctx, 5); // max attempts } } @@ -594,17 +697,27 @@ void retryWithContentOnUnprocessedException() { void differentBackoffBasedOnStatus(RetryRule retryRule) { final WebClient client = client(retryRule); + AggregatedHttpResponse res; final Stopwatch sw = Stopwatch.createStarted(); - AggregatedHttpResponse res = client.get("/503-then-success").aggregate().join(); + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + res = client.get("/503-then-success").aggregate().join(); + ctx = captor.get(); + } assertThat(res.contentUtf8()).isEqualTo("Succeeded after retry"); assertThat(sw.elapsed(TimeUnit.MILLISECONDS)).isBetween((long) (10 * 0.9), (long) (1000 * 1.1)); + awaitValidClientRequestContext(ctx, 2); + waitForEveryAttemptResponse(ctx); reqCount.set(0); sw.reset().start(); - res = client.get("/500-then-success").aggregate().join(); + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + res = client.get("/500-then-success").aggregate().join(); + ctx = captor.get(); + } assertThat(res.contentUtf8()).isEqualTo("Succeeded after retry"); assertThat(sw.elapsed(TimeUnit.MILLISECONDS)).isGreaterThanOrEqualTo((long) (1000 * 0.9)); + awaitValidClientRequestContext(ctx, 2); } @Test @@ -613,8 +726,13 @@ void retryWithRequestBody() { .onServerErrorStatus() .onException() .thenBackoff(Backoff.fixed(10))); - final AggregatedHttpResponse res = client.post("/post-ping-pong", "bar").aggregate().join(); + final AggregatedHttpResponse res; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + res = client.post("/post-ping-pong", "bar").aggregate().join(); + ctx = captor.get(); + } assertThat(res.contentUtf8()).isEqualTo("bar"); + awaitValidClientRequestContext(ctx, 2); } @Test @@ -652,7 +770,12 @@ void shouldGetExceptionWhenFactoryIsClosed() { // // Peel CompletionException first. - Throwable t = peel(catchThrowable(() -> client.get("/service-unavailable").aggregate().join())); + Throwable t; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + t = peel(catchThrowable(() -> client.get("/service-unavailable").aggregate().join())); + ctx = captor.get(); + } + awaitValidClientRequestContext(ctx, 1); // not able to schedule second retry. if (t instanceof UnprocessedRequestException) { final Throwable cause = t.getCause(); assertThat(cause).isInstanceOf(IllegalStateException.class); @@ -660,7 +783,8 @@ void shouldGetExceptionWhenFactoryIsClosed() { } assertThat(t).isInstanceOf(IllegalStateException.class) .satisfies(cause -> assertThat(cause.getMessage()).matches( - "(?i).*(factory has been closed|not accepting a task).*")); + "(?i).*(factory has been closed|not accepting a task" + + "|factory is closing or closed).*")); } @Test @@ -679,12 +803,17 @@ void doNotRetryWhenResponseIsAborted() throws Exception { .decorator(LoggingClient.newDecorator()) .build(); responseAbortServiceCallCounter.set(0); - final HttpResponse httpResponse = client.get("/response-abort"); + final HttpResponse httpResponse; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + httpResponse = client.get("/response-abort"); + ctx = captor.get(); + } if (abortCause == null) { httpResponse.abort(); } else { httpResponse.abort(abortCause); } + awaitValidClientRequestContext(ctx, 1); final RequestLog log = context.get().log().whenComplete().join(); final Throwable requestCause = log.requestCause(); @@ -712,23 +841,28 @@ void doNotRetryWhenResponseIsAborted() throws Exception { @Test void doNotRetryWhenSubscriberIsCancelled() throws Exception { final WebClient client = client(retryAlways); - client.get("/subscriber-cancel").subscribe( - new Subscriber() { - @Override - public void onSubscribe(Subscription s) { - s.cancel(); // Cancel as soon as getting the subscription. - } - @Override - public void onNext(HttpObject httpObject) {} + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + client.get("/subscriber-cancel").subscribe( + new Subscriber() { + @Override + public void onSubscribe(Subscription s) { + s.cancel(); // Cancel as soon as getting the subscription. + } + + @Override + public void onNext(HttpObject httpObject) {} - @Override - public void onError(Throwable t) {} + @Override + public void onError(Throwable t) {} - @Override - public void onComplete() {} - }); + @Override + public void onComplete() {} + }); + ctx = captor.get(); + } + awaitValidClientRequestContext(ctx, 1); TimeUnit.SECONDS.sleep(1L); // Sleep to check if there's a retry. assertThat(subscriberCancelServiceCallCounter.get()).isEqualTo(1); } @@ -755,11 +889,15 @@ void doNotRetryWhenRequestIsAborted() throws Exception { } else { req.abort(abortCause); } - client.execute(req).aggregate(); + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + client.execute(req).aggregate(); + ctx = captor.get(); + } TimeUnit.SECONDS.sleep(1); // No request is made. assertThat(responseAbortServiceCallCounter.get()).isZero(); + awaitValidClientRequestContext(ctx, 0); final RequestLog log = context.get().log().whenComplete().join(); if (abortCause == null) { assertThat(log.requestCause()).isExactlyInstanceOf(AbortedStreamException.class); @@ -785,9 +923,14 @@ void exceptionInDecorator() { .decorator(RetryingClient.newDecorator(strategy, 5)) .build(); - assertThatThrownBy(() -> client.get("/").aggregate().join()) - .hasCauseExactlyInstanceOf(AnticipatedException.class); + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + assertThatThrownBy(() -> client.get("/").aggregate().join()) + .isInstanceOf(CompletionException.class) + .hasCauseExactlyInstanceOf(AnticipatedException.class); + ctx = captor.get(); + } assertThat(retryCounter.get()).isEqualTo(5); + awaitValidClientRequestContext(ctx, 5); } @Test @@ -798,9 +941,13 @@ void exceptionInRule() { }; final WebClient client = client(rule); - assertThatThrownBy(client.get("/").aggregate()::join) - .isInstanceOf(CompletionException.class) - .hasCauseReference(exception); + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + assertThatThrownBy(client.get("/").aggregate()::join) + .isInstanceOf(CompletionException.class) + .hasCauseReference(exception); + ctx = captor.get(); + } + awaitValidClientRequestContext(ctx, 1); } @Test @@ -811,9 +958,13 @@ void exceptionInRuleWithContent() { }; final WebClient client = client(rule, 10000, 0, 100); - assertThatThrownBy(client.get("/").aggregate()::join) - .isInstanceOf(CompletionException.class) - .hasCauseReference(exception); + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + assertThatThrownBy(client.get("/").aggregate()::join) + .isInstanceOf(CompletionException.class) + .hasCauseReference(exception); + ctx = captor.get(); + } + awaitValidClientRequestContext(ctx, 1); } @Test @@ -828,11 +979,15 @@ void useSameEventLoopWhenAggregate() throws InterruptedException { }) .decorator(RetryingClient.newDecorator(RetryRule.failsafe(), 2)) .build(); - client.get("/503-then-success").aggregate().whenComplete((unused, cause) -> { - assertThat(eventLoop.get().inEventLoop()).isTrue(); - latch.countDown(); - }); + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + client.get("/503-then-success").aggregate().whenComplete((unused, cause) -> { + assertThat(eventLoop.get().inEventLoop()).isTrue(); + latch.countDown(); + }); + ctx = captor.get(); + } latch.await(); + awaitValidClientRequestContext(ctx, 2); } private WebClient client(RetryRule retryRule) { @@ -886,6 +1041,53 @@ private WebClient client(RetryRuleWithContent retryRuleWithContent .build(); } + private static void waitForEveryAttemptResponse(ClientRequestContext ctx) { + final State state = AbstractRetryingClient.state(ctx); + final int numStartedAttempts = state.startedAttempts().size(); + assertThat(state.isRetryingComplete()).isTrue(); + + await().untilAsserted(() -> { + final List> startedAttempts = state.startedAttempts(); + assertThat(startedAttempts).hasSize(numStartedAttempts); + for (Attempt attempt : startedAttempts) { + assertThat(attempt.attemptRes().isComplete()).isTrue(); + } + }); + } + + private static void assertValidRetryingState(ClientRequestContext ctx, int expectedAttempts) { + final State state = AbstractRetryingClient.state(ctx); + assertThat(state.isRetryingComplete()).isTrue(); + + final List> startedAttempts = state.startedAttempts(); + assertThat(startedAttempts).hasSize(expectedAttempts); + for (Attempt attempt : startedAttempts) { + assertThat(attempt.attemptRes().isComplete()).isTrue(); + } + } + + private static void awaitValidClientRequestContext(ClientRequestContext ctx) { + awaitValidClientRequestContext(ctx, ctx.log().children().size()); + } + + private static void awaitValidClientRequestContext(ClientRequestContext ctx, int expectedNumRequests) { + await().untilAsserted(() -> { + assertThat(ctx.log().isComplete()).isTrue(); + assertThat(ctx.log().children()).hasSize(expectedNumRequests); + ctx.log().children().forEach(childLogAccess -> { + try { + final RequestLog childLog = childLogAccess.whenComplete().get(); + assertThat(childLog).isNotNull(); + assertThat(childLog.isComplete()).isTrue(); + } catch (InterruptedException | ExecutionException e) { + fail(e); + } + }); + + assertValidRetryingState(ctx, expectedNumRequests); + }); + } + private static class RetryIfContentMatch implements RetryRuleWithContent { private final String retryContent; private final RetryDecision decision = RetryDecision.retry(Backoff.fixed(100)); diff --git a/core/src/test/java/com/linecorp/armeria/client/retry/RetryingClientWithHedgingTest.java b/core/src/test/java/com/linecorp/armeria/client/retry/RetryingClientWithHedgingTest.java new file mode 100644 index 00000000000..daa10ae73fe --- /dev/null +++ b/core/src/test/java/com/linecorp/armeria/client/retry/RetryingClientWithHedgingTest.java @@ -0,0 +1,842 @@ +/* + * 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.linecorp.armeria.client.retry.AbstractRetryingClient.ARMERIA_RETRY_COUNT; +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 static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.nio.charset.Charset; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BiFunction; +import java.util.function.Consumer; +import java.util.function.Function; + +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +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.ResponseCancellationException; +import com.linecorp.armeria.client.ResponseTimeoutException; +import com.linecorp.armeria.client.WebClient; +import com.linecorp.armeria.client.WebClientBuilder; +import com.linecorp.armeria.client.endpoint.EndpointGroup; +import com.linecorp.armeria.client.endpoint.EndpointSelectionStrategy; +import com.linecorp.armeria.client.logging.LoggingClient; +import com.linecorp.armeria.client.retry.AbstractRetryingClient.State; +import com.linecorp.armeria.client.retry.AbstractRetryingClient.State.Attempt; +import com.linecorp.armeria.common.AggregatedHttpResponse; +import com.linecorp.armeria.common.ExchangeType; +import com.linecorp.armeria.common.HttpRequest; +import com.linecorp.armeria.common.HttpResponse; +import com.linecorp.armeria.common.HttpStatus; +import com.linecorp.armeria.common.MediaType; +import com.linecorp.armeria.common.SessionProtocol; +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.RequestLogProperty; +import com.linecorp.armeria.common.util.Exceptions; +import com.linecorp.armeria.server.HttpService; +import com.linecorp.armeria.server.RoutingContext; +import com.linecorp.armeria.server.ServerBuilder; +import com.linecorp.armeria.server.ServiceRequestContext; +import com.linecorp.armeria.server.logging.LoggingService; +import com.linecorp.armeria.testing.junit5.server.ServerExtension; + +class RetryingClientWithHedgingTest { + private static class TestServer extends ServerExtension { + private CountDownLatch responseLatch = new CountDownLatch(1); + private final AtomicInteger numRequests = new AtomicInteger(); + private HttpService helloService = mock(HttpService.class); + + TestServer() { + super(true); + } + + @Override + protected void configure(ServerBuilder sb) throws Exception { + sb.decorator(LoggingService.newDecorator()); + sb.blockingTaskExecutor(5); + + sb.service("/hello", + new HttpService() { + @Override + public HttpResponse serve(ServiceRequestContext ctx, HttpRequest req) + throws Exception { + + final CompletableFuture responseFuture = new CompletableFuture<>(); + final HttpResponse res = HttpResponse.of(responseFuture); + + // We are using a blocking task executor to not block the event loop so we are + // able to receive request cancellations. + ctx.blockingTaskExecutor().execute(() -> { + if (numRequests.incrementAndGet() != 1) { + responseFuture.completeExceptionally(new IllegalStateException( + "Expected only one request, but got: " + numRequests.get())); + return; + } + + try { + responseLatch.await(); + } catch (InterruptedException e) { + responseFuture.completeExceptionally(e); + fail(e); + } + + try { + responseFuture.complete(helloService.serve(ctx, req)); + } catch (Exception e) { + responseFuture.completeExceptionally(e); + } + }); + + return res; + } + + @Override + public ExchangeType exchangeType(RoutingContext routingContext) { + return ExchangeType.UNARY; + } + }); + } + + private void reset() { + helloService = mock(HttpService.class); + responseLatch = new CountDownLatch(1); + numRequests.set(0); + } + + public HttpService getHelloService() { + return helloService; + } + + public int getNumRequests() { + return numRequests.get(); + } + + public void unlatchResponse() { + responseLatch.countDown(); + } + } + + private static final String SERVER1_RESPONSE = "s1"; + private static final String SERVER2_RESPONSE = "s2#"; + private static final String SERVER3_RESPONSE = "s3##"; + + @RegisterExtension + private static final TestServer server1 = new TestServer(); + @RegisterExtension + private static final TestServer server2 = new TestServer(); + @RegisterExtension + private static final TestServer server3 = new TestServer(); + + private @Nullable ClientRequestContext ctx; + private static ClientFactory clientFactory; + private static final RetryRule NO_RETRY_RULE = RetryRule.builder().thenNoRetry(); + + @BeforeAll + static void beforeAll() { + // use different eventLoop from server's so that clients don't hang when the eventLoop in server hangs + clientFactory = ClientFactory.builder() + .workerGroup(5).build(); + } + + @AfterAll + static void afterAll() { + clientFactory.closeAsync(); + } + + @BeforeEach + void beforeEach() { + server1.reset(); + server2.reset(); + server3.reset(); + } + + @AfterEach + void afterEach() { + server1.unlatchResponse(); + server2.unlatchResponse(); + server3.unlatchResponse(); + } + + @Test + void firstServerWins() throws Exception { + when(server1.getHelloService().serve(any(), any())).thenReturn(HttpResponse.of(SERVER1_RESPONSE)); + + final RetryConfig hedgingNoRetryConfig = RetryConfig + .builder(NO_RETRY_RULE) + .maxTotalAttempts(3) + .hedgingDelayMillis(500) + .build(); + + final WebClient client = client(hedgingNoRetryConfig); + + final CompletableFuture responseFuture; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + responseFuture = client.get("/hello").aggregate(); + ctx = captor.get(); + } + + server1.unlatchResponse(); + + await().untilAsserted(() -> { + assertValidAggregatedResponse(responseFuture, SERVER1_RESPONSE); + assertValidClientRequestContext( + ctx, GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER1_RESPONSE), + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER1_RESPONSE), null, null); + + assertValidServerRequestContext(server1, 1, false); + assertNoServerRequestContext(server2); + assertNoServerRequestContext(server3); + }); + } + + @Test + void secondServerWins() throws Exception { + when(server2.getHelloService().serve(any(), any())).thenReturn(HttpResponse.of(SERVER2_RESPONSE)); + + final RetryConfig hedgingNoRetryConfig = RetryConfig + .builder(NO_RETRY_RULE) + .maxTotalAttempts(3) + .hedgingDelayMillis(500) + .build(); + + final WebClient client = client(hedgingNoRetryConfig); + + final CompletableFuture responseFuture; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + responseFuture = client.get("/hello").aggregate(); + ctx = captor.get(); + } + + await().untilAsserted(() -> { + assertThat(server1.getNumRequests()).isOne(); + assertThat(server3.getNumRequests()).isOne(); + }); + + server2.unlatchResponse(); + + await() + .untilAsserted(() -> { + assertValidAggregatedResponse(responseFuture, SERVER2_RESPONSE); + + assertValidClientRequestContext( + ctx, GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER2_RESPONSE), + VERIFY_REQUEST_CANCELLED, + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER2_RESPONSE), + VERIFY_REQUEST_CANCELLED + ); + + assertValidServerRequestContext(server1, 1, true); + assertValidServerRequestContext(server2, 2, false); + assertValidServerRequestContext(server3, 3, true); + }); + } + + @Test + void thirdServerWins() throws Exception { + when(server3.getHelloService().serve(any(), any())).thenReturn(HttpResponse.of(SERVER3_RESPONSE)); + + final RetryConfig hedgingNoRetryConfig = RetryConfig + .builder(NO_RETRY_RULE) + .maxTotalAttempts(3) + .hedgingDelayMillis(500) + .build(); + + final WebClient client = client(hedgingNoRetryConfig); + + final CompletableFuture responseFuture; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + responseFuture = client.get("/hello").aggregate(); + ctx = captor.get(); + } + + await().untilAsserted(() -> { + assertThat(server1.getNumRequests()).isOne(); + assertThat(server2.getNumRequests()).isOne(); + }); + + server3.unlatchResponse(); + + await() + .untilAsserted(() -> { + assertValidAggregatedResponse(responseFuture, SERVER3_RESPONSE); + assertValidClientRequestContext( + ctx, GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER3_RESPONSE), + VERIFY_REQUEST_CANCELLED, + VERIFY_REQUEST_CANCELLED, + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER3_RESPONSE) + ); + + assertValidServerRequestContext(server1, 1, true); + assertValidServerRequestContext(server2, 2, true); + assertValidServerRequestContext(server3, 3, false); + }); + } + + @Test + void respectsShorterBackoffTriggeringFasterRetry() throws Exception { + when(server1.getHelloService().serve(any(), any())) + .thenReturn(HttpResponse.of(HttpStatus.TOO_MANY_REQUESTS, MediaType.PLAIN_TEXT, + SERVER1_RESPONSE)); + when(server2.getHelloService().serve(any(), any())) + .thenReturn(HttpResponse.of(HttpStatus.TOO_MANY_REQUESTS, MediaType.PLAIN_TEXT, + SERVER2_RESPONSE)); + when(server3.getHelloService().serve(any(), any())) + .thenReturn(HttpResponse.of(HttpStatus.TOO_MANY_REQUESTS, MediaType.PLAIN_TEXT, + SERVER3_RESPONSE)); + + class SpyableAttemptLimitingBackoff implements Backoff { + private final Backoff delegate = Backoff.withoutDelay().withMaxAttempts(2); + + @Override + public long nextDelayMillis(int numAttemptsSoFar) { + return delegate.nextDelayMillis(numAttemptsSoFar); + } + } + + final Backoff backoff = spy(new SpyableAttemptLimitingBackoff()); + + final RetryConfig config = RetryConfig + .builder(RetryRule.of( + RetryRule + .builder() + .onStatus(HttpStatus.TOO_MANY_REQUESTS) + .thenBackoff(backoff), + NO_RETRY_RULE + )) + .maxTotalAttempts(3) + .hedgingDelayMillis(500) + .build(); + + final WebClient client = client(config); + final CompletableFuture responseFuture; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + responseFuture = client.get("/hello").aggregate(); + ctx = captor.get(); + } + + server2.unlatchResponse(); + server3.unlatchResponse(); + + // 500ms for the hedging request to server 2 and 250ms tolerance. + await().atMost(500 + 250, TimeUnit.MILLISECONDS).untilAsserted(() -> { + assertThat(server1.getNumRequests()).isOne(); + // Server 1 is blocking, after 500ms hedging request to server 2 will come to the rescue. + assertThat(server2.getNumRequests()).isOne(); + // Server 2 responds immediately, backoff will order retry immediately. + // This cancels the hedging request that came with the request to server 2. + assertThat(server3.getNumRequests()).isOne(); + // The hedged request to server 3 will not be issued because we are out of total attempts. + // Server 3 responds immediately, but we will not continue as we are out of attempts. + + verify(backoff, times(1)).nextDelayMillis(1); + }); + + // Should be enough for the client to receive and process the two responses from server + // 2 and server 3. + Thread.sleep(500); + // Server 1 answers, again with TOO_MANY_REQUESTS, and it tries to retry. + // We exceeded the number of attempts _of the backoff_ so we should stop retrying and take the last + // response which is the response from server 1. + server1.unlatchResponse(); + + await().untilAsserted(() -> { + assertValidAggregatedResponse(responseFuture, HttpStatus.TOO_MANY_REQUESTS, SERVER1_RESPONSE); + assertValidClientRequestContext( + ctx, GET_VERIFY_RESPONSE_HAS_CONTENT_AND_STATUS.apply(SERVER1_RESPONSE, + HttpStatus.TOO_MANY_REQUESTS), + GET_VERIFY_RESPONSE_HAS_CONTENT_AND_STATUS.apply(SERVER1_RESPONSE, + HttpStatus.TOO_MANY_REQUESTS), + GET_VERIFY_RESPONSE_HAS_CONTENT_AND_STATUS.apply(SERVER2_RESPONSE, + HttpStatus.TOO_MANY_REQUESTS), + GET_VERIFY_RESPONSE_HAS_CONTENT_AND_STATUS.apply(SERVER3_RESPONSE, + HttpStatus.TOO_MANY_REQUESTS) + ); + + assertValidServerRequestContext(server1, 1, false); + assertValidServerRequestContext(server2, 2, false); + assertValidServerRequestContext(server3, 3, false); + }); + } + + @Test + void allServerLosePickLastResponse() throws Exception { + when(server1.getHelloService().serve(any(), any())) + .thenReturn(HttpResponse.of(HttpStatus.TOO_MANY_REQUESTS, MediaType.PLAIN_TEXT, + SERVER1_RESPONSE)); + when(server2.getHelloService().serve(any(), any())) + .thenReturn(HttpResponse.of(HttpStatus.TOO_MANY_REQUESTS, MediaType.PLAIN_TEXT, + SERVER2_RESPONSE)); + when(server3.getHelloService().serve(any(), any())) + .thenReturn(HttpResponse.of(HttpStatus.TOO_MANY_REQUESTS, MediaType.PLAIN_TEXT, + SERVER3_RESPONSE)); + + final RetryConfig config = RetryConfig + .builder(RetryRule.of( + RetryRule + .builder() + .onStatus(HttpStatus.TOO_MANY_REQUESTS) + .thenBackoff(Backoff.fixed(10_000)), + NO_RETRY_RULE + )) + .maxTotalAttempts(3) + .hedgingDelayMillis(500) + .build(); + + final WebClient client = client(config); + final CompletableFuture responseFuture; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + responseFuture = client.get("/hello").aggregate(); + ctx = captor.get(); + } + + server1.unlatchResponse(); + server2.unlatchResponse(); + server3.unlatchResponse(); + + await().untilAsserted(() -> { + assertValidAggregatedResponse(responseFuture, HttpStatus.TOO_MANY_REQUESTS, SERVER3_RESPONSE); + assertValidClientRequestContext(ctx, + GET_VERIFY_RESPONSE_HAS_CONTENT_AND_STATUS.apply( + SERVER3_RESPONSE, + HttpStatus.TOO_MANY_REQUESTS + ), + GET_VERIFY_RESPONSE_HAS_CONTENT_AND_STATUS.apply( + SERVER1_RESPONSE, + HttpStatus.TOO_MANY_REQUESTS + ), + GET_VERIFY_RESPONSE_HAS_CONTENT_AND_STATUS.apply( + SERVER2_RESPONSE, + HttpStatus.TOO_MANY_REQUESTS + ), + GET_VERIFY_RESPONSE_HAS_CONTENT_AND_STATUS.apply( + SERVER3_RESPONSE, + HttpStatus.TOO_MANY_REQUESTS + ) + ); + + assertValidServerRequestContext(server1, 1, false); + assertValidServerRequestContext(server2, 2, false); + assertValidServerRequestContext(server3, 3, false); + }); + } + + @Test + void thirdServerWinsEvenAfterPerAttemptTimeout() throws Exception { + when(server1.getHelloService().serve(any(), any())).thenReturn(HttpResponse.of(SERVER1_RESPONSE)); + when(server2.getHelloService().serve(any(), any())).thenReturn(HttpResponse.of(SERVER2_RESPONSE)); + when(server3.getHelloService().serve(any(), any())).thenReturn(HttpResponse.of(SERVER3_RESPONSE)); + + final RetryConfig hedgingNoRetryConfig = RetryConfig + .builder(RetryRule.builder().onTimeoutException().thenBackoff(Backoff.fixed(10_000))) // should + // be always overtaken by hedging task + .maxTotalAttempts(3) + .responseTimeoutMillisForEachAttempt(100) + .hedgingDelayMillis(200) + .build(); + + final WebClient client = clientBuilder() + .decorator( + RetryingClient.newDecorator(hedgingNoRetryConfig) + ) + .build(); + + final CompletableFuture responseFuture; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + responseFuture = client.get("/hello").aggregate(); + ctx = captor.get(); + } + + await().pollInterval(25, TimeUnit.MILLISECONDS).untilAsserted(() -> { + assertThat(server3.getNumRequests()).isOne(); + }); + + // As we know that the third server received a request, we know that + // we are at >= T + 400. This means that: + server1.unlatchResponse(); // issued at T (timed out) + server2.unlatchResponse(); // issued at T + 200 (timed out) + server3.unlatchResponse(); // issued at T + 400 (hopefully not timed out yet) + + await() + .untilAsserted(() -> { + assertValidServerRequestContext(server1, 1, true); + assertValidServerRequestContext(server2, 2, true); + assertValidServerRequestContext(server3, 3, false); + + assertValidAggregatedResponse(responseFuture, SERVER3_RESPONSE); + assertValidClientRequestContext( + ctx, GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER3_RESPONSE), + VERIFY_REQUEST_TIMED_OUT, + VERIFY_REQUEST_TIMED_OUT, GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER3_RESPONSE) + ); + }); + } + + @Test + void thirdServerWinsEvenAfterRetriableResponse() throws Exception { + when(server1.getHelloService().serve(any(), any())) + .thenReturn(HttpResponse.of(HttpStatus.TOO_MANY_REQUESTS, MediaType.PLAIN_TEXT, + SERVER1_RESPONSE)); + when(server2.getHelloService().serve(any(), any())) + .thenReturn(HttpResponse.of(HttpStatus.TOO_MANY_REQUESTS, MediaType.PLAIN_TEXT, + SERVER2_RESPONSE)); + when(server3.getHelloService().serve(any(), any())) + .thenReturn(HttpResponse.of(SERVER3_RESPONSE)); + + final RetryConfig config = RetryConfig + .builder(RetryRule.of( + RetryRule + .builder() + .onStatus(HttpStatus.TOO_MANY_REQUESTS) + .thenBackoff(Backoff.withoutDelay()), + NO_RETRY_RULE + )) + .maxTotalAttempts(3) + .hedgingDelayMillis(200) + .build(); + + final WebClient client = client(config); + + final CompletableFuture responseFuture; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + responseFuture = client.get("/hello").aggregate(); + ctx = captor.get(); + } + + server1.unlatchResponse(); + server2.unlatchResponse(); + server3.unlatchResponse(); + + await().untilAsserted(() -> { + assertThat(server1.getNumRequests()).isOne(); + assertThat(server2.getNumRequests()).isOne(); + assertThat(server3.getNumRequests()).isOne(); + }); + + await().untilAsserted(() -> { + assertValidServerRequestContext(server1, 1); + assertValidServerRequestContext(server2, 2); + assertValidServerRequestContext(server3, 3); + + assertValidAggregatedResponse(responseFuture, SERVER3_RESPONSE); + assertValidClientRequestContext( + ctx, GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER3_RESPONSE), + GET_VERIFY_RESPONSE_HAS_CONTENT_AND_STATUS.apply(SERVER1_RESPONSE, + HttpStatus.TOO_MANY_REQUESTS), + GET_VERIFY_RESPONSE_HAS_CONTENT_AND_STATUS.apply(SERVER2_RESPONSE, + HttpStatus.TOO_MANY_REQUESTS), + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER3_RESPONSE) + ); + }); + } + + @Test + void loosesAfterNonRetriableResponse() throws Exception { + when(server1.getHelloService().serve(any(), any())) + .thenReturn(HttpResponse.of(SERVER1_RESPONSE)); + when(server2.getHelloService().serve(any(), any())) + .thenReturn(HttpResponse.of(HttpStatus.INTERNAL_SERVER_ERROR, MediaType.PLAIN_TEXT, + SERVER2_RESPONSE)); + + final RetryConfig config = RetryConfig + .builder(NO_RETRY_RULE) + .maxTotalAttempts(3) + // Should be long enough so we can complete the second request before we continue issuing + // a third request to the third server. + .hedgingDelayMillis(500) + .build(); + + final WebClient client = client(config); + + final CompletableFuture responseFuture; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + responseFuture = client.get("/hello").aggregate(); + ctx = captor.get(); + } + + await().untilAsserted(() -> assertThat(server1.getNumRequests()).isEqualTo(1)); + + server2.unlatchResponse(); + + await().untilAsserted(() -> { + assertThat(server2.getNumRequests()).isOne(); + }); + + Thread.sleep(500); + server1.unlatchResponse(); + server2.unlatchResponse(); + + await().untilAsserted(() -> { + assertValidServerRequestContext(server1, 1, true); + assertValidServerRequestContext(server2, 2, false); + assertNoServerRequestContext(server3); + + assertValidAggregatedResponse(responseFuture, HttpStatus.INTERNAL_SERVER_ERROR, SERVER2_RESPONSE); + assertValidClientRequestContext( + ctx, GET_VERIFY_RESPONSE_HAS_CONTENT_AND_STATUS.apply(SERVER2_RESPONSE, + HttpStatus.INTERNAL_SERVER_ERROR), + VERIFY_REQUEST_CANCELLED, + GET_VERIFY_RESPONSE_HAS_CONTENT_AND_STATUS.apply(SERVER2_RESPONSE, + HttpStatus.INTERNAL_SERVER_ERROR), null + ); + }); + } + + @Test + void loosesAfterResponseTimeout() throws Exception { + final RetryConfig config = RetryConfig + .builder(NO_RETRY_RULE) + .maxTotalAttempts(3) + .hedgingDelayMillis(0) + .build(); + + final WebClient client = clientBuilder() + .responseTimeoutMillis(500) + .decorator( + RetryingClient.newDecorator(config) + ).build(); + + final CompletableFuture responseFuture; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + responseFuture = client.get("/hello").aggregate(); + ctx = captor.get(); + } + + await().untilAsserted(() -> { + assertThat(responseFuture).isCompletedExceptionally(); + assertThatThrownBy(responseFuture::get).satisfies(throwable -> { + final Throwable rootCause = Exceptions.peel(throwable); + assertThat(rootCause).isInstanceOf(ResponseTimeoutException.class); + }); + + final List childLogExceptions = new ArrayList<>(); + + final RequestLogVerifier catchException = log -> { + assertThat(log.responseCause()).isNotNull(); + childLogExceptions.add(log.responseCause()); + }; + + assertValidClientRequestContext(ctx, VERIFY_RESPONSE_TIMEOUT, catchException, catchException, + catchException); + + int numTimeouts = 0; + int numCancelled = 0; + + for (final @Nullable Throwable childException : childLogExceptions) { + if (childException instanceof ResponseTimeoutException) { + numTimeouts++; + } else if (childException instanceof ResponseCancellationException) { + numCancelled++; + } else { + fail("Unexpected exception: " + childException); + } + } + + assertThat(numTimeouts + numCancelled).isEqualTo(3); + // At least one attempt needs to time out. + assertThat(numTimeouts).isPositive(); + + assertValidServerRequestContext(server1, 1, true); + assertValidServerRequestContext(server2, 2, true); + assertValidServerRequestContext(server3, 3, true); + }); + } + + // todo(szymon): test being able to set different hedging delays for different servers + + private static WebClientBuilder clientBuilder() { + return WebClient.builder(SessionProtocol.H2C, + EndpointGroup.of(EndpointSelectionStrategy.roundRobin(), + server1.httpEndpoint(), + server2.httpEndpoint(), + server3.httpEndpoint())) + .requestAutoAbortDelayMillis(0) + .decorator(LoggingClient.newDecorator()) + .factory(clientFactory); + } + + private static WebClient client(RetryConfig config) { + return clientBuilder() + .decorator(RetryingClient.newDecorator(config)).build(); + } + + private static String getResponseContent(AggregatedHttpResponse response) { + return response.content().toString(Charset.defaultCharset()); + } + + private static void assertValidAggregatedResponse(CompletableFuture resFuture, + String expectedContent) { + assertThat(resFuture.isDone()).isTrue(); + final AggregatedHttpResponse res = resFuture.getNow(null); + assertThat(res.status()).isEqualTo(HttpStatus.OK); + assertThat(getResponseContent(res)).isEqualTo(expectedContent); + } + + private static void assertValidAggregatedResponse(CompletableFuture resFuture, + HttpStatus expectedStatus, String expectedContent) { + assertThat(resFuture.isDone()).isTrue(); + final AggregatedHttpResponse res = resFuture.getNow(null); + assertThat(getResponseContent(res)).isEqualTo(expectedContent); + assertThat(res.status()).isEqualTo(expectedStatus); + } + + private static void assertValidRootClientRequestContext(ClientRequestContext ctx, + RequestLogVerifier logVerifierCtx, + int expectedNumChildren) { + assertThat(ctx.log().isComplete()).isTrue(); + assertThat(ctx.log().children()).hasSize(expectedNumChildren); + final RequestLog log = ctx.log().getIfAvailable(RequestLogProperty.RESPONSE_CONTENT, + RequestLogProperty.RESPONSE_CAUSE, + RequestLogProperty.REQUEST_HEADERS); + assertThat(log).isNotNull(); + logVerifierCtx.accept(log); + } + + private static void assertValidRetryingState(ClientRequestContext ctx, int expectedAttempts) { + final State state = AbstractRetryingClient.state(ctx); + assertThat(state.isRetryingComplete()).isTrue(); + + final List> startedAttempts = state.startedAttempts(); + assertThat(startedAttempts).hasSize(expectedAttempts); + for (Attempt attempt : startedAttempts) { + assertThat(attempt.attemptRes().isComplete()).isTrue(); + } + } + + private static void assertValidClientRequestContext(ClientRequestContext ctx, + RequestLogVerifier logVerifierCtx, + RequestLogVerifier logVerifierServer1, + @Nullable RequestLogVerifier logVerifierServer2, + @Nullable RequestLogVerifier logVerifierServer3 + ) { + final int expectedAttempts = 1 + (logVerifierServer2 == null ? 0 : 1) + + (logVerifierServer3 == null ? 0 : 1); + + assertValidRootClientRequestContext(ctx, logVerifierCtx, expectedAttempts); + assertValidChildLog(ctx.log().children().get(0), 1, logVerifierServer1); + if (logVerifierServer2 != null) { + assertValidChildLog(ctx.log().children().get(1), 2, logVerifierServer2); + } + + if (logVerifierServer3 != null) { + assertValidChildLog(ctx.log().children().get(2), 3, logVerifierServer3); + } + + assertValidRetryingState(ctx, expectedAttempts); + } + + private static void assertValidChildLog(RequestLogAccess logAccess, int attemptNumber, + RequestLogVerifier requestLogVerifier) { + assertThat(logAccess.isComplete()).isTrue(); + // After the check right above, all properties of the RequestLog should be available. + final @Nullable RequestLog log = logAccess.getIfAvailable(RequestLogProperty.RESPONSE_CONTENT, + RequestLogProperty.RESPONSE_CAUSE, + RequestLogProperty.REQUEST_HEADERS); + assertThat(log).isNotNull(); + + if (attemptNumber > 1) { + assertThat(log.requestHeaders().getInt(ARMERIA_RETRY_COUNT)).isEqualTo(attemptNumber - 1); + } else { + assertThat(log.requestHeaders().contains(ARMERIA_RETRY_COUNT)).isFalse(); + } + + requestLogVerifier.accept(log); + } + + private static void assertValidServerRequestContext(ServerExtension server, int attemptNumber) { + assertValidServerRequestContext(server, attemptNumber, false); + } + + private static void assertValidServerRequestContext(ServerExtension server, int attemptNumber, + boolean expectCancelled) { + assertThat(server.requestContextCaptor().size()).isEqualTo(1); + + final ServiceRequestContext sctx = server.requestContextCaptor().peek(); + assertThat(sctx).isNotNull(); + assertThat(sctx.log().isComplete()).isTrue(); + + assertThat(sctx.isCancelled()).isEqualTo(expectCancelled); + + final RequestLog slog = sctx.log().getIfAvailable(RequestLogProperty.REQUEST_HEADERS, + RequestLogProperty.REQUEST_CONTENT); + + assertThat(slog).isNotNull(); + if (attemptNumber > 1) { + assertThat(slog.requestHeaders().getInt(ARMERIA_RETRY_COUNT)).isEqualTo(attemptNumber - 1); + } else { + assertThat(slog.requestHeaders().contains(ARMERIA_RETRY_COUNT)).isFalse(); + } + + assertThat(slog.requestHeaders().path()).contains("hello"); + } + + private static void assertNoServerRequestContext(ServerExtension server) { + assertThat(server.requestContextCaptor().size()).isEqualTo(0); + } + + @FunctionalInterface + private interface RequestLogVerifier extends Consumer {} + + private static final RequestLogVerifier VERIFY_REQUEST_CANCELLED = + log -> { + assertThat(log.responseCause()).isInstanceOf(ResponseCancellationException.class); + }; + + private static final RequestLogVerifier VERIFY_REQUEST_TIMED_OUT = + log -> { + assertThat(log.responseCause()).isInstanceOf(ResponseTimeoutException.class); + }; + // + private static final RequestLogVerifier VERIFY_RESPONSE_TIMEOUT = + log -> { + assertThat(log.responseCause()).isInstanceOf(ResponseTimeoutException.class); + }; + + private static final BiFunction + GET_VERIFY_RESPONSE_HAS_CONTENT_AND_STATUS = + (expectedResponseContent, expectedStatus) -> log -> { + assertThat(log.responseLength()).isEqualTo(expectedResponseContent.length()); + assertThat(log.responseStatus()).isEqualTo(expectedStatus); + }; + + private static final Function GET_VERIFY_RESPONSE_HAS_CONTENT = + expectedResponseContent -> GET_VERIFY_RESPONSE_HAS_CONTENT_AND_STATUS.apply(expectedResponseContent, + HttpStatus.OK); +} diff --git a/junit5/src/main/java/com/linecorp/armeria/testing/server/ServiceRequestContextCaptor.java b/junit5/src/main/java/com/linecorp/armeria/testing/server/ServiceRequestContextCaptor.java index de83161c4b7..f2511fb01c6 100644 --- a/junit5/src/main/java/com/linecorp/armeria/testing/server/ServiceRequestContextCaptor.java +++ b/junit5/src/main/java/com/linecorp/armeria/testing/server/ServiceRequestContextCaptor.java @@ -151,6 +151,15 @@ public ServiceRequestContext take() throws InterruptedException { return serviceContexts.take(); } + /** + * Retrieves, but does not remove, the first captured {@link ServiceRequestContext}, or returns + * {@code null} if there are no captured contexts. + */ + @Nullable + public ServiceRequestContext peek() { + return serviceContexts.peek(); + } + /** * Retrieves and removes the first captured {@link ServiceRequestContext}, waiting up to * {@value DEFAULT_TIMEOUT_IN_SECONDS} seconds if necessary until an element becomes available. diff --git a/thrift/thrift0.13/src/test/java/com/linecorp/armeria/it/client/retry/RetryingRpcClientWithHedgingTest.java b/thrift/thrift0.13/src/test/java/com/linecorp/armeria/it/client/retry/RetryingRpcClientWithHedgingTest.java new file mode 100644 index 00000000000..537586df49a --- /dev/null +++ b/thrift/thrift0.13/src/test/java/com/linecorp/armeria/it/client/retry/RetryingRpcClientWithHedgingTest.java @@ -0,0 +1,787 @@ +/* + * 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.it.client.retry; + +import static com.linecorp.armeria.client.retry.AbstractRetryingClient.ARMERIA_RETRY_COUNT; +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 static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BiFunction; +import java.util.function.Consumer; +import java.util.function.Function; + +import org.apache.thrift.TApplicationException; +import org.apache.thrift.TException; +import org.apache.thrift.async.AsyncMethodCallback; +import org.apache.thrift.transport.TTransportException; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; +import org.mockito.Mockito; + +import com.linecorp.armeria.client.ClientRequestContext; +import com.linecorp.armeria.client.ClientRequestContextCaptor; +import com.linecorp.armeria.client.Clients; +import com.linecorp.armeria.client.ResponseCancellationException; +import com.linecorp.armeria.client.ResponseTimeoutException; +import com.linecorp.armeria.client.endpoint.EndpointGroup; +import com.linecorp.armeria.client.endpoint.EndpointSelectionStrategy; +import com.linecorp.armeria.client.retry.AbstractRetryingClient; +import com.linecorp.armeria.client.retry.AbstractRetryingClient.State; +import com.linecorp.armeria.client.retry.AbstractRetryingClient.State.Attempt; +import com.linecorp.armeria.client.retry.Backoff; +import com.linecorp.armeria.client.retry.RetryConfig; +import com.linecorp.armeria.client.retry.RetryRule; +import com.linecorp.armeria.client.retry.RetryRuleWithContent; +import com.linecorp.armeria.client.retry.RetryingRpcClient; +import com.linecorp.armeria.client.thrift.ThriftClients; +import com.linecorp.armeria.common.RpcRequest; +import com.linecorp.armeria.common.RpcResponse; +import com.linecorp.armeria.common.SessionProtocol; +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.RequestLogProperty; +import com.linecorp.armeria.common.stream.ClosedStreamException; +import com.linecorp.armeria.common.util.Exceptions; +import com.linecorp.armeria.internal.testing.AnticipatedException; +import com.linecorp.armeria.server.ServerBuilder; +import com.linecorp.armeria.server.ServiceRequestContext; +import com.linecorp.armeria.server.logging.LoggingService; +import com.linecorp.armeria.server.thrift.THttpService; +import com.linecorp.armeria.testing.junit5.server.ServerExtension; + +import testing.thrift.main.HelloService; +import testing.thrift.main.HelloService.Iface; + +class RetryingRpcClientWithHedgingTest { + private static class TestServer extends ServerExtension { + private CountDownLatch responseLatch = new CountDownLatch(1); + private final AtomicInteger numRequests = new AtomicInteger(); + + private volatile HelloService.Iface serviceHandler; + + TestServer() { + super(true); + } + + @Override + protected void configure(ServerBuilder sb) throws Exception { + sb.blockingTaskExecutor(5); + sb.service("/thrift", THttpService.builder().useBlockingTaskExecutor(true).addService( + (Iface) name -> { + numRequests.incrementAndGet(); + try { + responseLatch.await(); + } catch (InterruptedException e) { + fail(e); + } + + return getServiceHandler().hello(name); + }).build() + .decorate(LoggingService.newDecorator()) + ); + } + + private void reset() { + requestContextCaptor().clear(); + serviceHandler = mock(HelloService.Iface.class); + responseLatch = new CountDownLatch(1); + numRequests.set(0); + } + + public HelloService.Iface getServiceHandler() { + return serviceHandler; + } + + public int getNumRequests() { + return numRequests.get(); + } + + public void unlatchResponse() { + responseLatch.countDown(); + } + } + + private static final String RETRIABLE_RESPONSE = "please-retry-thanks"; + private static final String SERVER1_RESPONSE = "s1"; + private static final String SERVER2_RESPONSE = "s2#"; + private static final String SERVER3_RESPONSE = "s3##"; + + private static final RetryRule NO_RETRY_RULE = RetryRule.builder().thenNoRetry(); + + @RegisterExtension + private static final TestServer server1 = new TestServer(); + @RegisterExtension + private static final TestServer server2 = new TestServer(); + @RegisterExtension + private static final TestServer server3 = new TestServer(); + + @BeforeEach + void beforeEach() { + server1.reset(); + server2.reset(); + server3.reset(); + } + + @AfterEach + void afterEach() { + // Unblock all servers. + server1.unlatchResponse(); + server2.unlatchResponse(); + server3.unlatchResponse(); + } + + @Test + void firstServerWins() throws Exception { + when(server1.getServiceHandler().hello(anyString())).thenReturn(SERVER1_RESPONSE); + + final HelloService.AsyncIface client = client( + RetryConfig + .builderForRpc(NO_RETRY_RULE) + .maxTotalAttempts(3) + .hedgingDelayMillis(500) + .build() + ); + + final CompletableFuture responseFuture; + final ClientRequestContext ctx; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + responseFuture = asyncHelloWith(client); + ctx = captor.get(); + } + + // Let server 1 win. + server1.unlatchResponse(); + + await().untilAsserted(() -> { + assertThat(responseFuture.get()).isEqualTo(SERVER1_RESPONSE); + + assertValidServerRequestContext(server1, 1, false); + assertNoServerRequestContext(server2); + assertNoServerRequestContext(server3); + + assertValidClientRequestContext( + ctx, GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER1_RESPONSE), + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER1_RESPONSE), + null, + null + ); + }); + } + + @Test + void secondServerWins() throws Exception { + when(server2.getServiceHandler().hello(anyString())).thenReturn(SERVER2_RESPONSE); + + final RetryConfig hedgingNoRetryConfig = RetryConfig + .builderForRpc(NO_RETRY_RULE) + .maxTotalAttempts(3) + .hedgingDelayMillis(100) + .build(); + + final HelloService.AsyncIface client = client(hedgingNoRetryConfig); + + final CompletableFuture responseFuture; + final ClientRequestContext ctx; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + responseFuture = asyncHelloWith(client); + ctx = captor.get(); + } + + await().untilAsserted(() -> { + assertThat(server1.getNumRequests()).isEqualTo(1); + assertThat(server3.getNumRequests()).isEqualTo(1); + }); + + server2.unlatchResponse(); + + await() + .untilAsserted(() -> { + assertThat(responseFuture.get()).isEqualTo(SERVER2_RESPONSE); + + assertValidClientRequestContext( + ctx, GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER2_RESPONSE), + VERIFY_REQUEST_CANCELLED, + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER2_RESPONSE), + VERIFY_REQUEST_CANCELLED + ); + + assertValidServerRequestContext(server1, 1, true); + assertValidServerRequestContext(server2, 2, false); + assertValidServerRequestContext(server3, 3, true); + }); + } + + @Test + void thirdServerWins() throws Exception { + when(server3.getServiceHandler().hello(anyString())).thenReturn(SERVER3_RESPONSE); + + final HelloService.AsyncIface client = client( + RetryConfig + .builderForRpc(NO_RETRY_RULE) + .maxTotalAttempts(3) + .hedgingDelayMillis(10) + .build() + ); + + final CompletableFuture responseFuture; + final ClientRequestContext ctx; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + responseFuture = asyncHelloWith(client); + ctx = captor.get(); + } + + // Let server 3 win. + server3.unlatchResponse(); + + await().untilAsserted(() -> { + assertThat(responseFuture.get()).isEqualTo(SERVER3_RESPONSE); + + assertValidServerRequestContext(server1, 1, true); + assertValidServerRequestContext(server2, 2, true); + assertValidServerRequestContext(server3, 3, false); + + assertValidClientRequestContext( + ctx, GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER3_RESPONSE), + VERIFY_REQUEST_CANCELLED, + VERIFY_REQUEST_CANCELLED, + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER3_RESPONSE) + ); + }); + } + + @Test + void respectsShorterBackoffTriggeringFasterRetry() throws Exception { + when(server1.getServiceHandler().hello(anyString())).thenReturn(RETRIABLE_RESPONSE); + when(server2.getServiceHandler().hello(anyString())).thenReturn(RETRIABLE_RESPONSE); + when(server3.getServiceHandler().hello(anyString())).thenReturn(RETRIABLE_RESPONSE); + + class SpyableAttemptLimitingBackoff implements Backoff { + private final Backoff delegate = Backoff.withoutDelay().withMaxAttempts(2); + + @Override + public long nextDelayMillis(int numAttemptsSoFar) { + return delegate.nextDelayMillis(numAttemptsSoFar); + } + } + + final Backoff backoff = Mockito.spy(new SpyableAttemptLimitingBackoff()); + + final RetryConfig config = RetryConfig.builderForRpc( + getRetryRetriableResponsesRule(backoff)) + .maxTotalAttempts(3) + .hedgingDelayMillis(500) + .build(); + + final HelloService.AsyncIface client = client(config); + final CompletableFuture responseFuture; + final ClientRequestContext ctx; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + responseFuture = asyncHelloWith(client); + ctx = captor.get(); + } + + server2.unlatchResponse(); + server3.unlatchResponse(); + + // 500ms for the hedging request to server 2 and 250ms tolerance. + await().atMost(500 + 250, TimeUnit.MILLISECONDS).untilAsserted(() -> { + assertThat(server1.getNumRequests()).isEqualTo(1); + assertThat(server2.getNumRequests()).isEqualTo(1); + assertThat(server3.getNumRequests()).isEqualTo(1); + Mockito.verify(backoff, Mockito.times(1)).nextDelayMillis(1); + }); + + Thread.sleep(500); + server1.unlatchResponse(); + + await().untilAsserted(() -> { + assertThat(responseFuture.get()).isEqualTo(RETRIABLE_RESPONSE); + assertValidClientRequestContext( + ctx, + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(RETRIABLE_RESPONSE), + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(RETRIABLE_RESPONSE), + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(RETRIABLE_RESPONSE), + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(RETRIABLE_RESPONSE) + ); + assertValidServerRequestContext(server1, 1, false); + assertValidServerRequestContext(server2, 2, false); + assertValidServerRequestContext(server3, 3, false); + }); + } + + @Test + void allServerLosePickLastResponse() throws Exception { + when(server1.getServiceHandler().hello(anyString())).thenReturn(RETRIABLE_RESPONSE); + when(server2.getServiceHandler().hello(anyString())).thenReturn(RETRIABLE_RESPONSE); + when(server3.getServiceHandler().hello(anyString())).thenReturn(RETRIABLE_RESPONSE); + + final RetryConfig config = RetryConfig.builderForRpc( + getRetryRetriableResponsesRule( + Backoff.fixed(10_000) + ) + ) + .maxTotalAttempts(3) + .hedgingDelayMillis(0) + .build(); + + final HelloService.AsyncIface client = client(config); + final CompletableFuture responseFuture; + final ClientRequestContext ctx; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + responseFuture = asyncHelloWith(client); + ctx = captor.get(); + } + + server1.unlatchResponse(); + server2.unlatchResponse(); + server3.unlatchResponse(); + + await().untilAsserted(() -> { + assertThat(responseFuture.get()).isEqualTo(RETRIABLE_RESPONSE); + + assertValidServerRequestContext(server1, 1, false); + assertValidServerRequestContext(server2, 2, false); + assertValidServerRequestContext(server3, 3, false); + + assertValidClientRequestContext( + ctx, + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(RETRIABLE_RESPONSE), + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(RETRIABLE_RESPONSE), + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(RETRIABLE_RESPONSE), + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(RETRIABLE_RESPONSE) + ); + }); + } + + @Test + void thirdServerWinsEvenAfterPerAttemptTimeout() throws Exception { + when(server3.getServiceHandler().hello(anyString())).thenReturn(SERVER3_RESPONSE); + + final RetryConfig hedgingNoRetryConfig = RetryConfig + .builderForRpc( + RetryRule.builder().onTimeoutException().thenBackoff(Backoff.fixed(10_000))) // should + // be always overtaken by hedging task + .maxTotalAttempts(3) + .responseTimeoutMillisForEachAttempt(100) + .hedgingDelayMillis(200) + .build(); + + final HelloService.AsyncIface client = client(hedgingNoRetryConfig); + + final CompletableFuture responseFuture; + final ClientRequestContext ctx; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + responseFuture = asyncHelloWith(client); + ctx = captor.get(); + } + + await().pollInterval(25, TimeUnit.MILLISECONDS).untilAsserted(() -> { + assertThat(server3.getNumRequests()).isOne(); + }); + + // As we know that the third server received a request, we know that + // we are at >= T + 400. This means that: + server1.unlatchResponse(); // issued at T (timed out) + server2.unlatchResponse(); // issued at T + 200 (timed out) + server3.unlatchResponse(); // issued at T + 400 (hopefully not timed out yet) + + await().untilAsserted(() -> { + assertThat(responseFuture.get()).isEqualTo(SERVER3_RESPONSE); + + assertValidServerRequestContext(server1, 1, true); + assertValidServerRequestContext(server2, 2, true); + assertValidServerRequestContext(server3, 3, false); + + assertValidClientRequestContext( + ctx, GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER3_RESPONSE), + VERIFY_RESPONSE_TIMEOUT, + VERIFY_RESPONSE_TIMEOUT, + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER3_RESPONSE) + ); + }); + } + + @Test + void thirdServerWinsEvenAfterRetriableResponse() throws Exception { + when(server1.getServiceHandler().hello(anyString())).thenReturn(RETRIABLE_RESPONSE); + when(server2.getServiceHandler().hello(anyString())).thenReturn(RETRIABLE_RESPONSE); + when(server3.getServiceHandler().hello(anyString())).thenReturn(SERVER3_RESPONSE); + + final RetryConfig config = RetryConfig.builderForRpc( + getRetryRetriableResponsesRule( + Backoff.withoutDelay() + ) + ) + .maxTotalAttempts(3) + .hedgingDelayMillis(10) + .build(); + + final HelloService.AsyncIface client = client(config); + + final CompletableFuture responseFuture; + final ClientRequestContext ctx; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + responseFuture = asyncHelloWith(client); + ctx = captor.get(); + } + + server1.unlatchResponse(); + server2.unlatchResponse(); + server3.unlatchResponse(); + + await().untilAsserted(() -> { + assertThat(responseFuture.get()).isEqualTo(SERVER3_RESPONSE); + + assertValidServerRequestContext(server1, 1, false); + assertValidServerRequestContext(server2, 2, false); + assertValidServerRequestContext(server3, 3, false); + + assertValidClientRequestContext( + ctx, GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER3_RESPONSE), + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(RETRIABLE_RESPONSE), + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(RETRIABLE_RESPONSE), + GET_VERIFY_RESPONSE_HAS_CONTENT.apply(SERVER3_RESPONSE) + ); + }); + } + + @Test + void loosesAfterNonRetriableResponse() throws TException { + when(server2.getServiceHandler().hello(anyString())).thenThrow(new AnticipatedException("Aachen")); + + final RetryConfig config = RetryConfig + .builderForRpc(NO_RETRY_RULE) + .maxTotalAttempts(3) + // Should be long enough so we can complete the second request before we continue issuing + // a third request to the third server. + .hedgingDelayMillis(500) + .build(); + + final HelloService.AsyncIface client = client(config); + + final CompletableFuture responseFuture; + final ClientRequestContext ctx; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + responseFuture = asyncHelloWith(client); + ctx = captor.get(); + } + + await().untilAsserted(() -> assertThat(server1.getNumRequests()).isEqualTo(1)); + + server2.unlatchResponse(); + + await().untilAsserted(() -> { + assertThatThrownBy(responseFuture::get) + .satisfies(throwable -> { + final Throwable peeledThrowable = Exceptions.peel(throwable); + assertThat(peeledThrowable).isInstanceOf(TApplicationException.class); + }); + + assertValidClientRequestContext( + ctx, + GET_VERIFY_RESPONSE_HAS_APPLICATION_EXCEPTION.apply(TApplicationException.INTERNAL_ERROR, + "Aachen"), + VERIFY_REQUEST_CANCELLED, + GET_VERIFY_RESPONSE_HAS_APPLICATION_EXCEPTION.apply(TApplicationException.INTERNAL_ERROR, + "Aachen"), + null + ); + + assertValidServerRequestContext(server1, 1, true); + assertValidServerRequestContext(server2, 2, false, AnticipatedException.class); + assertNoServerRequestContext(server3); + }); + } + + @Test + void loosesAfterResponseTimeout() { + final RetryConfig config = RetryConfig + .builderForRpc(NO_RETRY_RULE) + .maxTotalAttempts(3) + .hedgingDelayMillis(0) + .build(); + + final HelloService.AsyncIface client = client(config, 500); + + final CompletableFuture responseFuture; + final ClientRequestContext ctx; + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + responseFuture = asyncHelloWith(client); + ctx = captor.get(); + } + + await().atLeast(400, TimeUnit.MILLISECONDS).atMost(600, TimeUnit.MILLISECONDS).untilAsserted(() -> { + assertThat(responseFuture).isCompletedExceptionally(); + assertThatThrownBy(responseFuture::get).satisfies(throwable -> { + final Throwable peeledThrowable = Exceptions.peel(throwable); + assertThat(peeledThrowable).isInstanceOf(TTransportException.class); + assertThat(peeledThrowable.getCause()).isInstanceOf(ResponseTimeoutException.class); + }); + + final List childLogExceptions = new ArrayList<>(); + + final RequestLogVerifier catchException = log -> { + assertThat(log.responseCause()).isInstanceOf(TTransportException.class); + final TTransportException cause = (TTransportException) log.responseCause(); + assertThat(cause.getCause()).isNotNull(); + childLogExceptions.add(cause.getCause()); + }; + + assertValidClientRequestContext(ctx, + VERIFY_RESPONSE_TIMEOUT, + catchException, + catchException, + catchException); + + int numTimeouts = 0; + int numCancelled = 0; + + for (final @Nullable Throwable childException : childLogExceptions) { + if (childException instanceof ResponseTimeoutException) { + numTimeouts++; + } else if (childException instanceof ResponseCancellationException) { + numCancelled++; + } else { + fail("Unexpected exception: " + childException); + } + } + + assertThat(numTimeouts + numCancelled).isEqualTo(3); + // At least one attempt needs to time out. + assertThat(numTimeouts).isPositive(); + + assertValidServerRequestContext(server1, 1, true); + assertValidServerRequestContext(server2, 2, true); + assertValidServerRequestContext(server3, 3, true); + }); + } + + private CompletableFuture asyncHelloWith(HelloService.AsyncIface client) { + final CompletableFuture future = new CompletableFuture<>(); + try { + client.hello("hello", new AsyncMethodCallback() { + @Override + public void onComplete(String response) { + future.complete(response); + } + + @Override + public void onError(Exception exception) { + future.completeExceptionally(exception); + } + }); + } catch (TException e) { + future.completeExceptionally(e); + } + + return future; + } + + private static HelloService.AsyncIface client(RetryConfig config) { + return client(config, 10000); + } + + private static HelloService.AsyncIface client(RetryConfig config, + long responseTimeoutMillis) { + return ThriftClients.builder( + SessionProtocol.H2C, + EndpointGroup.of( + EndpointSelectionStrategy.roundRobin(), + server1.endpoint(SessionProtocol.H2C), + server2.endpoint(SessionProtocol.H2C), + server3.endpoint(SessionProtocol.H2C) + ) + ) + .responseTimeoutMillis(responseTimeoutMillis) + .path("/thrift") + .rpcDecorator(RetryingRpcClient.builder(config).newDecorator()) + .build(HelloService.AsyncIface.class); + } + + private static RetryRuleWithContent getRetryRetriableResponsesRule(Backoff backoff) { + return RetryRuleWithContent + .builder() + .onResponse((ctx, response) -> { + return response.whenComplete().handle( + (responseData, cause) -> { + assertThat(cause).isNull(); + assertThat(responseData).isInstanceOf(String.class); + return responseData.equals(RETRIABLE_RESPONSE); + } + ); + }) + .thenBackoff(backoff); + } + + private interface RequestLogVerifier extends Consumer {} + + private static final RequestLogVerifier VERIFY_REQUEST_CANCELLED = + log -> { + assertThat(log.responseCause()).isInstanceOf(TTransportException.class); + assertThat(log.responseCause().getCause()).isInstanceOf(ResponseCancellationException.class); + }; + + private static final RequestLogVerifier VERIFY_RESPONSE_TIMEOUT = + log -> { + assertThat(log.responseCause()).isInstanceOf(TTransportException.class); + assertThat(log.responseCause().getCause()).isInstanceOf(ResponseTimeoutException.class); + }; + + private static final Function GET_VERIFY_RESPONSE_HAS_CONTENT = + expectedResponseContent -> log -> { + assertThat(log.responseContent()).isInstanceOf(RpcResponse.class); + assertThat((CompletionStage) log.responseContent()) + .isCompletedWithValue(expectedResponseContent); + }; + + private static final BiFunction + GET_VERIFY_RESPONSE_HAS_APPLICATION_EXCEPTION = + (expectedType, expectedMessage) -> log -> { + assertThat(log.responseCause()).isInstanceOf(TApplicationException.class); + final TApplicationException cause = (TApplicationException) log.responseCause(); + assertThat(cause.getType()).isEqualTo(expectedType); + assertThat(cause.getMessage()).contains(expectedMessage); + }; + + private static void assertValidRetryingState(ClientRequestContext ctx, int expectedAttempts) { + final State state = AbstractRetryingClient.state(ctx); + assertThat(state.isRetryingComplete()).isTrue(); + + final List> startedAttempts = state.startedAttempts(); + assertThat(startedAttempts).hasSize(expectedAttempts); + for (Attempt attempt : startedAttempts) { + assertThat(attempt.attemptRes().whenComplete().isDone()).isTrue(); + } + } + + private static void assertValidRootClientRequestContext(ClientRequestContext ctx, + RequestLogVerifier logVerifierCtx, + int expectedNumChildren) { + assertThat(ctx.log().isComplete()).isTrue(); + assertThat(ctx.log().children()).hasSize(expectedNumChildren); + final RequestLog log = ctx.log().getIfAvailable(RequestLogProperty.RESPONSE_CONTENT, + RequestLogProperty.RESPONSE_CAUSE, + RequestLogProperty.REQUEST_HEADERS); + assertThat(log).isNotNull(); + logVerifierCtx.accept(log); + + assertValidRetryingState(ctx, expectedNumChildren); + } + + private void assertValidClientRequestContext(ClientRequestContext ctx, + RequestLogVerifier logVerifierCtx, + RequestLogVerifier logVerifierServer1, + @Nullable RequestLogVerifier logVerifierServer2, + @Nullable RequestLogVerifier logVerifierServer3 + ) { + final int expectedNumChildren = 1 + (logVerifierServer2 == null ? 0 : 1) + + (logVerifierServer3 == null ? 0 : 1); + + assertValidRootClientRequestContext(ctx, logVerifierCtx, expectedNumChildren); + assertValidChildLog(ctx.log().children().get(0), 1, logVerifierServer1); + if (logVerifierServer2 != null) { + assertValidChildLog(ctx.log().children().get(1), 2, logVerifierServer2); + } + + if (logVerifierServer3 != null) { + assertValidChildLog(ctx.log().children().get(2), 3, logVerifierServer3); + } + } + + void assertValidChildLog(RequestLogAccess logAccess, int attemptNumber, + RequestLogVerifier requestLogVerifier) { + assertThat(logAccess.isComplete()).isTrue(); + // After the check right above, all properties of the RequestLog should be available. + final @Nullable RequestLog log = logAccess.getIfAvailable(RequestLogProperty.RESPONSE_CONTENT, + RequestLogProperty.RESPONSE_CAUSE, + RequestLogProperty.REQUEST_HEADERS); + assertThat(log).isNotNull(); + + if (attemptNumber > 1) { + assertThat(log.requestHeaders().getInt(ARMERIA_RETRY_COUNT)).isEqualTo(attemptNumber - 1); + } else { + assertThat(log.requestHeaders().contains(ARMERIA_RETRY_COUNT)).isFalse(); + } + + requestLogVerifier.accept(log); + } + + private void assertValidServerRequestContext(ServerExtension server, int attemptNumber, + boolean expectCancelled) { + assertValidServerRequestContext(server, attemptNumber, expectCancelled, null); + } + + private void assertValidServerRequestContext(ServerExtension server, int attemptNumber, + boolean expectCancelled, + @Nullable Class expectedResponseException) { + assertThat(server.requestContextCaptor().size()).isEqualTo(1); + + final ServiceRequestContext sctx; + sctx = server.requestContextCaptor().peek(); + assertThat(sctx).isNotNull(); + assertThat(sctx.log().isComplete()).isTrue(); + + assertThat(sctx.isCancelled()).isEqualTo(expectCancelled); + + if (expectCancelled) { + assertThat(sctx.cancellationCause()).isInstanceOf(ClosedStreamException.class); + assertThat(sctx.cancellationCause().getMessage()) + .contains("received a RST_STREAM frame: CANCEL"); + } + + final RequestLog slog = sctx.log().getIfAvailable(RequestLogProperty.REQUEST_HEADERS, + RequestLogProperty.REQUEST_CONTENT); + + assertThat(slog).isNotNull(); + if (attemptNumber > 1) { + assertThat(slog.requestHeaders().getInt(ARMERIA_RETRY_COUNT)).isEqualTo(attemptNumber - 1); + } else { + assertThat(slog.requestHeaders().contains(ARMERIA_RETRY_COUNT)).isFalse(); + } + + assertThat(slog.requestContent()).isInstanceOf(RpcRequest.class); + assertThat(((RpcRequest) slog.requestContent()).params().get(0)).isEqualTo("hello"); + + if (expectedResponseException != null) { + assertThat(slog.responseCause()).isInstanceOf(expectedResponseException); + } else if (expectCancelled) { + assertThat(slog.responseCause()).isInstanceOf(ClosedStreamException.class); + assertThat(slog.responseCause().getMessage()) + .contains("received a RST_STREAM frame: CANCEL"); + } else { + assertThat(slog.responseCause()).isNull(); + } + } + + private void assertNoServerRequestContext(ServerExtension server) { + assertThat(server.requestContextCaptor().size()).isEqualTo(0); + } +}