combinator for zipping after broadcasting

[yurique]

Description

We have the .broadcastThrough method on streams, which runs each emitted element through N pipes, combining the output of those pipes into a single stream (without any guarantees about the order of the elements in the output).

Sometimes it is helpful to be able to do that, but also have the outputs of the pipes zipped together (in the correct order).

Example

import cats.effect.IO
import cats.effect.unsafe.implicits.global
import BroadcastZip._

val s: fs2.Stream[IO, Int] = fs2.Stream.emits(1 to 10)

val zipped = s.broadcastThroughAndZip[IO, Int](
  _.map(_ + 50),
  _.map(_ + 30),
  _.map(_ + 40),
  _.map(_ + 20)
)

zipped.compile.toList.unsafeRunSync()

// ---> 
// List(List(51, 31, 41, 21), List(52, 32, 42, 22), List(53, 33, 43, 23), List(54, 34, 44, 24), List(55, 35, 45, 25), List(56, 36, 46, 26), List(57, 37, 47, 27), List(58, 38, 48, 28), List(59, 39, 49, 29), List(60, 40, 50, 30))

The sub-streams should emit equal numbers of elements:

import cats.effect.IO
import cats.effect.unsafe.implicits.global
import BroadcastZip._

val s: fs2.Stream[IO, Int] = fs2.Stream.emits(1 to 10)

val zipped = s.broadcastThroughAndZip[IO, Int](
  _.map(_ + 50),
  _.map(_ + 30),
  _.map(_ + 40),
  _.map(_ + 20).dropRight(1)
)

zipped.compile.toList.unsafeRunSync()

// ---> 
// java.lang.IllegalStateException: sub-stream 3 ended, other streams emitted

The code

import cats.effect.Async
import cats.effect.Concurrent
import cats.effect.std.CyclicBarrier
import cats.syntax.all._
import fs2.Pipe

object BroadcastZip {

  implicit class StreamWithBroadcastZip[F[_], A](
    stream: fs2.Stream[F, A]
  )(implicit F: Async[F]) {

    def broadcastThroughAndZip[F2[x] >: F[x], B](
      pipes: Pipe[F2, A, B]*
    )(implicit F2: Concurrent[F2]): fs2.Stream[F2, List[B]] = {
      fs2.Stream
        .eval(
          (CyclicBarrier[F2](pipes.size), F2.ref(false), F2.ref(false)).tupled
        ).flatMap { case (barrier, someEmitted, someEnded) =>
          stream
            .broadcastThrough[F2, (Int, B)](
              pipes.zipWithIndex.map { case (pipe, index) =>
                (in: fs2.Stream[F2, A]) =>
                  pipe(in)
                    .evalMap { out =>
                      someEnded.get.flatMap[(Int, B)] { someEnded =>
                        if (someEnded) {
                          F2.raiseError[(Int, B)](
                            new IllegalStateException(s"sub-stream ${index} emitted, other streams ended")
                          )
                        } else {
                          someEmitted.set(true) >>
                            barrier.await >>
                            someEmitted.set(false).as(index -> out)
                        }
                      }
                    }
                    .onComplete {
                      fs2.Stream
                        .eval {
                          someEmitted.get.flatMap { someEmitted =>
                            if (someEmitted) {
                              F2.raiseError[Unit](
                                new IllegalStateException(s"sub-stream ${index} ended, other streams emitted")
                              )
                            } else {
                              someEnded.set(true)
                            }
                          }
                        }.flatMap(_ => fs2.Stream.empty)
                    }

              }: _*
            )
        }
        .chunkN(pipes.size, allowFewer = true)
        .flatMap {
          case chunk if chunk.size < pipes.size =>
            fs2.Stream.raiseError[F2](
              new IllegalStateException(s"failed to zip: ${chunk.size} sub-streams emitted (${pipes.size} expected)")
            )

          case chunk =>
            fs2.Stream.emit(chunk.toList.sortBy(_._1).map(_._2))
        }

    }
  }

}