Scala GraphDSL类代码示例

//设置package包名称以及导入依赖的类
package com.packt.chapter8

import akka.actor.ActorSystem
import akka.kafka.scaladsl.{Consumer, Producer}
import akka.kafka.{ConsumerSettings, ProducerSettings, Subscriptions}
import akka.stream.{ActorMaterializer, ClosedShape}
import akka.stream.scaladsl.{Flow, GraphDSL, RunnableGraph, Sink, Source}
import org.apache.kafka.clients.consumer.{ConsumerConfig, ConsumerRecord}
import org.apache.kafka.clients.producer.ProducerRecord
import org.apache.kafka.common.TopicPartition
import org.apache.kafka.common.serialization.{ByteArrayDeserializer, ByteArraySerializer, StringDeserializer, StringSerializer}

import scala.concurrent.duration._

object ProcessingKafkaApplication extends App {
  implicit val actorSystem = ActorSystem("SimpleStream")
  implicit val actorMaterializer = ActorMaterializer()

  val bootstrapServers = "localhost:9092"
  val kafkaTopic = "akka_streams_topic"
  val partition = 0
  val subscription = Subscriptions.assignment(new TopicPartition(kafkaTopic, partition))

  val consumerSettings = ConsumerSettings(actorSystem, new ByteArrayDeserializer, new StringDeserializer)
    .withBootstrapServers(bootstrapServers)
    .withGroupId("akka_streams_group")
    .withProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest")

  val producerSettings = ProducerSettings(actorSystem, new ByteArraySerializer, new StringSerializer)
    .withBootstrapServers(bootstrapServers)

  val runnableGraph = RunnableGraph.fromGraph(GraphDSL.create() { implicit builder =>
    import GraphDSL.Implicits._

    val tickSource = Source.tick(0 seconds, 5 seconds, "Hello from Akka Streams using Kafka!")
    val kafkaSource = Consumer.plainSource(consumerSettings, subscription)
    val kafkaSink = Producer.plainSink(producerSettings)
    val printlnSink = Sink.foreach(println)

    val mapToProducerRecord = Flow[String].map(elem => new ProducerRecord[Array[Byte], String](kafkaTopic, elem))
    val mapFromConsumerRecord = Flow[ConsumerRecord[Array[Byte], String]].map(record => record.value())

    tickSource  ~> mapToProducerRecord   ~> kafkaSink
    kafkaSource ~> mapFromConsumerRecord ~> printlnSink

    ClosedShape
  })

  runnableGraph.run()
} 

//设置package包名称以及导入依赖的类
package eu.svez.akka.stream.cats

import akka.NotUsed
import akka.stream.SinkShape
import akka.stream.scaladsl.{GraphDSL, Sink, Source}
import akka.stream.testkit.TestSubscriber
import cats.implicits._
import eu.svez.akka.stream.cats.Stages._

class PartitionValidatedSpec extends StageSpec {

  "PartitionValidated" should "partition a flow of Validation[E, A] in two flows of E and A" in new Test {
    val src = Source(List(
      1.valid[String],
      2.valid[String],
      "BOOM!".invalid[Int],
      3.valid[String],
      "BOOM 2!".invalid[Int]
    ))

    src.runWith(testSink)

    successProbe.request(3)
    failureProbe.request(2)

    successProbe.expectNext(1)
    successProbe.expectNext(2)
    successProbe.expectNext(3)
    failureProbe.expectNext("BOOM!")
    failureProbe.expectNext("BOOM 2!")
    successProbe.expectComplete()
    failureProbe.expectComplete()
  }

  trait Test {
    val failureProbe = TestSubscriber.probe[String]()
    val successProbe = TestSubscriber.probe[Int]()

    val testSink = Sink.fromGraph(GraphDSL.create() { implicit builder: GraphDSL.Builder[NotUsed] =>
      import GraphDSL.Implicits._

      val valStage = builder.add(PartitionValidated[String, Int]())

      valStage.invalid ~> Sink.fromSubscriber(failureProbe)
      valStage.valid ~> Sink.fromSubscriber(successProbe)

      SinkShape(valStage.in)
    })
  }
} 

//设置package包名称以及导入依赖的类
package eu.svez.akka.stream

