MaybeSource rewritten as GraphStage #22789

akka-stream/src/main/scala/akka/stream/impl/MaybeSource.scala

@@ -0,0 +1,84 @@
+/**
+ * Copyright (C) 2009-2017 Lightbend Inc. <http://www.lightbend.com>
+ */
+package akka.stream.impl
+
+import akka.annotation.InternalApi
+import akka.dispatch.ExecutionContexts
+import akka.stream.{ AbruptStageTerminationException, Attributes, Outlet, SourceShape }
+import akka.stream.impl.Stages.DefaultAttributes
+import akka.stream.stage.{ GraphStageLogic, GraphStageWithMaterializedValue, OutHandler }
+import akka.util.OptionVal
+
+import scala.concurrent.Promise
+import scala.util.Try
+
+/**
+ * INTERNAL API
+ */
+@InternalApi private[akka] object MaybeSource extends GraphStageWithMaterializedValue[SourceShape[AnyRef], Promise[Option[AnyRef]]] {
+  val out = Outlet[AnyRef]("MaybeSource.out")
+  override val shape = SourceShape(out)
+
+  override protected def initialAttributes = DefaultAttributes.maybeSource
+
+  override def createLogicAndMaterializedValue(inheritedAttributes: Attributes): (GraphStageLogic, Promise[Option[AnyRef]]) = {
+    import scala.util.{ Success ⇒ ScalaSuccess, Failure ⇒ ScalaFailure }
+    val promise = Promise[Option[AnyRef]]()
+    val logic = new GraphStageLogic(shape) with OutHandler {
+
+      private var arrivedEarly: OptionVal[AnyRef] = OptionVal.None
+
+      override def preStart(): Unit = {
+        promise.future.value match {
+          case Some(value) ⇒
+            // already completed, shortcut
+            handleCompletion(value)
+          case None ⇒
+            // callback on future completion
+            promise.future.onComplete(
+              getAsyncCallback(handleCompletion).invoke
+            )(ExecutionContexts.sameThreadExecutionContext)
+        }
+      }
+
+      override def onPull(): Unit = arrivedEarly match {
+        case OptionVal.Some(value) ⇒
+          push(out, value)
+          completeStage()
+        case OptionVal.None ⇒
+      }
+
+      private def handleCompletion(elem: Try[Option[AnyRef]]): Unit = {
+        elem match {
+          case ScalaSuccess(None) ⇒
+            completeStage()
+          case ScalaSuccess(Some(value)) ⇒
+            if (isAvailable(out)) {
+              push(out, value)
+              completeStage()
+            } else {
+              arrivedEarly = OptionVal.Some(value)
+            }
+          case ScalaFailure(ex) ⇒
+            failStage(ex)
+        }
+      }
+
+      override def onDownstreamFinish(): Unit = {
+        promise.tryComplete(ScalaSuccess(None))
+      }
+
+      override def postStop(): Unit = {
+        if (!promise.isCompleted)
+          promise.tryFailure(new AbruptStageTerminationException(this))
+      }
+
+      setHandler(out, this)
+
+    }
+    (logic, promise)
+  }
+
+  override def toString = "MaybeSource"
+}