@trpc/server 11.0.0-rc.630 → 11.0.0-rc.632

package/src/unstable-core-do-not-import/stream/jsonl.ts

@@ -1,12 +1,8 @@
-import { Unpromise } from '../../vendor/unpromise';
-import { getTRPCErrorFromUnknown } from '../error/TRPCError';
 import { isAsyncIterable, isFunction, isObject, run } from '../utils';
 import type { Deferred } from './utils/createDeferred';
 import { createDeferred } from './utils/createDeferred';
-import {
-  createReadableStream,
-  isCancelledStreamResult,
-} from './utils/createReadableStream';
+import { createReadableStream } from './utils/createReadableStream';
+import { withRefCount } from './utils/withRefCount';
 
 /**
  * A subset of the standard ReadableStream properties needed by tRPC internally.
@@ -58,31 +54,31 @@ type ChunkDefinition = [
   type: ChunkValueType,
   chunkId: ChunkIndex,
 ];
-type DehydratedValue = [
+type EncodedValue = [
   // data
   [unknown] | [],
   // chunk descriptions
   ...ChunkDefinition[],
 ];
 
-type Head = Record<string, DehydratedValue>;
+type Head = Record<string, EncodedValue>;
 type PromiseChunk =
   | [
       chunkIndex: ChunkIndex,
       status: PROMISE_STATUS_FULFILLED,
-      value: DehydratedValue,
+      value: EncodedValue,
     ]
   | [chunkIndex: ChunkIndex, status: PROMISE_STATUS_REJECTED, error: unknown];
 type IterableChunk =
   | [
       chunkIndex: ChunkIndex,
       status: ASYNC_ITERABLE_STATUS_RETURN,
-      value: DehydratedValue,
+      value: EncodedValue,
     ]
   | [
       chunkIndex: ChunkIndex,
       status: ASYNC_ITERABLE_STATUS_VALUE,
-      value: DehydratedValue,
+      value: EncodedValue,
     ]
   | [
       chunkIndex: ChunkIndex,
@@ -127,42 +123,38 @@ function createBatchStreamProducer(opts: ProducerOptions) {
   const placeholder = 0 as PlaceholderValue;
 
   const stream = createReadableStream<ChunkData>();
-  const pending = new Set<ChunkIndex>();
-
-  function maybeClose() {
-    if (pending.size === 0 && !stream.cancelled()) {
+  const pending = withRefCount(new Set<ChunkIndex>(), () => {
+    if (!stream.cancelled()) {
       stream.controller.close();
     }
-  }
-  function dehydratePromise(
-    promise: Promise<unknown>,
-    path: (string | number)[],
-  ) {
-    //
+  });
+
+  const maybeEnqueue = (chunk: ChunkData) => {
+    if (!stream.cancelled()) {
+      stream.controller.enqueue(chunk);
+    }
+  };
+
+  function encodePromise(promise: Promise<unknown>, path: (string | number)[]) {
     const error = checkMaxDepth(path);
     if (error) {
-      promise.catch(() => {
-        // ignore
+      // Catch any errors from the original promise to ensure they're reported
+      promise.catch((cause) => {
+        opts.onError?.({ error: cause, path });
       });
+      // Replace the promise with a rejected one containing the max depth error
       promise = Promise.reject(error);
     }
     const idx = counter++ as ChunkIndex;
     pending.add(idx);
 
-    Unpromise.race([promise, stream.cancelledPromise])
+    promise
       .then((it) => {
-        if (isCancelledStreamResult(it)) {
-          return;
-        }
-        stream.controller.enqueue([
-          idx,
-          PROMISE_STATUS_FULFILLED,
-          dehydrate(it, path),
-        ]);
+        maybeEnqueue([idx, PROMISE_STATUS_FULFILLED, encode(it, path)]);
       })
       .catch((cause) => {
         opts.onError?.({ error: cause, path });
-        stream.controller.enqueue([
+        maybeEnqueue([
           idx,
           PROMISE_STATUS_REJECTED,
           opts.formatError?.({ error: cause, path }),
@@ -170,78 +162,55 @@ function createBatchStreamProducer(opts: ProducerOptions) {
       })
       .finally(() => {
         pending.delete(idx);
-        maybeClose();
       });
     return idx;
   }
-  function dehydrateAsyncIterable(
+  function encodeAsyncIterable(
     iterable: AsyncIterable<unknown>,
     path: (string | number)[],
   ) {
-    const error = checkMaxDepth(path);
-    if (error) {
-      iterable = {
-        [Symbol.asyncIterator]() {
-          throw error;
-        },
-      };
-    }
     const idx = counter++ as ChunkIndex;
     pending.add(idx);
     run(async () => {
+      const error = checkMaxDepth(path);
+      if (error) {
+        throw error;
+      }
       const iterator = iterable[Symbol.asyncIterator]();
 
       while (true) {
-        const next = await Unpromise.race([
-          iterator.next().catch(getTRPCErrorFromUnknown),
-          stream.cancelledPromise,
-        ]);
-
-        if (next instanceof Error) {
-          opts.onError?.({ error: next, path });
-
-          stream.controller.enqueue([
-            idx,
-            ASYNC_ITERABLE_STATUS_ERROR,
-            opts.formatError?.({ error: next, path }),
-          ]);
-          return;
-        }
-        if (isCancelledStreamResult(next)) {
-          await iterator.return?.();
-          break;
+        if (stream.cancelled()) {
+          const res = await iterator.return?.();
+          return res?.value;
         }
+        const next = await iterator.next();
+
         if (next.done) {
-          stream.controller.enqueue([
+          maybeEnqueue([
             idx,
             ASYNC_ITERABLE_STATUS_RETURN,
-            dehydrate(next.value, path),
+            encode(next.value, path),
           ]);
           break;
         }
-        stream.controller.enqueue([
+        maybeEnqueue([
           idx,
           ASYNC_ITERABLE_STATUS_VALUE,
-          dehydrate(next.value, path),
+          encode(next.value, path),
         ]);
       }
-
-      pending.delete(idx);
-      maybeClose();
-    }).catch((cause) => {
-      // this shouldn't happen, but node crashes if we don't catch it
-      opts.onError?.({
-        error: new Error(
-          'You found a bug - please report it on https://github.com/trpc/trpc',
-          // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-          // @ts-ignore https://github.com/tc39/proposal-error-cause
-          {
-            cause,
-          },
-        ),
-        path,
+    })
+      .catch((cause) => {
+        opts.onError?.({ error: cause, path });
+        maybeEnqueue([
+          idx,
+          ASYNC_ITERABLE_STATUS_ERROR,
+          opts.formatError?.({ error: cause, path }),
+        ]);
+      })
+      .finally(() => {
+        pending.delete(idx);
       });
-    });
     return idx;
   }
   function checkMaxDepth(path: (string | number)[]) {
@@ -250,12 +219,12 @@ function createBatchStreamProducer(opts: ProducerOptions) {
     }
     return null;
   }
-  function dehydrateAsync(
+  function encodeAsync(
     value: unknown,
     path: (string | number)[],
   ): null | [type: ChunkValueType, chunkId: ChunkIndex] {
     if (isPromise(value)) {
-      return [CHUNK_VALUE_TYPE_PROMISE, dehydratePromise(value, path)];
+      return [CHUNK_VALUE_TYPE_PROMISE, encodePromise(value, path)];
     }
     if (isAsyncIterable(value)) {
       if (opts.maxDepth && path.length >= opts.maxDepth) {
@@ -263,29 +232,26 @@ function createBatchStreamProducer(opts: ProducerOptions) {
       }
       return [
         CHUNK_VALUE_TYPE_ASYNC_ITERABLE,
-        dehydrateAsyncIterable(value, path),
+        encodeAsyncIterable(value, path),
       ];
     }
     return null;
   }
-  function dehydrate(
-    value: unknown,
-    path: (string | number)[],
-  ): DehydratedValue {
+  function encode(value: unknown, path: (string | number)[]): EncodedValue {
     if (value === undefined) {
       return [[]];
     }
     if (!isObject(value)) {
       return [[value]];
     }
-    const reg = dehydrateAsync(value, path);
+    const reg = encodeAsync(value, path);
     if (reg) {
       return [[placeholder], [null, ...reg]];
     }
     const newObj = {} as Record<string, unknown>;
     const asyncValues: ChunkDefinition[] = [];
     for (const [key, item] of Object.entries(value)) {
-      const transformed = dehydrateAsync(item, [...path, key]);
+      const transformed = encodeAsync(item, [...path, key]);
       if (!transformed) {
         newObj[key] = item;
         continue;
@@ -298,8 +264,9 @@ function createBatchStreamProducer(opts: ProducerOptions) {
 
   const newHead: Head = {};
   for (const [key, item] of Object.entries(data)) {
-    newHead[key] = dehydrate(item, [key]);
+    newHead[key] = encode(item, [key]);
   }
+  pending.activate();
 
   return [newHead, stream.readable] as const;
 }
@@ -440,7 +407,7 @@ export async function jsonlStreamConsumer<THead>(opts: {
   /**
    * This `AbortController` will be triggered when there are no more listeners to the stream.
    */
