akka.stream.FlowShape 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.bg.imp

import akka.NotUsed
import akka.stream.FlowShape
import akka.stream.contrib.PartitionWith
import akka.stream.scaladsl.{Flow, GraphDSL, Merge, Partition}
import cmwell.bg.BGMetrics
import cmwell.common.formats.BGMessage
import cmwell.common._
import cmwell.zstore.ZStore
import com.typesafe.scalalogging.LazyLogging

import scala.concurrent.ExecutionContext

object RefsEnricher extends LazyLogging {

  def toSingle(bgm: BGMetrics, irwReadConcurrency: Int, zStore: ZStore)
              (implicit ec: ExecutionContext): Flow[BGMessage[Command], BGMessage[SingleCommand], NotUsed] = {

    Flow.fromGraph(GraphDSL.create() { implicit b =>
      import GraphDSL.Implicits._

      // CommandRef goes left, all rest go right
      // update metrics for each type of command
      val commandsPartitioner = b.add(PartitionWith[BGMessage[Command], BGMessage[CommandRef], BGMessage[Command]] {
        case bgm @ BGMessage(_, CommandRef(_)) => Left(bgm.asInstanceOf[BGMessage[CommandRef]])
        case bgm => Right(bgm)
      })

      val commandRefsFetcher = Flow[BGMessage[CommandRef]].mapAsync(irwReadConcurrency) {
        case bgMessage @ BGMessage(_, CommandRef(ref)) => {
          zStore.get(ref).map { payload =>
            bgMessage.copy(message = CommandSerializer.decode(payload))
          }
        }
      }

      val singleCommandsMerge = b.add(Merge[BGMessage[Command]](2))

      commandsPartitioner.out0 ~> commandRefsFetcher ~> singleCommandsMerge.in(0)

      commandsPartitioner.out1 ~> singleCommandsMerge.in(1)

      FlowShape(commandsPartitioner.in,singleCommandsMerge.out.map {
        bgMessage => {
          // cast to SingleCommand while updating metrics
          bgMessage.message match {
            case wc: WriteCommand           => bgm.writeCommandsCounter += 1
                                               bgm.infotonCommandWeightHist += wc.infoton.weight
            case oc: OverwriteCommand       => bgm.overrideCommandCounter += 1
                                               bgm.infotonCommandWeightHist += oc.infoton.weight
            case _: UpdatePathCommand       => bgm.updatePathCommandsCounter += 1
            case _: DeletePathCommand       => bgm.deletePathCommandsCounter += 1
            case _: DeleteAttributesCommand => bgm.deleteAttributesCommandsCounter += 1
            case unknown                    => logger.error(s"unknown command [$unknown]")
          }
          bgm.commandMeter.mark()
          bgMessage.copy(message = bgMessage.message.asInstanceOf[SingleCommand])
        }
      }.outlet)
    })
  }
} 

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 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 swave.docs

import scala.util.{Failure, Success}
import akka.NotUsed
import akka.actor.ActorSystem
import akka.stream.scaladsl._
import akka.stream.{ActorMaterializer, ActorMaterializerSettings, FlowShape}
import swave.core.util.XorShiftRandom

object AkkaPi extends App {
  implicit val system = ActorSystem("AkkaPi")
  import system.dispatcher

  private val settings =
    ActorMaterializerSettings(system).withSyncProcessingLimit(Int.MaxValue).withInputBuffer(256, 256)
  implicit val materializer = ActorMaterializer(settings)

  val random = XorShiftRandom()

  Source
    .fromIterator(() ⇒ Iterator.continually(random.nextDouble()))
    .grouped(2)
    .map { case x +: y +: Nil ⇒ Point(x, y) }
    .via(broadcastFilterMerge)
    .scan(State(0, 0)) { _ withNextSample _ }
    .splitWhen(_.totalSamples % 1000000 == 1)
    .drop(999999)
    .concatSubstreams
    .map(state ⇒ f"After ${state.totalSamples}%,10d samples π is approximated as ${state.π}%.6f")
    .take(30)
    .runForeach(println)
    .onComplete {
      case Success(_) =>
        val time = System.currentTimeMillis() - system.startTime
        println(f"Done. Total time: $time%,6d ms, Throughput: ${30000.0 / time}%.3fM samples/sec\n")
        system.terminate()

      case Failure(e) => println(e)
    }

