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 index 67d7fb2b4f3..cf05998f12b 100644 --- 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 findService(Class 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 index a34851350db..4c4435332da 100644 --- 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 serviceChain(ServiceRequestContext ctx) { - final ImmutableList.Builder builder = ImmutableList.builder(); - router.findAll(ctx.routingContext()).forEach(routed -> { + @Nullable + public T findService(ServiceRequestContext ctx, + Class serviceClass) { + for (Routed 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 index 0db89742c89..7f433976acc 100644 --- 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 @Nullable T findService(Class serviceClass) { + public T findService(Class serviceClass) { return unwrap().findService(serviceClass); }