akka.stream.Outlet 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 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 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 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 cmwell.tools.data.sparql

import akka.actor.ActorRef
import akka.stream.{Attributes, FlowShape, Inlet, Outlet}
import akka.pattern._
import akka.util.Timeout
import akka.stream.stage._
import cmwell.tools.data.ingester.Ingester.IngestEvent
import cmwell.tools.data.sparql.InfotonReporter.{RequestDownloadStats, RequestIngestStats, ResponseDownloadStats, ResponseIngestStats}
import cmwell.tools.data.utils.logging.DataToolsLogging

import scala.concurrent.ExecutionContext
import ExecutionContext.Implicits.global
import scala.concurrent.duration._

object StpPeriodicLogger {
  def apply(infotonReporter: ActorRef, logFrequency: FiniteDuration) =
    new StpPeriodicLogger(infotonReporter, logFrequency)
}

class StpPeriodicLogger(infotonReporter: ActorRef, logFrequency: FiniteDuration)
  extends GraphStage[FlowShape[IngestEvent, IngestEvent]] with DataToolsLogging {

  val in = Inlet[IngestEvent]("StpPeriodicLogger.in")
  val out = Outlet[IngestEvent]("StpPeriodicLogger.out")
  override val shape = FlowShape.of(in, out)

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

    override def preStart(): Unit = schedulePeriodically(None, logFrequency)

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

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

    override protected def onTimer(timerKey: Any): Unit = {

      implicit val timeout : akka.util.Timeout = 5.seconds

      for {
        downloadStatsMap <- (infotonReporter ? RequestDownloadStats).mapTo[ResponseDownloadStats]
        ingestStatsObj <- (infotonReporter ? RequestIngestStats).mapTo[ResponseIngestStats]
        ingestStatsOption = ingestStatsObj.stats
      } yield {
        downloadStatsMap.stats.get(SparqlTriggeredProcessor.sparqlMaterializerLabel).foreach(materialized => {
          ingestStatsOption.foreach(ingestStats => {
            logger.info(s" *** STP Agent ${materialized.label}" +
              s" *** Materialized ${materialized.receivedInfotons}, Ingested: ${ingestStats.ingestedInfotons}, Failed: ${ingestStats.failedInfotons}")
          })
        })
      }
    }

  }

} 

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 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 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 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 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 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.voice

import java.nio.ByteOrder

import scala.concurrent.{Future, Promise}

import akka.stream.scaladsl.Flow
import akka.stream.stage._
import akka.stream.{Attributes, FlowShape, Inlet, Outlet}
import akka.util.ByteString

class IPDiscoveryFlow(openValve: () => Unit)
    extends GraphStageWithMaterializedValue[FlowShape[ByteString, ByteString], Future[VoiceUDPFlow.FoundIP]] {

  val in: Inlet[ByteString]   = Inlet("IPDiscoveryFlow.in")
  val out: Outlet[ByteString] = Outlet("IPDiscoveryFlow.out")

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

  override def createLogicAndMaterializedValue(
      inheritedAttributes: Attributes
  ): (GraphStageLogic, Future[VoiceUDPFlow.FoundIP]) = {
    val promise = Promise[VoiceUDPFlow.FoundIP]
    val logic = new GraphStageLogicWithLogging(shape) with InHandler with OutHandler {

      override def onPush(): Unit = {
        val data = grab(in)
        log.debug(s"Grabbing data for IP discovery $data")
        val byteBuf = data.asByteBuffer.order(ByteOrder.BIG_ENDIAN)
        val tpe     = byteBuf.getShort

        require(tpe == 0x2, s"Was expecting IP discovery result, got $tpe")

        byteBuf.getShort //Length
        byteBuf.getInt   //SSRC
        val nullTermString = new Array[Byte](64)
        byteBuf.get(nullTermString)
        val address = new String(nullTermString, 0, nullTermString.iterator.takeWhile(_ != 0).length)
        val port    = byteBuf.getChar.toInt //Char is unsigned short

        promise.success(VoiceUDPFlow.FoundIP(address, port))
        log.debug("Success doing IP discovery")

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

        openValve()
      }

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

      override def onUpstreamFailure(ex: Throwable): Unit = {
        promise.tryFailure(new Exception("Connection failed.", ex))
        super.onUpstreamFailure(ex)
      }

      setHandlers(in, out, this)
    }

    (logic, promise.future)
  }
}
object IPDiscoveryFlow {
  def flow(openValve: () => Unit): Flow[ByteString, ByteString, Future[VoiceUDPFlow.FoundIP]] =
    Flow.fromGraph(new IPDiscoveryFlow(openValve))
} 

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 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)
      }
    }
  }
} 

// 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.services.commands.tracker

import akka.stream.{Inlet, Outlet, Shape}
import com.daml.ledger.api.v1.command_submission_service.SubmitRequest
import com.daml.ledger.api.v1.completion.Completion
import com.daml.ledger.api.v1.ledger_offset.LedgerOffset
import com.daml.ledger.client.services.commands.CompletionStreamElement
import com.daml.util.Ctx
import com.google.protobuf.empty.Empty

import scala.collection.immutable
import scala.util.Try

private[tracker] final case class CommandTrackerShape[Context](
    submitRequestIn: Inlet[Ctx[Context, SubmitRequest]],
    submitRequestOut: Outlet[Ctx[(Context, String), SubmitRequest]],
    commandResultIn: Inlet[Either[Ctx[(Context, String), Try[Empty]], CompletionStreamElement]],
    resultOut: Outlet[Ctx[Context, Completion]],
    offsetOut: Outlet[LedgerOffset])
    extends Shape {

  override def inlets: immutable.Seq[Inlet[_]] = Vector(submitRequestIn, commandResultIn)

  override def outlets: immutable.Seq[Outlet[_]] = Vector(submitRequestOut, resultOut, offsetOut)

  override def deepCopy(): Shape =
    CommandTrackerShape[Context](
      submitRequestIn.carbonCopy(),
      submitRequestOut.carbonCopy(),
      commandResultIn.carbonCopy(),
      resultOut.carbonCopy(),
      offsetOut.carbonCopy())
} 

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()
        }
      }
    )
  }
} 

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)
      })
    }
  }
} 