  println("Main thread exit.")

  ///////////////////////////////////////////

  def broadcastFilterMerge: Flow[Point, Sample, NotUsed] =
    Flow.fromGraph(GraphDSL.create() { implicit b =>
      import GraphDSL.Implicits._

      val broadcast   = b.add(Broadcast[Point](2)) // split one upstream into 2 downstreams
      val filterInner = b.add(Flow[Point].filter(_.isInner).map(_ => InnerSample))
      val filterOuter = b.add(Flow[Point].filterNot(_.isInner).map(_ => OuterSample))
      val merge       = b.add(Merge[Sample](2)) // merge 2 upstreams into one downstream

      broadcast.out(0) ~> filterInner ~> merge.in(0)
      broadcast.out(1) ~> filterOuter ~> merge.in(1)

      FlowShape(broadcast.in, merge.out)
    })

  //////////////// MODEL ///////////////

  case class Point(x: Double, y: Double) {
    def isInner: Boolean = x * x + y * y < 1.0
  }

  sealed trait Sample
  case object InnerSample extends Sample
  case object OuterSample extends Sample

  case class State(totalSamples: Long, inCircle: Long) {
    def π: Double = (inCircle.toDouble / totalSamples) * 4.0
    def withNextSample(sample: Sample) =
      State(totalSamples + 1, if (sample == InnerSample) inCircle + 1 else inCircle)
  }
} 

package swave.docs

import scala.util.{Failure, Success}
import scala.concurrent.duration._
import akka.NotUsed
import akka.actor.ActorSystem
import akka.stream.scaladsl._
import akka.stream.{ActorMaterializer, ActorMaterializerSettings, FlowShape, ThrottleMode}
import swave.core.util.XorShiftRandom

object AkkaPiThrottled extends App {
  implicit val system = ActorSystem("AkkaPi")
  import system.dispatcher

  private val settings = ActorMaterializerSettings(system)
  implicit val materializer = ActorMaterializer(settings)

  val random = XorShiftRandom()

  Source
    .fromIterator(() ⇒ Iterator.continually(random.nextDouble()))
    .grouped(2)
    .map { case x +: y +: Nil ⇒ Point(x, y) }
    .via(broadcastFilterMerge)
    .scan(State(0, 0)) { _ withNextSample _ }
    .conflate(Keep.right)
    .throttle(1, 1.second, 0, ThrottleMode.Shaping)
    .map { state ⇒ println(f"After ${state.totalSamples}%,10d samples π is approximated as ${state.π}%.6f"); state }
    .take(20)
    .runWith(Sink.last)
    .onComplete {
      case Success(State(totalSamples, _)) =>
        val time = System.currentTimeMillis() - system.startTime
        val throughput = totalSamples / 1000.0 / time
        println(f"Done. Total time: $time%,6d ms, Throughput: $throughput%.3fM samples/sec\n")
        system.terminate()

      case Failure(e) => println(e)
    }

  println("Main thread exit.")

  ///////////////////////////////////////////

  def broadcastFilterMerge: Flow[Point, Sample, NotUsed] =
    Flow.fromGraph(GraphDSL.create() { implicit b =>
      import GraphDSL.Implicits._

      val broadcast   = b.add(Broadcast[Point](2)) // split one upstream into 2 downstreams
    val filterInner = b.add(Flow[Point].filter(_.isInner).map(_ => InnerSample))
      val filterOuter = b.add(Flow[Point].filterNot(_.isInner).map(_ => OuterSample))
      val merge       = b.add(Merge[Sample](2)) // merge 2 upstreams into one downstream

      broadcast.out(0) ~> filterInner ~> merge.in(0)
      broadcast.out(1) ~> filterOuter ~> merge.in(1)

      FlowShape(broadcast.in, merge.out)
    })

  //////////////// MODEL ///////////////

