akka.stream.stage.GraphStageLogic Scala Examples

package io.opencensus.scala.akka.http.utils

import akka.NotUsed
import akka.http.scaladsl.model.{HttpEntity, HttpResponse}
import akka.stream.scaladsl.Flow
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
import akka.stream.{Attributes, FlowShape, Inlet, Outlet}

object ExecuteAfterResponse {

  private class AfterResponseFlow[Element](
      onFinish: () => Unit,
      onFailure: Throwable => Unit
  ) extends GraphStage[FlowShape[Element, Element]] {
    private val in  = Inlet[Element]("in")
    private val out = Outlet[Element]("out")

    override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
      new GraphStageLogic(shape) with InHandler with OutHandler {
        def onPush(): Unit = push(out, grab(in))
        def onPull(): Unit = pull(in)

        setHandler(in, this)
        setHandler(out, this)

        override def onUpstreamFinish(): Unit = {
          onFinish()
          super.onUpstreamFinish()
        }
        override def onUpstreamFailure(ex: Throwable): Unit = {
          onFailure(ex)
          super.onUpstreamFailure(ex)
        }
      }

    override val shape = FlowShape(in, out)
  }

  private object AfterResponseFlow {
    def apply[Element](
        onFinish: () => Unit,
        onFailure: Throwable => Unit
    ): Flow[Element, Element, NotUsed] =
      Flow.fromGraph(new AfterResponseFlow(onFinish, onFailure))
  }

  def onComplete(
      response: HttpResponse,
      onFinish: () => Unit,
      onFailure: Throwable => Unit
  ): HttpResponse = {

    response.copy(
      entity = if (response.status.allowsEntity) {
        response.entity.transformDataBytes(
          AfterResponseFlow(onFinish, onFailure)
        )
      } else {
        onFinish()
        HttpEntity.Empty
      }
    )
  }
} 

package de.knutwalker.akka.http

import de.knutwalker.akka.stream.JsonStreamParser

import akka.http.scaladsl.model.HttpEntity
import akka.http.scaladsl.model.MediaTypes.`application/json`
import akka.http.scaladsl.unmarshalling.{ FromEntityUnmarshaller, Unmarshaller }
import akka.http.scaladsl.util.FastFuture
import akka.stream.scaladsl.Sink
import akka.stream.stage.{ GraphStageLogic, GraphStageWithMaterializedValue, InHandler }
import akka.stream.{ AbruptStageTerminationException, Attributes, Inlet, SinkShape }

import jawn.Facade

import scala.concurrent.{ Future, Promise }
import java.util.NoSuchElementException

object JsonSupport extends JsonSupport {
  private def firstElementSink[J <: AnyRef]: Sink[J, Future[J]] =
    Sink.fromGraph(new FirstElementSinkStage[J])

  private final class FirstElementSinkStage[J <: AnyRef] extends GraphStageWithMaterializedValue[SinkShape[J], Future[J]] {
    private[this] val in: Inlet[J] = Inlet("firstElement.in")

    override val shape: SinkShape[J] = SinkShape.of(in)
    override protected def initialAttributes: Attributes = Attributes.name("firstElement")

    override def createLogicAndMaterializedValue(inheritedAttributes: Attributes): (GraphStageLogic, Future[J]) = {
      val p: Promise[J] = Promise()
      (new GraphStageLogic(shape) with InHandler {
        private[this] var element: J = null.asInstanceOf[J]

        override def preStart(): Unit = pull(in)

        def onPush(): Unit = {
          if (element eq null) {
            element = grab(in)
          }
          pull(in)
        }

        override def onUpstreamFinish(): Unit = {
          val el = element
          element = null.asInstanceOf[J]
          if (el ne null) {
            p.trySuccess(el)
          } else {
            p.tryFailure(new NoSuchElementException("No complete json entity consumed"))
          }
          completeStage()
        }

        override def onUpstreamFailure(ex: Throwable): Unit = {
          element = null.asInstanceOf[J]
          p.tryFailure(ex)
          failStage(ex)
        }

        override def postStop(): Unit = {
          if (!p.isCompleted) {
            p.failure(new AbruptStageTerminationException(this))
            ()
          }
        }

        setHandler(in, this)
      }, p.future)
    }

    override def toString: String = "FirstElementSinkStage"
  }
}

trait JsonSupport {

  implicit def jsonUnmarshaller[J <: AnyRef : Facade]: FromEntityUnmarshaller[J] =
    Unmarshaller.withMaterializer[HttpEntity, J](_ => implicit mat => {
      case HttpEntity.Strict(_, data) => FastFuture(JsonStreamParser.parse[J](data))
      case entity                     => entity.dataBytes.via(JsonStreamParser[J]).runWith(JsonSupport.firstElementSink[J])
    }).forContentTypes(`application/json`)
} 

package cmwell.crawler

import akka.stream.{Attributes, FlowShape, Inlet, Outlet}
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
import com.typesafe.scalalogging.Logger

object CrawlerRatePrinter {
  def apply(crawlerId: String,
            printFrequency: Int,
            maxPrintRateMillis: Int)(logger: Logger): CrawlerRatePrinter =
    new CrawlerRatePrinter(crawlerId, printFrequency, maxPrintRateMillis)(logger)
}

class CrawlerRatePrinter(crawlerId: String,
                         printFrequency: Int,
                         maxPrintRateMillis: Int)(logger: Logger) extends GraphStage[FlowShape[Long, Long]] {
  val in = Inlet[Long]("CrawlerRatePrinter.in")
  val out = Outlet[Long]("CrawlerRatePrinter.out")
  override val shape = FlowShape.of(in, out)

  override def createLogic(attr: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) {
      private val startTime: Double = System.currentTimeMillis
      private var totalElementsGot: Long = 0L
      private var printedAtElementNo: Long = 0L
      private var printedAtTime: Double = startTime
      private var localStartTime: Double = startTime
      private var localTotalElementsGot: Long = 0L
      private var localRate: Double = 0

      setHandler(
        in,
        new InHandler {
          override def onPush(): Unit = {
            val elem = grab(in)
            totalElementsGot += 1
            localTotalElementsGot += 1
            if (totalElementsGot - printedAtElementNo >= printFrequency) {
              val currentTime = System.currentTimeMillis
              if (currentTime - printedAtTime > maxPrintRateMillis) {
                val rate = totalElementsGot / (currentTime - startTime) * 1000
                if (currentTime - localStartTime > 15000) {
                  localRate = localTotalElementsGot / (currentTime - localStartTime) * 1000
                  localTotalElementsGot = 0
                  localStartTime = currentTime
                }
                logger.info(s"$crawlerId Current offset is $elem. Total $totalElementsGot offsets already processed. " +
                s"Read rate: avg: ${rate.formatted("%.2f")} current: ${localRate.formatted("%.2f")} offsets/second")
                printedAtElementNo = totalElementsGot
                printedAtTime = currentTime
              }
            }
            push(out, elem)
          }
        }
      )
      setHandler(out, new OutHandler {
        override def onPull(): Unit = {
          pull(in)
        }
      })
    }
} 

package akka.grpc.internal

import akka.stream.{ Attributes, FlowShape, Inlet, Outlet }
import akka.stream.stage.{ GraphStage, GraphStageLogic, InHandler, OutHandler }


class CancellationBarrierGraphStage[T] extends GraphStage[FlowShape[T, T]] {
  val in: Inlet[T] = Inlet("CancellationBarrier")
  val out: Outlet[T] = Outlet("CancellationBarrier")

  override val shape: FlowShape[T, T] = FlowShape(in, out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) {
      setHandler(
        in,
        new InHandler {
          override def onPush(): Unit = emit(out, grab(in))
        })

      setHandler(
        out,
        new OutHandler {
          override def onPull(): Unit = pull(in)

          override def onDownstreamFinish(): Unit = {
            if (!hasBeenPulled(in))
              pull(in)

            setHandler(
              in,
              new InHandler {
                override def onPush(): Unit = {
                  grab(in)
                  pull(in)
                }
              })
          }
        })
    }
} 

package example.akkawschat.cli

import akka.stream._
import akka.stream.scaladsl.{ Flow, Source, GraphDSL }
import akka.stream.stage.{ InHandler, GraphStageLogic, GraphStage }

import scala.concurrent.ExecutionContext

object Prompt {
   ⇒
              collectedString = collectedString.dropRight(1)
              prompt()
            case x ⇒
              //println(s"Got ${x.toInt}")
              collectedString += x
              print(x)
              pull(characterInput)
          }
        }
      })
      setHandler(outputLinesIn, new InHandler {
        def onPush(): Unit = {
          print(s"$RESTORE$ERASE_LINE${grab(outputLinesIn)}\n$SAVE$promptLine")
          pull(outputLinesIn)
        }
      })
      setHandler(readLinesOut, eagerTerminateOutput)

      override def preStart(): Unit = {
        pull(outputLinesIn)
        print(SAVE) // to make sure we don't jump back to former SAVE position in the terminal
        prompt()
      }

      def promptLine = s"$RESTORE$ERASE_LINE$SAVE> $collectedString"

      def prompt(): Unit = {
        print(promptLine)
        pull(characterInput)
      }
    }
} 

package example.akkawschat.cli

import akka.stream.stage.{ OutHandler, GraphStageLogic, GraphStage }
import akka.stream._

import scala.annotation.tailrec
import scala.concurrent.{ Future, ExecutionContext }
import scala.util.control.NoStackTrace

class ConsoleInput(implicit ec: ExecutionContext) extends GraphStage[SourceShape[Char]] {
  val out = Outlet[Char]("consoleOut")
  val shape: SourceShape[Char] = SourceShape(out)

  def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) {
      TTY.noEchoStty()

      @volatile var cancelled = false
      def getOne(): Unit = {
        val callback = getAsyncCallback[Char](push(out, _))

        Future {
          @tailrec def read(): Unit =
            if (cancelled) throw new Exception with NoStackTrace
            else if (System.in.available() > 0)
              callback.invoke(System.in.read().toChar)
            else {
              Thread.sleep(10)
              read()
            }

          read()
        }
      }

      setHandler(out, new OutHandler {
        def onPull(): Unit = getOne()

        override def onDownstreamFinish(cause: Throwable): Unit = {
          cancelled = true
          super.onDownstreamFinish(cause)
        }
      })

      override def postStop(): Unit =
        TTY.saneStty()
    }
} 

