Add flatMap method for collapsing streams into one #4227

sleeplesslord · 2022-04-25T23:28:19Z

Motivation:
StreamMessage has map, but it is inconvenient to use it when the mapping function itself returns a StreamMessage. This PR adds flatMap which allows applying a function that returns a StreamMessage without ending up with nested StreamMessages.

Modifications:

Add FlatMapStreamMessage, which uses a function that returns a StreamMessage to modify the stream.
Add StreamMessage.flatMap() to allow easily creating a FlatMapStreamMessage from an existing stream.

Result:

Closes Add StreamMessage.flatMap() #3966
You can now use StreamMessage.flatMap to modify a stream using a function that returns a StreamMessage.

codecov · 2022-04-26T00:35:16Z

Codecov Report

Attention: Patch coverage is 76.92308% with 48 lines in your changes are missing coverage. Please review.

Project coverage is 74.09%. Comparing base (8194436) to head (b7ac1c4).
Report is 45 commits behind head on main.

❗ Current head b7ac1c4 differs from pull request most recent head 3b35003. Consider uploading reports for the commit 3b35003 to get more accurate results

Files	Patch %	Lines
...rp/armeria/common/stream/FlatMapStreamMessage.java	76.84%	30 Missing and 17 partials ⚠️
.../linecorp/armeria/common/stream/StreamMessage.java	80.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #4227       +/-   ##
===========================================
+ Coverage        0   74.09%   +74.09%     
- Complexity      0    20908    +20908     
===========================================
  Files           0     1810     +1810     
  Lines           0    77148    +77148     
  Branches        0     9867     +9867     
===========================================
+ Hits            0    57161    +57161     
- Misses          0    15328    +15328     
- Partials        0     4659     +4659

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

ikhoon · 2022-05-04T06:19:17Z

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