  case class Point(x: Double, y: Double) {
    def isInner: Boolean = x * x + y * y < 1.0
  }

  sealed trait Sample
  case object InnerSample extends Sample
  case object OuterSample extends Sample

  case class State(totalSamples: Long, inCircle: Long) {
    def π: Double = (inCircle.toDouble / totalSamples) * 4.0
    def withNextSample(sample: Sample) =
      State(totalSamples + 1, if (sample == InnerSample) inCircle + 1 else inCircle)
  }
} 

package com.github.j5ik2o.bank.useCase

import akka.{ Done, NotUsed }
import akka.stream.{ FlowShape, QueueOfferResult }
import akka.stream.scaladsl.{ Flow, GraphDSL, Sink, SourceQueueWithComplete, Unzip, Zip }

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

object UseCaseSupport {
  implicit class FlowOps[A, B](val self: Flow[A, B, NotUsed]) extends AnyVal {
    def zipPromise: Flow[(A, Promise[B]), (B, Promise[B]), NotUsed] =
      Flow
        .fromGraph(GraphDSL.create() { implicit b =>
          import GraphDSL.Implicits._
          val unzip = b.add(Unzip[A, Promise[B]])
          val zip   = b.add(Zip[B, Promise[B]])
          unzip.out0 ~> self ~> zip.in0
          unzip.out1 ~> zip.in1
          FlowShape(unzip.in, zip.out)
        })
  }
}

trait UseCaseSupport {

  protected def offerToQueue[A, B](
      sourceQueue: SourceQueueWithComplete[(A, Promise[B])]
  )(request: A, promise: Promise[B])(implicit ec: ExecutionContext): Future[B] = {
    sourceQueue.offer((request, promise)).flatMap {
      case QueueOfferResult.Enqueued =>
        promise.future
      case QueueOfferResult.Failure(t) =>
        Future.failed(new Exception("Failed to offer request", t))
      case QueueOfferResult.Dropped =>
        Future.failed(
          new Exception(
            s"Failed to enqueue resolve request, the queue buffer was full, please check the bank.interface.buffer-size setting"
          )
        )
      case QueueOfferResult.QueueClosed =>
        Future.failed(new Exception("Failed to enqueue request batch write, the queue was closed"))
    }
  }

  protected def completePromiseSink[T]: Sink[(T, Promise[T]), Future[Done]] = Sink.foreach {
    case (response, promise) =>
      promise.success(response)
  }

} 

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.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 example.akkastream.basic

import akka.NotUsed
import akka.actor.ActorSystem
import akka.stream.scaladsl.{ Broadcast, Flow, GraphDSL, Sink, Source, Zip }
import akka.stream.{ ActorMaterializer, FlowShape, SourceShape }

import scala.io.StdIn

object PartialGraph2 extends App {
  implicit val system = ActorSystem()
  implicit val mat = ActorMaterializer()
  import system.dispatcher

  val pairs: Source[(Int, Int), NotUsed] = Source.fromGraph(GraphDSL.create() { implicit b =>
    import GraphDSL.Implicits._

    // prepare graph elements
    val zip = b.add(Zip[Int, Int]())

    def ints = Source.fromIterator(() => Iterator.from(1))

    // connect the graph
    ints.filter(_ % 2 != 0) ~> zip.in0
    ints.filter(_ % 2 == 0) ~> zip.in1

    // expose port
    SourceShape(zip.out)
  })

  val firstPair = pairs.runWith(Sink.head)
  firstPair.foreach(println)

  val pairUpWithToString = Flow.fromGraph(GraphDSL.create() { implicit b =>
    import GraphDSL.Implicits._
    val broadcast = b.add(Broadcast[Int](2))
    val zip = b.add(Zip[Int, String]())

    broadcast.out(0)  ~> zip.in0
    broadcast.out(1).map(_.toString) ~> zip.in1

    FlowShape(broadcast.in, zip.out)
  })

  Source(List(1)).via(pairUpWithToString).runWith(Sink.head).foreach(println)

  StdIn.readLine()
  system.terminate()
} 

package example.akkastream.graph