package com.sksamuel.pulsar4s.akka.streams

import akka.Done
import akka.stream.Attributes
import akka.stream.Outlet
import akka.stream.SourceShape
import akka.stream.stage.AsyncCallback
import akka.stream.stage.GraphStageLogic
import akka.stream.stage.GraphStageWithMaterializedValue
import akka.stream.stage.OutHandler
import com.sksamuel.exts.Logging
import com.sksamuel.pulsar4s.Consumer
import com.sksamuel.pulsar4s.ConsumerMessage
import com.sksamuel.pulsar4s.MessageId
import org.apache.pulsar.client.api.ConsumerStats

import scala.concurrent.ExecutionContext
import scala.concurrent.Future
import scala.util.Failure
import scala.util.Success

trait CommittableMessage[T] {
  def ack(cumulative: Boolean = false): Future[Done]
  def nack(): Future[Done]
  def message: ConsumerMessage[T]
}

class PulsarCommittableSourceGraphStage[T](create: () => Consumer[T], seek: Option[MessageId])
  extends GraphStageWithMaterializedValue[SourceShape[CommittableMessage[T]], Control]
    with Logging {

  private val out = Outlet[CommittableMessage[T]]("pulsar.out")
  override def shape: SourceShape[CommittableMessage[T]] = SourceShape(out)

  private class PulsarCommittableSourceLogic(shape: Shape) extends GraphStageLogic(shape) with OutHandler with Control {
    setHandler(out, this)

    var consumer: Consumer[T] = _
    var receiveCallback: AsyncCallback[CommittableMessage[T]] = _

    override def preStart(): Unit = {
      implicit val context: ExecutionContext = super.materializer.executionContext
      consumer = create()
      seek foreach consumer.seek
      receiveCallback = getAsyncCallback(push(out, _))
    }

    override def onPull(): Unit = {
      implicit val context: ExecutionContext = super.materializer.executionContext
      logger.debug("Pull received; asking consumer for message")

      consumer.receiveAsync.onComplete {
        case Success(msg) =>
          logger.debug(s"Message received: $msg")
          receiveCallback.invoke(new CommittableMessage[T] {
            override def message: ConsumerMessage[T] = msg
            override def ack(cumulative: Boolean): Future[Done] = {
              logger.debug(s"Acknowledging message: $msg")
              val ackFuture = if (cumulative) {
                consumer.acknowledgeCumulativeAsync(msg.messageId)
              } else {
                consumer.acknowledgeAsync(msg.messageId)
              }
              ackFuture.map(_ => Done)
            }
            override def nack(): Future[Done] = {
              logger.debug(s"Negatively acknowledging message: $msg")
              consumer.negativeAcknowledgeAsync(msg.messageId).map(_ => Done)
            }
          })
        case Failure(e) =>
          logger.warn("Error when receiving message", e)
          failStage(e)
      }
    }

    override def stop(): Unit = completeStage()

    override def shutdown()(implicit ec: ExecutionContext): Future[Done] = {
      completeStage()
      consumer.closeAsync.map(_ => Done)
    }

    def stats: ConsumerStats = consumer.stats
  }

  override def createLogicAndMaterializedValue(inheritedAttributes: Attributes): (GraphStageLogic, Control) = {
    val logic = new PulsarCommittableSourceLogic(shape)
    (logic, logic)
  }
} 

package com.sksamuel.pulsar4s.akka.streams

import akka.Done
import akka.stream.Attributes
import akka.stream.Outlet
import akka.stream.SourceShape
import akka.stream.stage.AsyncCallback
import akka.stream.stage.GraphStageLogic
import akka.stream.stage.GraphStageWithMaterializedValue
import akka.stream.stage.OutHandler
import com.sksamuel.exts.Logging
import com.sksamuel.pulsar4s.Consumer
import com.sksamuel.pulsar4s.ConsumerMessage
import com.sksamuel.pulsar4s.MessageId
import org.apache.pulsar.client.api.ConsumerStats

import scala.concurrent.ExecutionContext
import scala.concurrent.Future
import scala.util.Failure
import scala.util.Success

trait Control {

  
  def stats: ConsumerStats
}

class PulsarSourceGraphStage[T](create: () => Consumer[T], seek: Option[MessageId])
  extends GraphStageWithMaterializedValue[SourceShape[ConsumerMessage[T]], Control]
    with Logging {

  private val out = Outlet[ConsumerMessage[T]]("pulsar.out")
  override def shape: SourceShape[ConsumerMessage[T]] = SourceShape(out)

  override def createLogicAndMaterializedValue(inheritedAttributes: Attributes): (GraphStageLogic, Control) = {

    val logic: GraphStageLogic with Control = new GraphStageLogic(shape) with OutHandler with Control {
      setHandler(out, this)

      var consumer: Consumer[T] = _
      var callback: AsyncCallback[ConsumerMessage[T]] = _

      override def preStart(): Unit = {
        consumer = create()
        seek foreach consumer.seek
        callback = getAsyncCallback(msg => push(out, msg))
      }

      override def onPull(): Unit = {
        implicit val context: ExecutionContext = super.materializer.executionContext
        logger.debug("Pull received; asking consumer for message")
        consumer.receiveAsync.onComplete {
          case Success(msg) =>
            logger.debug(s"Msg received $msg")
            callback.invoke(msg)
            consumer.acknowledge(msg.messageId)
          case Failure(e) =>
            logger.warn("Error when receiving message", e)
            failStage(e)
        }
      }

      override def stop(): Unit = completeStage()

      override def shutdown()(implicit ec: ExecutionContext): Future[Done] = {
        completeStage()
        consumer.closeAsync.map(_ => Done)
      }

      override def stats: ConsumerStats = consumer.stats
    }

    (logic, logic)
  }
} 

package com.sksamuel.pulsar4s.akka.streams

import akka.Done
import akka.stream.stage.{AsyncCallback, GraphStageLogic, GraphStageWithMaterializedValue, InHandler}
import akka.stream.{Attributes, Inlet, SinkShape}
import com.sksamuel.exts.Logging
import com.sksamuel.pulsar4s.{Producer, ProducerMessage, Topic}

import scala.concurrent.duration._
import scala.concurrent.{Await, ExecutionContextExecutor, Future, Promise}
import scala.util.{Failure, Success}

class PulsarMultiSinkGraphStage[T](createFn: Topic => Producer[T], initTopics: Set[Topic] = Set.empty)
  extends GraphStageWithMaterializedValue[SinkShape[(Topic, ProducerMessage[T])], Future[Done]]
    with Logging {

  private val in = Inlet.create[(Topic, ProducerMessage[T])]("pulsar.in")

  override def shape: SinkShape[(Topic, ProducerMessage[T])] = SinkShape.of(in)

  override def createLogicAndMaterializedValue(inheritedAttributes: Attributes): (GraphStageLogic, Future[Done]) = {

    val promise = Promise[Done]()

    val logic: GraphStageLogic = new GraphStageLogic(shape) with InHandler {
      setHandler(in, this)

      implicit def context: ExecutionContextExecutor = super.materializer.executionContext

      var producers: Map[Topic, Producer[T]] = _
      var next: AsyncCallback[(Topic, ProducerMessage[T])] = _
      var error: Throwable = _

      override def preStart(): Unit = {
        producers = initTopics.map(t => t -> createFn(t)).toMap
        next = getAsyncCallback { _ => pull(in) }
        pull(in)
      }

      private def getProducer(topic: Topic): Producer[T] =
        producers.get(topic) match {
          case Some(p) => p
          case None =>
            logger.debug(s"creating new producer for topic $topic")
            val producer = createFn(topic)
            producers += topic -> producer
            producer
        }

      override def onPush(): Unit = {
        try {
          val (topic, message) = grab(in)
          logger.debug(s"Sending message $message to $topic")
          val producer = getProducer(topic)
          producer.sendAsync(message).onComplete {
            case Success(_) => next.invoke(topic -> message)
            case Failure(e) =>
              logger.error("Failing pulsar sink stage", e)
              failStage(e)
          }
        } catch {
          case e: Throwable =>
            logger.error("Failing pulsar sink stage", e)
            failStage(e)
        }
      }

      override def postStop(): Unit = {
        logger.debug("Graph stage stopping; closing producers")
        val fs = producers.flatMap { case (_, p) =>
          Seq(
            p.flushAsync,
            p.closeAsync
          )
        }
        Await.ready(Future.sequence(fs), 15.seconds)
      }

      override def onUpstreamFailure(ex: Throwable): Unit = {
        promise.tryFailure(ex)
      }

      override def onUpstreamFinish(): Unit = {
        promise.trySuccess(Done)
      }
    }

    (logic, promise.future)
  }

} 

package com.sksamuel.pulsar4s.akka.streams

import akka.Done
import akka.stream.stage.{AsyncCallback, GraphStageLogic, GraphStageWithMaterializedValue, InHandler}
import akka.stream.{Attributes, Inlet, SinkShape}
import com.sksamuel.exts.Logging
import com.sksamuel.pulsar4s.{Producer, ProducerMessage}

import scala.concurrent.{ExecutionContextExecutor, Future, Promise}
import scala.util.{Failure, Success}

