raboof/pekko
1
0
Fork 0
Code Issues Pull requests Projects Releases 6 Packages Wiki Activity Actions 11

=str Remove outHandler of FlattenMerge (#29890)

Browse source
This commit is contained in:
kerr 2020-12-15 18:18:48 +08:00 • committed by GitHub
parent 8bb74b592d
commit 8c0ffe3b2a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 70 additions and 71 deletions
Download patch file Download diff file Expand all files Collapse all files

141
akka-stream/src/main/scala/akka/stream/impl/fusing/StreamOfStreams.scala
View file

@ -39,97 +39,96 @@ import akka.util.ccompat.JavaConverters._
override def initialAttributes = DefaultAttributes.flattenMerge override def initialAttributes = DefaultAttributes.flattenMerge
override val shape = FlowShape(in, out) override val shape = FlowShape(in, out)
override def createLogic(enclosingAttributes: Attributes) = new GraphStageLogic(shape) { override def createLogic(enclosingAttributes: Attributes) =
var sources = Set.empty[SubSinkInlet[T]] new GraphStageLogic(shape) with OutHandler with InHandler {
var pendingSingleSources = 0 var sources = Set.empty[SubSinkInlet[T]]
def activeSources = sources.size + pendingSingleSources var pendingSingleSources = 0
def activeSources = sources.size + pendingSingleSources
// To be able to optimize for SingleSource without materializing them the queue may hold either // To be able to optimize for SingleSource without materializing them the queue may hold either
// SubSinkInlet[T] or SingleSource // SubSinkInlet[T] or SingleSource
var queue: BufferImpl[AnyRef] = _ var queue: BufferImpl[AnyRef] = _
override def preStart(): Unit = queue = BufferImpl(breadth, enclosingAttributes) override def preStart(): Unit = queue = BufferImpl(breadth, enclosingAttributes)
def pushOut(): Unit = { def pushOut(): Unit = {
queue.dequeue() match { queue.dequeue() match {
case src: SubSinkInlet[T] @unchecked => case src: SubSinkInlet[T] @unchecked =>
push(out, src.grab()) push(out, src.grab())
if (!src.isClosed) src.pull() if (!src.isClosed) src.pull()
else removeSource(src) else removeSource(src)
case single: SingleSource[T] @unchecked => case single: SingleSource[T] @unchecked =>
push(out, single.elem) push(out, single.elem)
removeSource(single) removeSource(single)
}
} }
}
setHandler(in, new InHandler {
override def onPush(): Unit = { override def onPush(): Unit = {
val source = grab(in) val source = grab(in)
addSource(source) addSource(source)
if (activeSources < breadth) tryPull(in) if (activeSources < breadth) tryPull(in)
} }
override def onUpstreamFinish(): Unit = if (activeSources == 0) completeStage()
})
setHandler(out, new OutHandler { override def onUpstreamFinish(): Unit = if (activeSources == 0) completeStage()
override def onPull(): Unit = { override def onPull(): Unit = {
pull(in) pull(in)
setHandler(out, outHandler) setHandler(out, new OutHandler {
} override def onPull(): Unit = {
}) // could be unavailable due to async input having been executed before this notification
if (queue.nonEmpty && isAvailable(out)) pushOut()
val outHandler = new OutHandler {
// could be unavailable due to async input having been executed before this notification
override def onPull(): Unit = if (queue.nonEmpty && isAvailable(out)) pushOut()
}
def addSource(source: Graph[SourceShape[T], M]): Unit = {
// If it's a SingleSource or wrapped such we can push the element directly instead of materializing it.
// Have to use AnyRef because of OptionVal null value.
TraversalBuilder.getSingleSource(source.asInstanceOf[Graph[SourceShape[AnyRef], M]]) match {
case OptionVal.Some(single) =>
if (isAvailable(out) && queue.isEmpty) {
push(out, single.elem.asInstanceOf[T])
} else {
queue.enqueue(single)
pendingSingleSources += 1
} }
case _ => })
val sinkIn = new SubSinkInlet[T]("FlattenMergeSink") }
sinkIn.setHandler(new InHandler {
override def onPush(): Unit = { setHandlers(in, out, this)
if (isAvailable(out)) {
push(out, sinkIn.grab()) def addSource(source: Graph[SourceShape[T], M]): Unit = {
sinkIn.pull() // If it's a SingleSource or wrapped such we can push the element directly instead of materializing it.
} else { // Have to use AnyRef because of OptionVal null value.
queue.enqueue(sinkIn) TraversalBuilder.getSingleSource(source.asInstanceOf[Graph[SourceShape[AnyRef], M]]) match {
} case OptionVal.Some(single) =>
if (isAvailable(out) && queue.isEmpty) {
push(out, single.elem.asInstanceOf[T])
} else {
queue.enqueue(single)
pendingSingleSources += 1
} }
override def onUpstreamFinish(): Unit = if (!sinkIn.isAvailable) removeSource(sinkIn) case _ =>
}) val sinkIn = new SubSinkInlet[T]("FlattenMergeSink")
sinkIn.pull() sinkIn.setHandler(new InHandler {
sources += sinkIn override def onPush(): Unit = {
val graph = Source.fromGraph(source).to(sinkIn.sink) if (isAvailable(out)) {
interpreter.subFusingMaterializer.materialize(graph, defaultAttributes = enclosingAttributes) push(out, sinkIn.grab())
sinkIn.pull()
} else {
queue.enqueue(sinkIn)
}
}
override def onUpstreamFinish(): Unit = if (!sinkIn.isAvailable) removeSource(sinkIn)
})
sinkIn.pull()
sources += sinkIn
val graph = Source.fromGraph(source).to(sinkIn.sink)
interpreter.subFusingMaterializer.materialize(graph, defaultAttributes = enclosingAttributes)
}
} }
}
def removeSource(src: AnyRef): Unit = { def removeSource(src: AnyRef): Unit = {
val pullSuppressed = activeSources == breadth val pullSuppressed = activeSources == breadth
src match { src match {
case sub: SubSinkInlet[T] @unchecked => case sub: SubSinkInlet[T] @unchecked =>
sources -= sub sources -= sub
case _: SingleSource[_] => case _: SingleSource[_] =>
pendingSingleSources -= 1 pendingSingleSources -= 1
}
if (pullSuppressed) tryPull(in)
if (activeSources == 0 && isClosed(in)) completeStage()
} }
if (pullSuppressed) tryPull(in)
if (activeSources == 0 && isClosed(in)) completeStage() override def postStop(): Unit = sources.foreach(_.cancel())
} }
override def postStop(): Unit = sources.foreach(_.cancel())
}
override def toString: String = s"FlattenMerge($breadth)" override def toString: String = s"FlattenMerge($breadth)"
} }