@trpc/server 11.0.0-alpha-tmp-app-router-example.388 → 11.0.0-alpha-tmp-issues-5851-take-two.451

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

@@ -1,19 +1,8 @@
-/* eslint-disable @typescript-eslint/no-non-null-assertion */
-import { isAsyncIterable, isFunction, isObject } from '../utils';
-
-// ---------- utils
-
-function createReadableStream<TValue = unknown>() {
-  let controller: ReadableStreamDefaultController<TValue> =
-    null as unknown as ReadableStreamDefaultController<TValue>;
-  const stream = new ReadableStream<TValue>({
-    start(c) {
-      controller = c;
-    },
-  });
-
-  return [stream, controller] as const;
-}
+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 } from './utils/createReadableStream';
 
 /**
  * A subset of the standard ReadableStream properties needed by tRPC internally.
@@ -65,19 +54,19 @@ type ChunkDefinition = [
   type: ChunkValueType,
   chunkId: ChunkIndex,
 ];
-type HydratedValue = [
+type DehydratedValue = [
   // data
   [unknown],
   // chunk descriptions
   ...ChunkDefinition[],
 ];
 
-type Head = Record<string, HydratedValue>;
+type Head = Record<string, DehydratedValue>;
 type PromiseChunk =
   | [
       chunkIndex: ChunkIndex,
       status: PROMISE_STATUS_FULFILLED,
-      value: HydratedValue,
+      value: DehydratedValue,
     ]
   | [chunkIndex: ChunkIndex, status: PROMISE_STATUS_REJECTED, error: unknown];
 type IterableChunk =
@@ -85,7 +74,7 @@ type IterableChunk =
   | [
       chunkIndex: ChunkIndex,
       status: ASYNC_ITERABLE_STATUS_VALUE,
-      value: HydratedValue,
+      value: DehydratedValue,
     ]
   | [
       chunkIndex: ChunkIndex,
@@ -105,18 +94,16 @@ export function isPromise(value: unknown): value is Promise<unknown> {
 type Serialize = (value: any) => any;
 type Deserialize = (value: any) => any;
 
+type PathArray = readonly (string | number)[];
 export type ProducerOnError = (opts: {
   error: unknown;
-  path: (string | number)[];
+  path: PathArray;
 }) => void;
 export interface ProducerOptions {
   serialize?: Serialize;
-  data: Record<number, unknown>;
+  data: Record<string, unknown> | unknown[];
   onError?: ProducerOnError;
-  formatError?: (opts: {
-    error: unknown;
-    path: (string | number)[];
-  }) => unknown;
+  formatError?: (opts: { error: unknown; path: PathArray }) => unknown;
   maxDepth?: number;
 }
 
@@ -131,14 +118,15 @@ function createBatchStreamProducer(opts: ProducerOptions) {
   let counter = 0 as ChunkIndex;
   const placeholder = 0 as PlaceholderValue;
 
-  const [stream, controller] = createReadableStream<ChunkData>();
+  const stream = createReadableStream<ChunkData>();
   const pending = new Set<ChunkIndex>();
+
   function maybeClose() {
-    if (pending.size === 0) {
-      controller.close();
+    if (pending.size === 0 && !stream.cancelled()) {
+      stream.controller.close();
     }
   }
-  function hydratePromise(
+  function dehydratePromise(
     promise: Promise<unknown>,
     path: (string | number)[],
   ) {
@@ -152,19 +140,24 @@ function createBatchStreamProducer(opts: ProducerOptions) {
     }
     const idx = counter++ as ChunkIndex;
     pending.add(idx);
-    const enqueue = (value: PromiseChunk) => {
-      controller.enqueue(value);
-    };
-    promise
+
+    Promise.race([promise, stream.cancelledPromise])
       .then((it) => {
-        enqueue([idx, PROMISE_STATUS_FULFILLED, hydrate(it, path)]);
+        if (it === null) {
+          return;
+        }
+        stream.controller.enqueue([
+          idx,
+          PROMISE_STATUS_FULFILLED,
+          dehydrate(it, path),
+        ]);
       })
-      .catch((error) => {
-        opts.onError?.({ error, path });
-        enqueue([
+      .catch((cause) => {
+        opts.onError?.({ error: cause, path });
+        stream.controller.enqueue([
           idx,
           PROMISE_STATUS_REJECTED,
-          opts.formatError?.({ error, path }),
+          opts.formatError?.({ error: cause, path }),
         ]);
       })
       .finally(() => {
@@ -173,7 +166,7 @@ function createBatchStreamProducer(opts: ProducerOptions) {
       });
     return idx;
   }
-  function hydrateAsyncIterable(
+  function dehydrateAsyncIterable(
     iterable: AsyncIterable<unknown>,
     path: (string | number)[],
   ) {
@@ -187,28 +180,56 @@ function createBatchStreamProducer(opts: ProducerOptions) {
     }
     const idx = counter++ as ChunkIndex;
     pending.add(idx);
-    void (async () => {
-      try {
-        for await (const item of iterable) {
-          controller.enqueue([
+    run(async () => {
+      const iterator = iterable[Symbol.asyncIterator]();
+
+      while (true) {
+        const next = await Promise.race([
+          iterator.next().catch(getTRPCErrorFromUnknown),
+          stream.cancelledPromise,
+        ]);
+
+        if (next instanceof Error) {
+          opts.onError?.({ error: next, path });
+
+          stream.controller.enqueue([
             idx,
-            ASYNC_ITERABLE_STATUS_VALUE,
-            hydrate(item, path),
+            ASYNC_ITERABLE_STATUS_ERROR,
+            opts.formatError?.({ error: next, path }),
           ]);
+          return;
+        }
+        if (next === 'cancelled') {
+          await iterator.return?.();
+          break;
+        }
+        if (next.done) {
+          stream.controller.enqueue([idx, ASYNC_ITERABLE_STATUS_DONE]);
+          break;
         }
-        controller.enqueue([idx, ASYNC_ITERABLE_STATUS_DONE]);
-      } catch (error) {
-        opts.onError?.({ error, path });
-        controller.enqueue([
+        stream.controller.enqueue([
           idx,
-          ASYNC_ITERABLE_STATUS_ERROR,
-          opts.formatError?.({ error, path }),
+          ASYNC_ITERABLE_STATUS_VALUE,
+          dehydrate(next.value, path),
         ]);
-      } finally {
-        pending.delete(idx);
-        maybeClose();
       }
-    })();
+
+      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,
+      });
+    });
     return idx;
   }
   function checkMaxDepth(path: (string | number)[]) {
@@ -217,12 +238,12 @@ function createBatchStreamProducer(opts: ProducerOptions) {
     }
     return null;
   }
-  function hydrateChunk(
+  function dehydrateChunk(
     value: unknown,
     path: (string | number)[],
   ): null | [type: ChunkValueType, chunkId: ChunkIndex] {
     if (isPromise(value)) {
-      return [CHUNK_VALUE_TYPE_PROMISE, hydratePromise(value, path)];
+      return [CHUNK_VALUE_TYPE_PROMISE, dehydratePromise(value, path)];
     }
     if (isAsyncIterable(value)) {
       if (opts.maxDepth && path.length >= opts.maxDepth) {
@@ -230,13 +251,16 @@ function createBatchStreamProducer(opts: ProducerOptions) {
       }
       return [
         CHUNK_VALUE_TYPE_ASYNC_ITERABLE,
-        hydrateAsyncIterable(value, path),
+        dehydrateAsyncIterable(value, path),
       ];
     }
     return null;
   }
-  function hydrate(value: unknown, path: (string | number)[]): HydratedValue {
-    const reg = hydrateChunk(value, path);
+  function dehydrate(
+    value: unknown,
+    path: (string | number)[],
+  ): DehydratedValue {
+    const reg = dehydrateChunk(value, path);
     if (reg) {
       return [[placeholder], [null, ...reg]];
     }
@@ -246,7 +270,7 @@ function createBatchStreamProducer(opts: ProducerOptions) {
     const newObj = {} as Record<string, unknown>;
     const asyncValues: ChunkDefinition[] = [];
     for (const [key, item] of Object.entries(value)) {
-      const transformed = hydrateChunk(item, [...path, key]);
+      const transformed = dehydrateChunk(item, [...path, key]);
       if (!transformed) {
         newObj[key] = item;
         continue;
@@ -259,10 +283,10 @@ function createBatchStreamProducer(opts: ProducerOptions) {
 
   const newHead: Head = {};
   for (const [key, item] of Object.entries(data)) {
-    newHead[key] = hydrate(item, [key]);
+    newHead[key] = dehydrate(item, [key]);
   }
 
-  return [newHead, stream] as const;
+  return [newHead, stream.readable] as const;
 }
 /**
  * JSON Lines stream producer
@@ -296,6 +320,7 @@ export function jsonlStreamProducer(opts: ProducerOptions) {
     )
     .pipeThrough(new TextEncoderStream());
 }
+
 class StreamInterruptedError extends Error {
   constructor(cause?: unknown) {
     // eslint-disable-next-line @typescript-eslint/ban-ts-comment
@@ -313,7 +338,7 @@ export type ConsumerOnError = (opts: { error: unknown }) => void;
 const nodeJsStreamToReaderEsque = (source: NodeJSReadableStreamEsque) => {
   return {
     getReader() {
-      const [stream, controller] = createReadableStream<Uint8Array>();
+      const { readable, controller } = createReadableStream<Uint8Array>();
       source.on('data', (chunk) => {
         controller.enqueue(chunk);
       });
@@ -323,7 +348,7 @@ const nodeJsStreamToReaderEsque = (source: NodeJSReadableStreamEsque) => {
       source.on('error', (error) => {
         controller.error(error);
       });
-      return stream.getReader();
+      return readable.getReader();
     },
   };
 };
@@ -388,19 +413,6 @@ function createConsumerStream<THead>(
   );
 }
 
-function createDeferred<TValue>() {
-  let resolve: (value: TValue) => void;
-  let reject: (error: unknown) => void;
-  const promise = new Promise<TValue>((res, rej) => {
-    resolve = res;
-    reject = rej;
-  });
-
-  return { promise, resolve: resolve!, reject: reject! };
-}
-
-type Deferred<TValue> = ReturnType<typeof createDeferred<TValue>>;
-
 /**
  * JSON Lines stream consumer
  * @see https://jsonlines.org/
@@ -412,6 +424,7 @@ export async function jsonlStreamConsumer<THead>(opts: {
   formatError?: (opts: { error: unknown }) => Error;
 }) {
   const { deserialize = (v) => v } = opts;
+  const streamAbortController = new AbortController();
 
   let source = createConsumerStream<Head>(opts.from);
   if (deserialize) {
@@ -427,19 +440,28 @@ export async function jsonlStreamConsumer<THead>(opts: {
 
   type ControllerChunk = ChunkData | StreamInterruptedError;
   type ChunkController = ReadableStreamDefaultController<ControllerChunk>;
-  const chunkDeferred = new Map<ChunkIndex, Deferred<ChunkController>>();
-  const controllers = new Map<ChunkIndex, ChunkController>();
+  type ControllerWrapper = {
+    controller: ChunkController;
+    abortController: AbortController;
+  };
+  const chunkDeferred = new Map<ChunkIndex, Deferred<ControllerWrapper>>();
+  const controllers = new Map<ChunkIndex, ControllerWrapper>();
 
-  function dehydrateChunkDefinition(value: ChunkDefinition) {
+  function hydrateChunkDefinition(value: ChunkDefinition, path: PathArray) {
     const [_path, type, chunkId] = value;
 
-    const [stream, controller] = createReadableStream<ChunkData>();
-    controllers.set(chunkId, controller);
+    const { readable, controller } = createReadableStream<ChunkData>();
+
+    const wrapper: ControllerWrapper = {
+      controller,
+      abortController: new AbortController(),
+    };
+    controllers.set(chunkId, wrapper);
 
     // resolve chunk deferred if it exists
     const deferred = chunkDeferred.get(chunkId);
     if (deferred) {
-      deferred.resolve(controller);
+      deferred.resolve(wrapper);
       chunkDeferred.delete(chunkId);
     }
 
@@ -447,7 +469,7 @@ export async function jsonlStreamConsumer<THead>(opts: {
       case CHUNK_VALUE_TYPE_PROMISE: {
         return new Promise((resolve, reject) => {
           // listen for next value in the stream
-          const reader = stream.getReader();
+          const reader = readable.getReader();
           reader
             .read()
             .then((it) => {
@@ -463,7 +485,7 @@ export async function jsonlStreamConsumer<THead>(opts: {
               const [_chunkId, status, data] = value as PromiseChunk;
               switch (status) {
                 case PROMISE_STATUS_FULFILLED:
-                  resolve(dehydrate(data));
+                  resolve(hydrate(data, path));
                   break;
                 case PROMISE_STATUS_REJECTED:
                   reject(
@@ -481,51 +503,83 @@ export async function jsonlStreamConsumer<THead>(opts: {
       }
       case CHUNK_VALUE_TYPE_ASYNC_ITERABLE: {
         return {
-          [Symbol.asyncIterator]: async function* () {
-            const reader = stream.getReader();
-            while (true) {
-              const { done, value } = await reader.read();
-              if (done) {
-                break;
-              }
-              if (value instanceof StreamInterruptedError) {
-                throw value;
-              }
-
-              const [_chunkId, status, data] = value as IterableChunk;
-
-              switch (status) {
-                case ASYNC_ITERABLE_STATUS_VALUE:
-                  yield dehydrate(data);
-                  break;
-                case ASYNC_ITERABLE_STATUS_DONE:
-                  controllers.delete(chunkId);
-                  return;
-                case ASYNC_ITERABLE_STATUS_ERROR:
+          [Symbol.asyncIterator]: () => {
+            const reader = readable.getReader();
+            const iterator: AsyncIterator<unknown> = {
+              next: async () => {
+                const { done, value } = await reader.read();
+                if (value instanceof StreamInterruptedError) {
+                  throw value;
+                }
+                if (done) {
                   controllers.delete(chunkId);
-                  throw (
-                    opts.formatError?.({ error: data }) ?? new AsyncError(data)
-                  );
-              }
-            }
+                  return {
+                    done: true,
+                    value: undefined,
+                  };
+                }
+
+                const [_chunkId, status, data] = value as IterableChunk;
+
+                switch (status) {
+                  case ASYNC_ITERABLE_STATUS_VALUE:
+                    return {
+                      done: false,
+                      value: hydrate(data, path),
+                    };
+                  case ASYNC_ITERABLE_STATUS_DONE:
+                    controllers.delete(chunkId);
+                    return {
+                      done: true,
+                      value: undefined,
+                    };
+                  case ASYNC_ITERABLE_STATUS_ERROR:
+                    controllers.delete(chunkId);
+                    throw (
+                      opts.formatError?.({ error: data }) ??
+                      new AsyncError(data)
+                    );
+                }
+              },
+              return: async () => {
+                wrapper.abortController.abort();
+                if (
+                  chunkDeferred.size === 0 &&
+                  Array.from(controllers.values()).every(
+                    (it) => it.abortController.signal.aborted,
+                  )
+                ) {
+                  // nothing is listening to the stream anymore
+                  streamAbortController.abort();
+                }
+                return {
+                  done: true,
+                  value: undefined,
+                };
+              },
+            };
+            return iterator;
           },
         };
       }
     }
   }
 
-  function dehydrate(value: HydratedValue): unknown {
+  function hydrate(value: DehydratedValue, path: PathArray): unknown {
     const [[data], ...asyncProps] = value;
 
     for (const value of asyncProps) {
-      const dehydrated = dehydrateChunkDefinition(value);
-
-      const [path] = value;
-      if (path === null) {
-        return dehydrated;
+      const [key] = value;
+      const hydrated = hydrateChunkDefinition(
+        value,
+        key === null ? path : [...path, key],
+      );
+
+      if (key === null) {
+        return hydrated;
       }
 
-      (data as any)[path] = dehydrated;
+      (data as any)[key] = hydrated;
     }
     return data;
   }
@@ -538,7 +592,7 @@ export async function jsonlStreamConsumer<THead>(opts: {
       deferred.reject(error);
     }
     chunkDeferred.clear();
-    for (const controller of controllers.values()) {
+    for (const { controller } of controllers.values()) {
       controller.enqueue(error);
       controller.close();
     }
@@ -547,12 +601,17 @@ export async function jsonlStreamConsumer<THead>(opts: {
   source
     .pipeTo(
       new WritableStream({
+        start(controller) {
+          streamAbortController.signal.addEventListener('abort', () => {
+            controller.error(streamAbortController.signal.reason);
+          });
+        },
         async write(chunkOrHead) {
           if (headDeferred) {
             const head = chunkOrHead as Record<number | string, unknown>;
 
             for (const [key, value] of Object.entries(chunkOrHead)) {
-              const parsed = dehydrate(value as any);
+              const parsed = hydrate(value as any, [key]);
               head[key] = parsed;
             }
             headDeferred.resolve(head as THead);
@@ -561,17 +620,17 @@ export async function jsonlStreamConsumer<THead>(opts: {
           }
           const chunk = chunkOrHead as ChunkData;
           const [idx] = chunk;
-          let controller = controllers.get(idx);
-          if (!controller) {
+          let wrapper = controllers.get(idx);
+          if (!wrapper) {
             let deferred = chunkDeferred.get(idx);
             if (!deferred) {
-              deferred = createDeferred<ChunkController>();
+              deferred = createDeferred();
               chunkDeferred.set(idx, deferred);
             }
 
-            controller = await deferred.promise;
+            wrapper = await deferred.promise;
           }
-          controller.enqueue(chunk);
+          wrapper.controller.enqueue(chunk);
         },
         close: closeOrAbort,
         abort: closeOrAbort,