import akka.actor.ActorSystem
import akka.stream.scaladsl.{ Balance, Broadcast, Flow, GraphDSL, Keep, Merge, RunnableGraph, Sink, Source }
import akka.stream.{ ActorMaterializer, FlowShape, SourceShape }

import scala.concurrent.Future
import scala.io.StdIn

object PartialGraph extends App {
  implicit val system = ActorSystem()
  implicit val mat = ActorMaterializer()
  import system.dispatcher

  def partial =
    GraphDSL
      .create() { implicit b =>
        import GraphDSL.Implicits._

        val B = b.add(Broadcast[Int](2))
        val C = b.add(Merge[Int](2))
        val D = Flow[Int].map(_ + 1)
        val E = b.add(Balance[Int](2))
        val F = b.add(Merge[Int](2))

        C <~ F
        B ~> C ~> F
        B ~> D ~> E ~> F

        FlowShape(B.in, E.out(1))
      }
      .named("partial")

  // 转换partial从FlowShape到Flow，可访问流DSL（比如：.filter() 函数）
  val flow = Flow.fromGraph(partial)

  val source = Source.fromGraph(GraphDSL.create() { implicit b =>
    import GraphDSL.Implicits._
    val merge = b.add(Merge[Int](2))
    Source.single(0) ~> merge
    Source(List(2, 3, 4)) ~> merge
    SourceShape(merge.out)
  })

  val sink: Sink[Int, Future[Int]] = Flow[Int].map(_ * 2).drop(10).named("nestedFlow").toMat(Sink.head)(Keep.right)

  val closed: RunnableGraph[Future[Int]] =
    source.via(flow.filter(_ > 1)).toMat(sink)(Keep.right)

  closed.run().foreach(println)

  StdIn.readLine()
  system.terminate()
} 

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

import ackcord.CacheSnapshot
import ackcord.data.raw.RawMessage
import ackcord.data.{Message, User}
import ackcord.requests.{CreateMessage, Request, Requests}
import ackcord.syntax._
import akka.NotUsed
import akka.stream.FlowShape
import akka.stream.scaladsl.{Broadcast, Flow, GraphDSL}

object CmdHelper {

  
  def isValidCommand[F[_]](needMention: Boolean, msg: Message)(
      implicit c: CacheSnapshot
  ): Option[List[String]] = {
    if (needMention) {
      val botUser = c.botUser
      //We do a quick check first before parsing the message
      val quickCheck = if (msg.mentions.contains(botUser.id)) Some(msg.content.split(" ").toList) else None

      quickCheck.flatMap { args =>
        MessageParser
          .parseEither(args, MessageParser[User])
          .toOption
          .flatMap {
            case (remaining, user) if user.id == botUser.id => Some(remaining)
            case (_, _)                                     => None
          }
      }
    } else Some(msg.content.split(" ").toList)
  }
} 

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

package com.github.dnvriend

import akka.NotUsed
import akka.event.LoggingAdapter
import akka.stream.FlowShape
import akka.stream.scaladsl.{ Broadcast, Flow, GraphDSL, Sink }

import scala.compat.Platform
import scala.collection.immutable._

object LogProgress {
  def flow[A](each: Long = 1000)(implicit log: LoggingAdapter = null): Flow[A, A, NotUsed] = Flow.fromGraph[A, A, NotUsed](GraphDSL.create() { implicit b =>
    import GraphDSL.Implicits._
    val logFlow = Flow[A].statefulMapConcat { () =>
      var last = Platform.currentTime
      var num = 0L
      (x: A) =>
        num += 1
        if (num % each == 0) {
          val duration = Platform.currentTime - last
          val logOpt = Option(log)
          Option(log).foreach(_.info("[{} ms / {}]: {}", duration, each, num))
          if (logOpt.isEmpty) println(s"[$duration ms / $each]: $num")
          last = Platform.currentTime
        }
        Iterable(x)
    }
    val bcast = b.add(Broadcast[A](2, eagerCancel = false))
    bcast ~> logFlow ~> Sink.ignore
    FlowShape.of(bcast.in, bcast.out(1))
  })
} 

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

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

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

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.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.mjakubowski84.parquet4s
import akka.stream.FlowShape
import akka.stream.scaladsl.{Broadcast, Flow, GraphDSL, Keep, Sink, ZipWith}
import com.github.mjakubowski84.parquet4s.ParquetWriter.ParquetWriterFactory
import org.apache.hadoop.fs.Path
import org.slf4j.{Logger, LoggerFactory}

