Skip to content

Commit

Permalink
feat: Add emitMulti with Spliterator support
Browse files Browse the repository at this point in the history
  • Loading branch information
He-Pin committed Feb 22, 2025
1 parent 22685b9 commit bc138b6
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,20 @@ class GraphStageLogicSpec extends StreamSpec with GraphInterpreterSpecKit with S
override def toString = "GraphStageLogicSpec.emitEmptyIterable"
}

object emitSplitIterator extends GraphStage[SourceShape[Int]] {
val out = Outlet[Int]("out")
override val shape = SourceShape(out)
override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) {
setHandler(out,
new OutHandler {
override def onPull(): Unit = emitMultiple(
out,
java.util.stream.Stream.of(1, 2, 3).spliterator(), () => emit(out, 42, () => completeStage()))
})
}
override def toString = "GraphStageLogicSpec.emitEmptyIterable"
}

private case class ReadNEmitN(n: Int) extends GraphStage[FlowShape[Int, Int]] {
override val shape = FlowShape(Inlet[Int]("readN.in"), Outlet[Int]("readN.out"))

Expand Down Expand Up @@ -196,6 +210,12 @@ class GraphStageLogicSpec extends StreamSpec with GraphInterpreterSpecKit with S

}

"emit properly when using split iterator" in {

Source.fromGraph(emitSplitIterator).runWith(Sink.seq).futureValue should ===(List(1, 2, 3, 42))

}

"invoke lifecycle hooks in the right order" in {
val g = new GraphStage[FlowShape[Int, Int]] {
val in = Inlet[Int]("in")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ import scala.annotation.tailrec
import scala.collection.{ immutable, mutable }
import scala.concurrent.{ Future, Promise }
import scala.concurrent.duration.FiniteDuration

import org.apache.pekko
import pekko.{ Done, NotUsed }
import pekko.actor._
Expand All @@ -37,6 +36,8 @@ import pekko.stream.stage.ConcurrentAsyncCallbackState.{ NoPendingEvents, State
import pekko.util.OptionVal
import pekko.util.unused

import java.util.Spliterator

/**
* Scala API: A GraphStage represents a reusable graph stream processing operator.
*
Expand Down Expand Up @@ -979,6 +980,26 @@ abstract class GraphStageLogic private[stream] (val inCount: Int, val outCount:
}
} else andThen()

/**
* Emit a sequence of elements through the given outlet and continue with the given thunk
* afterwards, suspending execution if necessary.
* This action replaces the [[OutHandler]] for the given outlet if suspension
* is needed and reinstalls the current handler upon receiving an `onPull()`
* signal (before invoking the `andThen` function).
*/
final protected def emitMultiple[T](out: Outlet[T], elems: Spliterator[T], andThen: () => Unit): Unit = {
val iter = new EmittingSpliterator[T](out, elems, getNonEmittingHandler(out), andThen)
if (isAvailable(out)) {
if (!iter.tryPush()) {
andThen()
} else {
setOrAddEmitting(out, iter)
}
} else {
setOrAddEmitting(out, iter)
}
}

/**
* Emit a sequence of elements through the given outlet, suspending execution if necessary.
* This action replaces the [[OutHandler]] for the given outlet if suspension
Expand Down Expand Up @@ -1118,6 +1139,19 @@ abstract class GraphStageLogic private[stream] (val inCount: Int, val outCount:
}
}

private final class EmittingSpliterator[T](_out: Outlet[T], elems: Spliterator[T], _previous: OutHandler,
_andThen: () => Unit)
extends Emitting[T](_out, _previous, _andThen) with java.util.function.Consumer[T] {

override def onPull(): Unit = if (!elems.tryAdvance(this)) {
followUp()
}

def tryPush(): Boolean = elems.tryAdvance(this)

override def accept(elem: T): Unit = push(out, elem)
}

private class EmittingCompletion[T](_out: Outlet[T], _previous: OutHandler)
extends Emitting[T](_out, _previous, DoNothing) {
override def onPull(): Unit = complete(out)
Expand Down

0 comments on commit bc138b6

Please sign in to comment.