Skip to content

Commit cfd75c9

Browse files
He-PinHe-Pin
committed
+str Add flatmapConcat with parallelism.
1 parent e778606 commit cfd75c9

File tree

12 files changed

+677
-21
lines changed

12 files changed

+677
-21
lines changed

‎akka-bench-jmh/src/main/scala/akka/stream/FlatMapConcatBenchmark.scala‎

Lines changed: 61 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import java.util.concurrent.TimeUnit
99

1010
import scala.concurrent.Await
1111
import scala.concurrent.duration._
12+
import scala.concurrent.Future
1213

1314
import com.typesafe.config.ConfigFactory
1415
import org.openjdk.jmh.annotations._
@@ -60,31 +61,87 @@ class FlatMapConcatBenchmark {
6061
@OperationsPerInvocation(OperationsPerInvocation)
6162
def sourceDotSingle(): Unit = {
6263
val latch = new CountDownLatch(1)
63-
6464
testSource.flatMapConcat(Source.single).runWith(new LatchSink(OperationsPerInvocation, latch))
65+
awaitLatch(latch)
66+
}
6567

68+
@Benchmark
69+
@OperationsPerInvocation(OperationsPerInvocation)
70+
def sourceDotSingleP1(): Unit = {
71+
val latch = new CountDownLatch(1)
72+
testSource.flatMapConcat(1, Source.single).runWith(new LatchSink(OperationsPerInvocation, latch))
6673
awaitLatch(latch)
6774
}
6875

6976
@Benchmark
7077
@OperationsPerInvocation(OperationsPerInvocation)
71-
def internalSingleSource(): Unit = {
78+
def sourceDotSingleP2(): Unit = {
7279
val latch = new CountDownLatch(1)
80+
testSource.flatMapConcat(2, Source.single).runWith(new LatchSink(OperationsPerInvocation, latch))
81+
awaitLatch(latch)
82+
}
7383

84+
@Benchmark
85+
@OperationsPerInvocation(OperationsPerInvocation)
86+
def internalSingleSource(): Unit = {
87+
val latch = new CountDownLatch(1)
7488
testSource
7589
.flatMapConcat(elem => new GraphStages.SingleSource(elem))
7690
.runWith(new LatchSink(OperationsPerInvocation, latch))
77-
7891
awaitLatch(latch)
7992
}
8093

8194
@Benchmark
8295
@OperationsPerInvocation(OperationsPerInvocation)
8396
def oneElementList(): Unit = {
8497
val latch = new CountDownLatch(1)
85-
8698
testSource.flatMapConcat(n => Source(n :: Nil)).runWith(new LatchSink(OperationsPerInvocation, latch))
99+
awaitLatch(latch)
100+
}
101+
102+
@Benchmark
103+
@OperationsPerInvocation(OperationsPerInvocation)
104+
def oneElementListP1(): Unit = {
105+
val latch = new CountDownLatch(1)
106+
testSource.flatMapConcat(1, n => Source(n :: Nil)).runWith(new LatchSink(OperationsPerInvocation, latch))
107+
awaitLatch(latch)
108+
}
109+
110+
@Benchmark
111+
@OperationsPerInvocation(OperationsPerInvocation)
112+
def oneElementListP2(): Unit = {
113+
val latch = new CountDownLatch(1)
114+
testSource.flatMapConcat(2, n => Source(n :: Nil)).runWith(new LatchSink(OperationsPerInvocation, latch))
115+
awaitLatch(latch)
116+
}
87117

118+
@Benchmark
119+
@OperationsPerInvocation(OperationsPerInvocation)
120+
def completedFuture(): Unit = {
121+
val latch = new CountDownLatch(1)
122+
testSource
123+
.flatMapConcat(n => Source.future(Future.successful(n)))
124+
.runWith(new LatchSink(OperationsPerInvocation, latch))
125+
awaitLatch(latch)
126+
}
127+
128+
@Benchmark
129+
@OperationsPerInvocation(OperationsPerInvocation)
130+
def completedFutureP1(): Unit = {
131+
val latch = new CountDownLatch(1)
132+
testSource
133+
.flatMapConcat(1, n => Source.future(Future.successful(n)))
134+
.runWith(new LatchSink(OperationsPerInvocation, latch))
135+
awaitLatch(latch)
136+
}
137+
138+
@Benchmark
139+
@OperationsPerInvocation(OperationsPerInvocation)
140+
def completedFutureP2(): Unit = {
141+
val latch = new CountDownLatch(1)
142+
testSource
143+
.flatMapConcat(2, n => Source.future(Future.successful(n)))
144+
.runWith(new LatchSink(OperationsPerInvocation, latch))
88145
awaitLatch(latch)
89146
}
90147

‎akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowFlatMapConcatParallelismSpec.scala‎

Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
/*
2+
* Copyright (C) 2014-2024 Lightbend Inc. <https://www.lightbend.com>
3+
*/
4+
5+
package akka.stream.scaladsl
6+
7+
import akka.NotUsed
8+
import akka.pattern.FutureTimeoutSupport
9+
import akka.stream.OverflowStrategy
10+
import akka.stream.testkit._
11+
import akka.stream.testkit.scaladsl.TestSink
12+
13+
import java.util.concurrent.ThreadLocalRandom
14+
import java.util.concurrent.atomic.AtomicInteger
15+
import scala.annotation.switch
16+
import scala.concurrent.{ ExecutionContext, Future }
17+
import scala.concurrent.duration.DurationInt
18+
import scala.util.control.NoStackTrace
19+
20+
class FlowFlatMapConcatParallelismSpec extends StreamSpec("""
21+
akka.stream.materializer.initial-input-buffer-size = 2
22+
""") with ScriptedTest with FutureTimeoutSupport {
23+
val toSeq = Flow[Int].grouped(1000).toMat(Sink.head)(Keep.right)
24+
25+
class BoomException extends RuntimeException("BOOM~~") with NoStackTrace
26+
"A flatMapConcat" must {
27+
28+
for (i <- 1 until 129) {
29+
s"work with value presented sources with parallelism: $i" in {
30+
Source(
31+
List(
32+
Source.empty[Int],
33+
Source.single(1),
34+
Source.empty[Int],
35+
Source(List(2, 3, 4)),
36+
Source.future(Future.successful(5)),
37+
Source.lazyFuture(() => Future.successful(6)),
38+
Source.future(after(1.millis)(Future.successful(7)))))
39+
.flatMapConcat(i, identity)
40+
.runWith(toSeq)
41+
.futureValue should ===(1 to 7)
42+
}
43+
}
44+
45+
def generateRandomValuePresentedSources(nums: Int): (Int, Seq[Source[Int, NotUsed]]) = {
46+
val seq = Seq.tabulate(nums) { _ =>
47+
val random = ThreadLocalRandom.current().nextInt(1, 10)
48+
(random: @switch) match {
49+
case 1 => Source.single(1)
50+
case 2 => Source(List(1))
51+
case 3 => Source.fromJavaStream(() => java.util.stream.Stream.of(1))
52+
case 4 => Source.future(Future.successful(1))
53+
case 5 => Source.future(after(1.millis)(Future.successful(1)))
54+
case _ => Source.empty[Int]
55+
}
56+
}
57+
val sum = seq.filterNot(_.eq(Source.empty[Int])).size
58+
(sum, seq)
59+
}
60+
61+
for (i <- 1 until 129) {
62+
s"work with generated value presented sources with parallelism: $i " in {
63+
val (sum, sources) = generateRandomValuePresentedSources(10000)
64+
Source(sources)
65+
.flatMapConcat(i, identity)
66+
.runWith(Sink.seq)
67+
.map(_.sum)(ExecutionContext.parasitic)
68+
.futureValue shouldBe sum
69+
}
70+
}
71+
72+
"work with value presented failed sources" in {
73+
val ex = new BoomException
74+
Source(
75+
List(
76+
Source.empty[Int],
77+
Source.single(1),
78+
Source.empty[Int],
79+
Source(List(2, 3, 4)),
80+
Source.future(Future.failed(ex)),
81+
Source.lazyFuture(() => Future.successful(5))))
82+
.flatMapConcat(ThreadLocalRandom.current().nextInt(1, 129), identity)
83+
.onErrorComplete[BoomException]()
84+
.runWith(toSeq)
85+
.futureValue should ===(1 to 4)
86+
}
87+
88+
"work with value presented sources when demands slow" in {
89+
val prob = Source(
90+
List(Source.empty[Int], Source.single(1), Source(List(2, 3, 4)), Source.lazyFuture(() => Future.successful(5))))
91+
.flatMapConcat(ThreadLocalRandom.current().nextInt(1, 129), identity)
92+
.runWith(TestSink())
93+
94+
prob.request(1)
95+
prob.expectNext(1)
96+
prob.expectNoMessage(1.seconds)
97+
prob.request(2)
98+
prob.expectNext(2, 3)
99+
prob.expectNoMessage(1.seconds)
100+
prob.request(2)
101+
prob.expectNext(4, 5)
102+
prob.expectComplete()
103+
}
104+
105+
"can do pre materialization when parallelism > 1" in {
106+
val materializationCounter = new AtomicInteger(0)
107+
val randomParallelism = ThreadLocalRandom.current().nextInt(4, 65)
108+
val prob = Source(1 to (randomParallelism * 3))
109+
.flatMapConcat(
110+
randomParallelism,
111+
value => {
112+
Source
113+
.lazySingle(() => {
114+
materializationCounter.incrementAndGet()
115+
value
116+
})
117+
.buffer(1, overflowStrategy = OverflowStrategy.backpressure)
118+
})
119+
.runWith(TestSink())
120+
121+
expectNoMessage(1.seconds)
122+
materializationCounter.get() shouldBe 0
123+
124+
prob.request(1)
125+
prob.expectNext(1.seconds, 1)
126+
expectNoMessage(1.seconds)
127+
materializationCounter.get() shouldBe (randomParallelism + 1)
128+
materializationCounter.set(0)
129+
130+
prob.request(2)
131+
prob.expectNextN(List(2, 3))
132+
expectNoMessage(1.seconds)
133+
materializationCounter.get() shouldBe 2
134+
materializationCounter.set(0)
135+
136+
prob.request(randomParallelism - 3)
137+
prob.expectNextN(4 to randomParallelism)
138+
expectNoMessage(1.seconds)
139+
materializationCounter.get() shouldBe (randomParallelism - 3)
140+
materializationCounter.set(0)
141+
142+
prob.request(randomParallelism)
143+
prob.expectNextN(randomParallelism + 1 to randomParallelism * 2)
144+
expectNoMessage(1.seconds)
145+
materializationCounter.get() shouldBe randomParallelism
146+
materializationCounter.set(0)
147+
148+
prob.request(randomParallelism)
149+
prob.expectNextN(randomParallelism * 2 + 1 to randomParallelism * 3)
150+
expectNoMessage(1.seconds)
151+
materializationCounter.get() shouldBe 0
152+
prob.expectComplete()
153+
}
154+
155+
}
156+
157+
}

‎akka-stream/src/main/scala/akka/stream/impl/JavaStreamSource.scala‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import java.util.function.Consumer
1212

1313
/** INTERNAL API */
1414
@InternalApi private[stream] final class JavaStreamSource[T, S <: java.util.stream.BaseStream[T, S]](
15-
open: () => java.util.stream.BaseStream[T, S])
15+
val open: () => java.util.stream.BaseStream[T, S])
1616
extends GraphStage[SourceShape[T]] {
1717

1818
val out: Outlet[T] = Outlet("JavaStreamSource")

‎akka-stream/src/main/scala/akka/stream/impl/Stages.scala‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ import akka.stream.Attributes._
7979
val mergePreferred = name("mergePreferred")
8080
val mergePrioritized = name("mergePrioritized")
8181
val flattenMerge = name("flattenMerge")
82+
val flattenConcat = name("flattenConcat")
8283
val recoverWith = name("recoverWith")
8384
val onErrorComplete = name("onErrorComplete")
8485
val broadcast = name("broadcast")

‎akka-stream/src/main/scala/akka/stream/impl/TraversalBuilder.scala‎

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,12 @@ package akka.stream.impl
66

77
import scala.collection.immutable.Map.Map1
88
import scala.language.existentials
9-
109
import akka.annotation.{ DoNotInherit, InternalApi }
1110
import akka.stream._
1211
import akka.stream.impl.StreamLayout.AtomicModule
1312
import akka.stream.impl.TraversalBuilder.{ AnyFunction1, AnyFunction2 }
1413
import akka.stream.impl.fusing.GraphStageModule
15-
import akka.stream.impl.fusing.GraphStages.SingleSource
14+
import akka.stream.impl.fusing.GraphStages.{ FutureSource, IterableSource, SingleSource }
1615
import akka.stream.scaladsl.Keep
1716
import akka.util.OptionVal
1817

@@ -369,6 +368,42 @@ import akka.util.OptionVal
369368
}
370369
}
371370

