ResponseTimeout can be scheduled at an earlier timing (#5793)

Motivation:

The original motivation of this PR stems from
#4591.
There has been requests for a timeout which spans over an entire
client's lifecycle. Following #5800, this can be achieved easily by
adjusting where to call `CancellationScheduler#start`.

I propose that we add an enum `ResponseTimeoutMode` which decides when a
`CancellationScheduler` will start to be scheduled. By doing so, we can
allow users to choose when a response timeout will start to be scheduled
per client.

Modifications:

- Introduced a new `ResponseTimeoutMode` enum and added options in
`AbstractClientOptionsBuilder` and `Flags` respectively
- Depending on the set `ResponseTimeoutMode`, the timeout is started on
1) context init 2) request start 3) or response read

Result:

- Closes #4591

<!--
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
-->

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

@@ -506,6 +506,18 @@ public AbstractClientOptionsBuilder contextCustomizer(
         return this;
     }
 
+    /**
+     * Sets the {@link ResponseTimeoutMode} which determines when a {@link #responseTimeout(Duration)}}
+     * will start to be scheduled.
+     *
+     * @see ResponseTimeoutMode
+     */
+    @UnstableApi
+    public AbstractClientOptionsBuilder responseTimeoutMode(ResponseTimeoutMode responseTimeoutMode) {
+        return option(ClientOptions.RESPONSE_TIMEOUT_MODE,
+                      requireNonNull(responseTimeoutMode, "responseTimeoutMode"));
+    }
+
     /**
      * Builds {@link ClientOptions} with the given options and the
      * {@linkplain ClientOptions#of() default options}.

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

@@ -198,6 +198,9 @@ final boolean tryInitialize() {
         final CancellationScheduler scheduler = cancellationScheduler();
         if (scheduler != null) {
             scheduler.updateTask(newCancellationTask());
+            if (ctx.responseTimeoutMode() == ResponseTimeoutMode.CONNECTION_ACQUIRED) {
+                scheduler.start();
+            }
         }
         if (ctx.isCancelled()) {
             // The previous cancellation task wraps the cause with an UnprocessedRequestException

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

@@ -223,6 +223,12 @@ public BlockingWebClientRequestPreparation exchangeType(ExchangeType exchangeTyp
         return this;
     }
 
+    @Override
+    public BlockingWebClientRequestPreparation responseTimeoutMode(ResponseTimeoutMode responseTimeoutMode) {
+        delegate.responseTimeoutMode(responseTimeoutMode);
+        return this;
+    }
+
     @Override
     public BlockingWebClientRequestPreparation requestOptions(RequestOptions requestOptions) {
         delegate.requestOptions(requestOptions);

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

@@ -295,4 +295,9 @@ public ClientBuilder contextCustomizer(
     public ClientBuilder contextHook(Supplier<? extends AutoCloseable> contextHook) {
         return (ClientBuilder) super.contextHook(contextHook);
     }
+
+    @Override
+    public ClientBuilder responseTimeoutMode(ResponseTimeoutMode responseTimeoutMode) {
+        return (ClientBuilder) super.responseTimeoutMode(responseTimeoutMode);
+    }
 }

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

@@ -156,6 +156,10 @@ public final class ClientOptions
     public static final ClientOption<Supplier<? extends AutoCloseable>> CONTEXT_HOOK =
             ClientOption.define("CONTEXT_HOOK", NOOP_CONTEXT_HOOK);
 
+    @UnstableApi
+    public static final ClientOption<ResponseTimeoutMode> RESPONSE_TIMEOUT_MODE =
+            ClientOption.define("RESPONSE_TIMEOUT_MODE", Flags.responseTimeoutMode());
+
     private static final List<AsciiString> PROHIBITED_HEADER_NAMES = ImmutableList.of(
             HttpHeaderNames.HTTP2_SETTINGS,
             HttpHeaderNames.METHOD,
@@ -395,6 +399,17 @@ public Supplier<AutoCloseable> contextHook() {
         return (Supplier<AutoCloseable>) get(CONTEXT_HOOK);
     }
 
+    /**
+     * Returns the {@link ResponseTimeoutMode} which determines when a {@link #responseTimeoutMillis()}
+     * will start to be scheduled.
+     *
+     * @see ResponseTimeoutMode
+     */
+    @UnstableApi
+    public ResponseTimeoutMode responseTimeoutMode() {
+        return get(RESPONSE_TIMEOUT_MODE);
+    }
+
     /**
      * Returns a new {@link ClientOptionsBuilder} created from this {@link ClientOptions}.
      */

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

