Provide a way to run service code out of I/O event loops (#5233)

Motivation: When service devs are not very familiar with asynchronous programming, it is very easy to drive Armeria's core event loops into havoc by blocking them. Framework devs may want to make sure Armeria at least handle I/O and function normally for a certain set of core services, such as `PrometheusExpositionService` or `HealthCheckService`, by isolating other non-core services from the I/O event loops. Modifications: * Add the `serviceWorkerGroup()` builder methods to `ServerBuilder` and `ServiceConfigSetters` so a user can specify the service worker groups as shown in the above example. * If `serviceWorkerGroup` is not specified, the `workerGroup` is used by default. * Change how Armeria assigns an event loop to a `ServiceRequestContext`. * If the `workerGroup` is different from `serviceWorkerGroup`, then an event loop from `serviceWorkerGroup` is used. * Otherwise, the IO event loop is used for executing services * Modified the constructor of `DefaultServiceRequestContext` so it accepts an `EventLoop`. * Modified `HttpServerHandler.handleRequest()`, so that `HttpService#serve` is executed from the `serviceWorkerGroup` * Modified so that we can guarantee that pending `RequestLogFuture`s are always scheduled from the context's event loop Result: - Closes #4099. - Users can add per-service/virtual host/server `serviceWorkerGroup` property that makes a service use a different `EventLoopGroup` than `ServerBuilder.workerGroup`. --------- Co-authored-by: kezhenxu94 <[email protected]>
line · Jan 26, 2024 · a54f418 · a54f418
1 parent d9703ed
commit a54f418
Show file tree

Hide file tree

Showing 32 changed files with 1,019 additions and 112 deletions.
diff --git a/benchmarks/jmh/src/jmh/java/com/linecorp/armeria/server/RoutersBenchmark.java b/benchmarks/jmh/src/jmh/java/com/linecorp/armeria/server/RoutersBenchmark.java
@@ -66,27 +66,27 @@ public class RoutersBenchmark {
                 new ServiceConfig(route1, route1,
                                   SERVICE, defaultLogName, defaultServiceName, defaultServiceNaming, 0, 0,
                                   false, AccessLogWriter.disabled(), CommonPools.blockingTaskExecutor(),
-                                  SuccessFunction.always(), 0, multipartUploadsLocation, ImmutableList.of(),
-                                  HttpHeaders.of(), ctx -> RequestId.random(), serviceErrorHandler,
-                                  NOOP_CONTEXT_HOOK),
+                                  SuccessFunction.always(), 0, multipartUploadsLocation,
+                                  CommonPools.workerGroup(), ImmutableList.of(), HttpHeaders.of(),
+                                  ctx -> RequestId.random(), serviceErrorHandler, NOOP_CONTEXT_HOOK),
                 new ServiceConfig(route2, route2,
                                   SERVICE, defaultLogName, defaultServiceName, defaultServiceNaming, 0, 0,
                                   false, AccessLogWriter.disabled(), CommonPools.blockingTaskExecutor(),
-                                  SuccessFunction.always(), 0, multipartUploadsLocation, ImmutableList.of(),
-                                  HttpHeaders.of(), ctx -> RequestId.random(), serviceErrorHandler,
-                                  NOOP_CONTEXT_HOOK));
+                                  SuccessFunction.always(), 0, multipartUploadsLocation,
+                                  CommonPools.workerGroup(), ImmutableList.of(), HttpHeaders.of(),
+                                  ctx -> RequestId.random(), serviceErrorHandler, NOOP_CONTEXT_HOOK));
         FALLBACK_SERVICE = new ServiceConfig(Route.ofCatchAll(), Route.ofCatchAll(), SERVICE,
                                              defaultLogName, defaultServiceName,
                                              defaultServiceNaming, 0, 0, false, AccessLogWriter.disabled(),
-                                             CommonPools.blockingTaskExecutor(),
-                                             SuccessFunction.always(), 0, multipartUploadsLocation,
+                                             CommonPools.blockingTaskExecutor(), SuccessFunction.always(), 0,
+                                             multipartUploadsLocation, CommonPools.workerGroup(),
                                              ImmutableList.of(), HttpHeaders.of(), ctx -> RequestId.random(),
                                              serviceErrorHandler, NOOP_CONTEXT_HOOK);
         HOST = new VirtualHost(
                 "localhost", "localhost", 0, null, SERVICES, FALLBACK_SERVICE, RejectedRouteHandler.DISABLED,
                 unused -> NOPLogger.NOP_LOGGER, defaultServiceNaming, defaultLogName, 0, 0, false,
                 AccessLogWriter.disabled(), CommonPools.blockingTaskExecutor(), 0, SuccessFunction.ofDefault(),
-                multipartUploadsLocation, ImmutableList.of(),
+                multipartUploadsLocation, CommonPools.workerGroup(), ImmutableList.of(),
                 ctx -> RequestId.random());
         ROUTER = Routers.ofVirtualHost(HOST, SERVICES, RejectedRouteHandler.DISABLED);
     }

