Add Publisher.flatMapConcatSingle

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/PublisherFlatMapConcatUtils.java

@@ -0,0 +1,164 @@
+/*
+ * Copyright © 2022 Apple Inc. and the ServiceTalk project authors
+ *
+ * Licensed 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
+ *
+ *   http://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 io.servicetalk.concurrent.api;
+
+import io.servicetalk.concurrent.Cancellable;
+import io.servicetalk.concurrent.SingleSource;
+
+import java.util.Queue;
+import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
+import java.util.function.Function;
+import javax.annotation.Nullable;
+
+import static io.servicetalk.concurrent.api.Publisher.defer;
+import static io.servicetalk.concurrent.api.SourceAdapters.toSource;
+import static io.servicetalk.concurrent.internal.ConcurrentUtils.releaseLock;
+import static io.servicetalk.concurrent.internal.ConcurrentUtils.tryAcquireLock;
+import static io.servicetalk.utils.internal.PlatformDependent.newUnboundedSpscQueue;
+import static java.lang.Math.min;
+
+final class PublisherFlatMapConcatUtils {
+    private PublisherFlatMapConcatUtils() {
+    }
+
+    static <T, R> Publisher<R> flatMapConcatSingle(final Publisher<T> publisher,
+                                                   final Function<? super T, ? extends Single<? extends R>> mapper) {
+        return defer(() -> publisher.flatMapMergeSingle(new OrderedMapper<>(mapper, newUnboundedSpscQueue(4)))
+                .shareContextOnSubscribe());
+    }
+
+    static <T, R> Publisher<R> flatMapConcatSingleDelayError(
+            final Publisher<T> publisher, final Function<? super T, ? extends Single<? extends R>> mapper) {
+        return defer(() -> publisher.flatMapMergeSingleDelayError(new OrderedMapper<>(mapper, newUnboundedSpscQueue(4)))
+                .shareContextOnSubscribe());
+    }
+
+    static <T, R> Publisher<R> flatMapConcatSingle(final Publisher<T> publisher,
+                                                   final Function<? super T, ? extends Single<? extends R>> mapper,
+                                                   final int maxConcurrency) {
+        return defer(() ->
+                publisher.flatMapMergeSingle(new OrderedMapper<>(mapper,
+                                newUnboundedSpscQueue(min(8, maxConcurrency))), maxConcurrency)
+                        .shareContextOnSubscribe());
+    }
+
+    static <T, R> Publisher<R> flatMapConcatSingleDelayError(
+            final Publisher<T> publisher, final Function<? super T, ? extends Single<? extends R>> mapper,
+            final int maxConcurrency) {
+        return defer(() ->
+                publisher.flatMapMergeSingleDelayError(new OrderedMapper<>(mapper,
+                                newUnboundedSpscQueue(min(8, maxConcurrency))), maxConcurrency)
+                        .shareContextOnSubscribe());
+    }
+
+    private static final class OrderedMapper<T, R> implements Function<T, Single<R>> {
+        @SuppressWarnings("rawtypes")
+        private static final AtomicIntegerFieldUpdater<OrderedMapper> consumerLockUpdater =
+                AtomicIntegerFieldUpdater.newUpdater(OrderedMapper.class, "consumerLock");
+        private final Function<? super T, ? extends Single<? extends R>> mapper;
+        private final Queue<Item<R>> results;
+        @SuppressWarnings("unused")
+        private volatile int consumerLock;
+
+        private OrderedMapper(final Function<? super T, ? extends Single<? extends R>> mapper,
+                              final Queue<Item<R>> results) {
+            this.mapper = mapper;
+            this.results = results;
+        }
+
+        @Override
+        public Single<R> apply(final T t) {
+            final Single<? extends R> single = mapper.apply(t);
+            final Item<R> item = new Item<>();
+            results.add(item);
+            return new Single<R>() {
+                @Override
+                protected void handleSubscribe(final SingleSource.Subscriber<? super R> subscriber) {
+                    assert item.subscriber == null; // flatMapMergeSingle only does a single subscribe.
+                    item.subscriber = subscriber;
+                    toSource(single).subscribe(new SingleSource.Subscriber<R>() {
+                        @Override
+                        public void onSubscribe(final Cancellable cancellable) {
+                            subscriber.onSubscribe(cancellable);
+                        }
+
+                        @Override
+                        public void onSuccess(@Nullable final R result) {
+                            item.onSuccess(result);
+                            tryPollQueue();
+                        }
+
+                        @Override
+                        public void onError(final Throwable t) {
+                            item.onError(t);
+                            tryPollQueue();
+                        }
+
+                        private void tryPollQueue() {
+                            boolean tryAcquire = true;
+                            while (tryAcquire && tryAcquireLock(consumerLockUpdater, OrderedMapper.this)) {
+                                try {
+                                    Item<R> i;
+                                    while ((i = results.peek()) != null && i.tryTerminate()) {
+                                        results.poll();
+                                    }
+                                    // flatMapMergeSingle takes care of exception propagation / cleanup
+                                } finally {
+                                    tryAcquire = !releaseLock(consumerLockUpdater, OrderedMapper.this);
+                                }
+                            }
+                        }
+                    });
+                }
+            }
+            // The inner Single will determine if a copy is justified when we subscribe to it.
+            .shareContextOnSubscribe();
+        }
+    }
+
+    private static final class Item<R> {
+        @Nullable
+        SingleSource.Subscriber<? super R> subscriber;
+        @Nullable
+        private Object result;
+        // 0 = not terminated, 1 = success, 2 = error
+        private byte terminalState;
+
+        void onError(Throwable cause) {
+            terminalState = 2;
+            result = cause;
+        }
+
+        void onSuccess(@Nullable R r) {
+            terminalState = 1;
+            result = r;
+        }
+
+        @SuppressWarnings("unchecked")
+        boolean tryTerminate() {
+            assert subscriber != null; // if terminated, must have a subscriber
+            if (terminalState == 1) {
+                subscriber.onSuccess((R) result);
+                return true;
+            } else if (terminalState == 2) {
+                assert result != null;
+                subscriber.onError((Throwable) result);
+                return true;
+            }
+            return false;
+        }
+    }
+}

servicetalk-concurrent-api/src/test/java/io/servicetalk/concurrent/api/PublisherFlatMapConcatSingleTest.java

@@ -0,0 +1,185 @@
+/*
+ * Copyright © 2022 Apple Inc. and the ServiceTalk project authors
+ *
+ * Licensed 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
+ *
+ *   http://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 io.servicetalk.concurrent.api;
+
+import io.servicetalk.concurrent.test.internal.TestPublisherSubscriber;
+import io.servicetalk.context.api.ContextMap;
+
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import java.time.Duration;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.Function;
+
+import static io.servicetalk.concurrent.api.Executors.newCachedThreadExecutor;
+import static io.servicetalk.concurrent.api.Single.succeeded;
+import static io.servicetalk.concurrent.api.SourceAdapters.toSource;
+import static io.servicetalk.concurrent.internal.DeliberateException.DELIBERATE_EXCEPTION;
+import static java.time.Duration.ofMillis;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.contains;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.sameInstance;
+import static org.junit.jupiter.api.Assumptions.assumeFalse;
+
+class PublisherFlatMapConcatSingleTest {
+    private static final ContextMap.Key<String> K1 = ContextMap.Key.newKey("k1", String.class);
+    private static final ContextMap.Key<String> K2 = ContextMap.Key.newKey("k2", String.class);
+    private final TestPublisherSubscriber<String> subscriber = new TestPublisherSubscriber<>();
+    private final TestSubscription subscription = new TestSubscription();
+    private static Executor executor;
+
+    @BeforeAll
+    static void beforeClass() {
+        executor = newCachedThreadExecutor();
+    }
+
+    @AfterAll
+    static void afterClass() throws Exception {
+        executor.closeAsync().toFuture().get();
+    }
+
+    @ParameterizedTest(name = "{displayName} [{index}] begin={0} end={1} maxConcurrency={2} delayError={3}")
+    @CsvSource(value = {"0,5,10,false", "0,25,10,false", "0,5,10,true", "0,25,10,true"})
+    void orderPreservedInRangeWithConcurrency(int begin, int end, int maxConcurrency, boolean delayError) {
+        toSource(concatSingle(Publisher.range(begin, end),
+                i -> executor.timer(getDuration()).toSingle().map(__ -> i + "x"), maxConcurrency, delayError)
+        ).subscribe(subscriber);
+        String[] expected = expected(begin, end);
+        subscriber.awaitSubscription().request(expected.length);
+        assertThat(subscriber.takeOnNext(expected.length), contains(expected));
+        subscriber.awaitOnComplete();
+    }
+
+    @ParameterizedTest(name = "{displayName} [{index}] begin={0} end={1} maxConcurrency={2} delayError={3}")
+    @CsvSource(value = {"0,5,10,false", "0,25,10,false", "0,5,10,true", "0,25,10,true"})
+    void errorPropagatedInOrderLast(int begin, int end, int maxConcurrency, boolean delayError) {
+        final int endLessOne = end - 1;
+        String[] expected = expected(begin, endLessOne);
+        toSource(concatSingle(Publisher.range(begin, end), i -> executor.timer(getDuration()).toSingle().map(__ -> {
+                    if (i == endLessOne) {
+                        throw DELIBERATE_EXCEPTION;
+                    }
+                    return i + "x";
+                }), maxConcurrency, delayError)
+        ).subscribe(subscriber);
+        subscriber.awaitSubscription().request(end - begin);
+        assertThat(subscriber.takeOnNext(expected.length), contains(expected));
+        assertThat(subscriber.awaitOnError(), sameInstance(DELIBERATE_EXCEPTION));
+    }
+
+    @ParameterizedTest(name = "{displayName} [{index}] begin={0} end={1} maxConcurrency={2} delayError={3}")
+    @CsvSource(value = {"0,5,10,false", "0,25,10,false", "0,5,10,true", "0,25,10,true"})
+    void errorPropagatedInOrderFirst(int begin, int end, int maxConcurrency, boolean delayError) {
+        toSource(concatSingle(Publisher.range(begin, end), i -> executor.timer(getDuration()).toSingle().map(__ -> {
+                    if (i == begin) {
+                        throw DELIBERATE_EXCEPTION;
+                    }
+                    return i + "x";
+                }), maxConcurrency, delayError)
+        ).subscribe(subscriber);
+        subscriber.awaitSubscription().request(end - begin);
+        if (delayError) {
+            String[] expected = expected(begin + 1, end);
+            assertThat(subscriber.takeOnNext(expected.length), contains(expected));
+        }
+        assertThat(subscriber.awaitOnError(), sameInstance(DELIBERATE_EXCEPTION));
+    }
+
+    @ParameterizedTest(name = "{displayName} [{index}] delayError={0}")
+    @ValueSource(booleans = {true, false})
+    void cancellationPropagated(boolean delayError) throws InterruptedException {
+        TestPublisher<Integer> source = new TestPublisher.Builder<Integer>().disableAutoOnSubscribe().build(
+                subscriber1 -> {
+            subscriber1.onSubscribe(subscription);
+            return subscriber1;
+        });
+        final TestCancellable cancellable = new TestCancellable();
+        final TestSingle<String> singleSource = new TestSingle.Builder<String>().disableAutoOnSubscribe()
+                .build(subscriber1 -> {
+            subscriber1.onSubscribe(cancellable);
+            return subscriber1;
+        });
+        toSource(concatSingle(source, i -> singleSource, 2, delayError))
+                .subscribe(subscriber);
+        subscriber.awaitSubscription().request(1);
+        subscription.awaitRequestN(1);
+        source.onNext(0);
+        singleSource.awaitSubscribed();
+        subscriber.awaitSubscription().cancel();
+        cancellable.awaitCancelled();
+        subscription.awaitCancelled();
+    }
+
+    @ParameterizedTest(name = "{displayName} [{index}] delayError={0}")
+    @ValueSource(booleans = {true, false})
+    void singleTerminalThrows(boolean delayError) {
+        toSource(concatSingle(Publisher.range(0, 2), i -> succeeded(i + "x"), 2, delayError)
+                .<String>map(x -> {
+                    throw DELIBERATE_EXCEPTION;
+                })
+        ).subscribe(subscriber);
+        subscriber.awaitSubscription().request(1);
+        assertThat(subscriber.awaitOnError(), sameInstance(DELIBERATE_EXCEPTION));
+    }
+
+    @ParameterizedTest(name = "{displayName} [{index}] delayError={0}")
+    @ValueSource(booleans = {true, false})
+    void asyncContext(boolean delayError) {
+        assumeFalse(AsyncContext.isDisabled());
+        final int begin = 0;
+        final int end = 2;
+        final int elements = end - begin;
+        AsyncContext.put(K1, "v1");
+        toSource(concatSingle(Publisher.range(begin, end), i -> {
+                    AsyncContext.put(K2, "v2");
+                    return succeeded(i + "x");
+                }, elements, delayError)
+                .map(x -> {
+                    assertThat(AsyncContext.get(K1), equalTo("v1"));
+                    assertThat(AsyncContext.get(K2), equalTo("v2"));
+                    return x;
+                })
+        ).subscribe(subscriber);
+        subscriber.awaitSubscription().request(elements);
+        assertThat(subscriber.takeOnNext(elements), contains(expected(begin, end)));
+        subscriber.awaitOnComplete();
+    }
+
+    private static Duration getDuration() {
+        // Introduce randomness to increase the likelihood of out of order task completion.
+        return ofMillis(ThreadLocalRandom.current().nextInt(5));
+    }
+
+    private static <T, R> Publisher<R> concatSingle(Publisher<T> publisher, Function<T, Single<R>> mapper,
+                                                    int maxConcurrency, boolean delayError) {
+        return delayError ?
+                publisher.flatMapConcatSingleDelayError(mapper, maxConcurrency) :
+                publisher.flatMapConcatSingle(mapper, maxConcurrency);
+    }
+
+    private static String[] expected(int begin, int end) {
+        String[] expected = new String[end - begin];
+        for (int i = begin; i < end; ++i) {
+            expected[i - begin] = i + "x";
+        }
+        return expected;
+    }
+}

servicetalk-concurrent-reactivestreams/src/test/java/io/servicetalk/concurrent/reactivestreams/tck/PublisherFlatMapConcatSingleDelayErrorTckTest.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright © 2022 Apple Inc. and the ServiceTalk project authors
+ *
+ * Licensed 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
+ *
+ *   http://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 io.servicetalk.concurrent.reactivestreams.tck;
+
+import io.servicetalk.concurrent.api.Publisher;
+import io.servicetalk.concurrent.api.Single;
+
+import org.testng.annotations.Test;
+
+@Test
+public class PublisherFlatMapConcatSingleDelayErrorTckTest extends AbstractPublisherOperatorTckTest<Integer> {
+    @Override
+    protected Publisher<Integer> composePublisher(Publisher<Integer> publisher, int elements) {
+        return publisher.flatMapConcatSingleDelayError(Single::succeeded);
+    }
+}

servicetalk-concurrent-reactivestreams/src/test/java/io/servicetalk/concurrent/reactivestreams/tck/PublisherFlatMapSingleTckTest.java → servicetalk-concurrent-reactivestreams/src/test/java/io/servicetalk/concurrent/reactivestreams/tck/PublisherFlatMapConcatSingleTckTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright © 2018 Apple Inc. and the ServiceTalk project authors
+ * Copyright © 2022 Apple Inc. and the ServiceTalk project authors
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -21,9 +21,9 @@
 import org.testng.annotations.Test;
 
 @Test
-public class PublisherFlatMapSingleTckTest extends AbstractPublisherOperatorTckTest<Integer> {
+public class PublisherFlatMapConcatSingleTckTest extends AbstractPublisherOperatorTckTest<Integer> {
     @Override
     protected Publisher<Integer> composePublisher(Publisher<Integer> publisher, int elements) {
-        return publisher.flatMapMergeSingle(Single::succeeded, 10);
+        return publisher.flatMapConcatSingle(Single::succeeded);
     }
 }