@@ -222,4 +222,9 @@ public ClientOptionsBuilder contextCustomizer(
     public ClientOptionsBuilder contextHook(Supplier<? extends AutoCloseable> contextHook) {
         return (ClientOptionsBuilder) super.contextHook(contextHook);
     }
+
+    @Override
+    public ClientOptionsBuilder responseTimeoutMode(ResponseTimeoutMode responseTimeoutMode) {
+        return (ClientOptionsBuilder) super.responseTimeoutMode(responseTimeoutMode);
+    }
 }

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

@@ -612,6 +612,15 @@ default boolean isTimedOut() {
     @UnstableApi
     ExchangeType exchangeType();
 
+    /**
+     * Returns the {@link ResponseTimeoutMode} which determines when a {@link #responseTimeoutMillis()}
+     * will start to be scheduled.
+     *
+     * @see ResponseTimeoutMode
+     */
+    @UnstableApi
+    ResponseTimeoutMode responseTimeoutMode();
+
     @Override
     default ClientRequestContext unwrap() {
         return (ClientRequestContext) RequestContext.super.unwrap();

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

@@ -167,6 +167,11 @@ public ExchangeType exchangeType() {
         return unwrap().exchangeType();
     }
 
+    @Override
+    public ResponseTimeoutMode responseTimeoutMode() {
+        return unwrap().responseTimeoutMode();
+    }
+
     @Override
     public void hook(Supplier<? extends AutoCloseable> contextHook) {
         unwrap().hook(contextHook);

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

@@ -30,7 +30,7 @@
 final class DefaultRequestOptions implements RequestOptions {
 
     static final DefaultRequestOptions EMPTY = new DefaultRequestOptions(-1, -1, -1, null,
-                                                                         ImmutableMap.of(), null);
+                                                                         ImmutableMap.of(), null, null);
 
     private final long responseTimeoutMillis;
     private final long writeTimeoutMillis;
@@ -40,17 +40,21 @@ final class DefaultRequestOptions implements RequestOptions {
     private final Map<AttributeKey<?>, Object> attributeMap;
     @Nullable
     private final ExchangeType exchangeType;
+    @Nullable
+    private final ResponseTimeoutMode responseTimeoutMode;
 
     DefaultRequestOptions(long responseTimeoutMillis, long writeTimeoutMillis,
                           long maxResponseLength, @Nullable Long requestAutoAbortDelayMillis,
                           Map<AttributeKey<?>, Object> attributeMap,
-                          @Nullable ExchangeType exchangeType) {
+                          @Nullable ExchangeType exchangeType,
+                          @Nullable ResponseTimeoutMode responseTimeoutMode) {
         this.responseTimeoutMillis = responseTimeoutMillis;
         this.writeTimeoutMillis = writeTimeoutMillis;
         this.maxResponseLength = maxResponseLength;
         this.requestAutoAbortDelayMillis = requestAutoAbortDelayMillis;
         this.attributeMap = attributeMap;
         this.exchangeType = exchangeType;
+        this.responseTimeoutMode = responseTimeoutMode;
     }
 
     @Override
@@ -85,6 +89,12 @@ public ExchangeType exchangeType() {
         return exchangeType;
     }
 
+    @Nullable
+    @Override
+    public ResponseTimeoutMode responseTimeoutMode() {
+        return responseTimeoutMode;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) {
@@ -101,13 +111,14 @@ public boolean equals(Object o) {
                writeTimeoutMillis == that.writeTimeoutMillis &&
                maxResponseLength == that.maxResponseLength &&
                attributeMap.equals(that.attributeMap) &&
-               exchangeType == that.exchangeType;
+               exchangeType == that.exchangeType &&
+               responseTimeoutMode == that.responseTimeoutMode;
     }
 
     @Override
     public int hashCode() {
         return Objects.hash(responseTimeoutMillis, writeTimeoutMillis, maxResponseLength,
-                            attributeMap, exchangeType);
+                            attributeMap, exchangeType, responseTimeoutMode);
     }
 
     @Override
@@ -119,6 +130,7 @@ public String toString() {
                           .add("requestAutoAbortDelayMillis", requestAutoAbortDelayMillis)
                           .add("attributeMap", attributeMap)
                           .add("exchangeType", exchangeType)
+                          .add("responseTimeoutMode", responseTimeoutMode)
                           .toString();
     }
 }

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

@@ -416,4 +416,11 @@ public FutureTransformingRequestPreparation<T> exchangeType(ExchangeType exchang
         delegate.exchangeType(exchangeType);
         return this;
     }
+
+    @Override
+    public FutureTransformingRequestPreparation<T> responseTimeoutMode(
+            ResponseTimeoutMode responseTimeoutMode) {
+        delegate.responseTimeoutMode(responseTimeoutMode);
+        return this;
+    }
 }

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

@@ -294,7 +294,9 @@ void initTimeout() {
             final CancellationScheduler responseCancellationScheduler =
                     ctxExtension.responseCancellationScheduler();
             responseCancellationScheduler.updateTask(newCancellationTask());
-            responseCancellationScheduler.start();
+            if (ctx.responseTimeoutMode() == ResponseTimeoutMode.REQUEST_SENT) {
+                responseCancellationScheduler.start();
+            }
         }
     }
 

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

@@ -112,4 +112,14 @@ default RequestOptionsBuilder toBuilder() {
      */
     @Nullable
     ExchangeType exchangeType();
+
+    /**
+     * Returns the {@link ResponseTimeoutMode} which determines when a {@link #responseTimeoutMillis()}
+     * will start to be scheduled.
+     *
+     * @see ResponseTimeoutMode
+     */
+    @Nullable
+    @UnstableApi
+    ResponseTimeoutMode responseTimeoutMode();
 }

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

@@ -28,6 +28,7 @@
 
 import com.linecorp.armeria.common.ExchangeType;
 import com.linecorp.armeria.common.annotation.Nullable;
+import com.linecorp.armeria.common.annotation.UnstableApi;
 
 import io.netty.util.AttributeKey;
 
@@ -41,6 +42,8 @@ public final class RequestOptionsBuilder implements RequestOptionsSetters {
     private long maxResponseLength = -1;
     @Nullable
     private Long requestAutoAbortDelayMillis;
+    @Nullable
+    private ResponseTimeoutMode responseTimeoutMode;
 
     @Nullable
     private Map<AttributeKey<?>, Object> attributes;
@@ -58,6 +61,7 @@ public final class RequestOptionsBuilder implements RequestOptionsSetters {
                 attributes = new HashMap<>(attrs);
             }
             exchangeType = options.exchangeType();
+            responseTimeoutMode = options.responseTimeoutMode();
         }
     }
 