diff --git a/...rc/test/java/com/linecorp/armeria/common/brave/RequestContextCurrentTraceContextTest.java b/...rc/test/java/com/linecorp/armeria/common/brave/RequestContextCurrentTraceContextTest.java
@@ -17,16 +17,13 @@
 package com.linecorp.armeria.common.brave;
 
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.when;
 
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoSettings;
 import org.mockito.quality.Strictness;
-import org.mockito.stubbing.Answer;
 
 import com.linecorp.armeria.common.HttpMethod;
 import com.linecorp.armeria.common.HttpRequest;
@@ -55,10 +52,7 @@ class RequestContextCurrentTraceContextTest {
     @BeforeEach
     void setUp() {
         when(eventLoop.inEventLoop()).thenReturn(true);
-        doAnswer((Answer<Void>) invocation -> {
-            invocation.<Runnable>getArgument(0).run();
-            return null;
-        }).when(eventLoop).execute(any());
+        when(eventLoop.next()).thenReturn(eventLoop);
 
         ctx = ServiceRequestContext.builder(HttpRequest.of(HttpMethod.GET, "/"))
                                    .eventLoop(eventLoop)

diff --git a/core/src/main/java/com/linecorp/armeria/common/HttpRequest.java b/core/src/main/java/com/linecorp/armeria/common/HttpRequest.java
@@ -810,4 +810,10 @@ default HttpRequest peekError(Consumer<? super Throwable> action) {
         requireNonNull(action, "action");
         return of(headers(), HttpMessage.super.peekError(action));
     }
+
+    @Override
+    default HttpRequest subscribeOn(EventExecutor eventExecutor) {
+        requireNonNull(eventExecutor, "eventExecutor");
+        return of(headers(), HttpMessage.super.subscribeOn(eventExecutor));
+    }
 }
diff --git a/core/src/main/java/com/linecorp/armeria/common/HttpResponse.java b/core/src/main/java/com/linecorp/armeria/common/HttpResponse.java
@@ -1178,4 +1178,9 @@ default <T extends Throwable> HttpResponse recover(Class<T> causeClass,
             }
         });
     }
+
+    @Override
+    default HttpResponse subscribeOn(EventExecutor eventExecutor) {
+        return of(HttpMessage.super.subscribeOn(eventExecutor));
+    }
 }
diff --git a/core/src/main/java/com/linecorp/armeria/common/stream/StreamMessage.java b/core/src/main/java/com/linecorp/armeria/common/stream/StreamMessage.java
@@ -1136,4 +1136,22 @@ default InputStream toInputStream(Function<? super T, ? extends HttpData> httpDa
     default StreamMessage<T> endWith(Function<@Nullable Throwable, ? extends @Nullable T> finalizer) {
         return new SurroundingPublisher<>(null, this, finalizer);
     }
+
+    /**
+     * Calls {@link #subscribe(Subscriber, EventExecutor)} to the upstream
+     * {@link StreamMessage} using the specified {@link EventExecutor} and relays the stream
+     * transparently downstream. This may be useful if one would like to hide an
+     * {@link EventExecutor} from an upstream {@link Publisher}.
+     *
+     * <p>For example:<pre>{@code
+     * Subscriber<Integer> mySubscriber = null;
+     * StreamMessage<Integer> upstream = ...; // publisher callbacks are invoked by eventLoop1
+     * upstream.subscribeOn(eventLoop1)
+     *         .subscribe(mySubscriber, eventLoop2); // mySubscriber callbacks are invoked with eventLoop2
+     * }</pre>
+     */
+    default StreamMessage<T> subscribeOn(EventExecutor eventExecutor) {
+        requireNonNull(eventExecutor, "eventExecutor");
+        return new SubscribeOnStreamMessage<>(this, eventExecutor);
+    }
 }
