From daaa1998837bdb6eaa42d9160292e781fadb3dc8 Mon Sep 17 00:00:00 2001
From: Daniel Choudhury <dannychoudhury@gmail.com>
Date: Fri, 20 Oct 2023 14:46:57 +0700
Subject: [PATCH] fix(stream-ssr): Cancel the timeout when the react stream has
 finished (#9317)

---
 .../streaming/createReactStreamingHandler.ts  |  7 +++
 packages/vite/src/streaming/streamHelpers.ts  | 60 ++++++++++++-------
 .../transforms/cancelTimeoutTransform.ts      |  7 +++
 3 files changed, 51 insertions(+), 23 deletions(-)
 create mode 100644 packages/vite/src/streaming/transforms/cancelTimeoutTransform.ts

diff --git a/packages/vite/src/streaming/createReactStreamingHandler.ts b/packages/vite/src/streaming/createReactStreamingHandler.ts
index a958e0487bc4..63de73f4a261 100644
--- a/packages/vite/src/streaming/createReactStreamingHandler.ts
+++ b/packages/vite/src/streaming/createReactStreamingHandler.ts
@@ -115,6 +115,13 @@ export const createReactStreamingHandler = async (
       },
       {
         waitForAllReady: isSeoCrawler,
+        onError: (err) => {
+          if (!isProd && viteDevServer) {
+            viteDevServer.ssrFixStacktrace(err)
+          }
+
+          console.error(err)
+        },
       }
     )
 
diff --git a/packages/vite/src/streaming/streamHelpers.ts b/packages/vite/src/streaming/streamHelpers.ts
index efcd01f3de9e..e92e9038c624 100644
--- a/packages/vite/src/streaming/streamHelpers.ts
+++ b/packages/vite/src/streaming/streamHelpers.ts
@@ -2,6 +2,11 @@ import path from 'node:path'
 
 import React from 'react'
 
+import type {
+  RenderToReadableStreamOptions,
+  ReactDOMServerReadableStream,
+} from 'react-dom/server'
+
 import type { TagDescriptor } from '@redwoodjs/web'
 // @TODO (ESM), use exports field. Cannot import from web because of index exports
 import {
@@ -10,6 +15,7 @@ import {
 } from '@redwoodjs/web/dist/components/ServerInject'
 
 import { createBufferedTransformStream } from './transforms/bufferedTransform'
+import { createTimeoutTransform } from './transforms/cancelTimeoutTransform'
 import { createServerInjectionTransform } from './transforms/serverInjectionTransform'
 
 interface RenderToStreamArgs {
@@ -24,6 +30,7 @@ interface RenderToStreamArgs {
 
 interface StreamOptions {
   waitForAllReady?: boolean
+  onError?: (err: Error) => void
 }
 
 export async function reactRenderToStreamResponse(
@@ -55,13 +62,22 @@ export async function reactRenderToStreamResponse(
   const { injectionState, injectToPage } = createInjector()
 
   // This makes it safe for us to inject at any point in the stream
-  const bufferedTransformStream = createBufferedTransformStream()
+  const bufferTransform = createBufferedTransformStream()
 
   // This is a transformer stream, that will inject all things called with useServerInsertedHtml
-  const serverInjectionTransformer = createServerInjectionTransform({
+  const serverInjectionTransform = createServerInjectionTransform({
     injectionState,
   })
 
+  // Timeout after 10 seconds
+  // @TODO make this configurable
+  const controller = new AbortController()
+  const timeoutHandle = setTimeout(() => {
+    controller.abort()
+  }, 10000)
+
+  const timeoutTransform = createTimeoutTransform(timeoutHandle)
+
   // @ts-expect-error Something in React's packages mean types dont come through
   // Possible that we need to upgrade the @types/* packages
   const { renderToReadableStream } = await import('react-dom/server.edge')
@@ -96,29 +112,27 @@ export async function reactRenderToStreamResponse(
     // This gets set if there are errors inside Suspense boundaries
     let didErrorOutsideShell = false
 
-    // Timeout after 10 seconds
-    // @TODO make this configurable
-    const controller = new AbortController()
-    setTimeout(() => {
-      controller.abort()
-    }, 10000)
+    // Assign here so we get types, the dynamic import messes types
+    const renderToStreamOptions: RenderToReadableStreamOptions = {
+      ...bootstrapOptions,
+      signal: controller.signal,
+      onError: (err: any) => {
+        didErrorOutsideShell = true
+        console.error('🔻 Caught error outside shell')
+        streamOptions.onError?.(err)
+      },
+    }
 
-    const reactStream = await renderToReadableStream(
-      renderRoot(currentPathName),
-      {
-        ...bootstrapOptions,
-        signal: controller.signal,
-        onError: (err: any) => {
-          didErrorOutsideShell = true
-          console.error('🔻 Caught error outside shell')
-          console.error(err)
-        },
-      }
-    )
+    const reactStream: ReactDOMServerReadableStream =
+      await renderToReadableStream(
+        renderRoot(currentPathName),
+        renderToStreamOptions
+      )
 
     const output = reactStream
-      .pipeThrough(bufferedTransformStream)
-      .pipeThrough(serverInjectionTransformer)
+      .pipeThrough(bufferTransform)
+      .pipeThrough(serverInjectionTransform)
+      .pipeThrough(timeoutTransform)
 
     if (waitForAllReady) {
       await reactStream.allReady
@@ -130,7 +144,7 @@ export async function reactRenderToStreamResponse(
     })
   } catch (e) {
     console.error('🔻 Failed to render shell')
-    console.error(e)
+    streamOptions.onError?.(e as Error)
 
     // @TODO Asking for clarification from React team. Their documentation on this is incomplete I think.
     // Having the Document (and bootstrap scripts) here allows client to recover from errors in the shell
diff --git a/packages/vite/src/streaming/transforms/cancelTimeoutTransform.ts b/packages/vite/src/streaming/transforms/cancelTimeoutTransform.ts
new file mode 100644
index 000000000000..f157ab7ca793
--- /dev/null
+++ b/packages/vite/src/streaming/transforms/cancelTimeoutTransform.ts
@@ -0,0 +1,7 @@
+export function createTimeoutTransform(timeoutHandle: NodeJS.Timeout) {
+  return new TransformStream({
+    flush() {
+      clearTimeout(timeoutHandle)
+    },
+  })
+}