address comments by @trustin

line · May 9, 2024 · b183f38 · b183f38
1 parent 4be1d4f
commit b183f38
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 22 deletions.
diff --git a/core/src/main/java/com/linecorp/armeria/internal/server/DefaultServiceRequestContext.java b/core/src/main/java/com/linecorp/armeria/internal/server/DefaultServiceRequestContext.java
@@ -266,21 +266,11 @@ public ServiceConfig config() {
  public <T extends HttpService> T findService(Class<? extends T> serviceClass) {
  requireNonNull(serviceClass, "serviceClass");
  final HttpService service = config().service();
- if (!(service instanceof InitialDispatcherService)) {
+ if (service instanceof InitialDispatcherService) {
+ return ((InitialDispatcherService) service).findService(this, serviceClass);
+ } else {
  return service.as(serviceClass);
  }
-
- if (serviceChain == null) {
- serviceChain = ((InitialDispatcherService) service).serviceChain(this);
- }
-
- for (HttpService service0 : serviceChain) {
- final T targetService = service0.as(serviceClass);
- if (targetService != null) {
- return targetService;
- }
- }
- return null;
  }
 
  @Override

diff --git a/core/src/main/java/com/linecorp/armeria/internal/server/RouteDecoratingService.java b/core/src/main/java/com/linecorp/armeria/internal/server/RouteDecoratingService.java
@@ -23,13 +23,13 @@
 import java.util.function.Function;
 
 import com.google.common.base.MoreObjects;
-import com.google.common.collect.ImmutableList;
 
 import com.linecorp.armeria.common.HttpRequest;
 import com.linecorp.armeria.common.HttpResponse;
 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;
@@ -202,15 +202,18 @@ public HttpResponse serve(ServiceRequestContext ctx, HttpRequest req) throws Exc
  return service.serve(ctx, req);
  }
 
- public List<HttpService> serviceChain(ServiceRequestContext ctx) {
- final ImmutableList.Builder<HttpService> builder = ImmutableList.builder();
- router.findAll(ctx.routingContext()).forEach(routed -> {
+ @Nullable
+ public <T extends HttpService> T findService(ServiceRequestContext ctx,
+ Class<? extends T> serviceClass) {
+ for (Routed<RouteDecoratingService> routed : router.findAll(ctx.routingContext())) {
  if (routed.isPresent()) {
- builder.add(routed.value().decorator());
+ final T service = routed.value().decorator().as(serviceClass);
+ if (service != null) {
+ return service;
+ }
  }
- });
- builder.add((HttpService) unwrap());
- return builder.build();
+ }
+ return unwrap().as(serviceClass);
  }
 
  @Override

diff --git a/core/src/main/java/com/linecorp/armeria/server/ServiceRequestContextWrapper.java b/core/src/main/java/com/linecorp/armeria/server/ServiceRequestContextWrapper.java
@@ -85,8 +85,9 @@ public ServiceConfig config() {
  return unwrap().config();
  }
 
+ @Nullable
  @Override
- public <T extends HttpService> @Nullable T findService(Class<? extends T> serviceClass) {
+ public <T extends HttpService> T findService(Class<? extends T> serviceClass) {
  return unwrap().findService(serviceClass);
  }