139 lines
4.9 KiB
TypeScript
139 lines
4.9 KiB
TypeScript
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);
|
|
}
|