diff --git a/core/src/main/java/com/linecorp/armeria/common/stream/SubscribeOnStreamMessage.java b/core/src/main/java/com/linecorp/armeria/common/stream/SubscribeOnStreamMessage.java
@@ -0,0 +1,117 @@
+/*
+ * Copyright 2023 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.common.stream;
+
+import java.util.concurrent.CompletableFuture;
+
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+
+import io.netty.util.concurrent.EventExecutor;
+
+final class SubscribeOnStreamMessage<T> implements StreamMessage<T> {
+
+    private final StreamMessage<T> upstream;
+    private final EventExecutor upstreamExecutor;
+
+    SubscribeOnStreamMessage(StreamMessage<T> upstream, EventExecutor upstreamExecutor) {
+        this.upstream = upstream;
+        this.upstreamExecutor = upstreamExecutor;
+    }
+
+    @Override
+    public boolean isOpen() {
+        return upstream.isOpen();
+    }
+
+    @Override
+    public boolean isEmpty() {
+        return upstream.isEmpty();
+    }
+
+    @Override
+    public long demand() {
+        return upstream.demand();
+    }
+
+    @Override
+    public CompletableFuture<Void> whenComplete() {
+        return upstream.whenComplete();
+    }
+
+    @Override
+    public EventExecutor defaultSubscriberExecutor() {
+        return upstreamExecutor;
+    }
+
+    @Override
+    public void subscribe(Subscriber<? super T> subscriber, EventExecutor downstreamExecutor,
+                          SubscriptionOption... options) {
+        final Subscriber<? super T> subscriber0;
+        if (upstreamExecutor == downstreamExecutor) {
+            subscriber0 = subscriber;
+        } else {
+            subscriber0 = new SchedulingSubscriber<>(downstreamExecutor, subscriber);
+        }
+        if (upstreamExecutor.inEventLoop()) {
+            upstream.subscribe(subscriber0, downstreamExecutor, options);
+        } else {
+            upstreamExecutor.execute(() -> upstream.subscribe(subscriber0, upstreamExecutor, options));
+        }
+    }
+
+    @Override
+    public void abort() {
+        upstream.abort();
+    }
+
+    @Override
+    public void abort(Throwable cause) {
+        upstream.abort(cause);
+    }
+
+    static class SchedulingSubscriber<T> implements Subscriber<T> {
+
+        private final Subscriber<? super T> downstream;
+        private final EventExecutor downstreamExecutor;
+
+        SchedulingSubscriber(EventExecutor downstreamExecutor, Subscriber<? super T> downstream) {
+            this.downstream = downstream;
+            this.downstreamExecutor = downstreamExecutor;
+        }
+
+        @Override
+        public void onSubscribe(Subscription s) {
+            downstreamExecutor.execute(() -> downstream.onSubscribe(s));
+        }
+
+        @Override
+        public void onNext(T t) {
+            downstreamExecutor.execute(() -> downstream.onNext(t));
+        }
+
+        @Override
+        public void onError(Throwable t) {
+            downstreamExecutor.execute(() -> downstream.onError(t));
+        }
+
+        @Override
+        public void onComplete() {
+            downstreamExecutor.execute(downstream::onComplete);
+        }
+    }
+}
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
@@ -72,6 +72,7 @@
 import io.micrometer.core.instrument.MeterRegistry;
 import io.netty.buffer.ByteBufAllocator;
 import io.netty.channel.Channel;
+import io.netty.channel.EventLoop;
 import io.netty.util.AttributeKey;
 
 /**
@@ -90,6 +91,7 @@ public final class DefaultServiceRequestContext
             DefaultServiceRequestContext.class, HttpHeaders.class, "additionalResponseTrailers");
 
     private final Channel ch;
+    private final EventLoop eventLoop;
     private final ServiceConfig cfg;
     private final RoutingContext routingContext;
     private final RoutingResult routingResult;
@@ -141,22 +143,24 @@ public final class DefaultServiceRequestContext
      *                               e.g. {@code System.currentTimeMillis() * 1000}.
      */
     public DefaultServiceRequestContext(
-            ServiceConfig cfg, Channel ch, MeterRegistry meterRegistry, SessionProtocol sessionProtocol,
-            RequestId id, RoutingContext routingContext, RoutingResult routingResult, ExchangeType exchangeType,
+            ServiceConfig cfg, Channel ch, EventLoop eventLoop, MeterRegistry meterRegistry,
+            SessionProtocol sessionProtocol, RequestId id, RoutingContext routingContext,
+            RoutingResult routingResult, ExchangeType exchangeType,
             HttpRequest req, @Nullable SSLSession sslSession, ProxiedAddresses proxiedAddresses,
             InetAddress clientAddress, InetSocketAddress remoteAddress, InetSocketAddress localAddress,
             long requestStartTimeNanos, long requestStartTimeMicros,
             Supplier<? extends AutoCloseable> contextHook) {
 
-        this(cfg, ch, meterRegistry, sessionProtocol, id, routingContext, routingResult, exchangeType,
-             req, sslSession, proxiedAddresses, clientAddress, remoteAddress, localAddress,
+        this(cfg, ch, eventLoop, meterRegistry, sessionProtocol, id, routingContext, routingResult,
+             exchangeType, req, sslSession, proxiedAddresses, clientAddress, remoteAddress, localAddress,
              null /* requestCancellationScheduler */, requestStartTimeNanos, requestStartTimeMicros,
              HttpHeaders.of(), HttpHeaders.of(), contextHook);
     }
 
     public DefaultServiceRequestContext(
-            ServiceConfig cfg, Channel ch, MeterRegistry meterRegistry, SessionProtocol sessionProtocol,
-            RequestId id, RoutingContext routingContext, RoutingResult routingResult, ExchangeType exchangeType,
+            ServiceConfig cfg, Channel ch, EventLoop eventLoop, MeterRegistry meterRegistry,
+            SessionProtocol sessionProtocol, RequestId id, RoutingContext routingContext,
+            RoutingResult routingResult, ExchangeType exchangeType,
             HttpRequest req, @Nullable SSLSession sslSession, ProxiedAddresses proxiedAddresses,
             InetAddress clientAddress, InetSocketAddress remoteAddress, InetSocketAddress localAddress,
             @Nullable CancellationScheduler requestCancellationScheduler,
@@ -170,6 +174,7 @@ public DefaultServiceRequestContext(
               requireNonNull(req, "req"), null, null, contextHook);
 
         this.ch = requireNonNull(ch, "ch");
+        this.eventLoop = requireNonNull(eventLoop, "eventLoop");
         this.cfg = requireNonNull(cfg, "cfg");
         this.routingContext = routingContext;
         this.routingResult = routingResult;
@@ -178,7 +183,9 @@ public DefaultServiceRequestContext(
         } else {
             this.requestCancellationScheduler =
                     CancellationScheduler.ofServer(TimeUnit.MILLISECONDS.toNanos(cfg.requestTimeoutMillis()));
-            this.requestCancellationScheduler.init(eventLoop());
+            // the cancellation scheduler uses channelEventLoop since #start is called
+            // from the netty pipeline logic
+            this.requestCancellationScheduler.init(ch.eventLoop());
         }
         this.sslSession = sslSession;
         this.proxiedAddresses = requireNonNull(proxiedAddresses, "proxiedAddresses");
@@ -301,7 +308,7 @@ public ContextAwareEventLoop eventLoop() {
         if (contextAwareEventLoop != null) {
             return contextAwareEventLoop;
         }
-        return contextAwareEventLoop = ContextAwareEventLoop.of(this, ch.eventLoop());
+        return contextAwareEventLoop = ContextAwareEventLoop.of(this, eventLoop);
     }
 
     @Override

diff --git a/core/src/main/java/com/linecorp/armeria/server/AbstractAnnotatedServiceConfigSetters.java b/core/src/main/java/com/linecorp/armeria/server/AbstractAnnotatedServiceConfigSetters.java
@@ -47,6 +47,8 @@
 import com.linecorp.armeria.server.annotation.ResponseConverterFunction;
 import com.linecorp.armeria.server.logging.AccessLogWriter;
 
+import io.netty.channel.EventLoopGroup;
+
 @UnstableApi
 abstract class AbstractAnnotatedServiceConfigSetters implements AnnotatedServiceConfigSetters {
 
@@ -280,6 +282,18 @@ public AbstractAnnotatedServiceConfigSetters multipartUploadsLocation(Path multi
         return this;
     }
 
+    @Override
+    public ServiceConfigSetters serviceWorkerGroup(EventLoopGroup serviceWorkerGroup, boolean shutdownOnStop) {
+        defaultServiceConfigSetters.serviceWorkerGroup(serviceWorkerGroup, shutdownOnStop);
+        return this;
+    }
+
+    @Override
+    public ServiceConfigSetters serviceWorkerGroup(int numThreads) {
+        defaultServiceConfigSetters.serviceWorkerGroup(numThreads);
+        return this;
+    }
+
     @Override
     public AbstractAnnotatedServiceConfigSetters requestIdGenerator(
             Function<? super RoutingContext, ? extends RequestId> requestIdGenerator) {

diff --git a/core/src/main/java/com/linecorp/armeria/server/AbstractServiceBindingBuilder.java b/core/src/main/java/com/linecorp/armeria/server/AbstractServiceBindingBuilder.java
@@ -35,6 +35,8 @@
 import com.linecorp.armeria.common.util.BlockingTaskExecutor;
 import com.linecorp.armeria.server.logging.AccessLogWriter;
 
+import io.netty.channel.EventLoopGroup;
+
 /**
  * A builder class for binding an {@link HttpService} fluently.
  *
@@ -175,6 +177,19 @@ public AbstractServiceBindingBuilder multipartUploadsLocation(Path multipartUplo
         return this;
     }
 
+    @Override
+    public AbstractServiceBindingBuilder serviceWorkerGroup(EventLoopGroup serviceWorkerGroup,
+                                                            boolean shutdownOnStop) {
+        defaultServiceConfigSetters.serviceWorkerGroup(serviceWorkerGroup, shutdownOnStop);
+        return this;
+    }
+
+    @Override
+    public AbstractServiceBindingBuilder serviceWorkerGroup(int numThreads) {
+        defaultServiceConfigSetters.serviceWorkerGroup(numThreads);
+        return this;
+    }
+
     @Override
     public AbstractServiceBindingBuilder requestIdGenerator(
             Function<? super RoutingContext, ? extends RequestId> requestIdGenerator) {

diff --git a/core/src/main/java/com/linecorp/armeria/server/AggregatedHttpResponseHandler.java b/core/src/main/java/com/linecorp/armeria/server/AggregatedHttpResponseHandler.java
@@ -57,7 +57,7 @@ final class AggregatedHttpResponseHandler extends AbstractHttpResponseHandler
 
     @Override
     public Void apply(@Nullable AggregatedHttpResponse response, @Nullable Throwable cause) {
-        final EventLoop eventLoop = reqCtx.eventLoop();
+        final EventLoop eventLoop = ctx.channel().eventLoop();
         if (eventLoop.inEventLoop()) {
             apply0(response, cause);
         } else {

diff --git a/core/src/main/java/com/linecorp/armeria/server/AnnotatedServiceBindingBuilder.java b/core/src/main/java/com/linecorp/armeria/server/AnnotatedServiceBindingBuilder.java
@@ -33,6 +33,8 @@
 import com.linecorp.armeria.server.annotation.ResponseConverterFunction;
 import com.linecorp.armeria.server.logging.AccessLogWriter;
 
+import io.netty.channel.EventLoopGroup;
+
 /**
  * A builder class for binding an {@link HttpService} fluently. This class can be instantiated through
  * {@link ServerBuilder#annotatedService()}.
@@ -260,6 +262,17 @@ public AnnotatedServiceBindingBuilder contextHook(Supplier<? extends AutoCloseab
         return (AnnotatedServiceBindingBuilder) super.contextHook(contextHook);
     }
 
+    @Override
+    public AnnotatedServiceBindingBuilder serviceWorkerGroup(EventLoopGroup serviceWorkerGroup,
+                                                             boolean shutdownOnStop) {
+        return (AnnotatedServiceBindingBuilder) super.serviceWorkerGroup(serviceWorkerGroup, shutdownOnStop);
+    }
+
+    @Override
+    public AnnotatedServiceBindingBuilder serviceWorkerGroup(int numThreads) {
+        return (AnnotatedServiceBindingBuilder) super.serviceWorkerGroup(numThreads);
+    }
+
     /**
      * Registers the given service to {@link ServerBuilder} and return {@link ServerBuilder}
      * to continue building {@link Server}.