所以我有这个小的自定义阶段,用于在Akka Streams中进行分区。
object CustomPartitioner {
/**
* Creates a Partition stage that, given a type A, makes a decision to whether to partition to subtype B or subtype C
*
* @param partitionF applies function, if true, route to B, otherwise route to C.
*
* @tparam A type of input
* @tparam B type of output on the first outlet.
* @tparam C type of output on the second outlet.
*
* @return A partition stage
*/
def apply[A, B, C](partitionF: A => Either[B, C]) =
new GraphStage[FanOutShape2[A, B, C]] {
private val in: Inlet[A] = Inlet[A]("in")
private val outB = Outlet[B]("outB")
private val outC = Outlet[C]("outC")
private val pendingB = MutableQueue.empty[B]
private val pendingC = MutableQueue.empty[C]
override def shape: FanOutShape2[A, B, C] = new FanOutShape2(in, outB, outC)
override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
new GraphStageLogic(shape) with InHandler with OutHandler {
setHandler(in, this)
setHandler(outB, this)
setHandler(outC, this)
override def onPush(): Unit = {
val elem = grab(in)
partitionF(elem) match {
case Left(b) =>
pendingB.enqueue(b)
tryPush(outB, pendingB, b)
case Right(c) =>
pendingC.enqueue(c)
tryPush(outC, pendingC, c)
}
}
override def onPull(): Unit = pull(in)
private def tryPush[T](out: Outlet[T], pending: MutableQueue[T]): Unit =
if (isAvailable(out) && pending.nonEmpty) push(out, pending.dequeue())
}
}
我已将其作为分区插入到流中,然后将其合并回接收器中。
当我尝试使用组件测试将消息推送通过流时
java.lang.IllegalArgumentException: Cannot pull port (in(256390569)) twice
然后测试失败
java.lang.AssertionError: assertion failed: expected: expecting request() signal but got unexpected message CancelSubscription(PublisherProbeSubscription(akka.stream.impl.fusing.ActorGraphInterpreter$BatchingActorInputBoundary$$anon$1@53c99b09,akka.testkit.TestProbe@2539cd1c))
我可以肯定我搞砸了setHandler调用,因为其中有两个可以同时处理outB和outC。但是我不知道如何解决它,以使整个系统只调用一次onPush和onPull。
答案 0 :(得分:0)
我设法让它起作用
override def onPull(): Unit =
if (!hasBeenPulled(in))
pull(in)