UNPKG

@trpc/client

Version:

The tRPC client library

869 lines (799 loc) • 28.4 kB
import type { AnyTRPCRouter } from '@trpc/server'; import type { BehaviorSubject } from '@trpc/server/observable'; import { behaviorSubject, observable } from '@trpc/server/observable'; import type { CombinedDataTransformer, TRPCClientIncomingMessage, TRPCClientIncomingRequest, TRPCClientOutgoingMessage, TRPCResponseMessage, } from '@trpc/server/unstable-core-do-not-import'; import { run, sleep, transformResult, } from '@trpc/server/unstable-core-do-not-import'; import { TRPCClientError } from '../../../TRPCClientError'; import type { TRPCConnectionState } from '../../internals/subscriptions'; import type { Operation, OperationResultEnvelope } from '../../types'; import type { Encoder } from './encoder'; import { jsonEncoder } from './encoder'; import type { WebSocketClientOptions } from './options'; import { exponentialBackoff, keepAliveDefaults, lazyDefaults } from './options'; import type { TCallbacks } from './requestManager'; import { RequestManager } from './requestManager'; import { callUserCallback, ResettableTimeout, TRPCWebSocketClosedError, } from './utils'; import type { WebSocketConnectionOptions } from './wsConnection'; import { backwardCompatibility, sendPing, WsConnection } from './wsConnection'; /** * A WebSocket client for managing TRPC operations, supporting lazy initialization, * reconnection, keep-alive, and request management. */ export class WsClient { /** * Observable tracking the current connection state, including errors. */ public readonly connectionState: BehaviorSubject< TRPCConnectionState<TRPCClientError<AnyTRPCRouter>> >; private allowReconnect = false; private requestManager = new RequestManager(); /** * The last id sent on the wire. Each request gets its own instead of its * `op.id`: a retry reuses the failed attempt's `op.id`, and the server's * late frames for that attempt must not reach the retry. A subscription * sent again after its socket was lost gets a new one for the same reason. */ private lastRequestId = 0; private readonly connectionOptions: WebSocketConnectionOptions; /** * The connection new requests are sent on. A reconnect replaces it only * once the new connection is fully open, `connectionParams` included, so * until then requests keep going to the old one. */ private activeConnection: WsConnection; private readonly reconnectRetryDelay: (attemptIndex: number) => number; private inactivityTimeout: ResettableTimeout; private readonly callbacks: Pick< WebSocketClientOptions, 'onOpen' | 'onClose' | 'onError' >; private readonly lazyMode: boolean; private readonly encoder: Encoder; constructor(opts: WebSocketClientOptions) { this.encoder = opts.experimental_encoder ?? jsonEncoder; // Initialize callbacks, connection parameters, and options. this.callbacks = { onOpen: opts.onOpen, onClose: opts.onClose, onError: opts.onError, }; const lazyOptions = { ...lazyDefaults, ...opts.lazy, }; // Set up inactivity timeout for lazy connections. this.inactivityTimeout = new ResettableTimeout(() => { if ( this.requestManager.hasOutgoingRequests() || this.requestManager.hasPendingRequests() ) { this.inactivityTimeout.reset(); return; } this.close().catch(() => null); }, lazyOptions.closeMs); // Initialize the WebSocket connection. this.connectionOptions = { WebSocketPonyfill: opts.WebSocket, urlOptions: opts, keepAlive: { ...keepAliveDefaults, ...opts.keepAlive, }, encoder: this.encoder, }; this.activeConnection = this.createConnection(); this.reconnectRetryDelay = opts.retryDelayMs ?? exponentialBackoff; this.lazyMode = lazyOptions.enabled; this.connectionState = behaviorSubject< TRPCConnectionState<TRPCClientError<AnyTRPCRouter>> >({ type: 'state', state: lazyOptions.enabled ? 'idle' : 'connecting', error: null, }); // Automatically open the connection if lazy mode is disabled. if (!this.lazyMode) { this.open().catch(() => null); } } private createConnection() { const connection = new WsConnection(this.connectionOptions); connection.wsObservable.subscribe({ next: (ws) => { if (!ws) return; this.setupWebSocketListeners(ws); }, }); return connection; } /** * Opens the WebSocket connection. Handles reconnection attempts and updates * the connection state accordingly. */ private async open() { this.allowReconnect = true; if (this.reconnecting) return this.reconnecting; this.connectionAttemptIndex = 0; if (this.connectionState.get().state === 'idle') { callUserCallback(() => this.connectionState.next({ type: 'state', state: 'connecting', error: null, }), ); } try { await this.activeConnection.open(); } catch (error) { this.reconnect( new TRPCWebSocketClosedError({ message: 'Initialization error', cause: error, }), ); // A socket whose `connectionParams` failed is open, but not usable. await this.activeConnection.close().catch(() => null); return this.reconnecting; } } /** * Closes the WebSocket connection and stops managing requests. * Ensures all outgoing and pending requests are properly finalized. */ public async close() { this.allowReconnect = false; this.replacementError = null; this.retiredConnection = null; this.inactivityTimeout.stop(); const requestsToAwait: Promise<void>[] = []; for (const request of this.requestManager.getRequests()) { if (request.message.method === 'subscription') { callUserCallback(() => request.callbacks.complete()); } else if (request.state === 'outgoing') { callUserCallback(() => request.callbacks.error( TRPCClientError.from( new TRPCWebSocketClosedError({ message: 'Closed before connection was established', }), ), ), ); } else { requestsToAwait.push(request.end); } } await Promise.all(requestsToAwait).catch(() => null); this.nextConnection?.ws?.close(); await this.activeConnection.close().catch(() => null); callUserCallback(() => this.connectionState.next({ type: 'state', state: 'idle', error: null, }), ); } /** * Method to request the server. * Handles data transformation, batching of requests, and subscription lifecycle. * * @param op - The operation details including type, path, input and signal. Each subscription to the returned observable is sent with a new id, not `op.id`. * @param transformer - Data transformer for serializing requests and deserializing responses * @param lastEventId - Optional ID of the last received event for subscriptions * * @returns An observable that emits operation results and handles cleanup */ public request({ op: { type, path, input, signal }, transformer, lastEventId, }: { op: Pick<Operation, 'id' | 'type' | 'path' | 'input' | 'signal'>; transformer: CombinedDataTransformer; lastEventId?: string; }) { return observable< OperationResultEnvelope<unknown, TRPCClientError<AnyTRPCRouter>>, TRPCClientError<AnyTRPCRouter> >((observer) => { const message: TRPCClientOutgoingMessage = { id: ++this.lastRequestId, method: type, params: { input: transformer.input.serialize(input), path, lastEventId, }, }; const abort = this.batchSend(message, { ...observer, next(event) { const transformed = transformResult(event, transformer.output); if (!transformed.ok) { observer.error(TRPCClientError.from(transformed.error)); return; } observer.next({ result: transformed.result, }); }, }); const onAbort = () => observer.complete(); if (signal?.aborted) { onAbort(); } else { signal?.addEventListener('abort', onAbort, { once: true }); } return () => { // A subscription sent again after its socket was lost has a new id. const { id } = message; const pending = this.requestManager.getPendingRequest(id); // Under its earlier ids, it may still be draining other sockets. const draining = this.requestManager .getDrainingRequests() .filter((request) => request.message === message); abort(); if (type === 'subscription' && pending) { this.stopSubscription(id, pending.ws); this.stopSubscription(id, pending.successor); } this.closeIfDrained(pending?.ws); for (const request of [pending, ...draining]) { for (const earlier of request?.movedFrom ?? []) { this.closeIfDrained(earlier); } } signal?.removeEventListener('abort', onAbort); }; }); } /** * The current WebSocket and its state, or `null` when there is no socket. */ public get connection() { return backwardCompatibility( this.activeConnection.ws ? this.activeConnection : (this.nextConnection ?? this.activeConnection), ); } /** * Manages the reconnection process for the WebSocket using retry logic. * Ensures that only one reconnection attempt is active at a time by tracking the current * reconnection state in the `reconnecting` promise. */ private reconnecting: Promise<void> | null = null; /** * The connection a reconnect is opening to replace the active one. */ private nextConnection: WsConnection | null = null; /** * The connection the active one replaced, while it's open. If the active * one is lost first, this one serves every request again until the next. */ private retiredConnection: WsConnection | null = null; /** * Why the latest attempt to replace a socket that's still open failed. * It's shown once that socket closes, since until then it serves every * request. */ private replacementError: TRPCWebSocketClosedError | null = null; /** * The retry attempt that opened the current connection. */ private connectionAttemptIndex = 0; /** * Set when the server rejects the current connection with an `id: null` * error, e.g. because `createContext` threw. The server closes the socket * right after, so the reconnect that follows carries on with the backoff * instead of starting over at attempt 0. Cleared again when the socket * receives a result, which proves the connection was accepted. * @see https://github.com/trpc/trpc/issues/4774 */ private connectionRejected = false; /** Subscriptions whose subscriber has seen `started`. */ private readonly started = new WeakSet<TRPCClientOutgoingMessage>(); /** * Opens a new connection, retrying until one opens. While the current * socket is open, it keeps serving every request, new ones included, until * the new one is fully open, `connectionParams` included. * * @param closedError - Why the current socket was lost, or `null` when it's * still open and serving requests, e.g. on a server-requested reconnect. */ private reconnect(closedError: TRPCWebSocketClosedError | null) { if (!this.allowReconnect) { return; } if (closedError) { const error = this.replacementError ?? closedError; this.replacementError = null; callUserCallback(() => this.connectionState.next({ type: 'state', state: 'connecting', error: TRPCClientError.from(error), }), ); } if (this.reconnecting) return; const tryReconnect = async (attemptIndex: number) => { try { await sleep(this.reconnectRetryDelay(attemptIndex)); if (this.allowReconnect) { const next = this.createConnection(); this.nextConnection = next; try { await next.open(); // `open()` resolves even if the socket closed while its // `connectionParams` were pending. if (!next.isOpen()) { throw new Error('WebSocket closed before it was ready'); } } catch (error) { await next.close().catch(() => null); throw error; } finally { this.nextConnection = null; } if (this.allowReconnect) { // Read only now: losing the active socket while `next` opened may // have restored the one it replaced. const previous = this.activeConnection; this.activeConnection = next; this.retiredConnection = previous.isOpen() ? previous : null; this.connectionAttemptIndex = attemptIndex; this.replacementError = null; this.handOff(previous.ws); } else { await next.close(); } } this.reconnecting = null; } catch (error) { const closedError = new TRPCWebSocketClosedError({ message: 'Initialization error', cause: error, }); if (this.allowReconnect && this.activeConnection.isOpen()) { // The old socket keeps serving every request, so this is shown only // once it closes. this.replacementError = closedError; } else if (this.allowReconnect) { callUserCallback(() => this.connectionState.next({ type: 'state', state: 'connecting', error: TRPCClientError.from(closedError), }), ); } await tryReconnect(attemptIndex + 1); } }; this.reconnecting = tryReconnect( this.connectionRejected ? this.connectionAttemptIndex + 1 : 0, ); } /** * Moves the pending subscriptions to the active socket, which has just * opened: each is resubscribed with the id of the last event it delivered. * A subscription whose old socket is still open keeps delivering from it * until the new socket confirms it with `started`; see `confirmMove()`. * Closes the previous socket once nothing is left on it. */ private handOff(previous: WebSocket | null) { if ( this.connectionState.get().state !== 'pending' && this.activeConnection.isOpen() ) { callUserCallback(() => this.connectionState.next({ type: 'state', state: 'pending', error: null, }), ); } const ws = this.activeConnection.ws; const subscriptions = this.requestManager .getPendingRequests() .filter( (request) => request.message.method === 'subscription' && request.ws !== ws, ); if (subscriptions.length) { this.send(subscriptions.map(({ message }) => message)); for (const request of subscriptions) { // The socket it was moving to hasn't started it, and won't be used. this.stopSubscription(request.message.id, request.successor); if (request.ws && request.ws.readyState === request.ws.OPEN) { request.successor = ws; request.overlapEventIds = new Set(); } else { // Its socket was lost, so it starts again. this.started.delete(request.message); request.ws = ws; request.successor = null; } } } this.closeIfDrained(previous); } /** * Completes a subscription's move once its new socket has confirmed it: * stops it on the old socket, whose later messages for it are dropped * apart from events without ids, and closes that socket once everything it * sent before the stop has arrived and nothing else is left on it. */ private confirmMove(request: { message: TRPCClientOutgoingMessage; ws: WebSocket | null; successor: WebSocket | null; movedFrom: Set<WebSocket>; }) { const previous = request.ws; request.ws = request.successor; request.successor = null; if (previous && this.stopSubscription(request.message.id, previous)) { request.movedFrom.add(previous); // `stopped` can wait on the procedure, but the "PONG" can't. sendPing(previous, () => this.releaseMovedFrom(request, previous)); } this.closeIfDrained(previous); } private releaseMovedFrom( request: { movedFrom: Set<WebSocket> }, ws: WebSocket, ) { if (!request.movedFrom.delete(ws)) return; this.requestManager.deleteDrained(); this.closeIfDrained(ws); } private hasRequestsOn(ws: WebSocket | null) { return [ ...this.requestManager.getPendingRequests(), ...this.requestManager.getDrainingRequests(), ].some( (request) => request.ws === ws || (ws !== null && request.movedFrom.has(ws)), ); } /** * Closes `ws` once it's been replaced and nothing is left on it. */ private closeIfDrained(ws: WebSocket | null | undefined) { if (!ws || ws === this.activeConnection.ws) return; if (this.hasRequestsOn(ws)) return; ws.close(); } /** * Stops a subscription on the socket it was sent on, if that's still open. * * @returns Whether the stop was sent. */ private stopSubscription( id: TRPCClientOutgoingMessage['id'], ws: WebSocket | null, ) { if (!ws || ws.readyState !== ws.OPEN) return false; ws.send(this.encoder.encode({ id, method: 'subscription.stop' })); return true; } /** * Fails the pending queries and mutations sent on `ws`, and returns its * pending subscriptions, which are kept to be sent again. */ private failPendingRequests(ws: WebSocket, cause: unknown) { const subscriptions = []; for (const request of this.requestManager.getPendingRequests()) { const { message, callbacks } = request; if (request.ws !== ws) continue; if (message.method === 'subscription') { subscriptions.push(request); continue; } callUserCallback(() => callbacks.error( TRPCClientError.from( cause ?? new TRPCWebSocketClosedError({ message: 'WebSocket closed', cause, }), ), ), ); this.requestManager.delete(message); } return subscriptions; } private setupWebSocketListeners(ws: WebSocket) { this.connectionRejected = false; /** * Fails the queries and mutations sent on `ws`. Its subscriptions are * sent again on the active socket if that's open, or else by * `reconnect()` once the next one opens. */ const handleCloseOrError = (cause: unknown) => { if (ws === this.retiredConnection?.ws) this.retiredConnection = null; const retired = this.retiredConnection; let restored = false; if (ws === this.activeConnection.ws && retired?.isOpen()) { restored = true; // The socket it replaced serves every request again, until the next // one opens. this.activeConnection = retired; this.retiredConnection = null; this.replacementError = new TRPCWebSocketClosedError({ message: 'WebSocket closed', cause, }); } for (const request of this.requestManager.getPendingRequests()) { if (request.successor === ws) { // Its old socket still has it. request.successor = null; request.overlapEventIds = null; } request.movedFrom.delete(ws); } for (const request of this.requestManager.getDrainingRequests()) { request.movedFrom.delete(ws); } this.requestManager.deleteDrained(); const lost = this.failPendingRequests(ws, cause).filter((request) => { if (!request.successor) return true; // Its new socket has it already. request.ws = request.successor; request.successor = null; return false; }); if ( ws === this.activeConnection.ws || (this.reconnecting && !this.activeConnection.isOpen()) ) { return; } for (const { message, callbacks } of lost) { this.started.delete(message); // The active socket may still have its old id, e.g. while stopping it. this.requestManager.requeue(message, ++this.lastRequestId); this.flushOutgoing(message, callbacks); } if (restored) this.reconnect(null); }; ws.addEventListener('open', () => { run(async () => { if (this.lazyMode) { this.inactivityTimeout.start(); } this.callbacks.onOpen?.(); // A socket opened to replace one that's still open changes nothing // the subscribers can see until `handOff()` puts it in use. if ( this.connectionState.get().state === 'pending' || (ws !== this.activeConnection.ws && this.activeConnection.isOpen()) ) { return; } callUserCallback(() => this.connectionState.next({ type: 'state', state: 'pending', error: null, }), ); }).catch((error) => { ws.close(3000); handleCloseOrError(error); }); }); const onMessage = ({ data }: MessageEvent) => { this.inactivityTimeout.reset(); // Handle PING/PONG as text regardless of encoder if (['PING', 'PONG'].includes(data)) return; const incomingMessage = this.encoder.decode( data, ) as TRPCClientIncomingMessage; if ('method' in incomingMessage) { if (ws === this.activeConnection.ws) { this.handleIncomingRequest(incomingMessage); } return; } // A replaced socket still answers its own queries and mutations, but // says nothing about the connection the client now uses. const isRetired = ws !== this.activeConnection.ws && ws !== this.nextConnection?.ws; if ( !isRetired && incomingMessage.id === null && 'error' in incomingMessage ) { this.connectionRejected = true; } else if (!isRetired && 'result' in incomingMessage) { // The server only sends results once `createContext` has succeeded, // so an earlier `id: null` error on this socket was not a rejection, // e.g. the server's reply to a malformed frame. this.connectionRejected = false; } this.handleResponseMessage(incomingMessage, ws); }; // With `ws`, a throw from a `message` listener stops the socket from // reading any further frames while it stays open. ws.addEventListener('message', (event) => { callUserCallback(() => onMessage(event)); }); ws.addEventListener('close', (event) => { handleCloseOrError(event); callUserCallback(() => this.callbacks.onClose?.(event)); if ( ws === this.activeConnection.ws && (!this.lazyMode || this.requestManager.hasPendingSubscriptions()) ) { this.reconnect( new TRPCWebSocketClosedError({ message: 'WebSocket closed', cause: event, }), ); } }); ws.addEventListener('error', (event) => { handleCloseOrError(event); callUserCallback(() => this.callbacks.onError?.(event)); if (ws !== this.activeConnection.ws) return; this.reconnect( new TRPCWebSocketClosedError({ message: 'WebSocket closed', cause: event, }), ); }); } private handleResponseMessage(message: TRPCResponseMessage, ws: WebSocket) { const request = this.requestManager.getPendingRequest(message.id) ?? this.requestManager.getDrainingRequest(message.id); if (!request) return; const result = 'result' in message ? message.result : null; if (request.successor === ws) { this.confirmMove(request); } else if (request.ws !== ws) { // A subscription moved to a new socket ignores the old one's frames, // apart from events without ids until it's stopped there: those can't // be replayed, and the new socket may have started listening after // the old one sent them. if (request.movedFrom.has(ws)) { if (result?.type === 'stopped') { this.releaseMovedFrom(request, ws); } else if (result?.type === 'data' && result.id === undefined) { callUserCallback(() => request.callbacks.next(message)); } } return; } // The subscriber doesn't see a move, even one its old socket didn't // live to finish. if (result?.type === 'started' && this.started.has(request.message)) { return; } if (result?.type === 'data' && result.id !== undefined) { const eventId = String(result.id); if (request.successor) { request.overlapEventIds?.add(eventId); } else if (request.overlapEventIds?.delete(eventId)) { // Already delivered by the old socket. return; } } callUserCallback(() => request.callbacks.next(message)); let completed = true; if ('result' in message && request.message.method === 'subscription') { if (message.result.type === 'data') { request.message.params.lastEventId = message.result.id; } else if (message.result.type === 'started') { this.started.add(request.message); } if (message.result.type !== 'stopped') { completed = false; } } if (completed) { callUserCallback(() => request.callbacks.complete()); this.requestManager.delete(request.message); } } private handleIncomingRequest(message: TRPCClientIncomingRequest) { if (message.method === 'reconnect') { this.reconnect(null); } } /** * Sends a message or batch of messages directly to the server. */ private send( messageOrMessages: TRPCClientOutgoingMessage | TRPCClientOutgoingMessage[], ) { if (!this.activeConnection.isOpen()) { throw new Error('Active connection is not open'); } const messages = messageOrMessages instanceof Array ? messageOrMessages : [messageOrMessages]; this.activeConnection.ws.send( this.encoder.encode(messages.length === 1 ? messages[0] : messages), ); } /** * Groups requests for batch sending. * * @returns A function to abort the batched request. */ private batchSend(message: TRPCClientOutgoingMessage, callbacks: TCallbacks) { this.inactivityTimeout.reset(); this.flushOutgoing(message, callbacks); return this.requestManager.register(message, callbacks); } /** * Sends `message` with the other outgoing requests once the connection is * open, batching those queued in the same tick. */ private flushOutgoing( message: TRPCClientOutgoingMessage, callbacks: TCallbacks, ) { run(async () => { // During a reconnect, requests go to the old socket while it's open. do { if (!this.activeConnection.isOpen()) { await this.open(); } await sleep(0); } while (!this.activeConnection.isOpen() && this.reconnecting); if (!this.requestManager.hasOutgoingRequests()) return; // Another run already sent this request. Each queued request's own run // sends it once the new socket opens. if ( !this.activeConnection.isOpen() && !this.requestManager.isOutgoing(message) ) { return; } const requests = this.requestManager.flush(this.activeConnection.ws); try { this.send(requests.map((request) => request.message)); } catch (err) { for (const request of requests) { this.requestManager.delete(request.message); callUserCallback(() => request.callbacks.error(TRPCClientError.from(err as Error)), ); } } }).catch((err) => { this.requestManager.delete(message); callUserCallback(() => callbacks.error(TRPCClientError.from(err))); }); } }