class PulsarSinkGraphStage[T](createFn: () => Producer[T])
  extends GraphStageWithMaterializedValue[SinkShape[ProducerMessage[T]], Future[Done]]
    with Logging {

  private val in = Inlet.create[ProducerMessage[T]]("pulsar.in")
  override def shape: SinkShape[ProducerMessage[T]] = SinkShape.of(in)

  override def createLogicAndMaterializedValue(inheritedAttributes: Attributes): (GraphStageLogic, Future[Done]) = {

    val promise = Promise[Done]()

    val logic: GraphStageLogic = new GraphStageLogic(shape) with InHandler {
      setHandler(in, this)

      implicit def context: ExecutionContextExecutor = super.materializer.executionContext

      var producer: Producer[T] = _
      var next: AsyncCallback[ProducerMessage[T]] = _
      var error: Throwable = _

      override def preStart(): Unit = {
        producer = createFn()
        next = getAsyncCallback { _ => pull(in) }
        pull(in)
      }

      override def onPush(): Unit = {
        try {
          val t = grab(in)
          logger.debug(s"Sending message $t")
          producer.sendAsync(t).onComplete {
            case Success(_) => next.invoke(t)
            case Failure(e) =>
              logger.error("Failing pulsar sink stage", e)
              failStage(e)
          }
        } catch {
          case e: Throwable =>
            logger.error("Failing pulsar sink stage", e)
            failStage(e)
        }
      }

      override def postStop(): Unit = {
        logger.debug("Graph stage stopping; closing producer")
        producer.flush()
        producer.close()
      }

      override def onUpstreamFailure(ex: Throwable): Unit = {
        promise.tryFailure(ex)
      }

      override def onUpstreamFinish(): Unit = {
        promise.trySuccess(Done)
      }
    }

    (logic, promise.future)
  }
} 

package org.squbs.pattern.stream

import java.io.File

import akka.actor.{ActorSystem, Props}
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
import akka.stream.{Attributes, FlowShape, Inlet, Outlet}
import com.typesafe.config.Config
import com.typesafe.scalalogging.Logger
import org.slf4j.LoggerFactory

abstract class PersistentBufferBase[T, S] (private[stream] val queue: PersistentQueue[T],
                                      onPushCallback: () => Unit = () => {})
                                     (implicit serializer: QueueSerializer[T],
                                  system: ActorSystem) extends GraphStage[FlowShape[T, S]] {

  def this(config: Config)(implicit serializer: QueueSerializer[T], system: ActorSystem) =
    this(new PersistentQueue[T](config))

  def this(persistDir: File)(implicit serializer: QueueSerializer[T], system: ActorSystem) =
    this(new PersistentQueue[T](persistDir))

  private[stream] val in = Inlet[T]("PersistentBuffer.in")
  private[stream] val out = Outlet[S]("PersistentBuffer.out")
  val shape: FlowShape[T, S] = FlowShape.of(in, out)
  val defaultOutputPort = 0
  @volatile protected var upstreamFailed = false
  @volatile protected var upstreamFinished = false
  protected val queueCloserActor = system.actorOf(Props(classOf[PersistentQueueCloserActor[T]], queue))

  protected def elementOut(e: Event[T]): S

  protected def autoCommit(index: Long) = {}

  def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) {

    private var lastPushed = 0L
    var downstreamWaiting = false

    override def preStart(): Unit = {
      // Start upstream demand
      pull(in)
    }

    setHandler(in, new InHandler {

      override def onPush(): Unit = {
        val element = grab(in)
        queue.enqueue(element)
        onPushCallback()
        if (downstreamWaiting) {
          queue.dequeue() foreach { element =>
            push(out, elementOut(element))
            downstreamWaiting = false
            lastPushed = element.index
            autoCommit(element.index)
          }
        }
        pull(in)
      }

      override def onUpstreamFinish(): Unit = {
        upstreamFinished = true

        if (downstreamWaiting) {
          queueCloserActor ! PushedAndCommitted(defaultOutputPort, lastPushed, queue.read(defaultOutputPort))
          queueCloserActor ! UpstreamFinished
          completeStage()
        }
      }

      override def onUpstreamFailure(ex: Throwable): Unit = {
        val logger = Logger(LoggerFactory.getLogger(this.getClass))
        logger.error("Received upstream failure signal: " + ex)
        upstreamFailed = true
        queueCloserActor ! UpstreamFailed
        completeStage()
      }
    })

    setHandler(out, new OutHandler {

      override def onPull(): Unit = {
        queue.dequeue() match {
          case Some(element) =>
            push(out, elementOut(element))
            lastPushed = element.index
            autoCommit(element.index)
          case None =>
            if (upstreamFinished) {
              queueCloserActor ! PushedAndCommitted(defaultOutputPort, lastPushed, queue.read(defaultOutputPort))
              queueCloserActor ! UpstreamFinished
              completeStage()
            } else downstreamWaiting = true
        }
      }
    })
  }
} 

package org.squbs.stream

import akka.stream._
import akka.stream.scaladsl.{GraphDSL, Source}
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}


final class TriggerMerge[T](eagerComplete: Boolean = false) extends GraphStage[FanInShape2[T, TriggerEvent, T]] {
  override val shape: FanInShape2[T, TriggerEvent, T] = new FanInShape2[T, TriggerEvent, T]("TriggerMerge")
  val in: Inlet[T] = shape.in0
  val trigger: Inlet[TriggerEvent] = shape.in1

  override def initialAttributes = Attributes.name("TriggerMerge")

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) {
    var flowEnabled, requestPending, triggerCompleted, willShutDown = false

    setHandler(in, new InHandler {
      override def onPush(): Unit = {
        push(out, grab(in))
        if (willShutDown) completeStage()
      }

      override def onUpstreamFinish(): Unit = {
        if (!isAvailable(in)) completeStage()
        willShutDown = true
      }
    })

    setHandler(trigger, new InHandler {
      override def onPush(): Unit = {
        val triggerEvent = grab(trigger)
        if (triggerEvent.value != flowEnabled) {
          if (triggerEvent.value && requestPending) {
            requestPending = false
            tryPull(in)
          }
          flowEnabled = triggerEvent.value
        }
        if (triggerCompleted) cancelTrigger()
        if (!hasBeenPulled(trigger)) tryPull(trigger)
      }

      override def onUpstreamFinish(): Unit = {
        if (!isAvailable(trigger)) cancelTrigger()
        triggerCompleted = true
      }

      def cancelTrigger(): Unit = if (eagerComplete) completeStage() else cancel(trigger)
    })

    setHandler(out, new OutHandler {
      override def onPull(): Unit = {
        if (willShutDown) completeStage()
        else {
          if (flowEnabled) tryPull(in)
          else {
            requestPending = true
            if (!hasBeenPulled(trigger)) tryPull(trigger)
          }
        }
      }
    })
  }

  def out: Outlet[T] = shape.out

  override def toString = "TriggerMerge"
}

class Trigger[T, M1, M2](eagerComplete: Boolean = false) {

  import GraphDSL.Implicits._

  private[stream] val source =
    (in: Graph[SourceShape[T], M1], trigger: Graph[SourceShape[TriggerEvent], M2]) => Source.fromGraph(
      GraphDSL.create(in, trigger)((_, _)) { implicit builder =>
        (sIn, sTrigger) =>
          val merge = builder.add(new TriggerMerge[T](eagerComplete))
          sIn ~> merge.in0
          sTrigger ~> merge.in1
          SourceShape(merge.out)
      })

  // for Java
  def source(in: javadsl.Source[T, M1], trigger: javadsl.Source[TriggerEvent, M2]):
      javadsl.Source[T, akka.japi.Pair[M1, M2]] = {
    source(in.asScala, trigger.asScala).mapMaterializedValue {
      case (m1, m2) => akka.japi.Pair(m1, m2)
    }.asJava
  }
} 

package org.squbs.streams

import akka.event.Logging
import akka.stream.ActorAttributes.SupervisionStrategy
import akka.stream._
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}

import scala.compat.java8.FunctionConverters._
import scala.util.control.NonFatal

object Deduplicate {

  def apply[T, U](key: T => U, duplicateCount: Long) =
    new Deduplicate[T, U](key, duplicateCount, new java.util.HashMap[U, MutableLong]())

  def apply[T, U](key: T => U, duplicateCount: Long, registry: java.util.Map[U, MutableLong]) =
    new Deduplicate[T, U](key, duplicateCount, registry)

  def apply[T](duplicateCount: Long = Long.MaxValue,
               registry: java.util.Map[T, MutableLong] = new java.util.HashMap[T, MutableLong]()): Deduplicate[T, T] =
    Deduplicate(t => t, duplicateCount, registry)
}


case class MutableLong(var value: Long = 0L) {
  def increment() = {
    value += 1
    value
  }
} 

package cmwell.dc.stream

import akka.stream.{Attributes, FlowShape, Inlet, Outlet}
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
import cmwell.dc.LazyLogging
import cmwell.dc.stream.MessagesTypesAndExceptions.{DcInfo, DcInfoKey}


object RatePrinter {
  def apply[A](dcKey: DcInfoKey,
               elementCount: A => Double,
               elementText: String,
               elementsText: String,
               printFrequency: Int): RatePrinter[A] =
    new RatePrinter(dcKey, elementCount, elementText, elementsText, printFrequency)
}

class RatePrinter[A](dcKey: DcInfoKey,
                     elementCount: A => Double,
                     elementText: String,
                     elementsText: String,
                     printFrequency: Int)
    extends GraphStage[FlowShape[A, A]]
    with LazyLogging {
  val in = Inlet[A]("RatePrinter.in")
  val out = Outlet[A]("RatePrinter.out")
  override val shape = FlowShape.of(in, out)

  override def createLogic(attr: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) {
      private val startTime: Double = System.currentTimeMillis
      private var totalElementsGot: Double = 0
      private var printedAtElementNo: Double = 0
      private var localStartTime: Double = startTime
      private var localTotalElementsGot: Double = 0
      private var localRate: Double = 0

      setHandler(
        in,
        new InHandler {
          override def onPush(): Unit = {
            val elem = grab(in)
            val currentTime = System.currentTimeMillis
            val currentElementsGot = elementCount(elem)
            totalElementsGot += currentElementsGot
            localTotalElementsGot += currentElementsGot
            if (totalElementsGot - printedAtElementNo >= printFrequency) {
              val rate = totalElementsGot / (currentTime - startTime) * 1000
              logger.info(
                s"Sync $dcKey: Got ${currentElementsGot.formatted("%.2f")} $elementText. Total ${totalElementsGot
                  .formatted("%.2f")} $elementsText. Read rate: avg: ${rate.formatted("%.2f")} current: ${localRate
                  .formatted("%.2f")} $elementText/second"
              )
              if (currentTime - localStartTime > 15000) {
                localRate = localTotalElementsGot / (currentTime - localStartTime) * 1000
                localTotalElementsGot = 0
                localStartTime = currentTime
              }
              printedAtElementNo = totalElementsGot
            }
            push(out, elem)
          }
        }
      )
      setHandler(out, new OutHandler {
        override def onPull(): Unit = {
          pull(in)
        }
      })
    }
} 