import akka.NotUsed
import akka.stream.FanOutShape2
import akka.stream.scaladsl.{Flow, GraphDSL, Partition}

import scala.util.{Failure, Success, Try}

object Stages {

  object PartitionEither {
    def apply[A, B]() = GraphDSL.create[FanOutShape2[Either[A, B], A, B]]() { implicit builder: GraphDSL.Builder[NotUsed] =>
      import GraphDSL.Implicits._

      val left = builder.add(Flow[Either[A, B]].map (_.left.get))
      val right = builder.add(Flow[Either[A, B]].map (_.right.get))
      val partition = builder.add(Partition[Either[A, B]](2, _.fold(_ ? 0, _ ? 1)))

      partition ~> left
      partition ~> right

      new FanOutShape2[Either[A, B], A, B](partition.in, left.out, right.out)
    }
  }

  implicit class EitherShape[A, B](val shape: FanOutShape2[Either[A, B], A, B]) extends AnyVal {
    def left = shape.out0
    def right = shape.out1
  }

  object PartitionTry {
    def apply[T]() = GraphDSL.create[FanOutShape2[Try[T], Throwable, T]]() { implicit builder ?
      import GraphDSL.Implicits._

      val success = builder.add(Flow[Try[T]].collect { case Success(a) ? a })
      val failure = builder.add(Flow[Try[T]].collect { case Failure(t) ? t })
      val partition = builder.add(Partition[Try[T]](2, _.map(_ ? 1).getOrElse(0)))

      partition ~> failure
      partition ~> success

      new FanOutShape2[Try[T], Throwable, T](partition.in, failure.out, success.out)
    }
  }

  implicit class TryShape[T](val shape: FanOutShape2[Try[T], Throwable, T]) extends AnyVal {
    def failure = shape.out0
    def success = shape.out1
  }

} 

//设置package包名称以及导入依赖的类
package eu.svez.akka.stream

import akka.NotUsed
import akka.stream.SinkShape
import akka.stream.scaladsl.{GraphDSL, Sink, Source}
import akka.stream.testkit.TestSubscriber
import Stages._

class PartitionEitherSpec extends StageSpec {

  "PartitionEither" should "partition a flow of Either[A, B] into two flows of A and B" in new Test {
    val src = Source(List(
      Right(1),
      Right(2),
      Left("One"),
      Right(3),
      Left("Two")
    ))

    src.runWith(testSink)

    rightProbe.request(4)
    leftProbe.request(3)

    rightProbe.expectNext(1)
    rightProbe.expectNext(2)
    rightProbe.expectNext(3)
    leftProbe.expectNext("One")
    leftProbe.expectNext("Two")
    rightProbe.expectComplete()
    leftProbe.expectComplete()
  }

  trait Test {
    val leftProbe = TestSubscriber.probe[String]()
    val rightProbe = TestSubscriber.probe[Int]()

    val testSink = Sink.fromGraph(GraphDSL.create() { implicit builder: GraphDSL.Builder[NotUsed] =>
      import GraphDSL.Implicits._

      val eitherStage = builder.add(PartitionEither[String, Int]())

      eitherStage.left ~> Sink.fromSubscriber(leftProbe)
      eitherStage.right ~> Sink.fromSubscriber(rightProbe)

      SinkShape(eitherStage.in)
    })
  }
} 

//设置package包名称以及导入依赖的类
package sample.graphDSL

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


object FlowFromGraph {

  def main(args: Array[String]): Unit = {
    implicit val system = ActorSystem("FlowFromGraph")
    implicit val ec = system.dispatcher
    implicit val materializer = ActorMaterializer()

    val processorFlow1: Flow[Int, Int, NotUsed] = Flow[Int].map(_ * 2)
    val processorFlow2: Flow[Int, Int, NotUsed] = Flow[Int].map(_ * 3)
    val listOfFlows = List(processorFlow1, processorFlow2)

    def compoundFlowFrom[T](indexFlows: Seq[Flow[T, T, NotUsed]]): Flow[T, T, NotUsed] = {
      require(indexFlows.nonEmpty, "Cannot create compound flow without any flows to combine")

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

        val broadcast: UniformFanOutShape[T, T] = b.add(Broadcast(indexFlows.size))
        val merge: UniformFanInShape[T, T] = b.add(Merge(indexFlows.size))

        indexFlows.foreach(broadcast ~> _ ~> merge)

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

    val compoundFlow = compoundFlowFrom(listOfFlows)

    Source(1 to 10)
      .via(compoundFlow)
      .runWith(Sink.foreach(println(_)))
      .onComplete(_ => system.terminate())
  }
} 

//设置package包名称以及导入依赖的类
package services

import javax.inject.{Inject, Singleton}

import akka.NotUsed
import akka.stream._
import akka.stream.scaladsl.{Broadcast, Flow, GraphDSL, Sink}
import play.api.Logger
import play.api.libs.concurrent.Execution.Implicits._

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


