implement request timeout annotation

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

@@ -27,6 +27,7 @@
 import com.linecorp.armeria.server.annotation.ServerSentEventResponseConverterFunction;
 import com.linecorp.armeria.server.annotation.decorator.LoggingDecoratorFactoryFunction;
 import com.linecorp.armeria.server.annotation.decorator.RateLimitingDecoratorFactoryFunction;
+import com.linecorp.armeria.server.annotation.decorator.RequestTimeoutDecoratorFunction;
 
 public enum BuiltInDependencyInjector implements DependencyInjector {
 
@@ -36,7 +37,8 @@ public enum BuiltInDependencyInjector implements DependencyInjector {
     private static final Set<Class<?>> builtInClasses =
             ImmutableSet.of(LoggingDecoratorFactoryFunction.class,
                             RateLimitingDecoratorFactoryFunction.class,
-                            ServerSentEventResponseConverterFunction.class);
+                            ServerSentEventResponseConverterFunction.class,
+                            RequestTimeoutDecoratorFunction.class);
 
     private static final Map<Class<?>, Object> instances = new ConcurrentHashMap<>();
 

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

@@ -201,7 +201,7 @@ private void setTimeoutNanosFromStart(long timeoutNanos) {
                 eventLoop.execute(() -> setTimeoutNanosFromStart0(timeoutNanos));
             }
         } else {
-            addPendingTimeoutNanos(timeoutNanos);
+            setPendingTimeoutNanos(timeoutNanos);
             addPendingTask(() -> setTimeoutNanosFromStart0(timeoutNanos));
         }
     }
@@ -260,18 +260,18 @@ private void setTimeoutNanosFromNow(long timeoutNanos) {
             if (eventLoop.inEventLoop()) {
                 setTimeoutNanosFromNow0(timeoutNanos);
             } else {
-                final long startTimeNanos = System.nanoTime();
+                final long eventLoopStartTimeNanos = System.nanoTime();
                 eventLoop.execute(() -> {
-                    final long passedTimeNanos0 = System.nanoTime() - startTimeNanos;
+                    final long passedTimeNanos0 = System.nanoTime() - eventLoopStartTimeNanos;
                     final long timeoutNanos0 = Math.max(1, timeoutNanos - passedTimeNanos0);
                     setTimeoutNanosFromNow0(timeoutNanos0);
                 });
             }
         } else {
-            final long startTimeNanos = System.nanoTime();
+            final long pendingTaskRegisterTimeNanos = System.nanoTime();
             setPendingTimeoutNanos(timeoutNanos);
             addPendingTask(() -> {
-                final long passedTimeNanos0 = System.nanoTime() - startTimeNanos;
+                final long passedTimeNanos0 = System.nanoTime() - pendingTaskRegisterTimeNanos;
                 final long timeoutNanos0 = Math.max(1, timeoutNanos - passedTimeNanos0);
                 setTimeoutNanosFromNow0(timeoutNanos0);
             });
@@ -411,7 +411,8 @@ public CompletableFuture<Void> whenTimedOut() {
         }
     }
 
-    private boolean isInitialized() {
+    @VisibleForTesting
+    public boolean isInitialized() {
         return pendingTask == noopPendingTask && eventLoop != null;
     }
 

core/src/main/java/com/linecorp/armeria/server/annotation/decorator/RequestTimeout.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright 2022 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.server.annotation.decorator;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+import java.util.concurrent.TimeUnit;
+
+import com.linecorp.armeria.server.annotation.DecoratorFactory;
+
+/**
+ * Annotation for request timeout.
+ */
+@DecoratorFactory(RequestTimeoutDecoratorFunction.class)
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ ElementType.TYPE, ElementType.METHOD })
+public @interface RequestTimeout {
+
+    /**
+     * Value of request timeout to set.
+     */
+    long value();
+
+    /**
+     * Time unit of request timeout to set.
+     */
+    TimeUnit unit() default TimeUnit.MILLISECONDS;
+}