// 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)
        }
      })
    }
  }
} 

// 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)
          }
        }
      })
    }
  }
} 

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)
} 

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.packt.chapter8

import akka.stream.{Attributes, Outlet, SourceShape}
import akka.stream.stage._

class HelloAkkaStreamsSource extends GraphStage[SourceShape[String]] {

  val out: Outlet[String] = Outlet("SystemInputSource")
  override val shape: SourceShape[String] = SourceShape(out)

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

      setHandler(out, new OutHandler {
        override def onPull() = {
          val line = "Hello World Akka Streams!"
          push(out, line)
        }
      })
    }
} 

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)
        }
      })
    }
} 

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 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 com.github.dnvriend.streams.customstage

import akka.stream.stage._
import akka.stream.testkit.scaladsl.TestSink
import akka.stream.{ Attributes, FlowShape, Inlet, Outlet }
import com.github.dnvriend.streams.TestSpec

class Ex3CustomFilterTest extends TestSpec {
  "CustomFilterStage" should "be implemented with a PushPullStage" in {

    
    withIterator(1) { src ⇒
      src.transform(() ⇒ new CustomFilterStage(_ % 2 == 0))
        .take(5)
        .runWith(TestSink.probe[Int])
        .request(Int.MaxValue)
        .expectNext(2, 4, 6, 8, 10)
        .expectComplete()
    }
  }

  it should "also be implemented as a GraphStage" in {
    class CustomFilterStage[A](p: A ⇒ Boolean) extends GraphStage[FlowShape[A, A]] {
      val in = Inlet[A]("Filter.in")
      val out = Outlet[A]("Filter.out")

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

      override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) {
        setHandler(in, new InHandler {
          override def onPush(): Unit = {
            val elem: A = grab(in)
            if (p(elem)) push(out, elem) else pull(in)
          }
        })

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

    withIterator(1) { src ⇒
      src.via(new CustomFilterStage(_ % 2 == 0))
        .take(5)
        .runWith(TestSink.probe[Int])
        .request(Int.MaxValue)
        .expectNext(2, 4, 6, 8, 10)
        .expectComplete()
    }
  }
} 

package com.github.dnvriend.streams.customstage

import akka.stream.stage._
import akka.stream.testkit.scaladsl.TestSink
import akka.stream.{ Attributes, FlowShape, Inlet, Outlet }
import com.github.dnvriend.streams.TestSpec

class Ex4StatefulStageTest extends TestSpec {

  
    withIterator(1) { src ⇒
      src.take(20)
        .transform(() ⇒ new SumEvenAndUnevenNumbersCollector(_ % 10 == 0)) // emit every 10 elements
        .runWith(TestSink.probe[(Int, Int)])
        .request(Int.MaxValue)
        .expectNext((20, 25), (60, 75))
        .expectComplete()
    }
  }

  it should "be implemented as a GraphShape" in {
    // as the StatefulStage will be deprecated, let's look at how to handle state in a GraphShape

    class CustomDuplicatorStage[A]() extends GraphStage[FlowShape[A, A]] {

      val in = Inlet[A]("Duplicator.in")
      val out = Outlet[A]("Duplicator.out")

      override def shape: FlowShape[A, A] = FlowShape.of(in, out)

      override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) {
        // note: all mutable state must be inside the GraphStageLogic
        var lastElem: Option[A] = None

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

          override def onUpstreamFinish(): Unit = {
            if (lastElem.isDefined) emit(out, lastElem.get)
            complete(out)
          }
        })

        setHandler(out, new OutHandler {
          override def onPull(): Unit = {
            if (lastElem.isDefined) {
              push(out, lastElem.get)
              lastElem = None
            } else {
              pull(in)
            }
          }
        })
      }
    }

    withIterator(1) { src ⇒
      src.take(2)
        .via(new CustomDuplicatorStage)
        .runWith(TestSink.probe[Int])
        .request(Int.MaxValue)
        .expectNext(1, 1, 2, 2)
        .expectComplete()
    }
  }
} 

package com.github.dnvriend.streams.customstage

import akka.stream.{ Outlet, Inlet, Attributes, FlowShape }
import akka.stream.stage._
import akka.stream.testkit.scaladsl.TestSink
import com.github.dnvriend.streams.TestSpec

class Ex2CustomMapTest extends TestSpec {

  "CustomMapStage" should "be implemented with a PushPullStage" in {

    
    withIterator(1) { src ⇒
      src.transform(() ⇒ new CustomMapStage(_ * 2))
        .take(2)
        .runWith(TestSink.probe[Int])
        .request(Int.MaxValue)
        .expectNext(2, 4)
        .expectComplete()
    }
  }

  it should "also be implemented as a GraphStage" in {
    class CustomMapStage[A, B](f: A ⇒ B) extends GraphStage[FlowShape[A, B]] {
      val in = Inlet[A]("Map.in")
      val out = Outlet[B]("Map.out")

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

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

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

    withIterator(1) { src ⇒
      src.via(new CustomMapStage(_ * 2))
        .take(2)
        .runWith(TestSink.probe[Int])
        .request(Int.MaxValue)
        .expectNext(2, 4)
        .expectComplete()
    }
  }
} 

// 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.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))
}