package uk.gov.hmrc.akka.xml

import akka.NotUsed
import akka.stream.scaladsl.Flow
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
import akka.stream.{Attributes, FlowShape, Inlet, Outlet}
import akka.util.ByteString


@deprecated("Use FastParsingStage instead","akka-xml-parser 1.0.0")
object MinimumChunk {

  def parser(minimumChunkSize: Int):
  Flow[ByteString, ByteString, NotUsed] = {
    Flow.fromGraph(new StreamingXmlParser(minimumChunkSize))
  }

  private class StreamingXmlParser(minimumChunkSize: Int)
    extends GraphStage[FlowShape[ByteString, ByteString]]
      with StreamHelper
      with ParsingDataFunctions {

    val in: Inlet[ByteString] = Inlet("Chunking.in")
    val out: Outlet[ByteString] = Outlet("Chunking.out")
    override val shape: FlowShape[ByteString, ByteString] = FlowShape(in, out)

    override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
      new GraphStageLogic(shape) {
        private var buffer = ByteString.empty

        setHandler(in, new InHandler {
          override def onPush(): Unit = {
            val elem = grab(in)
            buffer ++= elem
            emitChunk()
          }

          override def onUpstreamFinish(): Unit = {
            emit(out, buffer)
            completeStage()
          }
        })

        setHandler(out, new OutHandler {
          override def onPull(): Unit = {
            pull(in)
          }
        })

        private def emitChunk(): Unit = {
          if (buffer.size > minimumChunkSize) {
            push(out, buffer)
            buffer = ByteString.empty
          } else {
            pull(in)
          }
        }

      }
  }

} 

package io.vamp.common.http

import akka.stream.stage.{ GraphStage, GraphStageLogic, InHandler, OutHandler }
import akka.stream.{ Attributes, FlowShape, Inlet, Outlet }


class TerminateFlowStage[T](pred: T ⇒ Boolean, forwardTerminatingMessage: Boolean = false, terminate: Boolean = true) extends GraphStage[FlowShape[T, T]] {

  val in = Inlet[T]("TerminateFlowStage.in")
  val out = Outlet[T]("TerminateFlowStage.out")

  override val shape = FlowShape.of(in, out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) {

      setHandlers(in, out, new InHandler with OutHandler {
        override def onPull(): Unit = {
          pull(in)
        }

        override def onPush(): Unit = {
          val chunk = grab(in)

          if (pred(chunk)) {
            if (forwardTerminatingMessage)
              push(out, chunk)
            if (terminate)
              failStage(new RuntimeException("Flow terminated by TerminateFlowStage"))
            else
              completeStage()
          }
          else push(out, chunk)
        }
      })
    }
} 

package mass.connector.sql

import java.sql.{ Connection, PreparedStatement, ResultSet }

import akka.stream.stage.{ GraphStage, GraphStageLogic, OutHandler }
import akka.stream.{ Attributes, Outlet, SourceShape }
import javax.sql.DataSource
import fusion.jdbc.ConnectionPreparedStatementCreator
import fusion.jdbc.util.JdbcUtils

import scala.util.control.NonFatal

class JdbcSourceStage(dataSource: DataSource, creator: ConnectionPreparedStatementCreator, fetchRowSize: Int)
    extends GraphStage[SourceShape[ResultSet]] {
  private val out: Outlet[ResultSet] = Outlet("JdbcSource.out")

  override def shape: SourceShape[ResultSet] = SourceShape(out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) with OutHandler {
      var maybeConn =
        Option.empty[(Connection, Boolean, PreparedStatement, ResultSet)]

      setHandler(out, this)

      override def onPull(): Unit =
        maybeConn match {
          case Some((_, _, _, rs)) if rs.next() =>
            push(out, rs)
          case Some(_) =>
            completeStage()
          case None =>
            () // doing nothing, waiting for in preStart() to be completed
        }

      override def preStart(): Unit =
        try {
          val conn = dataSource.getConnection
          val autoCommit = conn.getAutoCommit
          conn.setAutoCommit(false)
          val stmt = creator(conn)
          val rs = stmt.executeQuery()
          //          rs.setFetchDirection(ResultSet.TYPE_FORWARD_ONLY)
          rs.setFetchSize(fetchRowSize)
          maybeConn = Option((conn, autoCommit, stmt, rs))
        } catch {
          case NonFatal(e) => failStage(e)
        }

      override def postStop(): Unit =
        for {
          (conn, autoCommit, stmt, rs) <- maybeConn
        } {
          JdbcUtils.closeResultSet(rs)
          JdbcUtils.closeStatement(stmt)
          conn.setAutoCommit(autoCommit)
          JdbcUtils.closeConnection(conn)
        }
    }
} 

package ackcord.util

import java.util.concurrent.atomic.AtomicBoolean

import scala.collection.immutable

import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
import akka.stream.{Attributes, FanInShape2, Inlet, Outlet}

//TODO: Maybe use a third inlet to determine where to listen to
class Switch[A](ref: AtomicBoolean, emitChangeTrue: immutable.Seq[A], emitChangeFalse: immutable.Seq[A])
    extends GraphStage[FanInShape2[A, A, A]] {
  override val shape: FanInShape2[A, A, A] = new FanInShape2[A, A, A]("Switch")

  val in1: Inlet[A]  = shape.in0
  val in2: Inlet[A]  = shape.in1
  val out: Outlet[A] = shape.out

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) with OutHandler {

      private var lastState: Boolean = ref.get()
      private var waitingOther: A    = _

      private def activeIn(): Inlet[A] = {
        val newState = ref.get()
        val newIn    = if (newState) in1 else in2

        if (lastState != newState) {
          lastState = newState

          emitMultiple(out, if (newState) emitChangeTrue else emitChangeFalse)

          if (waitingOther != null) {
            emit(out, waitingOther)
            waitingOther = null.asInstanceOf[A]
          }

          tryPull(newIn)
        }

        newIn
      }

      private def setInHandler(in: Inlet[A]): Unit = {
        setHandler(
          in,
          new InHandler {
            override def onPush(): Unit = {

              if (activeIn() == in) {
                emit(out, grab(in))
              } else {
                require(waitingOther == null, "Pushed other when a waiting other was already defined")
                waitingOther = grab(in)
              }
            }
          }
        )
      }

      setInHandler(in1)
      setInHandler(in2)

      setHandler(out, this)

      override def onPull(): Unit = pull(activeIn())
    }
} 

package ackcord.lavaplayer

import scala.concurrent.duration._

import akka.NotUsed
import akka.stream.scaladsl.Source
import akka.stream.stage.{GraphStage, GraphStageLogic, OutHandler, TimerGraphStageLogicWithLogging}
import akka.stream.{Attributes, Outlet, SourceShape}
import akka.util.ByteString
import com.sedmelluq.discord.lavaplayer.player.AudioPlayer

class LavaplayerSource(player: AudioPlayer) extends GraphStage[SourceShape[ByteString]] {
  val out: Outlet[ByteString] = Outlet("LavaplayerSource.out")

  override def shape: SourceShape[ByteString] = SourceShape(out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new TimerGraphStageLogicWithLogging(shape) with OutHandler {
      override def onPull(): Unit =
        tryPushFrame()

      override protected def onTimer(timerKey: Any): Unit = timerKey match {
        case "RetryProvide" => tryPushFrame()
      }

      def tryPushFrame(): Unit = {
        val frame = player.provide()
        if (frame != null) {
          push(out, ByteString.fromArray(frame.getData))
          //log.debug("Sending data")
        } else {
          //log.debug("Scheduling attempt to provide frame")
          scheduleOnce("RetryProvide", 20.millis)
        }
      }

      setHandler(out, this)
    }
}
object LavaplayerSource {

  def source(player: AudioPlayer): Source[ByteString, NotUsed] =
    Source.fromGraph(new LavaplayerSource(player))
} 

package ackcord.util

import akka.NotUsed
import akka.stream.scaladsl.Flow
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
import akka.stream.{Attributes, FlowShape, Inlet, Outlet}

class RepeatLast[A] extends GraphStage[FlowShape[A, A]] {
  val in: Inlet[A]                    = Inlet("RepeatLast.in")
  val out: Outlet[A]                  = Outlet("RepeatLast.out")
  override def shape: FlowShape[A, A] = FlowShape(in, out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) with InHandler with OutHandler {
      var elem: A = _

      override def onPush(): Unit = {
        elem = grab(in)
        push(out, elem)
      }

      override def onPull(): Unit = {
        if (elem != null && isAvailable(out)) push(out, elem)
        if (!hasBeenPulled(in)) pull(in)
      }

      setHandlers(in, out, this)
    }
}
object RepeatLast {

  
  def flow[A]: Flow[A, A, NotUsed] = Flow.fromGraph(new RepeatLast[A])
} 

package sangria.streaming

import scala.language.higherKinds
import akka.NotUsed
import akka.event.Logging
import akka.stream.ActorAttributes.SupervisionStrategy
import akka.stream._
import akka.stream.scaladsl.{Merge, Sink, Source}
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}

import scala.concurrent.Future

object akkaStreams {
  type AkkaSource[+T] = Source[T, NotUsed]

  abstract class SimpleLinearGraphStage[T] extends GraphStage[FlowShape[T, T]] {
    val in = Inlet[T](Logging.simpleName(this) + ".in")
    val out = Outlet[T](Logging.simpleName(this) + ".out")
    override val shape = FlowShape(in, out)
  }