  def stringPersister(pf: String => Future[Unit]): Flow[String, String, NotUsed] =
    Flow.fromGraph(GraphDSL.create() { implicit builder =>
      import GraphDSL.Implicits._

      val persistenceSink = Sink.foreach[String] { content =>
        val f = pf(content)
        f.onComplete {
          case Success(u) => Logger.debug(s"Persisted content: '$content'")
          case Failure(t) => Logger.error(s"Failed to persist content: '$content", t)
        }
      }

      val bcast = builder.add(Broadcast[String](2))
      bcast.out(1) ~> persistenceSink

      FlowShape(bcast.in, bcast.out(0))
    })
} 

//设置package包名称以及导入依赖的类
package org.hpi.esb.datavalidator.validation

import akka.NotUsed
import akka.stream._
import akka.stream.scaladsl.{Flow, GraphDSL}
import org.hpi.esb.datavalidator.configuration.Config
import org.hpi.esb.datavalidator.data.{SimpleRecord, Statistics}
import org.hpi.esb.datavalidator.kafka.TopicHandler
import org.hpi.esb.datavalidator.validation.graphstage.{AccumulateWhileUnchanged, IgnoreLastElements, ZipWhileEitherAvailable}

class StatisticsValidation(inTopicHandler: TopicHandler,
                           outTopicHandler: TopicHandler, windowSize: Long,
                           materializer: ActorMaterializer)
  extends Validation[Statistics](inTopicHandler, outTopicHandler, materializer) {

  override val valueName = "Statistics"
  override val queryName = "Statistics Query"


  val collectByWindow = new AccumulateWhileUnchanged[SimpleRecord, Long](r => windowStart(r.timestamp))
  val calculateStatistics = Flow[Seq[SimpleRecord]].map(s =>
    s.foldLeft(new Statistics()())((stats, record) => stats.getUpdatedWithValue(record.timestamp, record.value)))


  def createSource(): Graph[SourceShape[(Option[Statistics], Option[Statistics])], NotUsed] = {

    GraphDSL.create() { implicit builder =>
      import GraphDSL.Implicits._

      val zip = builder.add(ZipWhileEitherAvailable[Statistics]())
      val ignoreLastTwoElements = builder.add(new IgnoreLastElements[(Option[Statistics], Option[Statistics])](ignoreCount = 1))

      inTopicHandler.topicSource ~> take(inNumberOfMessages) ~> toSimpleRecords ~> collectByWindow ~> calculateStatistics ~> zip.in0
      outTopicHandler.topicSource ~> take(outNumberOfMessages) ~> toStatistics ~> zip.in1
      zip.out ~> ignoreLastTwoElements

      SourceShape(ignoreLastTwoElements.out)
    }
  }

  def windowStart(timestamp: Long): Long = {
    timestamp - (timestamp % windowSize)
  }
} 

//设置package包名称以及导入依赖的类
package org.hpi.esb.datavalidator.validation

import akka.NotUsed
import akka.stream.scaladsl.GraphDSL
import akka.stream.{ActorMaterializer, Graph, SourceShape}
import org.hpi.esb.datavalidator.configuration.Config
import org.hpi.esb.datavalidator.data.SimpleRecord
import org.hpi.esb.datavalidator.kafka.TopicHandler
import org.hpi.esb.datavalidator.validation.graphstage.ZipWhileEitherAvailable

