Something went wrong. Try again.
Something went wrong. Try again.
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101import { BlockingAtom } from "./blocking-atom.ts";import { BlockingQueue } from "./blocking-queue.ts";import { CancellationError } from "./errors.ts";import { makeGate } from "./gate.ts";
const STREAM_CONFIG_DEFAULT = { signal: undefined as AbortSignal | undefined, queueSize: 1000, threshPause: 800, threshResume: 200, onPause: () => {}, onResume: () => {},};
type StreamEvent = | ["yield", MessageEvent["data"]] | ["end"] | ["error", Error];
export async function* streamSocket( ws: WebSocket, config_: Partial<typeof STREAM_CONFIG_DEFAULT> = {},): AsyncGenerator<unknown> { const { signal, ...config } = { ...STREAM_CONFIG_DEFAULT, ...config_ }; signal?.throwIfAborted();
const queue = new BlockingQueue<StreamEvent>(config.queueSize); const gate = makeGate();
const onMessage = gate.many((m: MessageEvent) => { queue.enqueue(["yield", m.data]);
if (queue.size > config.threshPause && config.onPause) { queueMicrotask(config.onPause); } });
const onError = gate.once((e: Error) => { if (e.message === "Unexpected EOF") { // todo: why are we getting this on client shutdown instead of onClose? queue.enqueue(["end"]); } else { queue.enqueue(["error", e]); } });
const onClose = gate.once(() => queue.enqueue(["end"]));
try { ws.addEventListener("message", onMessage); ws.addEventListener("error", onError); ws.addEventListener("close", onClose);
while (true) { signal?.throwIfAborted();
const [event, value] = await queue.dequeue(signal); if (queue.size < config.threshResume && config.onResume) { queueMicrotask(config.onResume); }
switch (event) { case "yield": yield value; continue; case "error": throw value; case "end": return; } } } finally { ws.removeEventListener("message", onMessage); ws.removeEventListener("error", onError); ws.removeEventListener("close", onClose); }}
export async function takeSocket( ws: WebSocket, signal?: AbortSignal,): Promise<MessageEvent["data"]> { signal?.throwIfAborted();
const atom = new BlockingAtom<MessageEvent["data"]>(); const gate = makeGate();
const onMessage = gate.once((m: MessageEvent) => atom.set(m.data)); try { ws.addEventListener("message", onMessage); const data = await atom.get(signal); if (!data) { throw new CancellationError("socket read aborted"); }
return data; } finally { ws.removeEventListener("message", onMessage); }}