import scala.concurrent.duration.FiniteDuration


private[parquet4s] object IndefiniteStreamParquetSink extends IOOps {

  protected val logger: Logger = LoggerFactory.getLogger(this.getClass)

  def apply[In, ToWrite: ParquetWriterFactory, Mat](path: Path,
                                                    maxChunkSize: Int,
                                                    chunkWriteTimeWindow: FiniteDuration,
                                                    buildChunkPath: ChunkPathBuilder[In] = ChunkPathBuilder.default,
                                                    preWriteTransformation: In => ToWrite = identity[In] _,
                                                    postWriteSink: Sink[Seq[In], Mat] = Sink.ignore,
                                                    options: ParquetWriter.Options = ParquetWriter.Options()
                                            ): Sink[In, Mat] = {
    validateWritePath(path, options)

    val internalFlow = Flow.fromGraph(GraphDSL.create() { implicit b =>
      import GraphDSL.Implicits._
    
      val inChunkFlow = b.add(Flow[In].groupedWithin(maxChunkSize, chunkWriteTimeWindow))
      val broadcastChunks = b.add(Broadcast[Seq[In]](outputPorts = 2))
      val writeFlow = Flow[Seq[In]].map { chunk =>
        val toWrite = chunk.map(preWriteTransformation)
        val chunkPath = buildChunkPath(path, chunk)
        if (logger.isDebugEnabled()) logger.debug(s"Writing ${toWrite.size} records to $chunkPath")
        ParquetWriter.writeAndClose(chunkPath.toString, toWrite, options)
      }
      val zip = b.add(ZipWith[Seq[In], Unit, Seq[In]]((chunk, _) => chunk))
      
      inChunkFlow ~> broadcastChunks ~> writeFlow ~> zip.in1
                     broadcastChunks ~> zip.in0

      FlowShape(inChunkFlow.in, zip.out)               
    })

    internalFlow.toMat(postWriteSink)(Keep.right)
  }

} 

package sample.stream_shared_state

import akka.Done
import akka.actor.{ActorSystem, Cancellable}
import akka.stream.scaladsl.{Flow, GraphDSL, Keep, Sink, Source, SourceQueueWithComplete, Zip}
import akka.stream.{FlowShape, OverflowStrategy}

import scala.collection.immutable
import scala.concurrent.Future
import scala.concurrent.duration._
import scala.util.{Failure, Success}



object ParametrizedFlow extends App {
  val service = ParameterizedFlowService

  Thread.sleep(5000)
  service.update(1.0)

  Thread.sleep(2000)
  service.update(1.5)
  Thread.sleep(2000)
  service.cancel()
  Thread.sleep(2000)

  println(service.result())
}

object ParameterizedFlowService {
  implicit val system = ActorSystem("ParameterizedFlowService")
  implicit val executionContext = system.dispatcher

  def update(element: Double): Unit = flow._1._2.offer(element)

  def cancel(): Boolean = flow._1._1.cancel()

  def result(): Future[Seq[Double]] = flow._2

  val fun = (flowValue: Int, paramValue: Double) => flowValue * paramValue
  val flow: ((Cancellable, SourceQueueWithComplete[Double]), Future[immutable.Seq[Double]]) =
    Source.tick(0.seconds, 500.millis, 10)
      .viaMat(createParamFlow(1, OverflowStrategy.dropBuffer, 0.5)(fun))(Keep.both)
      .wireTap(x => println(x))
      .toMat(Sink.seq)(Keep.both)
      .run()

  val done: Future[Done] = flow._1._2.watchCompletion()
  terminateWhen(done)