-  abortController: AbortController | null;
+  abortController: AbortController;
 }) {
   const { deserialize = (v) => v } = opts;
 
@@ -458,46 +425,26 @@ export async function jsonlStreamConsumer<THead>(opts: {
 
   type ControllerChunk = ChunkData | StreamInterruptedError;
   type ChunkController = ReadableStreamDefaultController<ControllerChunk>;
-  type ControllerWrapper = {
-    controller: ChunkController;
-    returned: boolean;
-  };
-  const chunkDeferred = new Map<ChunkIndex, Deferred<ControllerWrapper>>();
-  const controllers = new Map<ChunkIndex, ControllerWrapper>();
-
-  const maybeAbort = () => {
-    if (
-      chunkDeferred.size === 0 &&
-      Array.from(controllers.values()).every((it) => it.returned)
-    ) {
-      // nothing is listening to the stream anymore
+
+  const controllers = withRefCount(
+    new Map<ChunkIndex, ChunkController>(),
+    () => {
       opts.abortController?.abort();
-    }
-  };
+    },
+  );
 
-  function hydrateChunkDefinition(value: ChunkDefinition) {
+  function decodeChunkDefinition(value: ChunkDefinition) {
     const [_path, type, chunkId] = value;
 
-    const { readable, controller } = createReadableStream<ChunkData>();
-
-    const wrapper: ControllerWrapper = {
-      controller,
-      returned: false,
-    };
-    controllers.set(chunkId, wrapper);
+    const stream = createReadableStream<ChunkData>();
 
-    // resolve chunk deferred if it exists
-    const deferred = chunkDeferred.get(chunkId);
-    if (deferred) {
-      deferred.resolve(wrapper);
-      chunkDeferred.delete(chunkId);
-    }
+    controllers.set(chunkId, stream.controller);
 
     switch (type) {
       case CHUNK_VALUE_TYPE_PROMISE: {
         return new Promise((resolve, reject) => {
           // listen for next value in the stream
-          const reader = readable.getReader();
+          const reader = stream.readable.getReader();
           reader
             .read()
             .then((it) => {
@@ -513,7 +460,7 @@ export async function jsonlStreamConsumer<THead>(opts: {
               const [_chunkId, status, data] = value as PromiseChunk;
               switch (status) {
                 case PROMISE_STATUS_FULFILLED:
-                  resolve(hydrate(data));
+                  resolve(decode(data));
 
                   break;
                 case PROMISE_STATUS_REJECTED:
@@ -525,17 +472,14 @@ export async function jsonlStreamConsumer<THead>(opts: {
             })
             .catch(reject)
             .finally(() => {
-              // reader.releaseLock();
               controllers.delete(chunkId);
-
-              maybeAbort();
             });
         });
       }
       case CHUNK_VALUE_TYPE_ASYNC_ITERABLE: {
         return {
           [Symbol.asyncIterator]: () => {
-            const reader = readable.getReader();
+            const reader = stream.readable.getReader();
             const iterator: AsyncIterator<unknown> = {
               next: async () => {
                 const { done, value } = await reader.read();
@@ -544,7 +488,6 @@ export async function jsonlStreamConsumer<THead>(opts: {
                 }
                 if (done) {
                   controllers.delete(chunkId);
-                  maybeAbort();
                   return {
                     done: true,
                     value: undefined,
@@ -557,18 +500,16 @@ export async function jsonlStreamConsumer<THead>(opts: {
                   case ASYNC_ITERABLE_STATUS_VALUE:
                     return {
                       done: false,
-                      value: hydrate(data),
+                      value: decode(data),
                     };
                   case ASYNC_ITERABLE_STATUS_RETURN:
                     controllers.delete(chunkId);
-                    maybeAbort();
                     return {
                       done: true,
-                      value: hydrate(data),
+                      value: decode(data),
                     };
                   case ASYNC_ITERABLE_STATUS_ERROR:
                     controllers.delete(chunkId);
-                    maybeAbort();
                     throw (
                       opts.formatError?.({ error: data }) ??
                       new AsyncError(data)
@@ -576,8 +517,7 @@ export async function jsonlStreamConsumer<THead>(opts: {
                 }
               },
               return: async () => {
-                wrapper.returned = true;
-                maybeAbort();
+                controllers.delete(chunkId);
                 return {
                   done: true,
                   value: undefined,
@@ -591,18 +531,18 @@ export async function jsonlStreamConsumer<THead>(opts: {
     }
   }
 
-  function hydrate(value: DehydratedValue): unknown {
+  function decode(value: EncodedValue): unknown {
     const [[data], ...asyncProps] = value;
 
     for (const value of asyncProps) {
       const [key] = value;
-      const hydrated = hydrateChunkDefinition(value);
+      const decoded = decodeChunkDefinition(value);
 
       if (key === null) {
-        return hydrated;
+        return decoded;
       }
 
-      (data as any)[key] = hydrated;
+      (data as any)[key] = decoded;
     }
     return data;
   }
@@ -611,11 +551,7 @@ export async function jsonlStreamConsumer<THead>(opts: {
     const error = new StreamInterruptedError(reason);
 
     headDeferred?.reject(error);
-    for (const deferred of chunkDeferred.values()) {
-      deferred.reject(error);
-    }
-    chunkDeferred.clear();
-    for (const { controller } of controllers.values()) {
+    for (const controller of controllers.values()) {
       controller.enqueue(error);
       controller.close();
     }
@@ -629,26 +565,21 @@ export async function jsonlStreamConsumer<THead>(opts: {
             const head = chunkOrHead as Record<number | string, unknown>;
 
             for (const [key, value] of Object.entries(chunkOrHead)) {
-              const parsed = hydrate(value as any);
+              const parsed = decode(value as any);
               head[key] = parsed;
             }
             headDeferred.resolve(head as THead);
             headDeferred = null;
+
+            controllers.activate();
             return;
           }
           const chunk = chunkOrHead as ChunkData;
           const [idx] = chunk;
-          let wrapper = controllers.get(idx);
-          if (!wrapper) {
-            let deferred = chunkDeferred.get(idx);
-            if (!deferred) {
-              deferred = createDeferred();
-              chunkDeferred.set(idx, deferred);
-            }
 
-            wrapper = await deferred.promise;
-          }
-          wrapper.controller.enqueue(chunk);
+          // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+          const controller = controllers.get(idx)!;
+          controller.enqueue(chunk);
         },
         close: closeOrAbort,
         abort: closeOrAbort,

package/src/unstable-core-do-not-import/stream/sse.ts

@@ -1,13 +1,12 @@
 import { getTRPCErrorFromUnknown } from '../error/TRPCError';
+import { isAbortError } from '../http/isAbortError';
 import type { MaybePromise } from '../types';
 import { identity, run } from '../utils';
 import type { EventSourceLike } from './sse.types';
 import type { inferTrackedOutput } from './tracked';
 import { isTrackedEnvelope } from './tracked';
-import { takeWithGrace, withCancel } from './utils/asyncIterable';
+import { takeWithGrace, withMaxDuration } from './utils/asyncIterable';
 import { createReadableStream } from './utils/createReadableStream';
-import type { PromiseTimer } from './utils/promiseTimer';
-import { createPromiseTimer } from './utils/promiseTimer';
 import { PING_SYM, withPing } from './utils/withPing';
 
 type Serialize = (value: any) => any;
@@ -80,27 +79,23 @@ export function sseStreamProducer<TValue = unknown>(
 
     let iterable: AsyncIterable<TValue | typeof PING_SYM> = opts.data;
 
-    iterable = withCancel(iterable, stream.cancelledPromise);
-
     if (opts.emitAndEndImmediately) {
       iterable = takeWithGrace(iterable, {
         count: 1,
         gracePeriodMs: 1,
-        onCancel: () => opts.abortCtrl.abort(),
+        abortCtrl: opts.abortCtrl,
       });
     }
 
-    let maxDurationTimer: PromiseTimer | null = null;
     if (
-      opts.maxDurationMs != null &&
+      opts.maxDurationMs &&
       opts.maxDurationMs > 0 &&
       opts.maxDurationMs !== Infinity
     ) {
-      maxDurationTimer = createPromiseTimer(opts.maxDurationMs).start();
-      iterable = withCancel(
-        iterable,
-        maxDurationTimer.promise.then(() => opts.abortCtrl.abort()),
-      );
+      iterable = withMaxDuration(iterable, {
+        maxDurationMs: opts.maxDurationMs,
+        abortCtrl: opts.abortCtrl,
+      });
     }
 
     if (ping.enabled && ping.intervalMs !== Infinity && ping.intervalMs > 0) {
@@ -133,20 +128,24 @@ export function sseStreamProducer<TValue = unknown>(
         chunk = null;
       }
     } catch (err) {
-      // ignore abort errors, send any other errors
-      if (!(err instanceof Error) || err.name !== 'AbortError') {
-        // `err` must be caused by `opts.data`, `JSON.stringify` or `serialize`.
-        // So, a user error in any case.
-        const error = getTRPCErrorFromUnknown(err);
-        const data = opts.formatError?.({ error }) ?? null;
-        stream.controller.enqueue({
-          event: SERIALIZED_ERROR_EVENT,
-          data: JSON.stringify(serialize(data)),
-        });
+      if (isAbortError(err)) {
+        // ignore abort errors, send any other errors
+        return;
       }
+      // `err` must be caused by `opts.data`, `JSON.stringify` or `serialize`.
+      // So, a user error in any case.
+      const error = getTRPCErrorFromUnknown(err);
+      const data = opts.formatError?.({ error }) ?? null;
+      stream.controller.enqueue({
+        event: SERIALIZED_ERROR_EVENT,
+        data: JSON.stringify(serialize(data)),
+      });
     } finally {
-      maxDurationTimer?.clear();
-      stream.controller.close();
+      try {
+        stream.controller.close();
+      } catch {
+        // ignore
+      }
     }
   }).catch((err) => {
     // should not be reached; just in case...

package/src/unstable-core-do-not-import/stream/utils/asyncIterable.ts

@@ -1,40 +1,47 @@
 import { Unpromise } from '../../../vendor/unpromise';
-import { noop } from '../../utils';
-import { createPromiseTimer } from './promiseTimer';
+import {
+  disposablePromiseTimer,
+  disposablePromiseTimerResult,
+} from './disposablePromiseTimer';
 
 /**
- * Derives a new {@link AsyncGenerator} based of {@link iterable}, that automatically stops with the
- * passed {@link cancel} promise.
+ * Derives a new {@link AsyncGenerator} based on {@link iterable}, that automatically stops after the specified duration.
  */
-export async function* withCancel<T>(
+export async function* withMaxDuration<T>(
   iterable: AsyncIterable<T>,
-  cancel: Promise<unknown>,
+  opts: { maxDurationMs: number; abortCtrl: AbortController },
 ): AsyncGenerator<T> {
-  const cancelPromise = cancel.then(noop);
   const iterator = iterable[Symbol.asyncIterator]();
-  // declaration outside the loop for garbage collection reasons
-  let result: null | IteratorResult<T> | void;
-  while (true) {
-    result = await Unpromise.race([iterator.next(), cancelPromise]);
-    if (result == null) {
-      await iterator.return?.();
-      break;
-    }
-    if (result.done) {
-      break;
+
+  const timer = disposablePromiseTimer(opts.maxDurationMs);
+  try {
+    const timerPromise = timer.start();
+
+    // declaration outside the loop for garbage collection reasons
+    let result: null | IteratorResult<T> | typeof disposablePromiseTimerResult;
+
+    while (true) {
+      result = await Unpromise.race([iterator.next(), timerPromise]);
+      if (result === disposablePromiseTimerResult) {
+        // cancelled due to timeout
+        opts.abortCtrl.abort();
+        const res = await iterator.return?.();
+        return res?.value;
+      }
+      if (result.done) {
+        return result;
+      }
+      yield result.value;
+      // free up reference for garbage collection
+      result = null;
     }
-    yield result.value;
-    // free up reference for garbage collection
-    result = null;
+  } finally {
+    // dispose timer
+    // Shouldn't be needed, but build breaks with `using` keyword
+    timer[Symbol.dispose]();
   }
 }
 
-interface TakeWithGraceOptions {
-  count: number;
-  gracePeriodMs: number;
-  onCancel: () => void;
-}
-
 /**
  * Derives a new {@link AsyncGenerator} based of {@link iterable}, that yields its first
  * {@link count} values. Then, a grace period of {@link gracePeriodMs} is started in which further
@@ -42,31 +49,45 @@ interface TakeWithGraceOptions {
  */
 export async function* takeWithGrace<T>(
   iterable: AsyncIterable<T>,
-  { count, gracePeriodMs, onCancel }: TakeWithGraceOptions,
+  opts: {
+    count: number;
+    gracePeriodMs: number;
+    abortCtrl: AbortController;
+  },
 ): AsyncGenerator<T> {
   const iterator = iterable[Symbol.asyncIterator]();
-  const timer = createPromiseTimer(gracePeriodMs);
+
+  // declaration outside the loop for garbage collection reasons
+  let result: null | IteratorResult<T> | typeof disposablePromiseTimerResult;
+
+  const timer = disposablePromiseTimer(opts.gracePeriodMs);
   try {
-    // declaration outside the loop for garbage collection reasons
-    let result: null | IteratorResult<T> | void;
+    let count = opts.count;
+
+    let timerPromise = new Promise<typeof disposablePromiseTimerResult>(() => {
+      // never resolves
+    });
+
     while (true) {
-      result = await Unpromise.race([iterator.next(), timer.promise]);
-      if (result == null) {
+      result = await Unpromise.race([iterator.next(), timerPromise]);
+      if (result === disposablePromiseTimerResult) {
         // cancelled
-        await iterator.return?.();
-        break;
+        const res = await iterator.return?.();
+        return res?.value;
       }
       if (result.done) {
-        break;
+        return result.value;
       }
       yield result.value;
       if (--count === 0) {
-        timer.start().promise.then(onCancel, noop);
+        timerPromise = timer.start();
+        // eslint-disable-next-line @typescript-eslint/no-floating-promises
+        timerPromise.then(() => opts.abortCtrl.abort());
       }
       // free up reference for garbage collection
       result = null;
     }
   } finally {
-    timer.clear();
+    timer[Symbol.dispose]();
   }
 }