Provide a way to complete a ClientRequestContext in Preprocessor

core/src/main/java/com/linecorp/armeria/client/retry/RetryingClient.java

@@ -316,7 +316,7 @@ private void doExecute0(ClientRequestContext ctx, HttpRequestDuplicator rootReqD
         assert ctxReq != null;
         final HttpResponse response;
         final ClientRequestContextExtension ctxExtension = derivedCtx.as(ClientRequestContextExtension.class);
-        if (!initialAttempt && ctxExtension != null && derivedCtx.endpoint() == null) {
+        if (!initialAttempt && ctxExtension != null && !ctxExtension.initializationTriggered()) {
             // clear the pending throwable to retry endpoint selection
             ClientPendingThrowableUtil.removePendingThrowable(derivedCtx);
             // if the endpoint hasn't been selected, try to initialize the ctx with a new endpoint/event loop

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

@@ -25,7 +25,6 @@
 import com.linecorp.armeria.client.ClientRequestContext;
 import com.linecorp.armeria.client.ResponseTimeoutException;
 import com.linecorp.armeria.client.RpcClient;
-import com.linecorp.armeria.client.endpoint.EndpointGroup;
 import com.linecorp.armeria.common.HttpRequest;
 import com.linecorp.armeria.common.Request;
 import com.linecorp.armeria.common.RpcRequest;
@@ -173,9 +172,7 @@ private void doExecute0(ClientRequestContext ctx, RpcRequest req,
         final RpcResponse res;
 
         final ClientRequestContextExtension ctxExtension = derivedCtx.as(ClientRequestContextExtension.class);
-        final EndpointGroup endpointGroup = derivedCtx.endpointGroup();
-        if (!initialAttempt && ctxExtension != null &&
-            endpointGroup != null && derivedCtx.endpoint() == null) {
+        if (!initialAttempt && ctxExtension != null && !ctxExtension.initializationTriggered()) {
             // clear the pending throwable to retry endpoint selection
             ClientPendingThrowableUtil.removePendingThrowable(derivedCtx);
             // if the endpoint hasn't been selected, try to initialize the ctx with a new endpoint/event loop

core/src/main/java/com/linecorp/armeria/internal/client/ClientRequestContextExtension.java

@@ -65,6 +65,18 @@ public interface ClientRequestContextExtension extends ClientRequestContext, Req
      */
     void finishInitialization(boolean success);
 
+    /**
+     * Initializes and fails the context immediately. This method can be thought of as a combination of
+     * {@link #init()} and {@link #finishInitialization(boolean)}.
+     * Returns true if this call initialized and failed the context.
+     */
+    boolean initAndFail(Throwable cause);
+
+    /**
+     * Unlike {@link #whenInitialized()}, returns whether an initialization is triggered.
+     */
+    boolean initializationTriggered();
+
     /**
      * A set of internal headers which are set by armeria internally.
      * These headers are merged with the lowest priority before getting sent over the wire.

core/src/main/java/com/linecorp/armeria/internal/client/ClientUtil.java

@@ -176,12 +176,19 @@ O executeWithFallback(U execution,
         if (ctxExt != null) {
             ctxExt.runContextCustomizer();
         }
+        O res;
         try {
-            return execution.execute(ctx, req);
+            res = execution.execute(ctx, req);
         } catch (Exception e) {
+            if (ctxExt != null && !ctxExt.initializationTriggered()) {
+                ctxExt.initAndFail(e);
+            }
             fail(ctx, e);
-            return errorResponseFactory.apply(ctx, e);
+            res = errorResponseFactory.apply(ctx, e);
         }
+
+        completeLogIfIncomplete(ctx, res);
+        return res;
     }
 
     private static <I extends Request, O extends Response, U extends Client<I, O>>

core/src/main/java/com/linecorp/armeria/internal/client/DefaultClientRequestContext.java

@@ -30,6 +30,7 @@
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -116,6 +117,10 @@ public final class DefaultClientRequestContext
             whenInitializedUpdater = AtomicReferenceFieldUpdater.newUpdater(
             DefaultClientRequestContext.class, CompletableFuture.class, "whenInitialized");
 
+    private static final AtomicIntegerFieldUpdater<DefaultClientRequestContext>
+            initializationTriggeredUpdater = AtomicIntegerFieldUpdater.newUpdater(
+            DefaultClientRequestContext.class, "initializationTriggered");
+
     private static SessionProtocol desiredSessionProtocol(SessionProtocol protocol, ClientOptions options) {
         if (!options.factory().options().preferHttp1()) {
             return protocol;
@@ -134,7 +139,8 @@ private static SessionProtocol desiredSessionProtocol(SessionProtocol protocol,
     private static final short STR_PARENT_LOG_AVAILABILITY = 1 << 1;
     private static boolean warnedNullRequestId;
 
-    private boolean initialized;
+    // 0: not initialized, 1: initialized
+    private volatile int initializationTriggered;
     @Nullable
     private EventLoop eventLoop;
     private EndpointGroup endpointGroup;
@@ -358,9 +364,10 @@ private static ServiceRequestContext serviceRequestContext() {
 
     @Override
     public CompletableFuture<Boolean> init() {
+        if (!initializationTriggeredUpdater.compareAndSet(this, 0, 1)) {
+            return whenInitialized();
+        }
         assert endpoint == null : endpoint;
-        assert !initialized;
-        initialized = true;
 
         final Throwable cancellationCause = cancellationCause();
         if (cancellationCause != null) {
@@ -383,6 +390,29 @@ public CompletableFuture<Boolean> init() {
         }
     }
 
+    @Override
+    public boolean initAndFail(Throwable cause) {
+        if (!initializationTriggeredUpdater.compareAndSet(this, 0, 1)) {
+            return false;
+        }
+        acquireEventLoop(endpointGroup);
+        failEarly(cause);
+        finishInitialization(false);
+        return true;
+    }
+
+    private void initNow() {
+        if (!initializationTriggeredUpdater.compareAndSet(this, 0, 1)) {
+            return;
+        }
+        finishInitialization(false);
+    }
+
+    @Override
+    public boolean initializationTriggered() {
+        return initializationTriggered == 1;
+    }
+
     private EndpointGroup mapEndpoint(EndpointGroup endpointGroup) {
         if (endpointGroup instanceof Endpoint) {
             return requireNonNull(options().endpointRemapper().apply((Endpoint) endpointGroup),
@@ -395,6 +425,7 @@ private EndpointGroup mapEndpoint(EndpointGroup endpointGroup) {
     private CompletableFuture<Boolean> initEndpoint(Endpoint endpoint) {
         updateEndpoint(endpoint);
         acquireEventLoop(endpoint);
+        maybeInitializeResponseCancellationScheduler();
         return initFuture(true, null);
     }
 
@@ -404,6 +435,7 @@ private CompletableFuture<Boolean> initEndpointGroup(EndpointGroup endpointGroup
         if (endpoint != null) {
             updateEndpoint(endpoint);
             acquireEventLoop(endpointGroup);
+            maybeInitializeResponseCancellationScheduler();
             return initFuture(true, null);
         }
 
@@ -412,6 +444,7 @@ private CompletableFuture<Boolean> initEndpointGroup(EndpointGroup endpointGroup
         return endpointGroup.select(this, temporaryEventLoop).handle((e, cause) -> {
             updateEndpoint(e);
             acquireEventLoop(endpointGroup);
+            maybeInitializeResponseCancellationScheduler();
 
             final boolean success;
             if (cause != null) {
@@ -462,13 +495,17 @@ public CompletableFuture<Boolean> whenInitialized() {
     public void finishInitialization(boolean success) {
         final CompletableFuture<Boolean> whenInitialized = this.whenInitialized;
         if (whenInitialized != null) {
-            whenInitialized.complete(success);
+            if (!whenInitialized.isDone()) {
+                whenInitialized.complete(success);
+            }
         } else {
             if (!whenInitializedUpdater.compareAndSet(this, null,
                                                       UnmodifiableFuture.completedFuture(success))) {
                 final CompletableFuture<Boolean> oldWhenInitialized = this.whenInitialized;
                 assert oldWhenInitialized != null;
-                oldWhenInitialized.complete(success);
+                if (!oldWhenInitialized.isDone()) {
+                    oldWhenInitialized.complete(success);
+                }
             }
         }
     }
@@ -484,7 +521,6 @@ private void acquireEventLoop(EndpointGroup endpointGroup) {
                     options().factory().acquireEventLoop(sessionProtocol(), endpointGroup, endpoint);
             eventLoop = releasableEventLoop.get();
             log.whenComplete().thenAccept(unused -> releasableEventLoop.release());
-            initializeResponseCancellationScheduler();
         }
     }
 
@@ -540,6 +576,7 @@ private void failEarly(Throwable cause) {
         final RequestLogBuilder logBuilder = logBuilder();
         logBuilder.endRequest(wrapped);
         logBuilder.endResponse(wrapped);
+        responseCancellationScheduler.finishNow(cause);
     }
 
     // TODO(ikhoon): Consider moving the logic for filling authority to `HttpClientDelegate.exceute()`.
@@ -616,17 +653,23 @@ private DefaultClientRequestContext(DefaultClientRequestContext ctx,
 
         this.endpointGroup = endpointGroup;
         updateEndpoint(endpoint);
+        if (endpoint != null) {
+            initNow();
+        }
         // We don't need to acquire an EventLoop for the initial attempt because it's already acquired by
         // the root context.
         if (endpoint == null || ctx.endpoint() == endpoint && ctx.log.children().isEmpty()) {
             eventLoop = ctx.eventLoop().withoutContext();
-            initializeResponseCancellationScheduler();
         } else {
             acquireEventLoop(endpoint);
         }
+        maybeInitializeResponseCancellationScheduler();
     }
 
-    private void initializeResponseCancellationScheduler() {
+    private void maybeInitializeResponseCancellationScheduler() {
+        if (responseCancellationScheduler.hasEventLoop()) {
+            return;
+        }
         final CancellationTask cancellationTask = cause -> {
             try (SafeCloseable ignored = RequestContextUtil.pop()) {
                 final HttpRequest request = request();
@@ -683,7 +726,7 @@ public SessionProtocol sessionProtocol() {
 
     @Override
     public void setSessionProtocol(SessionProtocol sessionProtocol) {
-        checkState(!initialized, "Cannot update sessionProtocol after initialization");
+        checkState(!initializationTriggered(), "Cannot update sessionProtocol after initialization");
         this.sessionProtocol = desiredSessionProtocol(requireNonNull(sessionProtocol, "sessionProtocol"),
                                                       options);
     }
@@ -789,10 +832,9 @@ public ContextAwareEventLoop eventLoop() {
 
     @Override
     public void setEventLoop(EventLoop eventLoop) {
-        checkState(!initialized, "Cannot update eventLoop after initialization");
+        checkState(!initializationTriggered(), "Cannot update eventLoop after initialization");
         checkState(this.eventLoop == null, "eventLoop can be updated only once");
         this.eventLoop = requireNonNull(eventLoop, "eventLoop");
-        initializeResponseCancellationScheduler();
     }
 
     @Override
@@ -820,7 +862,7 @@ public EndpointGroup endpointGroup() {
 
     @Override
     public void setEndpointGroup(EndpointGroup endpointGroup) {
-        checkState(!initialized, "Cannot update endpointGroup after initialization");
+        checkState(!initializationTriggered(), "Cannot update endpointGroup after initialization");
         this.endpointGroup = requireNonNull(endpointGroup, "endpointGroup");
     }
 
@@ -1035,7 +1077,14 @@ public CancellationScheduler responseCancellationScheduler() {
     @Override
     public void cancel(Throwable cause) {
         requireNonNull(cause, "cause");
-        responseCancellationScheduler.finishNow(cause);
+        if (initializationTriggered()) {
+            // happy path
+            responseCancellationScheduler.finishNow(cause);
+            return;
+        }
+        if (!initAndFail(cause)) {
+            responseCancellationScheduler.finishNow(cause);
+        }
     }
 
     @Nullable

core/src/main/java/com/linecorp/armeria/internal/common/CancellationScheduler.java

@@ -133,6 +133,8 @@ default void finishNow() {
     @VisibleForTesting
     State state();
 
+    boolean hasEventLoop();
+
     enum State {
         INIT,
         SCHEDULED,

core/src/main/java/com/linecorp/armeria/internal/common/DefaultCancellationScheduler.java

@@ -273,6 +273,11 @@ private ScheduleResult setTimeoutNanosFromNow0(long newTimeoutNanos) {
         return ScheduleResult.INVOKE_LATER;
     }
 
+    @Override
+    public boolean hasEventLoop() {
+        return eventLoop != null;
+    }
+
     private EventExecutor eventLoop() {
         assert eventLoop != null;
         return eventLoop;

core/src/main/java/com/linecorp/armeria/internal/common/NoopCancellationScheduler.java

@@ -116,4 +116,9 @@ public void updateTask(CancellationTask cancellationTask) {
     public State state() {
         return State.INIT;
     }
+
+    @Override
+    public boolean hasEventLoop() {
+        return false;
+    }
 }

core/src/test/java/com/linecorp/armeria/client/HttpPreprocessorTest.java

@@ -18,6 +18,7 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.awaitility.Awaitility.await;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -31,6 +32,7 @@
 import com.linecorp.armeria.common.HttpResponse;
 import com.linecorp.armeria.common.SessionProtocol;
 import com.linecorp.armeria.internal.client.ClientRequestContextExtension;
+import com.linecorp.armeria.internal.client.DefaultClientRequestContext;
 import com.linecorp.armeria.internal.common.CancellationScheduler.State;
 import com.linecorp.armeria.testing.junit5.common.EventLoopExtension;
 
@@ -112,6 +114,26 @@ void cancellationSchedulerIsInitializedCorrectly() {
         assertThat(client.get("/").status().code()).isEqualTo(200);
     }
 
+    @Test
+    void failureCompletesContext() {
+        final RuntimeException exception = new RuntimeException("test");
+        try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) {
+            final WebClient client = WebClient.of(new HttpPreprocessor() {
+                @Override
+                public HttpResponse execute(PreClient<HttpRequest, HttpResponse> delegate,
+                                            PreClientRequestContext ctx, HttpRequest req) throws Exception {
+                    throw exception;
+                }
+            });
+
+            assertThatThrownBy(() -> client.blocking().get("/")).isEqualTo(exception);
+            final DefaultClientRequestContext ctx = (DefaultClientRequestContext) captor.get();
+            await().untilAsserted(() -> assertThat(ctx.whenInitialized()).isDone());
+            await().untilAsserted(() -> assertThat(ctx.log().isComplete()).isTrue());
+            assertThat(ctx.eventLoop()).isNotNull();
+        }
+    }
+
     private static final class RunnablePreprocessor implements HttpPreprocessor {
 
         private static HttpPreprocessor of(Runnable runnable) {