如何根据另一个Akka流的元素聚合一个Akka流的元素?

时间:2016-03-20 01:22:15

标签: scala stream akka reactive-programming akka-stream

示例方案:将流的字节分组为由另一个(整数)流确定的大小的块。

def partition[A, B, C](
  first:Source[A, NotUsed],
  second:Source[B, NotUsed],
  aggregate:(Int => Seq[A], B) => C
):Source[C, NotUsed] = ???

val bytes:Source[Byte, NotUsed] = ???
val sizes:Source[Int, NotUsed] = ???

val chunks:Source[ByteString, NotUsed] =
  partition(bytes, sizes, (grab, count) => ByteString(grab(count)))

我的初步尝试包括Flow#scanFlow#prefixAndTail的组合,但感觉不太合适(见下文)。我还看了Framing,但它似乎不适用于上面的示例场景(也不足以适应非字节串流)。我猜我唯一的选择是使用Graphs(或者更通用的FlowOps#transform),但我还不够熟练(还)使用Akka流来尝试。

到目前为止,这是我能够提出的(具体到示例场景):

val chunks:Source[ByteString, NotUsed] = sizes
  .scan(bytes prefixAndTail 0) {
    (grouped, count) => grouped flatMapConcat {
      case (chunk, remainder) => remainder prefixAndTail count
    }
  }
  .flatMapConcat(identity)
  .collect { case (chunk, _) if chunk.nonEmpty => ByteString(chunk:_*) }

1 个答案:

答案 0 :(得分:4)

我认为您可以将处理实现为自定义GraphStage。该阶段将有两个Inlet元素。一个取字节,另一个取大小。它将有一个Outlet元素产生值。

考虑以下输入流。

def randomChars = Iterator.continually(Random.nextPrintableChar())
def randomNumbers = Iterator.continually(math.abs(Random.nextInt() % 50))

val bytes: Source[Char, NotUsed] =
  Source.fromIterator(() => randomChars)

val sizes: Source[Int, NotUsed] =
  Source.fromIterator(() => randomNumbers).filter(_ != 0)

然后使用描述自定义流处理(http://doc.akka.io/docs/akka/2.4.2/scala/stream/stream-customize.html)的信息,您可以构建GraphStage

case class ZipFraming() extends GraphStage[FanInShape2[Int, Char, (Int, ByteString)]] {

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

  override val shape: FanInShape2[Int, Char, (Int, ByteString)] =
    new FanInShape2[Int, Char, (Int, ByteString)]("ZipFraming")

  val inFrameSize: Inlet[Int] = shape.in0
  val inElements: Inlet[Char] = shape.in1

  def out: Outlet[(Int, ByteString)] = shape.out

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) {
      // we will buffer as much as 512 characters from the input
      val MaxBufferSize = 512
      // the buffer for the received chars
      var buffer = Vector.empty[Char]
      // the needed number of elements
      var needed: Int = -1
      // if the downstream is waiting
      var isDemanding = false

      override def preStart(): Unit = {
        pull(inFrameSize)
        pull(inElements)
      }

      setHandler(inElements, new InHandler {
        override def onPush(): Unit = {
          // we buffer elements as long as we can
          if (buffer.size < MaxBufferSize) {
            buffer = buffer :+ grab(inElements)
            pull(inElements)
          }
          emit()
        }
      })

      setHandler(inFrameSize, new InHandler {
        override def onPush(): Unit = {
          needed = grab(inFrameSize)
          emit()
        }
      })

      setHandler(out, new OutHandler {
        override def onPull(): Unit = {
          isDemanding = true
          emit()
        }
      })

      def emit(): Unit = {
        if (needed > 0 && buffer.length >= needed && isDemanding) {
          val (emit, reminder) = buffer.splitAt(needed)
          push(out, (needed, ByteString(emit.map(_.toByte).toArray)))
          buffer = reminder
          needed = -1
          isDemanding = false
          pull(inFrameSize)
          if (!hasBeenPulled(inElements)) pull(inElements)
        }
      }
    }
}

这就是你如何运行它。

RunnableGraph.fromGraph(GraphDSL.create(bytes, sizes)(Keep.none) { implicit b =>
  (bs, ss) =>
    import GraphDSL.Implicits._

    val zipFraming = b.add(ZipFraming())

    ss ~> zipFraming.in0
    bs ~> zipFraming.in1

    zipFraming.out ~> Sink.foreach[(Int, ByteString)](e => println((e._1, e._2.utf8String)))

    ClosedShape
}).run()