pekko/akka-stream/src/main/scala/akka/stream/impl/UnfoldResourceSourceAsync.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).recover {
          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"

}