@@ -135,13 +139,20 @@ public RequestOptionsBuilder exchangeType(ExchangeType exchangeType) {
         return this;
     }
 
+    @Override
+    @UnstableApi
+    public RequestOptionsBuilder responseTimeoutMode(ResponseTimeoutMode responseTimeoutMode) {
+        this.responseTimeoutMode = requireNonNull(responseTimeoutMode, "responseTimeoutMode");
+        return this;
+    }
+
     /**
      * Returns a newly created {@link RequestOptions} with the properties specified so far.
      */
     public RequestOptions build() {
         if (responseTimeoutMillis < 0 && writeTimeoutMillis < 0 &&
             maxResponseLength < 0 && requestAutoAbortDelayMillis == null && attributes == null &&
-            exchangeType == null) {
+            exchangeType == null && responseTimeoutMode == null) {
             return EMPTY;
         } else {
             final Map<AttributeKey<?>, Object> attributes;
@@ -152,7 +163,7 @@ public RequestOptions build() {
             }
             return new DefaultRequestOptions(responseTimeoutMillis, writeTimeoutMillis,
                                              maxResponseLength, requestAutoAbortDelayMillis,
-                                             attributes, exchangeType);
+                                             attributes, exchangeType, responseTimeoutMode);
         }
     }
 }

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

@@ -157,4 +157,12 @@ public interface RequestOptionsSetters {
      */
     @UnstableApi
     RequestOptionsSetters exchangeType(ExchangeType exchangeType);
+
+    /**
+     * Sets the {@link ResponseTimeoutMode} which determines when a {@link #responseTimeout(Duration)}}
+     * will start to be scheduled.
+     *
+     * @see ResponseTimeoutMode
+     */
+    RequestOptionsSetters responseTimeoutMode(ResponseTimeoutMode responseTimeoutMode);
 }

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

