@trpc/server 11.0.0-alpha-tmp-subscription-connection-state.489 → 11.0.0-alpha-tmp-12-06-react.667

package/dist/adapters/ws.mjs

@@ -4,23 +4,95 @@ import { run, isObject, isAsyncIterable } from '../unstable-core-do-not-import/u
 import { parseConnectionParamsFromUnknown } from '../unstable-core-do-not-import/http/parseConnectionParams.mjs';
 import { parseTRPCMessage } from '../unstable-core-do-not-import/rpc/parseTRPCMessage.mjs';
 import { isObservable, observableToAsyncIterable } from '../observable/observable.mjs';
+import { Unpromise } from '../vendor/unpromise/unpromise.mjs';
+import { iteratorResource } from '../unstable-core-do-not-import/stream/utils/asyncIterable.mjs';
+import '../unstable-core-do-not-import/stream/utils/disposable.mjs';
 import { isTrackedEnvelope } from '../unstable-core-do-not-import/stream/tracked.mjs';
 import { transformTRPCResponse } from '../unstable-core-do-not-import/transformer.mjs';
-import { toURL } from '../unstable-core-do-not-import/http/toURL.mjs';
 import '../unstable-core-do-not-import/rootConfig.mjs';
 import { callProcedure } from '../unstable-core-do-not-import/router.mjs';
+import { createURL } from './node-http/incomingMessageToRequest.mjs';
 
