Initial Commit

packages/foundation/flow/streams/src/pipeline/route/route.ts

@@ -0,0 +1,138 @@
+import { StreamRouteUnknownKeyError } from "../../errors/stream-error.js";
+import type { RouteOptions } from "./route.types.js";
+
+/**
+ * Internal: build a `Map<K, ReadableStream<T>>` of routed branches, single-pass
+ * (no chunk is duplicated). The reader drives the source and dispatches each
+ * chunk to the controller of exactly one branch — chosen by `keyFn`. On
+ * unknown key under `onUnknownKey: 'error'`, every branch errors with
+ * `StreamRouteUnknownKeyError` and the source is cancelled.
+ */
+function buildRoutedMap<T, K>(
+    source: ReadableStream<T>,
+    keyFn: (chunk: T) => K,
+    keys: readonly K[],
+    options: RouteOptions = {}
+): Map<K, ReadableStream<T>> {
+    const onUnknownKey = options.onUnknownKey ?? "drop";
+
+    const controllers = new Map<K, ReadableStreamDefaultController<T>>();
+    const knownKeys = new Set<K>(keys);
+    const result = new Map<K, ReadableStream<T>>();
+
+    // We need every branch's controller wired up before the reader starts so
+    // that early chunks are not lost when the source resolves quickly.
+    const controllerReady: Promise<void>[] = [];
+    for (const k of keys) {
+        let resolveReady!: () => void;
+        const ready = new Promise<void>((r) => {
+            resolveReady = r;
+        });
+        controllerReady.push(ready);
+
+        const branch = new ReadableStream<T>({
+            start(controller) {
+                controllers.set(k, controller);
+                resolveReady();
+            },
+            cancel() {
+                // Per-branch cancel is non-fatal; other branches keep flowing.
+                // Drop the controller so writes to it become no-ops.
+                controllers.delete(k);
+            },
+        });
+        result.set(k, branch);
+    }
+
+    const startReader = async (): Promise<void> => {
+        // Wait until every branch has registered its controller.
+        await Promise.all(controllerReady);
+
+        const reader = source.getReader();
+        try {
+            while (true) {
+                const { value, done } = await reader.read();
+                if (done) break;
+                const k = keyFn(value as T);
+                if (!knownKeys.has(k)) {
+                    if (onUnknownKey === "error") {
+                        const err = new StreamRouteUnknownKeyError({
+                            key: String(k),
+                        });
+                        for (const c of controllers.values()) c.error(err);
+                        controllers.clear();
+                        try {
+                            await reader.cancel(err);
+                        } catch {
+                            // ignore cancel failures
+                        }
+                        return;
+                    }
+                    // 'drop' — silently discard.
+                    continue;
+                }
+                const c = controllers.get(k);
+                if (c !== undefined) {
+                    c.enqueue(value as T);
+                }
+            }
+            for (const c of controllers.values()) c.close();
+            controllers.clear();
+        } catch (err) {
+            for (const c of controllers.values()) c.error(err);
+            controllers.clear();
+        } finally {
+            try {
+                reader.releaseLock();
+            } catch {
+                // already released
+            }
+        }
+    };
+
+    // Kick the reader on first pull from any branch. We bind a `pull` by
+    // wrapping each branch in a tee-style trigger: simplest correct approach
+    // is to start eagerly on the next microtask once all branches are
+    // constructed. Since branches are constructed synchronously above, this
+    // microtask fires after the caller has the Map in hand.
+    queueMicrotask(() => {
+        void startReader();
+    });
+
+    return result;
+}
+
+/**
+ * Routes each chunk of `source` to exactly one of N output branches, keyed by
+ * the result of `keyFn`. Returns a record keyed by the declared `keys`. Each
+ * chunk is delivered to a single branch (no duplication, single-pass).
+ *
+ * Per-key ordering is preserved. On unknown key, behaviour depends on
+ * `options.onUnknownKey` (`'drop'` default, or `'error'`).
+ */
+export function route<T, K extends string>(
+    source: ReadableStream<T>,
+    keyFn: (chunk: T) => K,
+    keys: readonly K[],
+    options?: RouteOptions
+): Record<K, ReadableStream<T>> {
+    const m = buildRoutedMap<T, K>(source, keyFn, keys, options);
+    const out = {} as Record<K, ReadableStream<T>>;
+    for (const k of keys) {
+        out[k] = m.get(k) as ReadableStream<T>;
+    }
+    return out;
+}
+
+/**
+ * Like `route`, but returns a `Map<K, ReadableStream<T>>` — supporting key
+ * types beyond `string` (numbers, symbols, objects).
+ */
+export function routeBy<T, K>(
+    source: ReadableStream<T>,
+    keyFn: (chunk: T) => K,
+    keys: readonly K[],
+    options?: RouteOptions
+): Map<K, ReadableStream<T>> {
+    return buildRoutedMap<T, K>(source, keyFn, keys, options);
+}