  class AkkaStreamsSubscriptionStream(implicit materializer: Materializer) extends SubscriptionStream[AkkaSource] {
    def supported[T[_]](other: SubscriptionStream[T]) = other.isInstanceOf[AkkaStreamsSubscriptionStream]

    def map[A, B](source: AkkaSource[A])(fn: A => B) = source.map(fn)

    def singleFuture[T](value: Future[T]) = Source.fromFuture(value)

    def single[T](value: T) = Source.single(value)

    def mapFuture[A, B](source: AkkaSource[A])(fn: A => Future[B]) =
      source.mapAsync(1)(fn)

    def first[T](s: AkkaSource[T]) = s.runWith(Sink.head)

    def failed[T](e: Throwable) = Source.failed(e).asInstanceOf[AkkaSource[T]]

    def onComplete[Ctx, Res](result: AkkaSource[Res])(op: => Unit) =
      result
        .via(OnComplete(() => op))
        .recover {case e => op; throw e}
        .asInstanceOf[AkkaSource[Res]]

    def flatMapFuture[Ctx, Res, T](future: Future[T])(resultFn: T => AkkaSource[Res]) =
      Source.fromFuture(future).flatMapMerge(1, resultFn)

    def merge[T](streams: Vector[AkkaSource[T]]) = {
      if (streams.size > 1)
        Source.combine(streams(0), streams(1), streams.drop(2): _*)(Merge(_))
      else if (streams.nonEmpty)
        streams.head
      else
        throw new IllegalStateException("No streams produced!")
    }

    def recover[T](stream: AkkaSource[T])(fn: Throwable => T) =
      stream recover {case e => fn(e)}
  }

  implicit def akkaSubscriptionStream(implicit materializer: Materializer): SubscriptionStream[AkkaSource] = new AkkaStreamsSubscriptionStream

  implicit def akkaStreamIsValidSubscriptionStream[A[_, _], Ctx, Res, Out](implicit materializer: Materializer, ev1: ValidOutStreamType[Res, Out]): SubscriptionStreamLike[Source[A[Ctx, Res], NotUsed], A, Ctx, Res, Out] =
    new SubscriptionStreamLike[Source[A[Ctx, Res], NotUsed], A, Ctx, Res, Out] {
      type StreamSource[X] = AkkaSource[X]
      val subscriptionStream = new AkkaStreamsSubscriptionStream
    }

  private final case class OnComplete[T](op: () => Unit) extends SimpleLinearGraphStage[T] {
    override def toString: String = "OnComplete"

    override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
      new GraphStageLogic(shape) with OutHandler with InHandler {
        def decider = inheritedAttributes.get[SupervisionStrategy].map(_.decider).getOrElse(Supervision.stoppingDecider)

        override def onPush(): Unit = {
          push(out, grab(in))
        }

        override def onPull(): Unit = pull(in)

        override def onDownstreamFinish() = {
          op()
          super.onDownstreamFinish()
        }

        override def onUpstreamFinish() = {
          op()
          super.onUpstreamFinish()
        }

        setHandlers(in, out, this)
      }
  }
} 

package com.netflix.iep.lwc

import akka.stream.Attributes
import akka.stream.FlowShape
import akka.stream.Inlet
import akka.stream.Outlet
import akka.stream.stage.GraphStage
import akka.stream.stage.GraphStageLogic
import akka.stream.stage.InHandler
import akka.stream.stage.OutHandler
import com.netflix.iep.lwc.ForwardingService.Message
import com.netflix.iep.lwc.fwd.cw.ClusterConfig
import com.typesafe.scalalogging.StrictLogging

