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

@@ -179,6 +179,9 @@ O executeWithFallback(U execution,
         try {
             return execution.execute(ctx, req);
         } catch (Exception e) {
+            if (ctxExt != null && !ctxExt.initializationTriggered()) {
+                ctxExt.initAndFail(e);
+            }
             fail(ctx, e);
             return errorResponseFactory.apply(ctx, e);
         }

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),
@@ -462,13 +492,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);
+                }
             }
         }
     }
@@ -540,6 +574,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,6 +651,9 @@ 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()) {
@@ -683,7 +721,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,7 +827,7 @@ 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();
@@ -820,7 +858,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 +1073,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/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) {

grpc/src/test/java/com/linecorp/armeria/client/grpc/GrpcPreprocessorTest.java

@@ -0,0 +1,78 @@
+/*
+ * 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.grpc;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.catchException;
+import static org.awaitility.Awaitility.await;
+
+import java.util.concurrent.CompletableFuture;
+
+import org.junit.jupiter.api.Test;
+
+import com.linecorp.armeria.client.ClientRequestContextCaptor;
+import com.linecorp.armeria.client.Clients;
+import com.linecorp.armeria.common.HttpResponse;
+import com.linecorp.armeria.internal.client.DefaultClientRequestContext;
+
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
+import testing.grpc.EmptyProtos.Empty;
+import testing.grpc.TestServiceGrpc.TestServiceBlockingStub;
+
+class GrpcPreprocessorTest {
+
+    @Test
+    void throwCompletesContext() {
+        final RuntimeException exception = new RuntimeException("test");
+        try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) {
+            final TestServiceBlockingStub stub = GrpcClients.builder((delegate, ctx, req) -> {
+                throw exception;
+            }).build(TestServiceBlockingStub.class);
+
+            final Exception thrown = catchException(() -> stub.emptyCall(Empty.getDefaultInstance()));
+            assertThat(thrown).isInstanceOf(StatusRuntimeException.class);
+            assertThat((StatusRuntimeException) thrown).hasCause(exception);
+            final DefaultClientRequestContext ctx = (DefaultClientRequestContext) captor.get();
+            await().untilAsserted(() -> assertThat(ctx.whenInitialized()).isDone());
+            await().untilAsserted(() -> assertThat(ctx.log().isComplete()).isTrue());
+            assertThat(ctx.eventLoop()).isNotNull();
+        }
+    }
+
+    @Test
+    void cancelCompletesContext() {
+        final RuntimeException exception = new RuntimeException("test");
+        try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) {
+            final TestServiceBlockingStub stub = GrpcClients.builder((delegate, ctx, req) -> {
+                return HttpResponse.of(CompletableFuture.supplyAsync(() -> {
+                    ctx.cancel();
+                    throw exception;
+                }));
+            }).build(TestServiceBlockingStub.class);
+
+            final Exception thrown = catchException(() -> stub.emptyCall(Empty.getDefaultInstance()));
+            assertThat(thrown).isInstanceOf(StatusRuntimeException.class);
+            final Status status = ((StatusRuntimeException) thrown).getStatus();
+            assertThat(status.getCause()).isSameAs(exception);
+            final DefaultClientRequestContext ctx = (DefaultClientRequestContext) captor.get();
+            await().untilAsserted(() -> assertThat(ctx.whenInitialized()).isDone());
+            await().untilAsserted(() -> assertThat(ctx.log().isComplete()).isTrue());
+            assertThat(ctx.eventLoop()).isNotNull();
+        }
+    }
+}