+
+    @Override
+    public CompletableFuture<Void> whenComplete() {
+        return source.whenComplete();


Should we complete this stream when all nested stream are completed?

ikhoon · 2022-05-04T06:21:15Z

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

+
+    @Override
+    public long demand() {
+        return source.demand();


I don't see we can provide the current demand with source.demand(). How about tracking the demand manually rather than delegating it?

ikhoon · 2022-05-04T06:24:19Z

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

+
+    @Override
+    public void abort() {
+        source.abort();


Should we also abort the nested StreamMessages in progress?

That sounds like a good idea. If I call source.abort() here, will that propagate to the Subscription so that I can just cancel my nested subscriptions in onError?

ikhoon · 2022-05-04T06:31:29Z

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

+
+        private void handleRequest(long n) {
+            requestedByDownstream = LongMath.saturatedAdd(requestedByDownstream, n);
+            upstream.request(maxConcurrency - sourceSubscriptions.size());


Are there any cases where sourceSubscriptions.size() is greater than maxConcurrency?

It should never happen

Let's add:

if (maxConcurrency - sourceSubscriptions.size() != 0) {...}

so that we don't request with 0.

Shouldn't we call flush before request? so that we don't call upstream.request(...) if we have enough elements to hand on to the downstream?

ikhoon · 2022-05-04T06:38:51Z

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

+                               .limit(available).forEach(sub -> sub.request(1));
+        }
+
+        private void dump() {


How about flush()? I found it not easy to understand this behavior with dump().

ikhoon · 2022-05-04T06:46:36Z

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

+            this.maxConcurrency = maxConcurrency;
+
+            sourceSubscriptions = new HashSet<>();
+            buffer = new LinkedList<>();


How about using ArrayDeque for less memory footprints?

ikhoon · 2022-05-04T06:48:20Z

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

+            }
+
+            final StreamMessage<U> newStreamMessage = function.apply(item);
+            newStreamMessage.subscribe(new FlatMapSubscriber<>(this));


Suggested change

newStreamMessage.subscribe(new FlatMapSubscriber<>(this));

newStreamMessage.subscribe(new FlatMapSubscriber<>(this), executor, options);

Is it OK to just pass executor here? I don't think there are any options that need to be applied

@ikhoon Gentle ping

I tried passing the executor here and I'm running into an issue I'm not sure how to address.

This code calls subscribe, which calls onSubscribe on FlatMapSubscriber. That will call subscribeChild on FlatMapAggregatingSubscriber, which will add the subscriber to sourceSubscriptions and then call request on FlatMapSubscriber. Then, if newStreamMessage is already completed it will call onComplete, which leads to the subscription being removed from sourceSubscriptions.

Essentially, in one call chain we go from onSubscribe to sub.request(1) all the way to onComplete, but since that removes the subscription from sourceSubscriptions I get a ConcurrentModificationException since the call to sub.request(1) is done in a forEach.

I guess I should just collect the subs into a new collection before calling forEach to avoid that? I'm not sure if there's a better approach

I see. How about always using executor to remove the child?
executor.execute(() -> handleCompleteChild(child));

If that works, we can remove all eventLoop.inEventLoop() check because we only use the executor.

ikhoon · 2022-05-04T06:49:25Z

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

+            }
+            canceled = true;
+
+            downstream.onError(cause);


Ditto: Cancel all inner StreamMessages?

ikhoon · 2022-05-04T06:59:38Z

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

+            final Optional<Long> requested = sourceSubscriptions.stream().map(FlatMapSubscriber::getRequested)
+                                                                .reduce(LongMath::saturatedAdd);


Suggested change

final Optional<Long> requested = sourceSubscriptions.stream().map(FlatMapSubscriber::getRequested)

.reduce(LongMath::saturatedAdd);

final long requested = sourceSubscriptions.stream().mapToLong(FlatMapSubscriber::getRequested).sum();

Will that handle overflowing in case getRequested returns Long.MAX_VALUE?

ikhoon · 2022-05-04T07:01:38Z

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

+                return;
+            }
+
+            final long available = getAvailableBufferSpace();


If available is Long.MAX_VALUE, we don't need backpressure with buffering.

trustin · 2022-06-08T13:23:19Z

Gentle ping, @sleeplesslord 🙇

sleeplesslord · 2022-11-26T13:33:50Z

Hey, sorry for letting this PR go stale for so long, I was addressing the comments but then I lost track of things 😅 I think I've covered most of the feedback now

minwoox · 2022-12-21T08:21:45Z

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

+     * As per
+     * <a href="https://github.com/reactive-streams/reactive-streams-jvm#2.13">
+     * Reactive Streams Specification 2.13</a>, the specified {@link Function} should not return
+     * a {@code null} value


Suggested change

* a {@code null} value

* a {@code null} value.

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

minwoox · 2022-12-21T08:29:35Z

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

+        this.source = (StreamMessage<T>) source;
+        this.function = (Function<T, StreamMessage<U>>) function;
+        this.maxConcurrency = maxConcurrency;
+        completionFuture = new CompletableFuture<>();


nit: We usually use EventLoopCheckingFuture for this.

minwoox · 2022-12-22T06:36:50Z

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

+            }
+
+            final StreamMessage<U> newStreamMessage = function.apply(item);
+            newStreamMessage.subscribe(new FlatMapSubscriber<>(this));


I see. How about always using executor to remove the child?
executor.execute(() -> handleCompleteChild(child));

minwoox · 2022-12-24T10:05:03Z

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

+                                                                               .filter(sub -> sub.getRequested()
+                                                                                              == 0)
+                                                                               .limit(available)
+                                                                               .collect(Collectors.toList());


nit:

Suggested change

.collect(Collectors.toList());

.collect(toImmutableList());

Yeah, creating a new list also works. 😄
So, Could you remove if (executor.inEventLoop()) {..} else {...} from this class where possible? 😄
e.g. completeChild, onNextChild?

Is that because the FlatMapSubscriber will run with the same executor, so we know that completeChild etc will always run in the same executor?

Yes exactly. 😄
All onX methods in FlatMapAggregatingSubscriber and FlatMapSubscriber are now called by the executor so we don't have to check it. 😄

