Fix to call the duplicator's child stream subscriber methods by the correct executor (#5783)

Motivation:
Subscriber methods must be called by the executor that is specified when subscribing to a `StreamMessage`. However, the methods of the duplicator's child stream subscriber are currently being called from the duplicator's executor, which is incorrect.

Modifications:
- Updated the implementation to ensure that the duplicator's child stream subscriber methods are called by the correct executor as specified during subscription.

Result:
- The duplicator's child stream subscriber methods are now called by the correct executor.

Author: minwoox

core/src/main/java/com/linecorp/armeria/common/stream/DefaultStreamMessageDuplicator.java

@@ -285,12 +285,7 @@ void subscribe(DownstreamSubscription<T> subscription) {
 
         private void doSubscribe(DownstreamSubscription<T> subscription) {
             if (state == State.ABORTED) {
-                final EventExecutor executor = subscription.executor;
-                if (executor.inEventLoop()) {
-                    failLateProcessorSubscriber(subscription);
-                } else {
-                    executor.execute(() -> failLateProcessorSubscriber(subscription));
-                }
+                subscription.failLateProcessorSubscriber();
                 return;
             }
 
@@ -301,63 +296,11 @@ private void doSubscribe(DownstreamSubscription<T> subscription) {
             }
         }
 
-        private static void failLateProcessorSubscriber(DownstreamSubscription<?> subscription) {
-            final Subscriber<?> lateSubscriber = subscription.subscriber();
-            try {
-                lateSubscriber.onSubscribe(NoopSubscription.get());
-                lateSubscriber.onError(
-                        new IllegalStateException("duplicator is closed or no more downstream can be added."));
-            } catch (Throwable t) {
-                throwIfFatal(t);
-                logger.warn("Subscriber should not throw an exception. subscriber: {}", lateSubscriber, t);
-            }
-        }
-
-        void unsubscribe(DownstreamSubscription<T> subscription, @Nullable Throwable cause) {
+        private void cleanupIfLastSubscription() {
             if (executor.inEventLoop()) {
-                doUnsubscribe(subscription, cause);
-            } else {
-                executor.execute(() -> doUnsubscribe(subscription, cause));
-            }
-        }
-
-        private void doUnsubscribe(DownstreamSubscription<T> subscription, @Nullable Throwable cause) {
-            if (!downstreamSubscriptions.remove(subscription)) {
-                return;
-            }
-
-            final Subscriber<? super T> subscriber = subscription.subscriber();
-            subscription.clearSubscriber();
-
-            final CompletableFuture<Void> completionFuture = subscription.whenComplete();
-            if (cause == null) {
-                try {
-                    subscriber.onComplete();
-                    completionFuture.complete(null);
-                } catch (Throwable t) {
-                    completionFuture.completeExceptionally(t);
-                    throwIfFatal(t);
-                    logger.warn("Subscriber.onComplete() should not raise an exception. subscriber: {}",
-                                subscriber, t);
-                } finally {
-                    doCleanupIfLastSubscription();
-                }
-                return;
-            }
-
-            try {
-                if (subscription.notifyCancellation || !(cause instanceof CancelledSubscriptionException)) {
-                    subscriber.onError(cause);
-                }
-                completionFuture.completeExceptionally(cause);
-            } catch (Throwable t) {
-                final Exception composite = new CompositeException(t, cause);
-                completionFuture.completeExceptionally(composite);
-                throwIfFatal(t);
-                logger.warn("Subscriber.onError() should not raise an exception. subscriber: {}",
-                            subscriber, composite);
-            } finally {
                 doCleanupIfLastSubscription();
+            } else {
+                executor.execute(this::doCleanupIfLastSubscription);
             }
         }
 
@@ -613,7 +556,7 @@ static final class DownstreamSubscription<T> implements Subscription {
         private final StreamMessage<T> streamMessage;
         private Subscriber<? super T> subscriber;
         private final StreamMessageProcessor<T> processor;
-        private final EventExecutor executor;
+        private final EventExecutor downstreamExecutor;
         private final boolean withPooledObjects;
         private final boolean notifyCancellation;
 
@@ -640,7 +583,7 @@ static final class DownstreamSubscription<T> implements Subscription {
             this.streamMessage = streamMessage;
             this.subscriber = subscriber;
             this.processor = processor;
-            this.executor = executor;
+            downstreamExecutor = executor;
             this.withPooledObjects = withPooledObjects;
             this.notifyCancellation = notifyCancellation;
         }
@@ -661,17 +604,36 @@ void clearSubscriber() {
             }
         }
 
+        void failLateProcessorSubscriber() {
+            if (downstreamExecutor.inEventLoop()) {
+                failLateProcessorSubscriber0();
+            } else {
+                downstreamExecutor.execute(this::failLateProcessorSubscriber0);
+            }
+        }
+
+        private void failLateProcessorSubscriber0() {
+            try {
+                subscriber.onSubscribe(NoopSubscription.get());
+                subscriber.onError(
+                        new IllegalStateException("duplicator is closed or no more downstream can be added."));
+            } catch (Throwable t) {
+                throwIfFatal(t);
+                logger.warn("Subscriber should not throw an exception. subscriber: {}", subscriber, t);
+            }
+        }
+
         // Called from processor.processorExecutor
         void invokeOnSubscribe() {
             if (invokedOnSubscribe) {
                 return;
             }
             invokedOnSubscribe = true;
 
-            if (executor.inEventLoop()) {
+            if (downstreamExecutor.inEventLoop()) {
                 invokeOnSubscribe0();
             } else {
-                executor.execute(this::invokeOnSubscribe0);
+                downstreamExecutor.execute(this::invokeOnSubscribe0);
             }
         }
 
@@ -680,7 +642,7 @@ void invokeOnSubscribe0() {
             try {
                 subscriber.onSubscribe(this);
             } catch (Throwable t) {
-                processor.unsubscribe(this, t);
+                unsubscribe(t);
                 throwIfFatal(t);
                 logger.warn("Subscriber.onSubscribe() should not raise an exception. subscriber: {}",
                             subscriber, t);
@@ -692,7 +654,7 @@ public void request(long n) {
             if (n <= 0) {
                 final Throwable cause = new IllegalArgumentException(
                         "n: " + n + " (expected: > 0, see Reactive Streams specification rule 3.9)");
-                processor.unsubscribe(this, cause);
+                unsubscribe(cause);
                 return;
             }
 
@@ -726,10 +688,10 @@ private void accumulateDemand(long n) {
         }
 
         void signal() {
-            if (executor.inEventLoop()) {
+            if (downstreamExecutor.inEventLoop()) {
                 doSignal();
             } else {
-                executor.execute(this::doSignal);
+                downstreamExecutor.execute(this::doSignal);
             }
         }
 
@@ -757,7 +719,7 @@ private boolean doSignalSingle(SignalQueue signals) {
 
             if (cancelledOrAborted != null) {
                 // Stream ended due to cancellation or abortion.
-                processor.unsubscribe(this, cancelledOrAborted);
+                unsubscribe(cancelledOrAborted);
                 return false;
             }
 
@@ -771,7 +733,7 @@ private boolean doSignalSingle(SignalQueue signals) {
             if (signal instanceof CloseEvent) {
                 // The stream has reached at its end.
                 offset++;
-                processor.unsubscribe(this, ((CloseEvent) signal).cause);
+                unsubscribe(((CloseEvent) signal).cause);
                 return false;
             }
 
@@ -812,7 +774,7 @@ private boolean doSignalSingle(SignalQueue signals) {
                     // If an exception such as IllegalReferenceCountException is raised while operating
                     // on the ByteBuf, catch it and notify the subscriber with it. So the
                     // subscriber does not hang forever.
-                    processor.unsubscribe(this, thrown);
+                    unsubscribe(thrown);
                     return false;
                 }
 
@@ -832,7 +794,7 @@ private boolean doSignalSingle(SignalQueue signals) {
                 try {
                     subscriber.onNext(obj);
                 } catch (Throwable t) {
-                    processor.unsubscribe(this, t);
+                    unsubscribe(t);
                     throwIfFatal(t);
                     logger.warn("Subscriber.onNext({}) should not raise an exception. subscriber: {}",
                                 obj, subscriber, t);
@@ -844,6 +806,54 @@ private boolean doSignalSingle(SignalQueue signals) {
             }
         }
 
+        void unsubscribe(@Nullable Throwable cause) {
+            if (downstreamExecutor.inEventLoop()) {
+                doUnsubscribe(cause);
+            } else {
+                downstreamExecutor.execute(() -> doUnsubscribe(cause));
+            }
+        }
+
+        private void doUnsubscribe(@Nullable Throwable cause) {
+            if (!processor.downstreamSubscriptions.remove(this)) {
+                return;
+            }
+
+            final Subscriber<? super T> subscriber = this.subscriber;
+            clearSubscriber();
+
+            final CompletableFuture<Void> completionFuture = whenComplete();
+            if (cause == null) {
+                try {
+                    subscriber.onComplete();
+                    completionFuture.complete(null);
+                } catch (Throwable t) {
+                    completionFuture.completeExceptionally(t);
+                    throwIfFatal(t);
+                    logger.warn("Subscriber.onComplete() should not raise an exception. subscriber: {}",
+                                subscriber, t);
+                } finally {
+                    processor.cleanupIfLastSubscription();
+                }
+                return;
+            }
+
+            try {
+                if (notifyCancellation || !(cause instanceof CancelledSubscriptionException)) {
+                    subscriber.onError(cause);
+                }
+                completionFuture.completeExceptionally(cause);
+            } catch (Throwable t) {
+                final Exception composite = new CompositeException(t, cause);
+                completionFuture.completeExceptionally(composite);
+                throwIfFatal(t);
+                logger.warn("Subscriber.onError() should not raise an exception. subscriber: {}",
+                            subscriber, composite);
+            } finally {
+                processor.cleanupIfLastSubscription();
+            }
+        }
+
         @Override
         public void cancel() {
             abort(subscriber instanceof AbortingSubscriber ? ((AbortingSubscriber<?>) subscriber).cause()

core/src/test/java/com/linecorp/armeria/common/stream/StreamMessageDuplicatorChildSubscriberTest.java

@@ -0,0 +1,104 @@
+/*
+ * Copyright 2021 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 static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.concurrent.CountDownLatch;
+
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+
+import com.linecorp.armeria.testing.junit5.common.EventLoopExtension;
+
+import io.netty.channel.EventLoop;
+
+class StreamMessageDuplicatorChildSubscriberTest {
+
+    @RegisterExtension
+    static final EventLoopExtension eventLoop1 = new EventLoopExtension();
+
+    @RegisterExtension
+    static final EventLoopExtension eventLoop2 = new EventLoopExtension();
+
+    @RegisterExtension
+    static final EventLoopExtension eventLoop3 = new EventLoopExtension();
+
+    @CsvSource({ "true", "false" })
+    @ParameterizedTest
+    void childSubscriberMethodsMustBeCalledByExecutors(boolean close) throws InterruptedException {
+        final StreamWriter<String> publisher = StreamMessage.streaming();
+        publisher.write("foo");
+        if (close) {
+            publisher.close();
+        } else {
+            publisher.abort();
+        }
+
+        final StreamMessageDuplicator<String> duplicator =
+                publisher.toDuplicator(eventLoop1.get());
+
+        final StreamMessage<String> first = duplicator.duplicate();
+        final StreamMessage<String> second = duplicator.duplicate();
+
+        duplicator.close();
+
+        final CountDownLatch latch = new CountDownLatch(2);
+        final EventLoop executor2 = eventLoop2.get();
+        first.subscribe(new ChildSubscriber(executor2, latch), executor2);
+
+        final EventLoop executor3 = eventLoop3.get();
+        second.subscribe(new ChildSubscriber(executor3, latch), executor3);
+        latch.await();
+    }
+
+    private static final class ChildSubscriber implements Subscriber<String> {
+
+        private final EventLoop eventLoop;
+        private final CountDownLatch latch;
+
+        ChildSubscriber(EventLoop eventLoop, CountDownLatch latch) {
+            this.eventLoop = eventLoop;
+            this.latch = latch;
+        }
+
+        @Override
+        public void onSubscribe(Subscription s) {
+            assertThat(eventLoop.inEventLoop()).isTrue();
+            s.request(Long.MAX_VALUE);
+        }
+
+        @Override
+        public void onNext(String data) {
+            assertThat(eventLoop.inEventLoop()).isTrue();
+        }
+
+        @Override
+        public void onError(Throwable t) {
+            assertThat(eventLoop.inEventLoop()).isTrue();
+            latch.countDown();
+        }
+
+        @Override
+        public void onComplete() {
+            assertThat(eventLoop.inEventLoop()).isTrue();
+            latch.countDown();
+        }
+    }
+}

core/src/test/java/com/linecorp/armeria/common/stream/StreamMessageDuplicatorTest.java

@@ -42,8 +42,6 @@
 import org.mockito.ArgumentCaptor;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 import com.google.common.base.Charsets;
 
@@ -61,8 +59,6 @@
 
 class StreamMessageDuplicatorTest {
 
-    private static final Logger logger = LoggerFactory.getLogger(StreamMessageDuplicatorTest.class);
-
     private static final List<ByteBuf> byteBufs = new ArrayList<>();
 
     @AfterEach