class ConfigManager
    extends GraphStage[FlowShape[Message, Map[String, ClusterConfig]]]
    with StrictLogging {

  private val in = Inlet[Message]("ConfigManager.in")
  private val out = Outlet[Map[String, ClusterConfig]]("ConfigManager.out")

  override val shape: FlowShape[Message, Map[String, ClusterConfig]] = FlowShape(in, out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = {
    new GraphStageLogic(shape) with InHandler with OutHandler {

      private val configs = scala.collection.mutable.AnyRefMap.empty[String, ClusterConfig]

      override def onPush(): Unit = {
        val msg = grab(in)
        if (msg.response.isUpdate) {
          val cluster = msg.cluster
          try {
            configs += cluster -> msg.response.clusterConfig
            logger.info(s"updated configuration for cluster $cluster")
          } catch {
            case e: Exception =>
              logger.warn(s"invalid config for cluster $cluster", e)
          }
        } else {
          configs -= msg.cluster
          logger.info(s"deleted configuration for cluster ${msg.cluster}")
        }
        push(out, configs.toMap)
      }

      override def onPull(): Unit = {
        pull(in)
      }

      override def onUpstreamFinish(): Unit = {
        completeStage()
      }

      setHandlers(in, out, this)
    }
  }
} 

package com.netflix.atlas.persistence

import java.nio.file.Files
import java.nio.file.Paths

import akka.NotUsed
import akka.stream.Attributes
import akka.stream.FlowShape
import akka.stream.Inlet
import akka.stream.Outlet
import akka.stream.stage.GraphStage
import akka.stream.stage.GraphStageLogic
import akka.stream.stage.InHandler
import akka.stream.stage.OutHandler
import akka.stream.stage.TimerGraphStageLogic
import com.netflix.atlas.core.model.Datapoint
import com.netflix.spectator.api.Registry
import com.typesafe.scalalogging.StrictLogging

import scala.concurrent.duration._

class RollingFileFlow(
  val dataDir: String,
  val rollingConf: RollingConfig,
  val registry: Registry
) extends GraphStage[FlowShape[Datapoint, NotUsed]]
    with StrictLogging {

  private val in = Inlet[Datapoint]("RollingFileSink.in")
  private val out = Outlet[NotUsed]("RollingFileSink.out")
  override val shape = FlowShape(in, out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = {

    new TimerGraphStageLogic(shape) with InHandler with OutHandler {

      private var hourlyWriter: HourlyRollingWriter = _

      override def preStart(): Unit = {
        logger.info(s"creating sink directory: $dataDir")
        Files.createDirectories(Paths.get(dataDir))

        hourlyWriter = new HourlyRollingWriter(dataDir, rollingConf, registry)
        hourlyWriter.initialize
        // This is to trigger rollover check when writer is idle for long time: e.g. in most cases
        // file writer will be idle while hour has ended but it is still waiting for late events
        schedulePeriodically(None, 5.seconds)
      }

      override def onPush(): Unit = {
        hourlyWriter.write(grab(in))
        pull(in)
      }

      override protected def onTimer(timerKey: Any): Unit = {
        hourlyWriter.write(RollingFileWriter.RolloverCheckDatapoint)
      }

      override def onUpstreamFinish(): Unit = {
        hourlyWriter.close()
        completeStage()
      }

      override def onUpstreamFailure(ex: Throwable): Unit = {
        hourlyWriter.close()
        failStage(ex)
      }

      setHandlers(in, out, this)

      override def onPull(): Unit = {
        // Nothing to emit
        pull(in)
      }
    }
  }
} 

package com.packt.chapter8

import akka.stream.{Attributes, FlowShape, Inlet, Outlet}
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
import com.packt.chapter8.WorkingWithGraphsApplication.GenericMsg

class StatefulCounterFlow extends GraphStage[FlowShape[Seq[GenericMsg], Int]] {

  val in: Inlet[Seq[GenericMsg]] = Inlet("IncomingGenericMsg")
  val out: Outlet[Int] = Outlet("OutgoingCount")
  override val shape: FlowShape[Seq[GenericMsg], Int] = FlowShape(in, out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) {
      var count = 0

      setHandler(in, new InHandler {
        override def onPush() = {
          val elem = grab(in)
          count += elem.size
          push(out, count)
        }
      })

      setHandler(out, new OutHandler {
        override def onPull() = {
          pull(in)
        }
      })
    }
} 

// Copyright (c) 2020 Digital Asset (Switzerland) GmbH and/or its affiliates. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

package com.daml.platform.server.api

import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
import akka.stream.{Attributes, FlowShape, Inlet, Outlet}

object DropRepeated {
  def apply[T](): GraphStage[FlowShape[T, T]] = new DropRepeated
}

final class DropRepeated[T] extends GraphStage[FlowShape[T, T]] {
  private val in = Inlet[T]("input")
  private val out = Outlet[T]("DropRepeated output")

  override def shape: FlowShape[T, T] = FlowShape(in, out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) {
      private var currentValue: Option[T] = None

      setHandler(
        in,
        new InHandler {
          override def onPush(): Unit = {
            val element = grab(in)
            if (currentValue.contains(element)) {
              pull(in)
            } else {
              currentValue = Some(element)
              push(out, element)
            }
          }
        }
      )

      setHandler(out, new OutHandler {
        override def onPull(): Unit = {
          pull(in)
        }
      })
    }
} 

// Copyright (c) 2020 Digital Asset (Switzerland) GmbH and/or its affiliates. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

package com.daml.ledger.client.binding.util

import akka.stream._
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
import org.slf4j.Logger

final case class Slf4JLogger[T, U](
    logger: Logger,
    prefix: String,
    project: T => U,
    logDemand: Boolean = false)
    extends GraphStage[FlowShape[T, T]] {

  override def toString = "Slf4JLog"

  val in: Inlet[T] = Inlet[T]("in")
  val out: Outlet[T] = Outlet[T]("out")

  override def shape: FlowShape[T, T] = FlowShape(in, out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) with OutHandler with InHandler {

      override def onPush(): Unit = {

        val elem = grab(in)
        if (logger.isDebugEnabled) logger.debug("[{}] Element: {}", prefix, project(elem))
        push(out, elem)
      }

      override def onPull(): Unit = {
        if (logDemand) logger.debug("[{}] Demand", prefix)
        pull(in)
      }

      override def onUpstreamFailure(cause: Throwable): Unit = {
        logger.warn(s"[$prefix] Upstream failed", cause)

        super.onUpstreamFailure(cause)
      }

      override def onUpstreamFinish(): Unit = {
        logger.debug("[{}] Upstream finished.", prefix)

        super.onUpstreamFinish()
      }

      override def onDownstreamFinish(cause: Throwable): Unit = {
        logger.debug("[{}] Downstream finished.", prefix)

        super.onDownstreamFinish(cause)
      }

      setHandlers(in, out, this)
    }
}

object Slf4JLogger {
  def apply[T](logger: Logger, prefix: String): Slf4JLogger[T, T] =
    new Slf4JLogger(logger, prefix, identity)
} 

package me.snov.akka.sqs.shape

import akka.Done
import akka.stream._
import akka.stream.stage.{GraphStageLogic, InHandler}
import com.amazonaws.handlers.AsyncHandler
import com.amazonaws.services.sqs.model._
import me.snov.akka.sqs.client.SqsClient

import scala.concurrent.Promise

private[sqs] class SqsPublishSinkGraphStageLogic(
                                                  client: SqsClient,
                                                  in: Inlet[SendMessageRequest],
                                                  shape: SinkShape[SendMessageRequest],
                                                  promise: Promise[Done]
                                                ) extends GraphStageLogic(shape) with StageLogging {

  private val MaxInFlight = 10
  private var inFlight = 0
  private var isShutdownInProgress = false
  private var amazonSendMessageHandler: AsyncHandler[SendMessageRequest, SendMessageResult] = _

  setHandler(in, new InHandler {
    override def onPush(): Unit = {
      inFlight += 1
      client.sendMessageAsync(grab(in), amazonSendMessageHandler)

      tryPull()
    }

    @scala.throws[Exception](classOf[Exception])
    override def onUpstreamFailure(exception: Throwable): Unit = {
      log.error(exception, "Upstream failure: {}", exception.getMessage)
      failStage(exception)
      promise.tryFailure(exception)
    }

    @scala.throws[Exception](classOf[Exception])
    override def onUpstreamFinish(): Unit = {
      log.debug("Upstream finish")
      isShutdownInProgress = true
      tryShutdown()
    }
  })

  override def preStart(): Unit = {
    setKeepGoing(true)

    val failureCallback = getAsyncCallback[Throwable](handleFailure)
    val sendCallback = getAsyncCallback[SendMessageResult](handleResult)

    amazonSendMessageHandler = new AsyncHandler[SendMessageRequest, SendMessageResult] {
      override def onError(exception: Exception): Unit =
        failureCallback.invoke(exception)

      override def onSuccess(request: SendMessageRequest, result: SendMessageResult): Unit =
        sendCallback.invoke(result)
    }

    // This requests one element at the Sink startup.
    pull(in)
  }

  private def tryShutdown(): Unit =
    if (isShutdownInProgress && inFlight <= 0) {
      completeStage()
      promise.trySuccess(Done)
    }

  private def tryPull(): Unit =
    if (inFlight < MaxInFlight && !isClosed(in) && !hasBeenPulled(in)) {
      pull(in)
    }

  private def handleFailure(exception: Throwable): Unit = {
    log.error(exception, "Client failure: {}", exception.getMessage)
    inFlight -= 1
    failStage(exception)
    promise.tryFailure(exception)
  }

  private def handleResult(result: SendMessageResult): Unit = {
    log.debug(s"Sent message {}", result.getMessageId)
    inFlight -= 1
    tryShutdown()
    tryPull()
  }
} 

package me.snov.akka.sqs.shape

import akka.event.{LoggingAdapter, NoLogging}
import akka.stream.ActorMaterializer
import akka.stream.stage.GraphStageLogic

private[sqs] trait StageLogging { self: GraphStageLogic =>

  private var loggingAdapter: LoggingAdapter = _

  def log: LoggingAdapter = {
    if (loggingAdapter eq null) {
      materializer match {
        case actorMaterializer: ActorMaterializer =>
          loggingAdapter = akka.event.Logging(actorMaterializer.system, self.getClass)
        case _ =>
          loggingAdapter = NoLogging
      }
    }

    loggingAdapter
  }
} 

package me.snov.akka.sqs.shape

import java.util

import akka.stream._
import akka.stream.stage.{AsyncCallback, GraphStageLogic, OutHandler}
import com.amazonaws.handlers.AsyncHandler
import com.amazonaws.services.sqs.model.{Message, ReceiveMessageRequest, ReceiveMessageResult}
import me.snov.akka.sqs.client.SqsClient

import scala.concurrent.duration._

private[sqs] class SqsSourceGraphStageLogic(client: SqsClient, out: Outlet[Message], shape: SourceShape[Message])
  extends GraphStageLogic(shape) with StageLogging {

  private val buffer: util.List[Message] = new util.ArrayList[Message]()
  private var handleMessagesCallback: AsyncCallback[util.List[Message]] = _
  private var awsReceiveMessagesHandler: AsyncHandler[ReceiveMessageRequest, ReceiveMessageResult] = _
  private var asyncReceiveMessagesIsInProgress = false
  private val errorCooldown = 5.seconds

  override def preStart(): Unit = {
    handleMessagesCallback = getAsyncCallback[util.List[Message]](handleMessages)
    awsReceiveMessagesHandler = new AsyncHandler[ReceiveMessageRequest, ReceiveMessageResult] {

      override def onError(exception: Exception): Unit = {
        log.error(exception, exception.getMessage)
        materializer.scheduleOnce(errorCooldown, new Runnable {
          override def run(): Unit = handleMessagesCallback.invoke(new util.ArrayList[Message]())
        })
      }

      override def onSuccess(request: ReceiveMessageRequest, result: ReceiveMessageResult): Unit =
        handleMessagesCallback.invoke(result.getMessages)
    }
  }

  private def handleMessages(messages: util.List[Message]): Unit = {
    asyncReceiveMessagesIsInProgress = false
    buffer.addAll(messages)
    getHandler(out).onPull()
  }

  private def loadMessagesAsync() = {
    if (!asyncReceiveMessagesIsInProgress) {
      asyncReceiveMessagesIsInProgress = true
      client.receiveMessageAsync(awsReceiveMessagesHandler)
    }
  }

  setHandler(out, new OutHandler {
    override def onPull(): Unit = {
      if (!buffer.isEmpty && isAvailable(out)) {
        push(shape.out, buffer.remove(0))
      }

      if (buffer.isEmpty) {
        loadMessagesAsync()
      }
    }
  })
} 

package me.snov.akka.sqs.shape

import akka.Done
import akka.stream._
import akka.stream.stage.{GraphStageLogic, GraphStageWithMaterializedValue}
import com.amazonaws.services.sqs.model.{SendMessageRequest, SendMessageResult}
import me.snov.akka.sqs.client.{SqsClient, SqsSettings}

import scala.concurrent.{Future, Promise}

object SqsPublishSinkShape {
  def apply(settings: SqsSettings): SqsPublishSinkShape = apply(SqsClient(settings))

  def apply(client: SqsClient): SqsPublishSinkShape = new SqsPublishSinkShape(client)
}

class SqsPublishSinkShape(client: SqsClient)
  extends GraphStageWithMaterializedValue[SinkShape[SendMessageRequest], Future[Done]] {

  val in: Inlet[SendMessageRequest] = Inlet("SqsPublishSinkShape.in")

  override val shape: SinkShape[SendMessageRequest] = SinkShape(in)

  override def createLogicAndMaterializedValue(inheritedAttributes: Attributes):
  (GraphStageLogic, Future[Done]) = {
    val promise = Promise[Done]()
    val logic = new SqsPublishSinkGraphStageLogic(client, in, shape, promise)

    (logic, promise.future)
  }
} 

package akka.stream.integration
package activemq

import akka.actor.ActorRef
import akka.camel.CamelMessage
import akka.stream._
import akka.stream.stage.{ GraphStage, GraphStageLogic, InHandler, OutHandler }

import scala.concurrent.{ ExecutionContext, Future, Promise }

private[activemq] class AckedFlow[A, B](implicit ec: ExecutionContext) extends GraphStage[FlowShape[(ActorRef, B), AckTup[A, B]]] {
  val in = Inlet[(ActorRef, B)]("AckedFlow.in")
  val out = Outlet[AckTup[A, B]]("AckedFlow.out")

  override val shape: FlowShape[(ActorRef, B), AckTup[A, B]] = FlowShape.of(in, out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) {
    var promises = Vector.empty[(Promise[A], Future[A])]
    setHandler(in, new InHandler {
      override def onPush(): Unit = {
        val (ref, b) = grab(in)
        val p = Promise[A]()
        val eventualResponse = p.future
        eventualResponse.onSuccess(successResponse(ref))
        eventualResponse.onFailure {
          case cause: Throwable =>
            ref ! akka.actor.Status.Failure(cause)
        }
        promises = promises.filterNot(_._1.isCompleted) :+ (p -> eventualResponse)
        push(out, p -> b)
      }
    })

    setHandler(out, new OutHandler {
      override def onPull(): Unit = {
        pull(in)
      }
    })
  }

  
  def successResponse(source: ActorRef): PartialFunction[A, Unit] = {
    case _ => source ! akka.camel.Ack
  }
}

class AckedResponseFlow[A, B](implicit ec: ExecutionContext, builder: MessageBuilder[A, CamelMessage]) extends AckedFlow[A, B] {
  override def successResponse(source: ActorRef): PartialFunction[A, Unit] = {
    case msg => source ! builder.build(msg)
  }
} 

package grpc.akkastreams

import akka.stream.{Attributes, Inlet, SinkShape}
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler}
import io.grpc.stub.CallStreamObserver

class GrpcSinkStage[I](observer: CallStreamObserver[I]) extends GraphStage[SinkShape[I]] {
  val in = Inlet[I]("grpc.in")
  override val shape: SinkShape[I] = SinkShape.of(in)
  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) with InHandler with Runnable {
      var element: Option[I] = None

      override def run(): Unit = getAsyncCallback((_: Unit) => {
        element match {
          case Some(value) if observer.isReady =>
            observer.onNext(value)
            tryPull(in)
          case _ => ()
        }
      }).invoke(())

      override def onPush(): Unit = {
        val value = grab(in)
        if (observer.isReady) {
          observer.onNext(value)
          pull(in)
        } else element = Some(value)
      }

      override def onUpstreamFinish(): Unit = observer.onCompleted()

      override def onUpstreamFailure(t: Throwable): Unit = observer.onError(t)

      override def preStart(): Unit = pull(in)

      observer.setOnReadyHandler(this)
      setHandler(in, this)
    }
} 