minwoox

Thanks! Left some more comments. 😉

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

minwoox · 2022-12-26T07:33:11Z

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

+
+        private void handleRequest(long n) {
+            requestedByDownstream = LongMath.saturatedAdd(requestedByDownstream, n);
+            upstream.request(maxConcurrency - sourceSubscriptions.size());


Let's add:

if (maxConcurrency - sourceSubscriptions.size() != 0) {...}

so that we don't request with 0.

minwoox · 2022-12-26T07:44:01Z

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

+
+        private void handleRequest(long n) {
+            requestedByDownstream = LongMath.saturatedAdd(requestedByDownstream, n);
+            upstream.request(maxConcurrency - sourceSubscriptions.size());


Shouldn't we call flush before request? so that we don't call upstream.request(...) if we have enough elements to hand on to the downstream?

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

jrhee17

Almost there 🙏

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

github-actions · 2024-03-21T07:18:01Z

🔍 Build Scan® (commit: `3b35003`)

Job name	Status	Build Scan®
build-windows-latest-jdk-19	✅	https://ge.armeria.dev/s/id5t4amg23how
build-self-hosted-unsafe-jdk-8	✅	https://ge.armeria.dev/s/k26t662pm7evg
build-self-hosted-unsafe-jdk-19-snapshot-blockhound	✅	https://ge.armeria.dev/s/ke77vevpynjoe
build-self-hosted-unsafe-jdk-17-min-java-17-coverage	✅	https://ge.armeria.dev/s/776rfk2oqpqd4
build-self-hosted-unsafe-jdk-17-min-java-11	✅	https://ge.armeria.dev/s/pfie27yozxykw
build-self-hosted-unsafe-jdk-17-leak	✅	https://ge.armeria.dev/s/ycdztikwwjyhm
build-self-hosted-unsafe-jdk-11	✅	https://ge.armeria.dev/s/tgbsf3duysqeq
build-macos-12-jdk-19	❌ (failure)	https://ge.armeria.dev/s/il6hokuioxqje

jrhee17

I've added a small commit which:

Addresses a ConcurrentModificationException thrown when the forEach call inside #requestAllAvailable calls childSubscribers#remove in #completeChild. childSubscribers is copied to mitigate this issue.
FlatMapStreamMessage could be closed, whereas the inner FlatMapSubscriber may not be closed. This is possible if abort is called before FlatMapSubscriber#onSubscribe is invoked. I've changed so that FlatMapSubscriber is added to childSubscribers immediately.
Miscellaneous possible NPEs were dealt with
Flaky tests were dealt with

Let me know if any changes don't make sense.

Lastly, sorry this PR took so long 😅 Looks good to me for merging 👍

minwoox · 2024-04-01T09:47:51Z

Thanks! I also put some miscellaneous changes.

minwoox

Thanks a lot, @sleeplesslord and @jrhee17!

ikhoon

Looks great! 🚀 💯

sleeplesslord · 2024-04-09T11:52:17Z

Thanks for the fixes and the merge, happy to see this make it in 🙏

Add flatMap method for collapsing streams into one

d95055c

sleeplesslord requested review from ikhoon, jrhee17, minwoox and trustin as code owners April 25, 2022 23:28

ikhoon added the new feature label Apr 26, 2022

ikhoon reviewed May 4, 2022

View reviewed changes

sleeplesslord added 4 commits May 5, 2022 09:48

Manually handle completion tracking for flatMap

d3fba09

Cancel inner subscriptions when FlatMap is cancelled

49fbc2e

Get demand from inner subscription

761da4c

Rename method for clarity

74bf0b7

sleeplesslord added 3 commits November 26, 2022 20:48

Use ArrayDeque for smaller memory footprint

80141c3

Request Long.MAX_VALUE from source when downstream requests it

0082442

Improve testing and add associativity test

3219a78

minwoox added this to the 1.22.0 milestone Dec 19, 2022

minwoox reviewed Dec 22, 2022

View reviewed changes

sleeplesslord added 4 commits December 24, 2022 10:02

Improve documentation for flatMap

589296b

Add check for maxConcurrency > 0 in flatMap

8d0df8e

Use EventLoopCheckingFuture

a9f6ab4

Use same executor for FlatMapSubscribers

9a6d5fc

minwoox reviewed Dec 24, 2022

View reviewed changes

sleeplesslord added 2 commits December 25, 2022 10:46

Remove unnecessary executor checks

9bfcb8a

Use toImmutableList

87ec0fb

minwoox reviewed Dec 26, 2022

View reviewed changes

sleeplesslord added 3 commits December 29, 2022 10:07

Pass SubscriptionOptions into FlatMapSubscriber

dc7c776

Cancel properly when getting invalid request

400d092

Prevent requesting 0 from upstream

84c4556

Fix lint

ca99d0c

jrhee17 reviewed Mar 8, 2023

View reviewed changes

ikhoon modified the milestones: 1.23.0, 1.24.0 Mar 15, 2023

sleeplesslord added 5 commits April 2, 2023 10:31

Fix implementations of isOpen and isEmpty

a69185c

Remove unnecessary subscription cancel

bc960cd

Handle future and downstream notification on cancel

54a77e2

Rename variable to a more fitting name

1cc7db7

Move check into event loop to remove volatile

80cebda

ikhoon modified the milestones: 1.24.0, 1.25.0 Jun 2, 2023

ikhoon modified the milestones: 1.25.0, 1.26.0 Jul 31, 2023

minwoox modified the milestones: 1.26.0, 1.27.0 Oct 11, 2023

ikhoon modified the milestones: 1.27.0, 1.28.0 Jan 5, 2024

Merge branch 'main' into flatmap

8b2b98e

handle leaks and concurrentmodificationexception

eeb2085

jrhee17 approved these changes Mar 25, 2024

View reviewed changes

minwoox added 2 commits April 1, 2024 18:40

Return early and decrease the demand before calling onNext

e7b3d7f

Call completionFuture.completeExceptionally after subscriber.onError

b7ac1c4

minwoox approved these changes Apr 1, 2024

View reviewed changes

minor clean up

3b35003

ikhoon approved these changes Apr 9, 2024

View reviewed changes

ikhoon merged commit 211a9e5 into line:main Apr 9, 2024
15 of 16 checks passed

sleeplesslord deleted the flatmap branch April 9, 2024 11:52

	newStreamMessage.subscribe(new FlatMapSubscriber<>(this));
	newStreamMessage.subscribe(new FlatMapSubscriber<>(this), executor, options);

		final Optional<Long> requested = sourceSubscriptions.stream().map(FlatMapSubscriber::getRequested)
		.reduce(LongMath::saturatedAdd);

	final Optional<Long> requested = sourceSubscriptions.stream().map(FlatMapSubscriber::getRequested)
	.reduce(LongMath::saturatedAdd);
	final long requested = sourceSubscriptions.stream().mapToLong(FlatMapSubscriber::getRequested).sum();

Add flatMap method for collapsing streams into one #4227

Add flatMap method for collapsing streams into one #4227

Conversation

sleeplesslord commented Apr 25, 2022

codecov bot commented Apr 26, 2022 • edited

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

trustin commented Jun 8, 2022

sleeplesslord commented Nov 26, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

minwoox Dec 24, 2022 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

minwoox left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jrhee17 left a comment

Choose a reason for hiding this comment

github-actions bot commented Mar 21, 2024 • edited

🔍 Build Scan® (commit: 3b35003)

jrhee17 left a comment

Choose a reason for hiding this comment

minwoox commented Apr 1, 2024

minwoox left a comment

Choose a reason for hiding this comment

ikhoon left a comment

Choose a reason for hiding this comment

sleeplesslord commented Apr 9, 2024

codecov bot commented Apr 26, 2022 •

edited

minwoox Dec 24, 2022 •

edited

github-actions bot commented Mar 21, 2024 •

edited

🔍 Build Scan® (commit: `3b35003`)