Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Add emitMulti with Spliterator support #1776

Merged
merged 1 commit into from
Feb 27, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  • the advantage of a Spliterator over an iterator is that you can use trySplit() to partition the workload so you can get multiple threads to work on the data
  • if you just use tryAdvance, you might as well be just using the standard iterator]
  • so what do we gain here?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you can avoid the hasNext and next call when you already have a Spliterator

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure this will make a big diff to performance. I'll approve this but without jmh bechmarks, I am very sceptical about whether this extra code is worth the extra maintenance overhead.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks

_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