I created a simple Sink based on akka.stream.impl.LazySink
. I have only tested it with a single element in the successful case so feel free to comment on here or the GitHub Gist.
import akka.NotUsed
import akka.stream.{Attributes, Inlet, SinkShape}
import akka.stream.scaladsl.{Sink, Source}
import akka.stream.stage._
class OneToOneOnDemandSink[T, +M](sink: T => Sink[T, M]) extends GraphStage[SinkShape[T]] {
val in: Inlet[T] = Inlet("OneToOneOnDemandSink.in")
override val shape = SinkShape(in)
override def createLogic(inheritedAttributes: Attributes) = new GraphStageLogic(shape) {
override def preStart(): Unit = pull(in)
val awaitingElementHandler = new InHandler {
override def onPush(): Unit = {
val element = grab(in)
val innerSource = createInnerSource(element)
val innerSink = sink(element)
Source.fromGraph(innerSource.source).runWith(innerSink)(subFusingMaterializer)
}
override def onUpstreamFinish(): Unit = completeStage()
override def onUpstreamFailure(ex: Throwable): Unit = failStage(ex)
}
setHandler(in, awaitingElementHandler)
def createInnerSource(element: T): SubSourceOutlet[T] = {
val innerSource = new SubSourceOutlet[T]("OneToOneOnDemandSink.innerSource")
innerSource.setHandler(new OutHandler {
override def onPull(): Unit = {
innerSource.push(element)
innerSource.complete()
if (isClosed(in)) {
completeStage()
} else {
pull(in)
setHandler(in, awaitingElementHandler)
}
}
override def onDownstreamFinish(): Unit = {
innerSource.complete()
if (isClosed(in)) {
completeStage()
}
}
})
setHandler(in, new InHandler {
override def onPush(): Unit = {
val illegalStateException = new IllegalStateException("Got a push that we weren't expecting")
innerSource.fail(illegalStateException)
failStage(illegalStateException)
}
override def onUpstreamFinish(): Unit = {
// We don't stop until the inner stream stops.
setKeepGoing(true)
}
override def onUpstreamFailure(ex: Throwable): Unit = {
innerSource.fail(ex)
failStage(ex)
}
})
innerSource
}
}
}
object OneToOneOnDemandSink {
def apply[T, M](sink: T => Sink[T, M]): Sink[T, NotUsed] = Sink.fromGraph(new OneToOneOnDemandSink(sink))
}
This will create a new Sink for each element so it avoids a whole lot of the complexity that LazySink
has and there is also no sensible materialized value to return.