  private def createParamFlow[A, P, O](bufferSize: Int, overflowStrategy: OverflowStrategy, initialParam: P)(fun: (A, P) => O) =
    Flow.fromGraph(GraphDSL.create(Source.queue[P](bufferSize, overflowStrategy)) { implicit builder =>
      queue =>
        import GraphDSL.Implicits._
        val zip = builder.add(Zip[A, P]())
        //Interesting use of the extrapolate operator
        //based on https://doc.akka.io/docs/akka/current/stream/stream-rate.html#understanding-extrapolate-and-expand
        val extra = builder.add(Flow[P].extrapolate(Iterator.continually(_), Some(initialParam)))
        val map = builder.add(Flow[(A, P)].map(r => fun(r._1, r._2)))

        queue ~> extra ~> zip.in1
        zip.out ~> map
        FlowShape(zip.in0, map.out)
    })

  private def terminateWhen(done: Future[_]) = {
    done.onComplete {
      case Success(_) =>
        println("Flow Success. About to terminate...")
        system.terminate()
      case Failure(e) =>
        println(s"Flow Failure: $e. About to terminate...")
        system.terminate()
    }
  }
} 

package sample.stream_actor

import akka.actor.{ActorRef, ActorSystem}
import akka.http.scaladsl.Http
import akka.http.scaladsl.model.StatusCodes
import akka.http.scaladsl.model.ws._
import akka.stream.scaladsl.{Flow, GraphDSL, Keep, Sink, Source}
import akka.stream.{FlowShape, SourceShape}
import sample.stream_actor.WindTurbineSimulator._

import scala.concurrent.duration._
import scala.concurrent.{ExecutionContext, Future}
import scala.util.{Failure, Success}

object WebSocketClient {
  def apply(id: String, endpoint: String, windTurbineSimulator: ActorRef)
           (implicit
            system: ActorSystem,
            executionContext: ExecutionContext) = {
    new WebSocketClient(id, endpoint, windTurbineSimulator)(system, executionContext)
  }
}

class WebSocketClient(id: String, endpoint: String, windTurbineSimulator: ActorRef)
                     (implicit
                      system: ActorSystem,
                      executionContext: ExecutionContext) {


  val webSocketFlow: Flow[Message, Message, Future[WebSocketUpgradeResponse]] = {
    val websocketUri = s"$endpoint/measurements/$id"
    Http().webSocketClientFlow(WebSocketRequest(websocketUri))
  }

  val outgoing = GraphDSL.create() { implicit builder =>
    val data = WindTurbineData(id)

    val flow = builder.add {
      Source.tick(1.second, 100.millis,())  //valve for the WindTurbineData frequency
        .map(_ => TextMessage(data.getNext))
    }

    SourceShape(flow.out)
  }

  val incoming = GraphDSL.create() { implicit builder =>
    val flow = builder.add {
      Flow[Message]
        .collect {
          case TextMessage.Strict(text) =>
            Future.successful(text)
          case TextMessage.Streamed(textStream) =>
            textStream.runFold("")(_ + _)
              .flatMap(Future.successful)
        }
        .mapAsync(1)(identity)
        .map(each => println(s"Client received msg: $each"))
    }

    FlowShape(flow.in, flow.out)
  }

  val (upgradeResponse, closed) = Source.fromGraph(outgoing)
    .viaMat(webSocketFlow)(Keep.right) // keep the materialized Future[WebSocketUpgradeResponse]
    .via(incoming)
    .toMat(Sink.ignore)(Keep.both) // also keep the Future[Done]
    .run()


  val connected =
    upgradeResponse.map { upgrade =>
      upgrade.response.status match {
        case StatusCodes.SwitchingProtocols => windTurbineSimulator ! Upgraded
        case statusCode => windTurbineSimulator ! FailedUpgrade(statusCode)
      }
    }

  connected.onComplete {
    case Success(_) => windTurbineSimulator ! Connected
    case Failure(ex) => windTurbineSimulator ! ConnectionFailure(ex)
  }

  closed.map { _ =>
    windTurbineSimulator ! Terminated
  }
  closed.onComplete {
    case Success(_)  => windTurbineSimulator ! Connected
    case Failure(ex) => windTurbineSimulator ! ConnectionFailure(ex)
  }
} 