@@ -0,0 +1,90 @@
+/*
+ * Copyright 2024 LINE Corporation
+ *
+ * LINE 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;
+
+import java.time.Duration;
+
+import com.linecorp.armeria.common.annotation.UnstableApi;
+import com.linecorp.armeria.common.util.TimeoutMode;
+
+/**
+ * Specifies when to start scheduling a response timeout.
+ *
+ * <p>For example:
+ * <pre>{@code
+ * |---request start---decorators(3s)---connection acquisition(2s)---request write(5s)---response read(4s)---|
+ * }</pre>
+ * <ul>
+ *     <li>
+ *         If {@link ResponseTimeoutMode#FROM_START} is used, the timeout task will be scheduled
+ *         immediately on request start. If the responseTimeout is less than 3 seconds, the timeout task
+ *         will trigger while the request goes through the decorator, and the request will fail before
+ *         acquiring a new connection. If the responseTimeout is greater than (3s + 2s + 5s + 4s) 14 seconds
+ *         the request will complete successfully.
+ *     </li>
+ *     <li>
+ *         If {@link ResponseTimeoutMode#CONNECTION_ACQUIRED} is used, the timeout task will be scheduled
+ *         after connection acquisition. If the responseTimeout is greater than (5s + 4s) 9 seconds the
+ *         request will complete successfully.
+ *     </li>
+ *     <li>
+ *         If {@link ResponseTimeoutMode#REQUEST_SENT} is used, the the timeout task will be scheduled after
+ *         the request is fully written. If the responseTimeout is greater than 4 seconds the request will
+ *         complete successfully.
+ *     </li>
+ * </ul>
+ *
+ * <p>When {@link TimeoutMode#SET_FROM_NOW} is used, the timeout is assumed to be scheduled when
+ * {@link ClientRequestContext#setResponseTimeout(TimeoutMode, Duration)} was called. If the request
+ * has already reached {@link ResponseTimeoutMode}, then the timeout is scheduled normally.
+ * If the request didn't reach {@link ResponseTimeoutMode} yet, the elapsed time is computed once
+ * {@link ResponseTimeoutMode} is reached and the timeout is scheduled accordingly.
+ * <pre>{@code
+ * |---request start---decorators(3s)---connection acquisition(2s)---request write(5s)---response read(4s)---|
+ * }</pre>
+ *  Assume {@link ResponseTimeoutMode#FROM_START} is set, and {@link TimeoutMode#SET_FROM_NOW}
+ *  with 1 second is called in the decorators. Then the timeout task will be triggered 1 second into connection
+ *  acquisition.
+ * <pre>{@code
+ * |---request start---decorators(3s)---connection acquisition(2s)---request write(5s)---response read(4s)---|
+ * }</pre>
+ *  Assume {@link ResponseTimeoutMode#REQUEST_SENT} is set, and {@link TimeoutMode#SET_FROM_NOW}
+ *  with 1 second is called in the decorators. The request will continue until the request is fully sent.
+ *  Since (2s + 5s) 7 seconds have elapsed which is greater than the 1-second timeout, the timeout task will be
+ *  invoked immediately before the response read starts.
+ */
+@UnstableApi
+public enum ResponseTimeoutMode {
+
+    /**
+     * The response timeout is scheduled when the request first starts to execute. More specifically,
+     * the scheduling will take place as soon as an endpoint is acquired but before the decorator chain
+     * is traversed.
+     */
+    FROM_START,
+
+    /**
+     * The response timeout is scheduled after the connection is acquired.
+     */
+    CONNECTION_ACQUIRED,
+
+    /**
+     * The response timeout is scheduled either after the client fully writes the request
+     * or when the response bytes are first read.
+     */
+    REQUEST_SENT,
+}

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

@@ -254,4 +254,9 @@ public RestClientBuilder contextCustomizer(
     public RestClientBuilder contextHook(Supplier<? extends AutoCloseable> contextHook) {
         return (RestClientBuilder) super.contextHook(contextHook);
     }
+
+    @Override
+    public RestClientBuilder responseTimeoutMode(ResponseTimeoutMode responseTimeoutMode) {
+        return (RestClientBuilder) super.responseTimeoutMode(responseTimeoutMode);
+    }
 }