class IdentityValidation(inTopicHandler: TopicHandler,
                         outTopicHandler: TopicHandler,
                         materializer: ActorMaterializer)
  extends Validation[SimpleRecord](inTopicHandler, outTopicHandler, materializer) {

  override val valueName = "SimpleRecords"
  override val queryName = "Identity Query"

  def createSource(): Graph[SourceShape[(Option[SimpleRecord], Option[SimpleRecord])], NotUsed] = {

    GraphDSL.create() { implicit builder =>
      import GraphDSL.Implicits._

      val zip = builder.add(ZipWhileEitherAvailable[SimpleRecord]())

      inTopicHandler.topicSource ~> take(inNumberOfMessages) ~> toSimpleRecords ~> zip.in0
      outTopicHandler.topicSource ~> take(outNumberOfMessages) ~> toSimpleRecords ~> zip.in1

      SourceShape(zip.out)
    }
  }
} 

//设置package包名称以及导入依赖的类
package akka_in_action.streams

import akka.actor.ActorSystem
import akka.stream.{ActorMaterializer, SourceShape}
import akka.stream.scaladsl.{GraphDSL, Sink, Source, Zip}
import GraphDSL.Implicits._
import scala.concurrent.duration._
import scala.concurrent.{Await, Future}

object PairsShape_ extends App {
  implicit val system = ActorSystem()
  implicit val ec = system.dispatcher
  implicit val materializer = ActorMaterializer()

  val pairs = Source.fromGraph(GraphDSL.create() { implicit builder =>
    val zip = builder.add(Zip[Int, Int]())

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

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

    SourceShape(zip.out)
  })

  val firstPair: Future[(Int, Int)] = pairs.runWith(Sink.head)
  val r = Await.result(firstPair, 300 millis)
  println(r)

  system.terminate()
} 

//设置package包名称以及导入依赖的类
package akka_in_action.streams

import akka.actor.ActorSystem
import akka.stream.{ActorMaterializer, ClosedShape}
import akka.stream.scaladsl.{Balance, Broadcast, Flow, GraphDSL, RunnableGraph, Sink, Source}

import scala.concurrent.Await
import scala.concurrent.duration._

object GraphBuilder extends App {
  implicit val system = ActorSystem("actor-system")
  implicit val materializer = ActorMaterializer()

  val topHeadSink = Sink.head[Int]
  val middleSink = Sink.head[Int]
  val bottomHeadSink = Sink.seq[Int]
  val sharedDoubler = Flow[Int].map(_ * 2)

  val results = RunnableGraph.fromGraph(GraphDSL.create(topHeadSink, middleSink, bottomHeadSink)
  ((_, _, _)) { implicit builder =>
    (topHS, midHS, bottomHS) =>
      import GraphDSL.Implicits._
      val broadcast = builder.add(Broadcast[Int](2))
      val source = builder.add(Source(1 to 10))

      Source.repeat(2) ~> broadcast.in
      broadcast.out(0) ~> sharedDoubler ~> topHS.in
      broadcast.out(1) ~> sharedDoubler ~> midHS.in
      source ~> bottomHS.in

      ClosedShape
  }).run()

  val r1 = Await.result(results._1, 300 millis)
  val r2 = Await.result(results._2, 300 millis)
  val r3 = Await.result(results._3, 300 millis)

  println(r1, r2, r3)
  system.terminate()
} 

//设置package包名称以及导入依赖的类
package akka_in_action.streams

import akka.actor.ActorSystem
import akka.stream.{ActorMaterializer, FlowShape, SourceShape}
import akka.stream.scaladsl.{Flow, GraphDSL, Sink, Source}
import GraphDSL.Implicits._

import scala.concurrent.duration._
import scala.concurrent.{Await, Future}

object MaterializedValue_ extends App {
  implicit val system = ActorSystem()
  implicit val ec = system.dispatcher
  implicit val materializer = ActorMaterializer()

  val foldFlow: Flow[Int, Int, Future[Int]] = Flow.fromGraph(
    GraphDSL.create(Sink.fold[Int, Int](0)(_ + _)) { implicit builder => fold =>
      FlowShape(fold.in, builder.materializedValue.mapAsync(4)(identity).outlet)
    })
  val r = Source(1 to 10)
    .via(foldFlow)
    .runWith(Sink.head)
  println(Await.result(r, 200 millis))

