Something went wrong. Try again.
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210import WebSocket, {ErrorEvent, MessageEvent} from 'isomorphic-ws'
import {combineSignals} from '#common/async/aborts'import {BlockingAtom} from '#common/async/blocking-atom'import {BlockingQueue} from '#common/async/blocking-queue'import {Breaker} from '#common/breaker'import {normalizeError, ProtocolError} from '#common/errors'import {z} from 'zod/v4'
import {parseJson} from './protocol'
/** * given a websocket, wait and take a single message off and return it * * @example * ``` * const ws = new WebSocket("wss://example.com/stream") * const timeout = timeoutSignal(5000) * * try { * const msg = await takeSocket(ws, timeout.signal) * doWhatever(msg) * } finally { * timeout.cleanup() * if (ws.readyState !== ws.CLOSED) * ws.close(); * } * ``` */export async function takeSocket(ws: WebSocket, signal?: AbortSignal): Promise<unknown> { signal?.throwIfAborted()
const atom = new BlockingAtom() const breaker = new Breaker()
const error = new AbortController() const multisignal = combineSignals(error.signal, signal)
const onMessage = breaker.tripThen((m: MessageEvent) => { atom.set(m.data) }) const onError = breaker.tripThen((e: unknown) => { error.abort(e) }) const onClose = breaker.tripThen(() => { error.abort('closed') })
try { ws.addEventListener('message', onMessage) ws.addEventListener('error', onError) ws.addEventListener('close', onClose)
const data = await atom.get(multisignal) if (!data) { const cause = normalizeError(multisignal.reason) throw new ProtocolError('socket read aborted', 408, {cause}) }
return data } finally { ws.removeEventListener('message', onMessage) ws.removeEventListener('error', onError) ws.removeEventListener('close', onClose) }}
/** exactly take socket, but will additionally apply a json decoding */export async function takeSocketJson<T>(ws: WebSocket, schema: z.ZodType<T>, signal?: AbortSignal): Promise<T> { const data = await takeSocket(ws, signal) return parseJson.pipe(schema).parseAsync(data)}
/** stream configuration options */
interface ConfigProps { maxDepth: number signal?: AbortSignal}
export const STREAM_CONFIG_DEFAULT: ConfigProps = { maxDepth: 1000,}
// symbols for iteration protocol
const yield$ = Symbol('yield$')const error$ = Symbol('error$')const end$ = Symbol('end$')
type StreamYield = [typeof yield$, unknown] | [typeof error$, Error] | [typeof end$]
/** * given a websocket, stream messages off the socket as an async generator * * @example * ``` * const ws = new WebSocket("wss://example.com/stream") * const timeout = timeoutSignal(5000) * * try { * // signal will fire in 5s, so we'll take as many messages as we get until then * for await (const msg of streamSocket(ws, { signal: timeout.signal })) { * doWhatever(msg) * } * } finally { * timeout.cleanup() * if (ws.readyState !== ws.CLOSED) * ws.close(); * } * ``` */export async function* streamSocket(ws: WebSocket, config_?: Partial<ConfigProps>) { const {signal, ...config} = {...STREAM_CONFIG_DEFAULT, ...(config_ || {})} signal?.throwIfAborted()
// await incoming messages without blocking const queue = new BlockingQueue<StreamYield>(config.maxDepth)
// if true, we're ignoring incoming messages until we drop the queue let inBackoffMode: boolean = false const backoffThresh = Math.floor(config.maxDepth * 0.9)
// we don't want to keep processing after we've been closed const breaker = new Breaker()
// define callback functions (need to be able to reference for removing them)
const onMessage = breaker.untilTripped((m: MessageEvent) => { if (inBackoffMode) { console.warn('ignoring incoming message due to backpressure protection!') return }
queue.enqueue([yield$, m.data]) inBackoffMode = queue.depth > backoffThresh if (inBackoffMode) { console.warn('message stream will start dropping messages due to backpressure!') } })
// todo: why are we getting this on client shutdown instead of onClose? const onError = breaker.tripThen((e: ErrorEvent) => { queue.enqueue([error$, normalizeError(e.error)]) })
const onClose = breaker.tripThen(() => { queue.enqueue([end$]) })
// finally get into our loop try { ws.addEventListener('message', onMessage) ws.addEventListener('error', onError) ws.addEventListener('close', onClose)
while (true) { signal?.throwIfAborted()
const [event, value] = await queue.dequeue(signal)
// TODO: eslint thinks inBackoffMode is always falsey... // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition if (inBackoffMode && queue.depth < backoffThresh) { console.debug('message stream will stop dropping messages due to eased backpressure') inBackoffMode = false }
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) }}
/** * exactly stream socket, but will additionally apply a json decoding * messages not validating will end the stream with an error */export async function* streamSocketJson(ws: WebSocket, config?: Partial<ConfigProps>): AsyncGenerator { for await (const message of streamSocket(ws, config)) { yield parseJson.parseAsync(message) }}
/** * exactly stream socket, but will additionally apply a json decoding * messages not validating will end the stream with an error */export async function* streamSocketSchema<T>( ws: WebSocket, schema: z.ZodType<T>, config?: Partial<ConfigProps>,): AsyncGenerator<T> { const parser = parseJson.pipe(schema) for await (const message of streamSocket(ws, config)) { yield await parser.parseAsync(message) }}