package de.nierbeck.floating.data.server

import akka.actor.{ActorRef, Props}
import akka.http.scaladsl.model.ws.{Message, TextMessage}
import akka.stream.FlowShape
import akka.stream.scaladsl.{Flow, GraphDSL, Merge, Source}
import de.nierbeck.floating.data.domain.Vehicle
import GraphDSL.Implicits._
import de.nierbeck.floating.data.server._
import de.nierbeck.floating.data.server.actors.websocket._


object Flows {

  def graphFlowWithStats(router: ActorRef): Flow[Message, Message, _] = {
    Flow.fromGraph(GraphDSL.create() { implicit builder =>


      // create an actor source
      val source = Source.actorPublisher[String](VehiclePublisher.props(router))

      // Graph elements we'll use
      val merge = builder.add(Merge[String](2))
      val filter = builder.add(Flow[String].filter(_ => false))

      // get BBox from request and send it to route, return nothing ...
      val mapMsgToString = builder.add(Flow[Message].map[String] {
        case TextMessage.Strict(msg) => {
          println(s"received message: $msg")
          if (msg.contains("close")) {
            router ! msg
          } else if (msg.contains("spark")) {
            router ! SPARK
          } else if (msg.contains("flink")) {
            router ! FLINK
          } else {
            val bbox = toBoundingBox(msg)
            println(s"transformedt to bbox: $bbox")
            router ! bbox
          }
          ""
        }
      })
      //outgoing message ...
      val mapStringToMsg = builder.add(Flow[String].map[Message](x => TextMessage.Strict(x)))

      //add source to flow
      val vehiclesSource = builder.add(source)

      // connect the graph
      mapMsgToString ~> filter ~> merge // this part of the merge will never provide msgs
      vehiclesSource ~> merge ~> mapStringToMsg

      // expose ports
      FlowShape(mapMsgToString.in, mapStringToMsg.out)
    })
  }

} 

package ch13

import akka.NotUsed
import akka.stream.FlowShape
import akka.stream.scaladsl.{Balance, Flow, GraphDSL, Merge}


object Balancer {
  def apply[In, Out](subFlow: Flow[In, Out, Any],
                     count: Int): Flow[In, Out, NotUsed] = {

    Flow.fromGraph(createGraph(subFlow, count))
  }

  import akka.stream.scaladsl.GraphDSL
  import GraphDSL.Implicits._

  def createGraph[Out, In](subFlow: Flow[In, Out, Any], count: Int) = {
    val balanceBlock  = Balance[In](count, waitForAllDownstreams = false)
    val mergeBlock = Merge[Out](count, eagerComplete = false)
    GraphDSL.create() { implicit builder ⇒
      val balancer = builder.add(balanceBlock)
      val merge = builder.add(mergeBlock)

      for (_ ← 1 to count) balancer ~> subFlow ~> merge

      FlowShape(balancer.in, merge.out)
    }
  }
} 

package cmwell.dc.stream.akkautils

import akka.NotUsed
import akka.stream.{FlowShape, Graph}
import akka.stream.scaladsl.{Balance, Flow, GraphDSL, Merge}


object ConcurrentFlow {

  def apply[I, O](parallelism: Int)(flow: Graph[FlowShape[I, O], NotUsed]): Graph[FlowShape[I, O], NotUsed] =
    GraphDSL.create() { implicit builder =>
      import GraphDSL.Implicits._
      val balancer = builder.add(Balance[I](parallelism))
      val merger = builder.add(Merge[O](parallelism))
      for (i <- 0 until parallelism) {
        balancer.out(i) ~> flow.async ~> merger.in(i)
      }
      FlowShape(balancer.in, merger.out)
    }
} 

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

import akka.NotUsed
import akka.actor.ActorSystem
import akka.stream.{ActorMaterializer, FlowShape}
import akka.stream.scaladsl.{Balance, Flow, GraphDSL, Merge, Sink, Source}

import scala.util.Random

trait PipeliningParallelizing extends App {

  implicit val actorSystem = ActorSystem("PipeliningParallelizing")
  implicit val actorMaterializer = ActorMaterializer()

