Provide a way to find dynamic decorators handling the current request (…

…line#5670)

Motivation:

`HttpService.as(Class)` can be used to unwrap and find an instance of a decorator or a service.
`ServiceRequestContext.config().service().as(Class)` can't find all decorators because decorators set with `ServerBuilder.decorator()` don't statically wrap the services.

I propose to add `ServiceRequestContext.findService(Class)` for finding both dynamic and static decorators handling a request.

Modifications:

- Add `ServiceRequestContext.findService(Class)` that gets all service chain from `InitialDispatcherService` and finds the specific service.

Result:

You can now easily find both dynamic and static decorators that handle a request by using `ServiceRequestContext.findService(Class)`.

core/src/main/java/com/linecorp/armeria/internal/server/DefaultServiceRequestContext.java

@@ -62,6 +62,8 @@
 import com.linecorp.armeria.internal.common.InitiateConnectionShutdown;
 import com.linecorp.armeria.internal.common.NonWrappingRequestContext;
 import com.linecorp.armeria.internal.common.util.TemporaryThreadLocals;
+import com.linecorp.armeria.internal.server.RouteDecoratingService.InitialDispatcherService;
+import com.linecorp.armeria.server.HttpService;
 import com.linecorp.armeria.server.ProxiedAddresses;
 import com.linecorp.armeria.server.Route;
 import com.linecorp.armeria.server.RoutingContext;
@@ -255,6 +257,18 @@ public ServiceConfig config() {
         return cfg;
     }
 
+    @Nullable
+    @Override
+    public <T extends HttpService> T findService(Class<? extends T> serviceClass) {
+        requireNonNull(serviceClass, "serviceClass");
+        final HttpService service = config().service();
+        if (service instanceof InitialDispatcherService) {
+            return ((InitialDispatcherService) service).findService(this, serviceClass);
+        } else {
+            return service.as(serviceClass);
+        }
+    }
+
     @Override
     public RoutingContext routingContext() {
         return routingContext;

core/src/main/java/com/linecorp/armeria/internal/server/RouteDecoratingService.java

@@ -29,6 +29,7 @@
 import com.linecorp.armeria.common.annotation.Nullable;
 import com.linecorp.armeria.server.HttpService;
 import com.linecorp.armeria.server.Route;
+import com.linecorp.armeria.server.Routed;
 import com.linecorp.armeria.server.Router;
 import com.linecorp.armeria.server.RoutingContext;
 import com.linecorp.armeria.server.ServiceConfig;
@@ -162,7 +163,7 @@ public String toString() {
                           .toString();
     }
 
-    private static class InitialDispatcherService extends SimpleDecoratingHttpService {
+    public static final class InitialDispatcherService extends SimpleDecoratingHttpService {
 
         private final Router<RouteDecoratingService> router;
         private final List<RouteDecoratingService> routeDecoratingServices;
@@ -201,6 +202,20 @@ public HttpResponse serve(ServiceRequestContext ctx, HttpRequest req) throws Exc
             return service.serve(ctx, req);
         }
 
+        @Nullable
+        public <T extends HttpService> T findService(ServiceRequestContext ctx,
+                                                     Class<? extends T> serviceClass) {
+            for (Routed<RouteDecoratingService> routed : router.findAll(ctx.routingContext())) {
+                if (routed.isPresent()) {
+                    final T service = routed.value().decorator().as(serviceClass);
+                    if (service != null) {
+                        return service;
+                    }
+                }
+            }
+            return unwrap().as(serviceClass);
+        }
+
         @Override
         public String toString() {
             return MoreObjects.toStringHelper(this)

core/src/main/java/com/linecorp/armeria/server/ServiceRequestContext.java

@@ -252,6 +252,19 @@ default SafeCloseable push() {
      */
     ServiceConfig config();
 
+    /**
+     * Finds the {@link HttpService} of the specified {@link Class} within the decorators or the service
+     * handling the current {@link Request}. This method is particularly useful when attempting to finding a
+     * dynamic decorator set using {@link ServerBuilder#decorator(DecoratingHttpServiceFunction)} or
+     * {@link ServerBuilder#decorator(Route, DecoratingHttpServiceFunction)}.
+     *
+     * <p>Note that {@link HttpService#as(Class)} by itself cannot find a dynamic decorator added directly
+     * to a {@link ServerBuilder}.
+     */
+    @UnstableApi
+    @Nullable
+    <T extends HttpService> T findService(Class<? extends T> serviceClass);
+
     /**
      * Returns the {@link RoutingContext} used to find the {@link Service}.
      */

core/src/main/java/com/linecorp/armeria/server/ServiceRequestContextWrapper.java

@@ -85,6 +85,12 @@ public ServiceConfig config() {
         return unwrap().config();
     }
 
+    @Nullable
+    @Override
+    public <T extends HttpService> T findService(Class<? extends T> serviceClass) {
+        return unwrap().findService(serviceClass);
+    }
+
     @Override
     public RoutingContext routingContext() {
         return unwrap().routingContext();

core/src/main/java/com/linecorp/armeria/server/VirtualHost.java

@@ -498,6 +498,10 @@ public Routed<ServiceConfig> findServiceConfig(RoutingContext routingCtx, boolea
                     // CorsService will handle the preflight request
                     // even if the service does not handle an OPTIONS method.
                     return routed;
+                } else {
+                    // `handlesCorsPreflight()` is false if `CorsService` is set as a route decorator.
+                    // However, this is not a problem because the CorsService is chosen and applied by
+                    // `InitialDispatcherService` regardless of the target service.
                 }
                 break;
             default:

core/src/test/java/com/linecorp/armeria/ServiceFinderTest.java

@@ -0,0 +1,115 @@
+/*
+ * 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;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+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.HttpRequest;
+import com.linecorp.armeria.common.HttpResponse;
+import com.linecorp.armeria.common.HttpStatus;
+import com.linecorp.armeria.common.annotation.Nullable;
+import com.linecorp.armeria.common.metric.MeterIdPrefixFunction;
+import com.linecorp.armeria.server.HttpService;
+import com.linecorp.armeria.server.Server;
+import com.linecorp.armeria.server.ServerBuilder;
+import com.linecorp.armeria.server.ServiceRequestContext;
+import com.linecorp.armeria.server.SimpleDecoratingHttpService;
+import com.linecorp.armeria.server.cors.CorsService;
+import com.linecorp.armeria.server.encoding.EncodingService;
+import com.linecorp.armeria.server.logging.LoggingService;
+import com.linecorp.armeria.server.metric.MetricCollectingService;
+import com.linecorp.armeria.testing.junit5.server.ServerExtension;
+
+class ServiceFinderTest {
+    @RegisterExtension
+    static final ServerExtension server = new ServerExtension() {
+        @Override
+        protected void configure(ServerBuilder sb) {
+            sb.decorator(CorsService.builderForAnyOrigin().newDecorator());
+            sb.decoratorUnder("/prefix", LoggingService.newDecorator());
+            sb.decoratorUnder("/prefix/nested",
+                              MetricCollectingService.newDecorator(MeterIdPrefixFunction.ofDefault("nested")));
+            sb.decoratorUnder("/unrelated", EncodingService.newDecorator());
+            sb.service("/prefix/nested/service", new MyService().decorate(MyDecorator::new));
+        }
+    };
+
+    @Test
+    void shouldFindService() throws InterruptedException {
+        final BlockingWebClient client = server.blockingWebClient();
+        final AggregatedHttpResponse res = client.get("/prefix/nested/service");
+        assertThat(res.status()).isEqualTo(HttpStatus.OK);
+        final ServiceRequestContext ctx = server.requestContextCaptor().take();
+        assertThat(ctx.findService(CorsService.class)).isNotNull();
+        assertThat(ctx.findService(LoggingService.class)).isNotNull();
+        assertThat(ctx.findService(MyDecorator.class)).isNotNull();
+        assertThat(ctx.findService(MyService.class)).isNotNull();
+        // Should not find the unrelated service.
+        assertThat(ctx.findService(EncodingService.class)).isNull();
+    }
+
+    @Test
+    void shouldFindServiceWithoutRouteDecorators() throws InterruptedException {
+        final MyService myService = new MyService();
+        final Server server = Server.builder()
+                                    .service("/prefix/nested/service",
+                                             myService.decorate(MyDecorator::new))
+                                    .build();
+        server.start().join();
+        final BlockingWebClient client = BlockingWebClient.of("http://127.0.0.1:" + server.activeLocalPort());
+        final AggregatedHttpResponse res = client.get("/prefix/nested/service");
+        assertThat(res.status()).isEqualTo(HttpStatus.OK);
+
+        final ServiceRequestContext ctx = myService.lastCtx;
+        assertThat(ctx).isNotNull();
+        assertThat(ctx.findService(MyDecorator.class)).isNotNull();
+        assertThat(ctx.findService(MyService.class)).isNotNull();
+
+        assertThat(ctx.findService(CorsService.class)).isNull();
+        server.stop();
+    }
+
+    private static class MyService implements HttpService {
+        @Nullable
+        private volatile ServiceRequestContext lastCtx;
+
+        @Override
+        public HttpResponse serve(ServiceRequestContext ctx, HttpRequest req) throws Exception {
+            lastCtx = ctx;
+            return HttpResponse.of(HttpStatus.OK);
+        }
+    }
+
+    private static class MyDecorator extends SimpleDecoratingHttpService {
+        /**
+         * Creates a new instance that decorates the specified {@link HttpService}.
+         */
+        protected MyDecorator(HttpService delegate) {
+            super(delegate);
+        }
+
+        @Override
+        public HttpResponse serve(ServiceRequestContext ctx, HttpRequest req) throws Exception {
+            return unwrap().serve(ctx, req);
+        }
+    }
+}