  val cyclicFold: Source[Int, Future[Int]] = Source.fromGraph(
    GraphDSL.create(Sink.fold[Int, Int](0)(_ + _)) { implicit builder => fold =>
      Source(1 to 10) ~> fold
      SourceShape(builder.materializedValue.mapAsync(4)(identity).outlet)
    })
  val r2 = cyclicFold
    .via(foldFlow)
    .runWith(Sink.head)
  println(Await.result(r2, 200 millis))

  system.terminate()
} 

//设置package包名称以及导入依赖的类
import java.nio.file.FileSystems

import akka.NotUsed
import akka.actor.ActorSystem
import akka.stream.{ActorMaterializer, ClosedShape, ThrottleMode}
import akka.stream.scaladsl.{GraphDSL, Merge, RunnableGraph, Sink, Source}
import akka.stream.alpakka.file.scaladsl
import org.apache.spark.streamdm.core.ExampleParser
import org.apache.spark.streamdm.core.specification.SpecificationParser
import pl.gosub.akka.online.{HoeffdingTreeProcessor, LearnerQuery}

import scala.concurrent.Await
import scala.concurrent.duration.Duration
import scala.concurrent.duration._


object HoeffdingTreeWithAlpakka extends App {
  implicit val system = ActorSystem()
  implicit val mat = ActorMaterializer()

  val specParser = new SpecificationParser

  val arffPath = this.getClass.getResource("/elecNormNew.arff").getPath // add for Windows .replaceFirst("^/(.:/)", "$1")

  val exampleSpec = specParser.fromArff(arffPath)

  val fsPath = this.getClass.getResource("/elecNormData.txt").getPath // add for Windows .replaceFirst("^/(.:/)", "$1")
  println(fsPath)


  val fs = FileSystems.getDefault
  val lines: Source[String, NotUsed] = scaladsl.FileTailSource.lines(
    path = fs.getPath(fsPath),
    maxLineSize = 8192,
    pollingInterval = 250.millis
  )

//  if the lines below do not work, please make sure that you got the linefeed character right wrt your operating system (LF vs CRLF)
//  lines.map(line => LearnerQuery(line.split(";").apply(0), ExampleParser.fromArff(line.split(";").apply(1), exampleSpec)))
//    .runForeach(line => System.out.println(line))

  val masterControlProgram = RunnableGraph.fromGraph(GraphDSL.create(Sink.foreach(print)) { implicit builder =>
    outMatches =>
      import GraphDSL.Implicits._
      val taggedInput = lines.map(line => LearnerQuery(line.split(";").apply(0), ExampleParser.fromArff(line.split(";").apply(1), exampleSpec)))

      taggedInput.statefulMapConcat(() => {
        val proc = new HoeffdingTreeProcessor(exampleSpec)
        proc.process(_)
      }) ~> outMatches
      ClosedShape
  }).run()

  import scala.concurrent.ExecutionContext.Implicits.global

  masterControlProgram.onComplete(_ => system.terminate())
  Await.ready(system.whenTerminated, Duration.Inf)

} 

//设置package包名称以及导入依赖的类
package eu.svez.akka.stream.cats

import akka.NotUsed
import akka.stream.SinkShape
import akka.stream.scaladsl.{GraphDSL, Sink, Source}
import akka.stream.testkit.TestSubscriber
import cats.data.NonEmptyList
import cats.implicits._
import eu.svez.akka.stream.cats.Stages._

class PartitionValidatedNelSpec extends StageSpec {

  "PartitionValidatedNel" should "partition a flow of Validation[E, A] in two flows of E and A" in new Test {
    val src = Source(List(
      1.valid[NonEmptyList[String]],
      2.valid[NonEmptyList[String]],
      NonEmptyList.of("BOOM!", "KABOOM!").invalid[Int],
      3.valid[NonEmptyList[String]],
      NonEmptyList.of("BOOM 2!").invalid[Int]
    ))

    src.runWith(testSink)

    successProbe.request(3)
    failureProbe.request(3)

    successProbe.expectNext(1)
    successProbe.expectNext(2)
    successProbe.expectNext(3)
    failureProbe.expectNext("BOOM!")
    failureProbe.expectNext("KABOOM!")
    failureProbe.expectNext("BOOM 2!")
    successProbe.expectComplete()
    failureProbe.expectComplete()
  }