core/src/main/java/com/linecorp/armeria/server/annotation/decorator/RequestTimeoutDecoratorFunction.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2022 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.server.annotation.decorator;
+
+import java.util.function.Function;
+
+import com.linecorp.armeria.common.HttpRequest;
+import com.linecorp.armeria.common.HttpResponse;
+import com.linecorp.armeria.common.util.TimeoutMode;
+import com.linecorp.armeria.server.HttpService;
+import com.linecorp.armeria.server.ServiceRequestContext;
+import com.linecorp.armeria.server.SimpleDecoratingHttpService;
+import com.linecorp.armeria.server.annotation.DecoratorFactoryFunction;
+
+/**
+ * A factory which creates a decorator that sets request timeout to current {@link ServiceRequestContext}.
+ */
+public final class RequestTimeoutDecoratorFunction implements DecoratorFactoryFunction<RequestTimeout> {
+
+    /**
+     * Creates a new decorator with the specified {@code parameter}.
+     */
+    @Override
+    public Function<? super HttpService, ? extends HttpService> newDecorator(RequestTimeout parameter) {
+        final long timeoutMillis = parameter.unit().toMillis(parameter.value());
+        return delegate -> new SimpleDecoratingHttpService(delegate) {
+            @Override
+            public HttpResponse serve(ServiceRequestContext ctx, HttpRequest req) throws Exception {
+                ServiceRequestContext.current()
+                                     .setRequestTimeoutMillis(TimeoutMode.SET_FROM_START, timeoutMillis);
+                return delegate.serve(ctx, req);
+            }
+        };
+    }
+}

core/src/test/java/com/linecorp/armeria/internal/server/annotation/RequestTimeoutAnnotationTest.java

@@ -0,0 +1,98 @@
+/*
+ * Copyright 2022 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.internal.server.annotation;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.concurrent.TimeUnit;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import com.linecorp.armeria.client.BlockingWebClient;
+import com.linecorp.armeria.common.AggregatedHttpResponse;
+import com.linecorp.armeria.common.HttpMethod;
+import com.linecorp.armeria.common.HttpRequest;
+import com.linecorp.armeria.common.HttpStatus;
+import com.linecorp.armeria.common.RequestHeaders;
+import com.linecorp.armeria.internal.server.DefaultServiceRequestContext;
+import com.linecorp.armeria.server.ServerBuilder;
+import com.linecorp.armeria.server.ServiceRequestContext;
+import com.linecorp.armeria.server.annotation.Get;
+import com.linecorp.armeria.server.annotation.decorator.RequestTimeout;
+import com.linecorp.armeria.testing.junit5.server.ServerExtension;
+
+public class RequestTimeoutAnnotationTest {
+
+    @RegisterExtension
+    static ServerExtension server = new ServerExtension() {
+        @Override
+        protected void configure(ServerBuilder sb) throws Exception {
+            sb.annotatedService("/myService", new MyAnnotationService());
+        }
+    };
+
+    static final long timeoutMillis = 1230;
+    static final long timeoutSeconds = 4560;
+
+    public static class MyAnnotationService {
+        @Get("/timeoutMillis")
+        @RequestTimeout(timeoutMillis)
+        public String timeoutMillis(ServiceRequestContext ctx, HttpRequest req) {
+            AnnotatedServiceTest.validateContextAndRequest(ctx, req);
+            return Long.toString(ctx.requestTimeoutMillis());
+        }
+
+        @Get("/timeoutSeconds")
+        @RequestTimeout(value = timeoutSeconds, unit = TimeUnit.SECONDS)
+        public String timeoutSeconds(ServiceRequestContext ctx, HttpRequest req) {
+            AnnotatedServiceTest.validateContextAndRequest(ctx, req);
+            return Long.toString(ctx.requestTimeoutMillis());
+        }
+
+        @Get("/subscriberIsInitialized")
+        public boolean subscriberIsInitialized(ServiceRequestContext ctx, HttpRequest req) {
+            AnnotatedServiceTest.validateContextAndRequest(ctx, req);
+            return ((DefaultServiceRequestContext) ctx).requestCancellationScheduler().isInitialized();
+        }
+    }
+
+    @Test
+    public void testRequestTimeoutSet() {
+        final BlockingWebClient client = BlockingWebClient.of(server.httpUri());
+
+        AggregatedHttpResponse response;
+
+        response = client.execute(RequestHeaders.of(HttpMethod.GET, "/myService/timeoutMillis"));
+        assertThat(response.status()).isEqualTo(HttpStatus.OK);
+        assertThat(Long.parseLong(response.contentUtf8())).isEqualTo(timeoutMillis);
+
+        response = client.execute(RequestHeaders.of(HttpMethod.GET, "/myService/timeoutSeconds"));
+        assertThat(response.status()).isEqualTo(HttpStatus.OK);
+        assertThat(Long.parseLong(response.contentUtf8()))
+                .isEqualTo(TimeUnit.SECONDS.toMillis(timeoutSeconds));
+    }
+
+    @Test
+    public void testCancellationSchedulerInit() {
+        final BlockingWebClient client = BlockingWebClient.of(server.httpUri());
+
+        final AggregatedHttpResponse response = client.execute(
+                RequestHeaders.of(HttpMethod.GET, "/myService/timeoutMillis"));
+        assertThat(response.status()).isEqualTo(HttpStatus.OK);
+        assertThat(Boolean.parseBoolean(response.contentUtf8())).isEqualTo(true);
+    }
+}