akka.stream.stage.InHandler Scala Examples

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

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

object ExecuteAfterResponse {

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

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

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

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

    override val shape = FlowShape(in, out)
  }

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

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

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

package de.knutwalker.akka.http

import de.knutwalker.akka.stream.JsonStreamParser

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

import jawn.Facade

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

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

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

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

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

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

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

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

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

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

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

    override def toString: String = "FirstElementSinkStage"
  }
}

trait JsonSupport {

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

package cmwell.crawler

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

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

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

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

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

package akka.grpc.internal

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


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

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

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

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

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

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

package example.akkawschat.cli

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

import scala.concurrent.ExecutionContext

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

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

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

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

package com.sksamuel.pulsar4s.akka.streams

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

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

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

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

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

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

    val promise = Promise[Done]()

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

      implicit def context: ExecutionContextExecutor = super.materializer.executionContext

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

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

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

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

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

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

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

    (logic, promise.future)
  }

} 

package com.sksamuel.pulsar4s.akka.streams

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

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

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

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

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

    val promise = Promise[Done]()

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

      implicit def context: ExecutionContextExecutor = super.materializer.executionContext

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

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

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

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

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

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

    (logic, promise.future)
  }
} 

package org.squbs.pattern.stream

import java.io.File

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

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

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

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

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

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

  protected def autoCommit(index: Long) = {}

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

    private var lastPushed = 0L
    var downstreamWaiting = false

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

    setHandler(in, new InHandler {

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

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

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

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

    setHandler(out, new OutHandler {

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

package org.squbs.stream

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


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

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

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

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

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

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

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

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

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

  def out: Outlet[T] = shape.out

  override def toString = "TriggerMerge"
}

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

  import GraphDSL.Implicits._

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

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

package org.squbs.streams

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

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

object Deduplicate {

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

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

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


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

package cmwell.dc.stream.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 uk.gov.hmrc.akka.xml

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


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

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

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

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

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

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

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

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

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

      }
  }

} 

package io.vamp.common.http

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


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

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

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

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

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

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

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

package 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.util

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

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

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

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

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

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

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

package sangria.streaming

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

import scala.concurrent.Future

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        setHandlers(in, out, this)
      }
  }
} 

package com.netflix.iep.lwc

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

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

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

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

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

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

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

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

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

      setHandlers(in, out, this)
    }
  }
} 

package com.netflix.atlas.persistence

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

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

import scala.concurrent.duration._

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

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

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

    new TimerGraphStageLogic(shape) with InHandler with OutHandler {

      private var hourlyWriter: HourlyRollingWriter = _

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

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

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

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

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

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

      setHandlers(in, out, this)

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

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

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 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.util.stream

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

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

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

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

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

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

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

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

          pending = elem
        }

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

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

package cmwell.util.stream

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

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

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

      private[this] var sum = 0L

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

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

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

      setHandlers(in, out, this)
    }
} 

package cool.graph.akkautil.stream

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

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

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

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

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

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

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

      setHandlers(in, out, this)
    }
} 

// 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.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, Inlet, SinkShape}
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler}
import io.grpc.stub.CallStreamObserver

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

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

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

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

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

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

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

package com.comcast.money.akka.acceptance.stream

import akka.stream.Attributes
import akka.stream.scaladsl.{ Keep, Sink, Source }
import akka.stream.stage.{ InHandler, OutHandler }
import com.comcast.money.akka.Blocking.RichFuture
import com.comcast.money.akka.SpanHandlerMatchers.{ haveSomeSpanNames, maybeCollectingSpanHandler }
import com.comcast.money.akka.stream.{ TracedFlow, TracedFlowLogic }
import com.comcast.money.akka.{ AkkaMoneyScope, MoneyExtension, SpanContextWithStack }
import org.scalatest.Ignore

class TracedFlowSpec extends AkkaMoneyScope {

  "MoneyExtension should pass a span through an Akka Stream" in {
    implicit val moneyExtension: MoneyExtension = MoneyExtension(actorSystem)
    implicit val spanContextWithStack: SpanContextWithStack = new SpanContextWithStack

    testStream().get()

    maybeCollectingSpanHandler should haveSomeSpanNames(testSpanNames)
  }

  "MoneyExtension should pass a span through an asynchronous Akka Stream" in {
    implicit val moneyExtension: MoneyExtension = MoneyExtension(actorSystem)
    implicit val spanContextWithStack: SpanContextWithStack = new SpanContextWithStack

    multithreadedTestStream().get()

    maybeCollectingSpanHandler should haveSomeSpanNames(testSpanNames)
  }

  val testSpanNames = Seq("flow-3", "flow-2", "flow-1")

  def testStream()(implicit spanContextWithStack: SpanContextWithStack, moneyExtension: MoneyExtension) =
    Source[(String, SpanContextWithStack)](List(("", spanContextWithStack)))
      .via(new TestFlowShape("flow-1"))
      .via(new TestFlowShape("flow-2"))
      .via(new TestFlowShape("flow-3", isFinalFlow = true))
      .runWith(Sink.seq)

  def multithreadedTestStream()(implicit spanContextWithStack: SpanContextWithStack, moneyExtension: MoneyExtension) =
    Source[(String, SpanContextWithStack)](List(("", spanContextWithStack)))
      .via(new TestFlowShape("flow-1").async)
      .via(new TestFlowShape("flow-2").async)
      .via(new TestFlowShape("flow-3", isFinalFlow = true).async)
      .runWith(Sink.seq)

  class TestFlowShape(id: String, isFinalFlow: Boolean = false)(implicit moneyExtension: MoneyExtension) extends TracedFlow[String, String] {

    override val inletName: String = "testin"
    override val outletName: String = "testout"

    override def createLogic(inheritedAttributes: Attributes) =
      new TracedFlowLogic {
        setHandler(in, new InHandler {
          override def onPush(): Unit = {
            val logic = (msg: String) => s"$msg$id"
            if (isFinalFlow) stopTracePush(key = id, stageLogic = logic)
            else tracedPush(id, logic)
          }
        })

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

} 

package akka.stream.integration
package activemq

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

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

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

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

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

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

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

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

package me.snov.akka.sqs.shape

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

import scala.concurrent.Promise

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

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

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

      tryPull()
    }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  override def toString = "Slf4JLog"

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

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

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

      override def onPush(): Unit = {

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

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

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

        super.onUpstreamFailure(cause)
      }

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

        super.onUpstreamFinish()
      }

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

        super.onDownstreamFinish(cause)
      }

      setHandlers(in, out, this)
    }
}

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

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