Skip to content

Commit 05d5b17

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

File tree

12 files changed

+680
-23
lines changed

12 files changed

+680
-23
lines changed

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

Lines changed: 64 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,8 @@ package akka.stream
66

77
import java.util.concurrent.CountDownLatch
88
import java.util.concurrent.TimeUnit
9-
10-
import scala.concurrent.Await
9+
import scala.concurrent.{ Await, Future }
1110
import scala.concurrent.duration._
12-
1311
import com.typesafe.config.ConfigFactory
1412
import org.openjdk.jmh.annotations._
1513

@@ -60,31 +58,91 @@ class FlatMapConcatBenchmark {
6058
@OperationsPerInvocation(OperationsPerInvocation)
6159
def sourceDotSingle(): Unit = {
6260
val latch = new CountDownLatch(1)
63-
6461
testSource.flatMapConcat(Source.single).runWith(new LatchSink(OperationsPerInvocation, latch))
62+
awaitLatch(latch)
63+
}
6564

65+
@Benchmark
66+
@OperationsPerInvocation(OperationsPerInvocation)
67+
def sourceDotSingleP1(): Unit = {
68+
val latch = new CountDownLatch(1)
69+
testSource.flatMapConcat(1, Source.single).runWith(new LatchSink(OperationsPerInvocation, latch))
6670
awaitLatch(latch)
6771
}
6872

6973
@Benchmark
7074
@OperationsPerInvocation(OperationsPerInvocation)
7175
def internalSingleSource(): Unit = {
7276
val latch = new CountDownLatch(1)
73-
7477
testSource
7578
.flatMapConcat(elem => new GraphStages.SingleSource(elem))
7679
.runWith(new LatchSink(OperationsPerInvocation, latch))
80+
awaitLatch(latch)
81+
}
7782

83+
@Benchmark
84+
@OperationsPerInvocation(OperationsPerInvocation)
85+
def internalSingleSourceP1(): Unit = {
86+
val latch = new CountDownLatch(1)
87+
testSource
88+
.flatMapConcat(1, elem => new GraphStages.SingleSource(elem))
89+
.runWith(new LatchSink(OperationsPerInvocation, latch))
7890
awaitLatch(latch)
7991
}
8092

8193
@Benchmark
8294
@OperationsPerInvocation(OperationsPerInvocation)
8395
def oneElementList(): Unit = {
8496
val latch = new CountDownLatch(1)
85-
8697
testSource.flatMapConcat(n => Source(n :: Nil)).runWith(new LatchSink(OperationsPerInvocation, latch))
98+
awaitLatch(latch)
99+
}
100+
101+
@Benchmark
102+
@OperationsPerInvocation(OperationsPerInvocation)
103+
def oneElementListP1(): Unit = {
104+
val latch = new CountDownLatch(1)
105+
testSource.flatMapConcat(1, n => Source(n :: Nil)).runWith(new LatchSink(OperationsPerInvocation, latch))
106+
awaitLatch(latch)
107+
}
108+
109+
@Benchmark
110+
@OperationsPerInvocation(OperationsPerInvocation)
111+
def completedFuture(): Unit = {
112+
val latch = new CountDownLatch(1)
113+
testSource
114+
.flatMapConcat(n => Source.future(Future.successful(n)))
115+
.runWith(new LatchSink(OperationsPerInvocation, latch))
116+
awaitLatch(latch)
117+
}
118+
119+
@Benchmark
120+
@OperationsPerInvocation(OperationsPerInvocation)
121+
def completedFutureP1(): Unit = {
122+
val latch = new CountDownLatch(1)
123+
testSource
124+
.flatMapConcat(1, n => Source.future(Future.successful(n)))
125+
.runWith(new LatchSink(OperationsPerInvocation, latch))
126+
awaitLatch(latch)
127+
}
87128

129+
@Benchmark
130+
@OperationsPerInvocation(OperationsPerInvocation)
131+
def normalFuture(): Unit = {
132+
val latch = new CountDownLatch(1)
133+
testSource
134+
.flatMapConcat(n => Source.future(Future(n)(system.dispatcher)))
135+
.runWith(new LatchSink(OperationsPerInvocation, latch))
136+
awaitLatch(latch)
137+
}
138+
139+
@Benchmark
140+
@OperationsPerInvocation(OperationsPerInvocation)
141+
def normalFutureP1(): Unit = {
142+
val latch = new CountDownLatch(1)
143+
testSource
144+
.flatMapConcat(1, n => Source.future(Future(n)(system.dispatcher)))
145+
.runWith(new LatchSink(OperationsPerInvocation, latch))
88146
awaitLatch(latch)
89147
}
90148

‎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[_] | _: FutureSource[_] | _: IterableSource[_] | _: JavaStreamSource[_, _] | EmptySource =>
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)