371+
/**
372+
* Try to find `SingleSource` or wrapped such. This is used as a
373+
* performance optimization in FlattenConcat and possibly other places.
374+
*/
375+
def getValuePresentedSource[A >: Null](graph: Graph[SourceShape[A], _]): OptionVal[Graph[SourceShape[A], _]] = {
376+
def isValuePresentedSource(graph: Graph[SourceShape[_ <: A], _]): Boolean = graph match {
377+
case _: SingleSource[_] | _: IterableSource[_] | _: JavaStreamSource[_, _] | EmptySource | _: FutureSource[_] =>
378+
true
379+
case _ => false
380+
}
381+
graph match {
382+
case _ if isValuePresentedSource(graph) => OptionVal.Some(graph)
383+
case _ =>
384+
graph.traversalBuilder match {
385+
case l: LinearTraversalBuilder =>
386+
l.pendingBuilder match {
387+
case OptionVal.Some(a: AtomicTraversalBuilder) =>
388+
a.module match {
389+
case m: GraphStageModule[_, _] =>
390+
m.stage match {
391+
case _ if isValuePresentedSource(m.stage.asInstanceOf[Graph[SourceShape[A], _]]) =>
392+
// It would be != EmptyTraversal if mapMaterializedValue was used and then we can't optimize.
393+
if ((l.traversalSoFar eq EmptyTraversal) && !l.attributes.isAsync)
394+
OptionVal.Some(m.stage.asInstanceOf[Graph[SourceShape[A], _]])
395+
else OptionVal.None
396+
case _ => OptionVal.None
397+
}
398+
case _ => OptionVal.None
399+
}
400+
case _ => OptionVal.None
401+
}
402+
case _ => OptionVal.None
403+
}
404+
}
405+
}
406+
372407
/**
373408
* Test if a Graph is an empty Source.
374409
* */

‎akka-stream/src/main/scala/akka/stream/impl/fusing/Ops.scala‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1268,6 +1268,7 @@ private[stream] object Collect {
12681268
*/
12691269
@InternalApi private[akka] final case class MapAsync[In, Out](parallelism: Int, f: In => Future[Out])
12701270
extends GraphStage[FlowShape[In, Out]] {
1271+
require(parallelism >= 1, "parallelism should >= 1")
12711272

12721273
import MapAsync._
12731274

0 commit comments

Comments
 (0)