如何基于一个 Akka 流的元素聚合另一个 Akka 流的元素

Question

示例场景：将一个流的字节分组为由另一个流（整数）确定的大小块。

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#scan和Flow#prefixAndTail的组合，但感觉不太对（见下文）。我也看了一下 Frameming，但它似乎不适用于上面的示例场景（也不足以容纳非字节串流）。我猜我唯一的选择是使用图形（或更通用的 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:_*) }

Answer 1

我认为您可以将处理实现为自定义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()