127 lines
4 KiB
Scala
127 lines
4 KiB
Scala
/*
|
|
* Copyright (C) 2015-2019 Lightbend Inc. <https://www.lightbend.com>
|
|
*/
|
|
|
|
package akka.stream.impl
|
|
|
|
import akka.Done
|
|
import akka.annotation.InternalApi
|
|
import akka.dispatch.ExecutionContexts.sameThreadExecutionContext
|
|
import akka.stream.ActorAttributes.SupervisionStrategy
|
|
import akka.stream._
|
|
import akka.stream.impl.Stages.DefaultAttributes
|
|
import akka.stream.stage._
|
|
|
|
import scala.concurrent.Future
|
|
import scala.util.{ Failure, Success, Try }
|
|
import scala.util.control.NonFatal
|
|
|
|
/**
|
|
* INTERNAL API
|
|
*/
|
|
@InternalApi private[akka] final class UnfoldResourceSourceAsync[T, S](
|
|
create: () => Future[S],
|
|
readData: (S) => Future[Option[T]],
|
|
close: (S) => Future[Done])
|
|
extends GraphStage[SourceShape[T]] {
|
|
val out = Outlet[T]("UnfoldResourceSourceAsync.out")
|
|
override val shape = SourceShape(out)
|
|
override def initialAttributes: Attributes = DefaultAttributes.unfoldResourceSourceAsync
|
|
|
|
def createLogic(inheritedAttributes: Attributes) = new GraphStageLogic(shape) with OutHandler {
|
|
lazy val decider = inheritedAttributes.mandatoryAttribute[SupervisionStrategy].decider
|
|
private implicit def ec = materializer.executionContext
|
|
private var state: Option[S] = None
|
|
|
|
private val createdCallback = getAsyncCallback[Try[S]] {
|
|
case Success(resource) =>
|
|
state = Some(resource)
|
|
if (isAvailable(out)) onPull()
|
|
case Failure(t) => failStage(t)
|
|
}.invokeWithFeedback _
|
|
|
|
private val errorHandler: PartialFunction[Throwable, Unit] = {
|
|
case NonFatal(ex) =>
|
|
decider(ex) match {
|
|
case Supervision.Stop =>
|
|
failStage(ex)
|
|
case Supervision.Restart => restartResource()
|
|
case Supervision.Resume => onPull()
|
|
}
|
|
}
|
|
|
|
private val readCallback = getAsyncCallback[Try[Option[T]]] {
|
|
case Success(data) =>
|
|
data match {
|
|
case Some(d) => push(out, d)
|
|
case None =>
|
|
// end of resource reached, lets close it
|
|
state match {
|
|
case Some(resource) =>
|
|
close(resource).onComplete(getAsyncCallback[Try[Done]] {
|
|
case Success(Done) => completeStage()
|
|
case Failure(ex) => failStage(ex)
|
|
}.invoke)
|
|
state = None
|
|
|
|
case None =>
|
|
// cannot happen, but for good measure
|
|
throw new IllegalStateException("Reached end of data but there is no open resource")
|
|
}
|
|
}
|
|
case Failure(t) => errorHandler(t)
|
|
}.invoke _
|
|
|
|
override def preStart(): Unit = createResource()
|
|
|
|
override def onPull(): Unit =
|
|
state match {
|
|
case Some(resource) =>
|
|
try {
|
|
readData(resource).onComplete(readCallback)(sameThreadExecutionContext)
|
|
} catch errorHandler
|
|
case None =>
|
|
// we got a pull but there is no open resource, we are either
|
|
// currently creating/restarting then the read will be triggered when creating the
|
|
// resource completes, or shutting down and then the pull does not matter anyway
|
|
}
|
|
|
|
override def postStop(): Unit = {
|
|
state.foreach(r => close(r))
|
|
}
|
|
|
|
private def restartResource(): Unit = {
|
|
state match {
|
|
case Some(resource) =>
|
|
// wait for the resource to close before restarting
|
|
close(resource).onComplete(getAsyncCallback[Try[Done]] {
|
|
case Success(Done) =>
|
|
createResource()
|
|
case Failure(ex) => failStage(ex)
|
|
}.invoke)
|
|
state = None
|
|
case None =>
|
|
createResource()
|
|
}
|
|
}
|
|
|
|
private def createResource(): Unit = {
|
|
create().onComplete { resource =>
|
|
createdCallback(resource).failed.foreach {
|
|
case _: StreamDetachedException =>
|
|
// stream stopped
|
|
resource match {
|
|
case Success(r) =>
|
|
close(r)
|
|
case Failure(ex) => throw ex // failed to open but stream is stopped already
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
setHandler(out, this)
|
|
|
|
}
|
|
override def toString = "UnfoldResourceSourceAsync"
|
|
|
|
}
|