  case class Wash(id: Int)
  case class Dry(id: Int)
  case class Done(id: Int)

  val tasks = (1 to 5).map(Wash)

  def washStage = Flow[Wash].map(wash => {
    val sleepTime = Random.nextInt(3) * 1000
    println(s"Washing ${wash.id}. It will take $sleepTime milliseconds.")
    Thread.sleep(sleepTime)
    Dry(wash.id)
  })

  def dryStage = Flow[Dry].map(dry => {
    val sleepTime = Random.nextInt(3) * 1000
    println(s"Drying ${dry.id}. It will take $sleepTime milliseconds.")
    Thread.sleep(sleepTime)
    Done(dry.id)
  })

  val parallelStage = Flow.fromGraph(GraphDSL.create() { implicit builder =>
    import GraphDSL.Implicits._

    val dispatchLaundry = builder.add(Balance[Wash](3))
    val mergeLaundry = builder.add(Merge[Done](3))

    dispatchLaundry.out(0) ~> washStage.async ~> dryStage.async ~> mergeLaundry.in(0)
    dispatchLaundry.out(1) ~> washStage.async ~> dryStage.async ~> mergeLaundry.in(1)
    dispatchLaundry.out(2) ~> washStage.async ~> dryStage.async ~> mergeLaundry.in(2)

    FlowShape(dispatchLaundry.in, mergeLaundry.out)
  })

  def runGraph(testingFlow: Flow[Wash, Done, NotUsed]) = Source(tasks).via(testingFlow).to(Sink.foreach(println)).run()
} 

package docs.scaladsl.mb

import akka.Done
import akka.NotUsed
import akka.stream.FlowShape
import akka.stream.scaladsl.Flow
import akka.stream.scaladsl.GraphDSL
import akka.stream.scaladsl.GraphDSL.Implicits._
import akka.stream.scaladsl.Merge
import akka.stream.scaladsl.Partition
import com.lightbend.lagom.scaladsl.api.ServiceCall
import com.lightbend.lagom.scaladsl.api.broker.Message

//#inject-service
class AnotherServiceImpl(helloService: HelloService) extends AnotherService {
  //#inject-service

  //#subscribe-to-topic
  helloService
    .greetingsTopic()
    .subscribe // <-- you get back a Subscriber instance
    .atLeastOnce(
      Flow.fromFunction(doSomethingWithTheMessage)
    )
  //#subscribe-to-topic

  var lastObservedMessage: String = _

  private def doSomethingWithTheMessage(greetingMessage: GreetingMessage): Done = {
    lastObservedMessage = greetingMessage.message
    Done
  }

  import scala.concurrent.ExecutionContext.Implicits.global

  override def foo: ServiceCall[NotUsed, String] = ServiceCall { req =>
    scala.concurrent.Future.successful(lastObservedMessage)
  }

  def subscribeWithMetadata = {
    //#subscribe-to-topic-with-metadata
    import com.lightbend.lagom.scaladsl.api.broker.Message
    import com.lightbend.lagom.scaladsl.broker.kafka.KafkaMetadataKeys

    helloService
      .greetingsTopic()
      .subscribe
      .withMetadata
      .atLeastOnce(
        Flow[Message[GreetingMessage]].map { msg =>
          val greetingMessage = msg.payload
          val messageKey      = msg.messageKeyAsString
          val kafkaHeaders    = msg.get(KafkaMetadataKeys.Headers)
          println(s"Message: $greetingMessage Key: $messageKey Headers: $kafkaHeaders")
          Done
        }
      )
    //#subscribe-to-topic-with-metadata
  }

  def skipMessages = {
    //#subscribe-to-topic-skip-messages
    helloService
      .greetingsTopic()
      .subscribe
      .atLeastOnce(
        Flow[GreetingMessage].map {
          case msg @ GreetingMessage("Kia ora") => doSomethingWithTheMessage(msg)
          case _                                => Done // Skip all messages where the message is not "Kia ora".
        }
      )
    //#subscribe-to-topic-skip-messages
  }
} 

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

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