package grpc.akkastreams

import akka.stream.{Attributes, Outlet, SourceShape}
import akka.stream.stage.{GraphStageLogic, GraphStageWithMaterializedValue, OutHandler}
import io.grpc.stub.{CallStreamObserver, StreamObserver}

import scala.concurrent.{Future, Promise}

class GrpcSourceStage[I, O](requestStream: CallStreamObserver[O])
  extends GraphStageWithMaterializedValue[SourceShape[I], Future[StreamObserver[I]]] {
  val out = Outlet[I]("grpc.out")
  override val shape: SourceShape[I] = SourceShape.of(out)

  override def createLogicAndMaterializedValue(
    inheritedAttributes: Attributes
  ): (GraphStageLogic, Future[StreamObserver[I]]) = {
    val promise: Promise[StreamObserver[I]] = Promise()

    val logic = new GraphStageLogic(shape) with OutHandler {
      val inObs = new StreamObserver[I] {
        override def onError(t: Throwable) =
          getAsyncCallback((t: Throwable) => fail(out, t)).invoke(t)

        override def onCompleted() =
          getAsyncCallback((_: Unit) => complete(out)).invoke(())

        override def onNext(value: I) =
          getAsyncCallback((value: I) => push(out, value)).invoke(value)
      }

      override def onPull(): Unit = requestStream.request(1)

      override def preStart(): Unit = {
        requestStream.disableAutoInboundFlowControl()
        promise.success(inObs)
      }

      setHandler(out, this)
    }

    (logic, promise.future)
  }
} 

package grpc.akkastreams

import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}

import akka.stream.{Attributes, FlowShape, Inlet, Outlet}
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
import io.grpc.stub.{ClientCallStreamObserver, ClientResponseObserver}

class GrpcGraphStage[I, O](operator: GrpcOperator[I, O]) extends GraphStage[FlowShape[I, O]] {
  val in = Inlet[I]("grpc.in")
  val out = Outlet[O]("grpc.out")

  override val shape: FlowShape[I, O] = FlowShape.of(in, out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) with InHandler with OutHandler {

      var requestStream = new AtomicReference[Option[ClientCallStreamObserver[I]]](None)
      val element = new AtomicReference[Option[I]](None)
      val requested = new AtomicBoolean(false)

      val outObs = new ClientResponseObserver[I, O] with Runnable {
        override def beforeStart(reqStream: ClientCallStreamObserver[I]): Unit = {
          requestStream.set(Some(reqStream))
          reqStream.disableAutoInboundFlowControl()
          reqStream.setOnReadyHandler(this)
        }

        override def onError(t: Throwable) =
          getAsyncCallback((t: Throwable) => fail(out, t)).invoke(t)

        override def onCompleted() =
          getAsyncCallback((_: Unit) => complete(out)).invoke(())

        override def onNext(value: O) =
          getAsyncCallback((value: O) => push(out, value)).invoke(value)

        override def run(): Unit = requestStream.get().foreach { reqStream =>
          if (requested.compareAndSet(true, false)) reqStream.request(1)
          if (reqStream.isReady) {
            element.getAndSet(None).foreach { value =>
              reqStream.onNext(value)
              tryPull(in)
            }
          }
        }
      }

      val inObs = operator(outObs)

      override def onPush(): Unit = {
        val value = grab(in)
        requestStream.get() match {
          case Some(reqStream) if reqStream.isReady() =>
            reqStream.onNext(value)
            pull(in)
          case _ => element.compareAndSet(None, Some(value))
        }
      }

      override def onUpstreamFinish(): Unit = inObs.onCompleted()

      override def onUpstreamFailure(t: Throwable): Unit = inObs.onError(t)

      override def onPull(): Unit =
        requestStream.get() match {
          case Some(reqStream) => reqStream.request(1)
          case _ => requested.compareAndSet(false, true)
        }

      override def preStart(): Unit = pull(in)

      setHandler(in, this)
      setHandler(out, this)
    }
} 

package cmwell.dc.stream.akkautils

import akka.stream.{Attributes, FlowShape, Inlet, Outlet}
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
import com.typesafe.scalalogging.LazyLogging


object DebugStage {
  def apply[A](name: String): DebugStage[A] = new DebugStage(name)
}

class DebugStage[A](name: String) extends GraphStage[FlowShape[A, A]] with LazyLogging {
  val in = Inlet[A]("DebugStage.in")
  val out = Outlet[A]("DebugStage.out")
  override val shape = FlowShape.of(in, out)
  override def createLogic(attr: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) {
      setHandler(
        in,
        new InHandler {
          override def onPush(): Unit = {
            logger.info(s"[$name]: grabbing element")
            val elem = grab(in)
            logger.info(s"[$name]: pushing the grabbed element $elem")
            push(out, elem)
          }
          override def onUpstreamFinish(): Unit = {
            logger.info(s"[$name]: onUpstreamFinish")
            super.onUpstreamFinish()
          }
          override def onUpstreamFailure(ex: Throwable): Unit = {
            logger.info(s"[$name]: onUpstreamFailure")
            super.onUpstreamFailure(ex)
          }
        }
      )
      setHandler(
        out,
        new OutHandler {
          override def onPull(): Unit = {
            logger.info(s"[$name]: pulling element")
            pull(in)
          }
          override def onDownstreamFinish(): Unit = {
            logger.info(s"[$name]: onDownstreamFinish")
            super.onDownstreamFinish()
          }
        }
      )
    }
} 

package com.olegych.scastie.util

import akka.actor.ActorRef
import akka.stream.{Outlet, SourceShape}
import akka.stream.stage.{GraphStageLogic, OutHandler}

import scala.collection.mutable.{Queue => MQueue}
import scala.reflect.runtime.universe._

class GraphStageLogicForwarder[T: TypeTag, U: TypeTag](out: Outlet[T], shape: SourceShape[T], coordinator: ActorRef, graphId: U)
    extends GraphStageLogic(shape) {

  setHandler(
    out,
    new OutHandler {
      override def onPull(): Unit = {
        deliver()
      }
    }
  )

  override def preStart(): Unit = {
    val thisGraphStageActorRef = getStageActor(bufferElement).ref
    coordinator ! ((graphId, thisGraphStageActorRef))
  }

  private val buffer = MQueue.empty[T]

  private def deliver(): Unit =
    if (isAvailable(out) && buffer.nonEmpty)
      push[T](out, buffer.dequeue)

  private def bufferElement(receive: (ActorRef, Any)): Unit =
    receive match {
      case (_, element: T @unchecked) =>
        buffer.enqueue(element)
        deliver()
    }

} 

package com.olegych.scastie.util

import akka.actor.ActorRef
import akka.stream.{Attributes, Outlet, SourceShape}
import akka.stream.stage.{GraphStage, GraphStageLogic}

import scala.reflect.runtime.universe._

class GraphStageForwarder[T: TypeTag, U: TypeTag](
    outletName: String,
    coordinator: ActorRef,
    graphId: U
) extends GraphStage[SourceShape[T]] {
  val out: Outlet[T] = Outlet(outletName)
  override val shape = SourceShape[T](out)
  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new GraphStageLogicForwarder(out, shape, coordinator, graphId)
} 

// Copyright (c) Microsoft. All rights reserved.

package com.microsoft.azure.iot.iothubreact

import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
import akka.stream.{Attributes, FlowShape, Inlet, Outlet}

private[iothubreact] class StreamManager
  extends GraphStage[FlowShape[MessageFromDevice, MessageFromDevice]] {

  private[this] val in          = Inlet[MessageFromDevice]("StreamCanceller.Flow.in")
  private[this] val out         = Outlet[MessageFromDevice]("StreamCanceller.Flow.out")
  private[this] var closeSignal = false

  override val shape = FlowShape.of(in, out)

  def close(): Unit = closeSignal = true

  override def createLogic(attr: Attributes): GraphStageLogic = {
    new GraphStageLogic(shape) {

      setHandler(in, new InHandler {
        override def onPush(): Unit = {
          val message: MessageFromDevice = grab(in)
          push(out, message)
        }
      })

      setHandler(out, new OutHandler {
        override def onPull(): Unit = {
          if (closeSignal) {
            cancel(in)
          } else {
            pull(in)
          }
        }
      })
    }
  }
} 

// Copyright (c) Microsoft. All rights reserved.

package com.microsoft.azure.iot.iothubreact.checkpointing

import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
import akka.stream.{Attributes, FlowShape, Inlet, Outlet}
import com.microsoft.azure.iot.iothubreact.MessageFromDevice
import com.microsoft.azure.iot.iothubreact.checkpointing.CheckpointService.UpdateOffset


private[iothubreact] class SaveOffsetOnPull(cpconfig: ICPConfiguration, partition: Int)
  extends GraphStage[FlowShape[MessageFromDevice, MessageFromDevice]] {

  val in   = Inlet[MessageFromDevice]("Checkpoint.Flow.in")
  val out  = Outlet[MessageFromDevice]("Checkpoint.Flow.out")
  val none = ""

  override val shape = FlowShape.of(in, out)

  // All state MUST be inside the GraphStageLogic, never inside the enclosing
  // GraphStage. This state is safe to read/write from all the callbacks
  // provided by GraphStageLogic and the registered handlers.
  override def createLogic(attr: Attributes): GraphStageLogic = {
    new GraphStageLogic(shape) {

      val checkpointService = CheckpointActorSystem(cpconfig).getCheckpointService(partition)
      var lastOffsetSent    = none

      // when a message enters the stage we safe its offset
      setHandler(in, new InHandler {
        override def onPush(): Unit = {
          val message: MessageFromDevice = grab(in)
          if (!message.isKeepAlive) lastOffsetSent = message.offset
          push(out, message)
        }
      })

      // when asked for more data we consider the saved offset processed and save it
      setHandler(out, new OutHandler {
        override def onPull(): Unit = {
          if (lastOffsetSent != none) checkpointService ! UpdateOffset(lastOffsetSent)
          pull(in)
        }
      })
    }
  }
} 

