Provide a way to complete a ClientRequestContext in Preprocessor (line#6466)

Motivation:

This issue was found while working on `GrpcServicesPreprocessor`
refactoring.

Currently, there is a bug where exceptions/response failures in
`Preprocessor` do not complete `ctx.whenInitialized`.
e.g.
```
WebClient.of((delegate, ctx, req) -> {
    // ctx.whenInitialized will not be completed
    throw new RuntimeException("error");
}).get("/");
```

While normally this may not be an issue, gRPC relies on completion of
this future to signal errors:
ref:
https://github.com/line/armeria/blob/1a31f3e8c9aba91fb663bfe5c4fc3eaa87c9d13f/grpc/src/main/java/com/linecorp/armeria/internal/client/grpc/ArmeriaClientCall.java#L180

I propose two modifications to resolve this issue:
- When a Preprocessor throws an exception from the `PreClient#execute`
calling thread, it should complete the ctx
```
WebClient.of((delegate, ctx, req) -> {
    throw new RuntimeException("error");
}).get("/");
```
- When `ctx.cancel` is called (even before `ctx.init`) is called, it
should complete the ctx
```
WebClient.of((delegate, ctx, req) -> {
    ctx.cancel(e);
    return HttpResponse.of(400);
}).get("/");
```

Modifications:

- Introduced an `initialized` flag to guard against concurrent
initialization attempts
- `finishInitialization` is also modified to guard against concurrent
calls
- A new `initAndFail` method is introduced which acquires an event loop,
initializes the cancellation scheduler, and completes the ctx
- `ctx.cancel` will trigger `initAndFail` if the ctx is not initialized
yet
  - If `PreClient#execute` throws an exception, `initAndFail` is called
- For derived contexts, if an endpoint does not exist, it means
`ClientUtil#initContextAnd*` needs to be called. Hence, initialization
is set only if an endpoint exists.
- `responseCancellationScheduler.finishNow` is called in
`DefaultClientRequestContext#failEarly` to record `cancellationCause`
consistently.
- `XdsPreprocessor` and `RouterFilter` now calls `ctx.cancel` if a
request is short-circuited before `ctx.init` is called.

Result:

- Requests failed at the `XdsPreprocessor`-level are propagated to the
user correctly when using gRPC.

<!--
Visit this URL to learn more about how to write a pull request
description:

https://armeria.dev/community/developer-guide#how-to-write-pull-request-description
-->

---------

Co-authored-by: Ikhun Um <ikhun.um@linecorp.com>

Author: jrhee17

core/src/main/java/com/linecorp/armeria/client/DefaultWebClient.java

@@ -126,7 +126,7 @@ public HttpResponse execute(HttpRequest req, RequestOptions requestOptions) {
         final DefaultClientRequestContext ctx = new DefaultClientRequestContext(
                 protocol, newReq, newReq.method(), null, reqTarget, endpointGroup, requestOptions, options(),
                 meterRegistry());
-        return ClientUtil.executeWithFallback(preClient, ctx, newReq, errorResponseFactory());
+        return ClientUtil.executePreClientWithFallback(preClient, ctx, newReq, errorResponseFactory());
     }
 
     private static HttpResponse abortRequestAndReturnFailureResponse(

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

@@ -325,14 +325,16 @@ 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
+            // tryCompleteLog is false because we handle it in completeLogIfBytesNotTransferred.
             response = initContextAndExecuteWithFallback(
                     unwrap(), ctxExtension, HttpResponse::of,
                     (context, cause) -> HttpResponse.ofFailure(cause), ctxReq, false);
         } else {
+            // tryCompleteLog is false because we handle it in completeLogIfBytesNotTransferred.
             response = executeWithFallback(unwrap(), derivedCtx,
                                            (context, cause) -> HttpResponse.ofFailure(cause), ctxReq, false);
         }

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

@@ -26,7 +26,6 @@
 import com.linecorp.armeria.client.ResponseTimeoutException;
 import com.linecorp.armeria.client.RpcClient;
 import com.linecorp.armeria.client.UnprocessedRequestException;
-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;
@@ -184,9 +183,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/common/CoreBlockHoundIntegration.java

@@ -71,5 +71,6 @@ public void applyTo(Builder builder) {
         builder.allowBlockingCallsInside("io.netty.handler.ssl.SslContext", "buildKeyStore");
         // StampedLock.writeLock() can be called
         builder.allowBlockingCallsInside("io.netty.buffer.AdaptivePoolingAllocator", "allocate");
+        builder.allowBlockingCallsInside("com.linecorp.armeria.common.util.Version", "getAll");
     }
 }

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

@@ -169,19 +169,26 @@ O executeWithFallback(U delegate, ClientRequestContext ctx,
     }
 
     public static <I extends Request, O extends Response, U extends PreClient<I, O>>
-    O executeWithFallback(U execution,
-                          PreClientRequestContext ctx, I req,
-                          BiFunction<ClientRequestContext, Throwable, O> errorResponseFactory) {
+    O executePreClientWithFallback(U execution,
+                                   PreClientRequestContext ctx, I req,
+                                   BiFunction<ClientRequestContext, Throwable, O> errorResponseFactory) {
         final ClientRequestContextExtension ctxExt = ctx.as(ClientRequestContextExtension.class);
         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/client/TailPreClient.java

@@ -79,6 +79,6 @@ public O execute(PreClientRequestContext ctx, I req) {
         assert ctxExt != null;
         final Client<I, O> delegate0 = clientCustomizer.apply(ctxExt, delegate);
         return ClientUtil.initContextAndExecuteWithFallback(delegate0, ctxExt,
-                                                            futureConverter, errorResponseFactory, req, true);
+                                                            futureConverter, errorResponseFactory, req, false);
     }
 }

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;