+function _ts_add_disposable_resource(env, value, async) {
+    if (value !== null && value !== void 0) {
+        if (typeof value !== "object" && typeof value !== "function") throw new TypeError("Object expected.");
+        var dispose, inner;
+        {
+            if (!Symbol.asyncDispose) throw new TypeError("Symbol.asyncDispose is not defined.");
+            dispose = value[Symbol.asyncDispose];
+        }
+        if (dispose === void 0) {
+            if (!Symbol.dispose) throw new TypeError("Symbol.dispose is not defined.");
+            dispose = value[Symbol.dispose];
+            inner = dispose;
+        }
+        if (typeof dispose !== "function") throw new TypeError("Object not disposable.");
+        if (inner) dispose = function() {
+            try {
+                inner.call(this);
+            } catch (e) {
+                return Promise.reject(e);
+            }
+        };
+        env.stack.push({
+            value: value,
+            dispose: dispose,
+            async: async
+        });
+    } else {
+        env.stack.push({
+            async: true
+        });
+    }
+    return value;
+}
+function _ts_dispose_resources(env) {
+    var _SuppressedError = typeof SuppressedError === "function" ? SuppressedError : function(error, suppressed, message) {
+        var e = new Error(message);
+        return e.name = "SuppressedError", e.error = error, e.suppressed = suppressed, e;
+    };
+    return (_ts_dispose_resources = function _ts_dispose_resources(env) {
+        function fail(e) {
+            env.error = env.hasError ? new _SuppressedError(e, env.error, "An error was suppressed during disposal.") : e;
+            env.hasError = true;
+        }
+        var r, s = 0;
+        function next() {
+            while(r = env.stack.pop()){
+                try {
+                    if (!r.async && s === 1) return s = 0, env.stack.push(r), Promise.resolve().then(next);
+                    if (r.dispose) {
+                        var result = r.dispose.call(r.value);
+                        if (r.async) return s |= 2, Promise.resolve(result).then(next, function(e) {
+                            fail(e);
+                            return next();
+                        });
+                    } else s |= 1;
+                } catch (e) {
+                    fail(e);
+                }
+            }
+            if (s === 1) return env.hasError ? Promise.reject(env.error) : Promise.resolve();
+            if (env.hasError) throw env.error;
+        }
+        return next();
+    })(env);
+}
 /**
  * Importing ws causes a build error
- * @link https://github.com/trpc/trpc/pull/5279
+ * @see https://github.com/trpc/trpc/pull/5279
  */ const WEBSOCKET_OPEN = 1; /* ws.WebSocket.OPEN */ 
 const unsetContextPromiseSymbol = Symbol('unsetContextPromise');
 function getWSConnectionHandler(opts) {
-    const { createContext , router  } = opts;
-    const { transformer  } = router._def._config;
+    const { createContext, router } = opts;
+    const { transformer } = router._def._config;
     return async (client, req)=>{
         const clientSubscriptions = new Map();
         const abortController = new AbortController();
+        if (opts.keepAlive?.enabled) {
+            const { pingMs, pongWaitMs } = opts.keepAlive;
+            handleKeepAlive(client, pingMs, pongWaitMs);
+        }
         function respond(untransformedJSON) {
             client.send(JSON.stringify(transformTRPCResponse(router._def._config, untransformedJSON)));
         }
@@ -73,9 +145,9 @@ function getWSConnectionHandler(opts) {
      *
      * - the context promise will be created immediately on connection if no connectionParams are expected
      * - if connection params are expected, they will be created once received
-     */ let ctxPromise = toURL(req.url ?? '').searchParams.get('connectionParams') === '1' ? unsetContextPromiseSymbol : createCtxPromise(()=>null);
+     */ let ctxPromise = createURL(req).searchParams.get('connectionParams') === '1' ? unsetContextPromiseSymbol : createCtxPromise(()=>null);
         async function handleRequest(msg) {
-            const { id , jsonrpc  } = msg;
+            const { id, jsonrpc } = msg;
             /* istanbul ignore next -- @preserve */ if (id === null) {
                 throw new TRPCError({
                     code: 'BAD_REQUEST',
@@ -86,8 +158,8 @@ function getWSConnectionHandler(opts) {
                 clientSubscriptions.get(id)?.abort();
                 return;
             }
-            const { path , lastEventId  } = msg.params;
-            let { input  } = msg.params;
+            const { path, lastEventId } = msg.params;
+            let { input } = msg.params;
             const type = msg.method;
             try {
                 if (lastEventId !== undefined) {
@@ -103,12 +175,14 @@ function getWSConnectionHandler(opts) {
                     }
                 }
                 await ctxPromise; // asserts context has been set
+                const abortController = new AbortController();
                 const result = await callProcedure({
                     procedures: router._def.procedures,
                     path,
                     getRawInput: async ()=>input,
                     ctx,
-                    type
+                    type,
+                    signal: abortController.signal
                 });
                 const isIterableResult = isAsyncIterable(result) || isObservable(result);
                 if (type !== 'subscription') {
@@ -147,76 +221,95 @@ function getWSConnectionHandler(opts) {
                         code: 'BAD_REQUEST'
                     });
                 }
-                const iterable = isObservable(result) ? observableToAsyncIterable(result) : result;
-                const iterator = iterable[Symbol.asyncIterator]();
-                const abortController = new AbortController();
-                const abortPromise = new Promise((resolve)=>{
-                    abortController.signal.onabort = ()=>resolve('abort');
-                });
+                const iterable = isObservable(result) ? observableToAsyncIterable(result, abortController.signal) : result;
                 run(async ()=>{
-                    while(true){
-                        const next = await Promise.race([
-                            iterator.next().catch(getTRPCErrorFromUnknown),
-                            abortPromise
-                        ]);
-                        if (next === 'abort') {
-                            await iterator.return?.();
-                            break;
-                        }
-                        if (next instanceof Error) {
-                            const error = getTRPCErrorFromUnknown(next);
-                            opts.onError?.({
-                                error,
-                                path,
-                                type,
-                                ctx,
-                                req,
-                                input
-                            });
+                    const env = {
+                        stack: [],
+                        error: void 0,
+                        hasError: false
+                    };
+                    try {
+                        const iterator = _ts_add_disposable_resource(env, iteratorResource(iterable), true);
+                        ;
+                        const abortPromise = new Promise((resolve)=>{
+                            abortController.signal.onabort = ()=>resolve('abort');
+                        });
+                        // We need those declarations outside the loop for garbage collection reasons. If they
+                        // were declared inside, they would not be freed until the next value is present.
+                        let next;
+                        let result;
+                        while(true){
+                            next = await Unpromise.race([
+                                iterator.next().catch(getTRPCErrorFromUnknown),
+                                abortPromise
+                            ]);
+                            if (next === 'abort') {
+                                await iterator.return?.();
+                                break;
+                            }
+                            if (next instanceof Error) {
+                                const error = getTRPCErrorFromUnknown(next);
+                                opts.onError?.({
+                                    error,
+                                    path,
+                                    type,
+                                    ctx,
+                                    req,
+                                    input
+                                });
+                                respond({
+                                    id,
+                                    jsonrpc,
+                                    error: getErrorShape({
+                                        config: router._def._config,
+                                        error,
+                                        type,
+                                        path,
+                                        input,
+                                        ctx
+                                    })
+                                });
+                                break;
+                            }
+                            if (next.done) {
+                                break;
+                            }
+                            result = {
+                                type: 'data',
+                                data: next.value
+                            };
+                            if (isTrackedEnvelope(next.value)) {
+                                const [id, data] = next.value;
+                                result.id = id;
+                                result.data = {
+                                    id,
+                                    data
+                                };
+                            }
                             respond({
                                 id,
                                 jsonrpc,
-                                error: getErrorShape({
-                                    config: router._def._config,
-                                    error,
-                                    type,
-                                    path,
-                                    input,
-                                    ctx
-                                })
+                                result
                             });
-                            break;
-                        }
-                        if (next.done) {
-                            break;
-                        }
-                        const result = {
-                            type: 'data',
-                            data: next.value
-                        };
-                        if (isTrackedEnvelope(next.value)) {
-                            const [id1, data] = next.value;
-                            result.id = id1;
-                            result.data = {
-                                id: id1,
-                                data
-                            };
+                            // free up references for garbage collection
+                            next = null;
+                            result = null;
                         }
                         respond({
                             id,
                             jsonrpc,
-                            result
+                            result: {
+                                type: 'stopped'
+                            }
                         });
+                        clientSubscriptions.delete(id);
+                    } catch (e) {
+                        env.error = e;
+                        env.hasError = true;
+                    } finally{
+                        const result = _ts_dispose_resources(env);
+                        if (result) await result;
                     }
-                    await iterator.return?.();
-                    respond({
-                        id,
-                        jsonrpc,
-                        result: {
-                            type: 'stopped'
-                        }
-                    });
-                    clientSubscriptions.delete(id);
                 }).catch((cause)=>{
                     const error = getTRPCErrorFromUnknown(cause);
                     opts.onError?.({
@@ -274,13 +367,24 @@ function getWSConnectionHandler(opts) {
                 });
             }
         }
-        client.on('message', async (message)=>{
+        client.on('message', async (rawData)=>{
+            // eslint-disable-next-line @typescript-eslint/no-base-to-string
+            const msgStr = rawData.toString();
+            if (msgStr === 'PONG') {
+                return;
+            }
+            if (msgStr === 'PING') {
+                if (!opts.dangerouslyDisablePong) {
+                    client.send('PONG');
+                }
+                return;
+            }
             if (ctxPromise === unsetContextPromiseSymbol) {
                 // If the ctxPromise wasn't created immediately, we're expecting the first message to be a TRPCConnectionParamsMessage
                 ctxPromise = createCtxPromise(()=>{
                     let msg;
                     try {
-                        msg = JSON.parse(message.toString());
+                        msg = JSON.parse(msgStr);
                         if (!isObject(msg)) {
                             throw new Error('Message was not an object');
                         }
@@ -297,7 +401,7 @@ function getWSConnectionHandler(opts) {
                 return;
             }
             try {
-                const msgJSON = JSON.parse(message.toString());
+                const msgJSON = JSON.parse(msgStr);
                 const msgs = Array.isArray(msgJSON) ? msgJSON : [
                     msgJSON
                 ];
@@ -343,48 +447,59 @@ function getWSConnectionHandler(opts) {
             abortController.abort();
         });
         if (ctxPromise !== unsetContextPromiseSymbol) {
-            // prevent unhandled promise rejection errors
-            await ctxPromise.catch(()=>null);
+            await ctxPromise;
         }
     };
 }
 /**
  * Handle WebSocket keep-alive messages
  */ function handleKeepAlive(client, pingMs = 30000, pongWaitMs = 5000) {
-    let heartbeatTimeout;
-    const heartbeatInterval = setInterval(()=>{
-        if (client.readyState !== WEBSOCKET_OPEN) {
-            return;
-        }
-        // First we send a ping message and wait for a pong
-        client.ping();
-        // We set a timeout to close the connection if the pong is not received
-        heartbeatTimeout = setTimeout(()=>{
-            client.terminate();
-            clearInterval(heartbeatInterval);
-        }, pongWaitMs);
-    }, pingMs).unref();
-    // When we receive a pong message, we clear the timeout
-    client.on('pong', ()=>{
-        heartbeatTimeout && clearTimeout(heartbeatTimeout);
-    });
-    // If the connection is closed, we clear the interval
+    let timeout = undefined;
+    let ping = undefined;
+    const schedulePing = ()=>{
+        const scheduleTimeout = ()=>{
+            timeout = setTimeout(()=>{
+                client.terminate();
+            }, pongWaitMs);
+        };
+        ping = setTimeout(()=>{
+            client.send('PING');
+            scheduleTimeout();
+        }, pingMs);
+    };
+    const onMessage = ()=>{
+        clearTimeout(ping);
+        clearTimeout(timeout);
+        schedulePing();
+    };
+    client.on('message', onMessage);
     client.on('close', ()=>{
-        clearInterval(heartbeatInterval);
+        clearTimeout(ping);
+        clearTimeout(timeout);
     });
+    schedulePing();
 }
 function applyWSSHandler(opts) {
-    const { wss , prefix , keepAlive  } = opts;
     const onConnection = getWSConnectionHandler(opts);
-    wss.on('connection', async (client, req)=>{
-        if (prefix && !req.url?.startsWith(prefix)) {
+    opts.wss.on('connection', (client, req)=>{
+        if (opts.prefix && !req.url?.startsWith(opts.prefix)) {
             return;
         }
-        await onConnection(client, req);
-        if (keepAlive?.enabled) {
-            const { pingMs , pongWaitMs  } = keepAlive;
-            handleKeepAlive(client, pingMs, pongWaitMs);
-        }
+        onConnection(client, req).catch((cause)=>{
+            opts.onError?.({
+                error: new TRPCError({
+                    code: 'INTERNAL_SERVER_ERROR',
+                    cause,
+                    message: 'Failed to handle WebSocket connection'
+                }),
+                req: req,
+                path: undefined,
+                type: 'unknown',
+                ctx: undefined,
+                input: undefined
+            });
+            client.close();
+        });
     });
     return {
         broadcastReconnectNotification: ()=>{
@@ -393,7 +508,7 @@ function applyWSSHandler(opts) {
                 method: 'reconnect'
             };
             const data = JSON.stringify(response);
-            for (const client of wss.clients){
+            for (const client of opts.wss.clients){
                 if (client.readyState === WEBSOCKET_OPEN) {
                     client.send(data);
                 }
@@ -402,4 +517,4 @@ function applyWSSHandler(opts) {
     };
 }
 
-export { applyWSSHandler, getWSConnectionHandler };
+export { applyWSSHandler, getWSConnectionHandler, handleKeepAlive };