EFF-807 Add Stream.broadcastN

.changeset/add-stream-broadcastn.md

@@ -0,0 +1,5 @@
+---
+"effect": minor
+---
+
+Add Stream.broadcastN for fixed-size stream broadcasts.

packages/effect/src/Stream.ts

@@ -67,7 +67,7 @@ import * as Option from "./Option.ts"
 import type { Pipeable } from "./Pipeable.ts"
 import type { Predicate, Refinement } from "./Predicate.ts"
 import { hasProperty, isNotUndefined, isTagged } from "./Predicate.ts"
-import type * as PubSub from "./PubSub.ts"
+import * as PubSub from "./PubSub.ts"
 import * as Pull from "./Pull.ts"
 import * as Queue from "./Queue.ts"
 import * as RcMap from "./RcMap.ts"
@@ -90,6 +90,7 @@ import type {
   OmitReason,
   ReasonTags,
   Tags,
+  TupleOf,
   unassigned
 } from "./Types.ts"
 import type * as Unify from "./Unify.ts"
@@ -8486,6 +8487,126 @@ export const aggregateWithin: {
     )
   }))))
 
+/**
+ * Fan out the stream, producing a fixed-size tuple of streams that each emit
+ * the same elements as the source stream.
+ *
+ * The source stream starts after all downstream streams have been subscribed.
+ * With the default suspend strategy, the source can only advance `capacity`
+ * chunks ahead of the slowest downstream stream. If a downstream stream is
+ * interrupted, it unsubscribes from the broadcast so it no longer contributes
+ * backpressure.
+ *
+ * **Example** (Broadcasting to two consumers)
+ *
+ * ```ts
+ * import { Console, Effect, Stream } from "effect"
+ *
+ * const program = Effect.scoped(
+ *   Effect.gen(function*() {
+ *     const [left, right] = yield* Stream.make(1, 2, 3).pipe(
+ *       Stream.broadcastN({ n: 2, capacity: 8 })
+ *     )
+ *
+ *     const values = yield* Effect.all([
+ *       Stream.runCollect(left),
+ *       Stream.runCollect(right)
+ *     ], { concurrency: "unbounded" })
+ *
+ *     yield* Console.log(values)
+ *   })
+ * )
+ *
+ * Effect.runPromise(program)
+ * // Output: [[1, 2, 3], [1, 2, 3]]
+ * ```
+ *
+ * @category Broadcast
+ * @since 4.0.0
+ */
+export const broadcastN: {
+  <const N extends number>(
+    options: {
+      readonly n: N
+      readonly capacity: "unbounded"
+      readonly replay?: number | undefined
+    } | {
+      readonly n: N
+      readonly capacity: number
+      readonly strategy?: "sliding" | "dropping" | "suspend" | undefined
+      readonly replay?: number | undefined
+    }
+  ): <A, E, R>(self: Stream<A, E, R>) => Effect.Effect<TupleOf<N, Stream<A, E>>, never, Scope.Scope | R>
+  <A, E, R, const N extends number>(
+    self: Stream<A, E, R>,
+    options: {
+      readonly n: N
+      readonly capacity: "unbounded"
+      readonly replay?: number | undefined
+    } | {
+      readonly capacity: number
+      readonly n: N
+      readonly strategy?: "sliding" | "dropping" | "suspend" | undefined
+      readonly replay?: number | undefined
+    }
+  ): Effect.Effect<TupleOf<N, Stream<A, E>>, never, Scope.Scope | R>
+} = dual(
+  2,
+  Effect.fnUntraced(function*<A, E, R, const N extends number>(
+    self: Stream<A, E, R>,
+    options: {
+      readonly n: N
+      readonly capacity: "unbounded"
+      readonly replay?: number | undefined
+    } | {
+      readonly n: N
+      readonly capacity: number
+      readonly strategy?: "sliding" | "dropping" | "suspend" | undefined
+      readonly replay?: number | undefined
+    }
+  ) {
+    const pubsub = yield* makePubSub<Take.Take<A, E>>(options)
+    const streams = new Array(options.n)
+    const parentScope = yield* Scope.Scope
+    for (let i = 0; i < options.n; i++) {
+      const scope = Scope.forkUnsafe(parentScope)
+      const subscription = yield* PubSub.subscribe(pubsub).pipe(
+        Effect.provideService(Scope.Scope, scope)
+      )
+      streams[i] = Channel.fromEffectTake(PubSub.take(subscription)).pipe(
+        Channel.onExit((exit) => Scope.close(scope, exit)),
+        fromChannel
+      )
+    }
+    yield* Channel.runForEach(self.channel, (value) => PubSub.publish(pubsub, value)).pipe(
+      Effect.onExit((exit) => PubSub.publish(pubsub, exit)),
+      Effect.forkScoped
+    )
+    return streams as TupleOf<N, Stream<A, E>>
+  })
+)
+
+const makePubSub = <A>(
+  options: {
+    readonly capacity: "unbounded"
+    readonly replay?: number | undefined
+  } | {
+    readonly capacity: number
+    readonly strategy?: "dropping" | "sliding" | "suspend" | undefined
+    readonly replay?: number | undefined
+  }
+) =>
+  Effect.acquireRelease(
+    options.capacity === "unbounded"
+      ? PubSub.unbounded<A>(options)
+      : options.strategy === "dropping"
+      ? PubSub.dropping<A>(options)
+      : options.strategy === "sliding"
+      ? PubSub.sliding<A>(options)
+      : PubSub.bounded<A>(options),
+    PubSub.shutdown
+  )
+
 /**
  * Creates a PubSub-backed stream that multicasts the source to all subscribers.
  *

packages/effect/test/Stream.test.ts

@@ -4587,6 +4587,36 @@ describe("Stream", () => {
         deepStrictEqual(result1, result2)
       })))
   })
+
+  describe("broadcastN", () => {
+    it.effect("fans out to a fixed number of streams", () =>
+      Effect.scoped(Effect.gen(function*() {
+        const [left, right] = yield* Stream.make(1, 2, 3).pipe(
+          Stream.broadcastN({ n: 2, capacity: 4 })
+        )
+
+        const result = yield* Effect.all([
+          Stream.runCollect(left),
+          Stream.runCollect(right)
+        ], { concurrency: "unbounded" })
+
+        assert.deepStrictEqual(result, [[1, 2, 3], [1, 2, 3]])
+      })))
+
+    it.effect("propagates failures to all downstream streams", () =>
+      Effect.scoped(Effect.gen(function*() {
+        const [left, right] = yield* Stream.fail("boom").pipe(
+          Stream.broadcastN({ n: 2, capacity: 4 })
+        )
+
+        const result = yield* Effect.all([
+          Stream.runCollect(left).pipe(Effect.exit),
+          Stream.runCollect(right).pipe(Effect.exit)
+        ], { concurrency: "unbounded" })
+
+        assert.deepStrictEqual(result, [Exit.fail("boom"), Exit.fail("boom")])
+      })))
+  })
 })
 
 const grouped = <A>(arr: Array<A>, size: number): Array<NonEmptyArray<A>> => {