  trait Test {
    val failureProbe = TestSubscriber.probe[String]()
    val successProbe = TestSubscriber.probe[Int]()

    val testSink = Sink.fromGraph(GraphDSL.create() { implicit builder: GraphDSL.Builder[NotUsed] =>
      import GraphDSL.Implicits._

      val valStage = builder.add(PartitionValidatedNel[String, Int]())

      valStage.invalid ~> Sink.fromSubscriber(failureProbe)
      valStage.valid ~> Sink.fromSubscriber(successProbe)

      SinkShape(valStage.in)
    })
  }
} 

//设置package包名称以及导入依赖的类
package eu.svez.akka.stream.cats

import akka.NotUsed
import akka.stream.SinkShape
import akka.stream.scaladsl.{GraphDSL, Sink, Source}
import akka.stream.testkit.TestSubscriber
import cats.data.Ior
import eu.svez.akka.stream.cats.Stages._

class PartitionIorSpec extends StageSpec {

  "PartitionIor" should "partition a flow of Ior[A, B] into two flows of A and B" in new Test {
    val src = Source(List(
      Ior.Right(1),
      Ior.Right(2),
      Ior.Left("One"),
      Ior.Right(3),
      Ior.Left("Two"),
      Ior.Both("Three", 4)
    ))

    src.runWith(testSink)

    rightProbe.request(4)
    leftProbe.request(3)

    rightProbe.expectNext(1)
    rightProbe.expectNext(2)
    rightProbe.expectNext(3)
    rightProbe.expectNext(4)
    leftProbe.expectNext("One")
    leftProbe.expectNext("Two")
    leftProbe.expectNext("Three")
    rightProbe.expectComplete()
    leftProbe.expectComplete()
  }

  trait Test {
    val leftProbe = TestSubscriber.probe[String]()
    val rightProbe = TestSubscriber.probe[Int]()

    val testSink = Sink.fromGraph(GraphDSL.create() { implicit builder: GraphDSL.Builder[NotUsed] =>
      import GraphDSL.Implicits._

      val iorStage = builder.add(PartitionIor[String, Int]())

      iorStage.left ~> Sink.fromSubscriber(leftProbe)
      iorStage.right ~> Sink.fromSubscriber(rightProbe)

      SinkShape(iorStage.in)
    })
  }
} 

//设置package包名称以及导入依赖的类
package eu.svez.akka.stream

import akka.NotUsed
import akka.stream.SinkShape
import akka.stream.scaladsl.{GraphDSL, Sink, Source}
import akka.stream.testkit.TestSubscriber
import eu.svez.akka.stream.Stages._

import scala.util.{Failure, Success}

class PartitionTrySpec extends StageSpec {

  "PartitionTry" should "partition a flow of Try[T] into two flows of Throwable and T" in new Test {
    val src = Source(List(
      Success(1),
      Success(2),
      Failure(new IllegalArgumentException("error 1")),
      Success(3),
      Failure(new ArrayIndexOutOfBoundsException("error 2"))
    ))

    src.runWith(testSink)

    successProbe.request(4)
    failureProbe.request(3)

    successProbe.expectNext(1)
    successProbe.expectNext(2)
    successProbe.expectNext(3)

    val t1 = failureProbe.expectNext()
    t1 shouldBe an[IllegalArgumentException]
    t1 should have message "error 1"

    val t2 = failureProbe.expectNext()
    t2 shouldBe an[ArrayIndexOutOfBoundsException]
    t2 should have message "error 2"

    successProbe.expectComplete()
    failureProbe.expectComplete()
  }

  trait Test {
    val failureProbe = TestSubscriber.probe[Throwable]()
    val successProbe = TestSubscriber.probe[Int]()

    val testSink = Sink.fromGraph(GraphDSL.create() { implicit builder: GraphDSL.Builder[NotUsed] =>
      import GraphDSL.Implicits._

      val tryStage = builder.add(PartitionTry[Int]())

      tryStage.failure ~> Sink.fromSubscriber(failureProbe)
      tryStage.success ~> Sink.fromSubscriber(successProbe)

      SinkShape(tryStage.in)
    })
  }
}