package org.hatdex.hat.utils

import akka.stream.stage.{ GraphStage, GraphStageLogic }
import akka.stream.{ Attributes, FanInShape2, Inlet, Outlet }

final class AugmentWith[T, U](augmentFunction: (T, U) ⇒ Either[T, U]) extends GraphStage[FanInShape2[T, U, T]] {
  private val left = Inlet[T]("left")
  private val right = Inlet[U]("right")
  private val out = Outlet[T]("out")

  override val shape = new FanInShape2(left, right, out)

  override def createLogic(attr: Attributes) = new GraphStageLogic(shape) {
    setHandler(left, eagerTerminateInput)
    setHandler(right, ignoreTerminateInput)
    setHandler(out, eagerTerminateOutput)

    var retainedL: T = _
    var retainedR: U = _

    def dispatch(l: T, r: U): Unit = {
      augmentFunction(l, r).fold(
        augmented ⇒ {
          retainedR = r
          emit(out, augmented, readL)
        },
        _ ⇒ {
          retainedL = l
          readR()
        })

    }

    val dispatchR = dispatch(retainedL, _: U)
    val dispatchL = dispatch(_: T, retainedR)
    val passL = () ⇒ emit(out, retainedL, () ⇒ { passAlong(left, out, doPull = true) })
    val readR = () ⇒ read(right)(dispatchR, passL)
    val readL = () ⇒ read(left)(dispatchL, readR)

    override def preStart(): Unit = {
      // all fan-in stages need to eagerly pull all inputs to get cycles started
      pull(right)
      read(left)(l ⇒ {
        retainedL = l
        readR()
      }, () ⇒ {
        abortReading(right)
      })
    }
  }
} 

package cool.graph.akkautil.stream

import akka.stream.ActorAttributes.SupervisionStrategy
import akka.stream.impl.fusing.GraphStages.SimpleLinearGraphStage
import akka.stream.stage.{GraphStageLogic, InHandler, OutHandler}
import akka.stream.{Attributes, Supervision}

case class OnCompleteStage[T](op: () ⇒ Unit) extends SimpleLinearGraphStage[T] {
  override def toString: String = "OnComplete"

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) with OutHandler with InHandler {
      def decider =
        inheritedAttributes
          .get[SupervisionStrategy]
          .map(_.decider)
          .getOrElse(Supervision.stoppingDecider)

      override def onPush(): Unit = {
        push(out, grab(in))
      }

      override def onPull(): Unit = pull(in)

      override def onDownstreamFinish() = {
        op()
        super.onDownstreamFinish()
      }

      override def onUpstreamFinish() = {
        op()
        super.onUpstreamFinish()
      }

      setHandlers(in, out, this)
    }
} 

package cmwell.util.stream

import akka.stream.ActorAttributes.SupervisionStrategy
import akka.stream._
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
import scala.util.control.NonFatal

case class TakeWeighted[T](maxWeight: Long, inclusive: Boolean = true, costFn: T => Long)
    extends GraphStage[FlowShape[T, T]] {
  val in: Inlet[T] = Inlet[T]("TakeWeighted.in")
  val out: Outlet[T] = Outlet[T]("TakeWeighted.out")
  override val shape = FlowShape(in, out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) with OutHandler with InHandler {

      private[this] var sum = 0L

      private def decider =
        inheritedAttributes.get[SupervisionStrategy].map(_.decider).getOrElse(Supervision.stoppingDecider)

      override def onPull(): Unit = pull(in)

      override def onPush(): Unit = {
        try {
          val elem = grab(in)
          val cost = costFn(elem)
          if (cost + sum < maxWeight) {
            sum += cost
            push(out, elem)
          } else {
            if (inclusive || cost + sum == maxWeight) {
              push(out, elem)
            }
            completeStage()
          }
        } catch {
          case NonFatal(ex) ⇒
            decider(ex) match {
              case Supervision.Stop ⇒ failStage(ex)
              case _ ⇒ pull(in)
            }
        }
      }

      setHandlers(in, out, this)
    }
} 

package cmwell.util.stream

import akka.stream.ActorAttributes.SupervisionStrategy
import akka.stream._
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
import scala.util.control.NonFatal

class MapInitAndLast[In, Out](init: In ⇒ Out, last: In ⇒ Out) extends GraphStage[FlowShape[In, Out]] {
  val in: Inlet[In] = Inlet[In]("MapInitAndLast.in")
  val out: Outlet[Out] = Outlet[Out]("MapInitAndLast.out")
  override val shape = FlowShape(in, out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) {

    private var pending: In = null.asInstanceOf[In]

    private def decider =
      inheritedAttributes.get[SupervisionStrategy].map(_.decider).getOrElse(Supervision.stoppingDecider)

    override def preStart(): Unit = tryPull(in)

    def pushWith(elem: In, f: In ⇒ Out): Unit = try { push(out, f(elem)) } catch {
      case NonFatal(ex) ⇒
        decider(ex) match {
          case Supervision.Stop ⇒ failStage(ex)
          case _ ⇒ pull(in)
        }
    }

    setHandler(
      in,
      new InHandler {
        override def onPush(): Unit = {
          val elem = grab(in)
          if (pending != null) pushWith(pending, init)
          else if (isAvailable(out)) pull(in)

          pending = elem
        }

        override def onUpstreamFinish(): Unit = {
          if (pending == null) completeStage()
          else {
            if (isAvailable(out)) {
              pushWith(pending, last)
              completeStage()
            }
          }
        }
      }
    )

    setHandler(
      out,
      new OutHandler {
        override def onPull(): Unit = {
          if (!isClosed(in)) {
            if (!hasBeenPulled(in)) {
              pull(in)
            }
          } else {
            if (pending != null) {
              pushWith(pending, last)
            }
            completeStage()
          }
        }
      }
    )
  }
} 

package cmwell.util.stream

import akka.stream.{Attributes, FlowShape, Inlet, Outlet}
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
// format: off
class StreamEventInspector[Elem](onUpstreamFinishInspection:   ()        => Unit = () => {},
                                 onUpstreamFailureInspection:  Throwable => Unit = _  => {},
                                 onDownstreamFinishInspection: ()        => Unit = () => {},
                                 onPushInspection:             Elem      => Unit = (_: Elem)  => {},
                                 onPullInspection:             ()        => Unit = () => {}) extends GraphStage[FlowShape[Elem, Elem]] {
  // format: on
  private val in = Inlet[Elem]("StreamEventInspector.in")
  private val out = Outlet[Elem]("StreamEventInspector.out")

  override val shape = FlowShape(in, out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) {

    setHandler(
      in,
      new InHandler {
        override def onPush(): Unit = {
          val elem = grab(in)
          onPushInspection(elem)
          push(out, elem)
        }

        override def onUpstreamFailure(ex: Throwable): Unit = {
          onUpstreamFailureInspection(ex)
          super.onUpstreamFailure(ex)
        }

        override def onUpstreamFinish(): Unit = {
          onUpstreamFinishInspection()
          super.onUpstreamFinish()
        }
      }
    )

    setHandler(
      out,
      new OutHandler {
        override def onPull(): Unit = {
          onPullInspection()
          pull(in)
        }

        override def onDownstreamFinish(): Unit = {
          onDownstreamFinishInspection()
          super.onDownstreamFinish()
        }
      }
    )
  }
} 

// Copyright (c) 2020 Digital Asset (Switzerland) GmbH and/or its affiliates. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

package com.daml.util.akkastreams

import akka.stream.scaladsl.Flow
import akka.stream.{Attributes, FlowShape, Inlet, Outlet}
import akka.stream.stage.{GraphStageLogic, GraphStageWithMaterializedValue, InHandler, OutHandler}

import scala.concurrent.{Future, Promise}


class ExtractMaterializedValue[T, Mat](toMaterialized: T => Option[Mat])
    extends GraphStageWithMaterializedValue[FlowShape[T, T], Future[Mat]] {

  val inlet: Inlet[T] = Inlet[T]("in")
  val outlet: Outlet[T] = Outlet[T]("out")

  override def createLogicAndMaterializedValue(
      inheritedAttributes: Attributes): (GraphStageLogic, Future[Mat]) = {
    val promise = Promise[Mat]()

    val logic = new GraphStageLogic(shape) {

      setHandler(
        inlet,
        new InHandler {
          override def onPush(): Unit = {
            val input = grab(inlet)
            push(outlet, input)
            toMaterialized(input).foreach { materialized =>
              promise.trySuccess(materialized)
              setSimplerHandler()
            }
          }

          private def setSimplerHandler(): Unit = {
            setHandler(inlet, new InHandler {
              override def onPush(): Unit =
                push(outlet, grab(inlet))
            })
          }

          override def onUpstreamFailure(ex: Throwable): Unit = {
            promise.tryFailure(ex)
            super.onUpstreamFailure(ex)
          }

          override def onUpstreamFinish(): Unit = {
            promise.tryFailure(
              new RuntimeException("Upstream completed before matching element arrived."))
            super.onUpstreamFinish()
          }
        }
      )

      setHandler(
        outlet,
        new OutHandler {
          override def onPull(): Unit = pull(inlet)

          override def onDownstreamFinish(cause: Throwable): Unit = {
            promise.tryFailure(
              new RuntimeException("Downstream completed before matching element arrived."))
            super.onDownstreamFinish(cause)
          }
        }
      )

    }

    logic -> promise.future
  }

  override def shape: FlowShape[T, T] = FlowShape(inlet, outlet)
}

object ExtractMaterializedValue {
  def apply[T, Mat](toOutputOrMaterialized: T => Option[Mat]): Flow[T, T, Future[Mat]] =
    Flow.fromGraph(new ExtractMaterializedValue[T, Mat](toOutputOrMaterialized))
}