From 4865619f64540581ae5f36c42267b6375a60f83d Mon Sep 17 00:00:00 2001
From: Josh Story <story@hey.com>
Date: Thu, 10 Oct 2024 08:03:47 -0700
Subject: [PATCH] [dynamicIO] track Prerender environment name during dev

When rendering in dev we track environment name to pass along server logs to the browser. Today the environment name is the default "Server" however when `dynamicIO` is enabled we want to have logs in the first Task where rendering starts be tagged with the environment "Prerender". This will help us understand what work is happening in the "Prerender" phase and what is happening in the dynamic "Server" phase while developing in dev mode.v
---
 .../app-render/app-render-render-utils.ts     | 31 +++++++++++
 .../next/src/server/app-render/app-render.tsx | 51 +++++++++++++++----
 packages/next/types/$$compiled.internal.d.ts  |  2 +-
 3 files changed, 72 insertions(+), 12 deletions(-)
 create mode 100644 packages/next/src/server/app-render/app-render-render-utils.ts

diff --git a/packages/next/src/server/app-render/app-render-render-utils.ts b/packages/next/src/server/app-render/app-render-render-utils.ts
new file mode 100644
index 00000000000000..caa5d31f2711a7
--- /dev/null
+++ b/packages/next/src/server/app-render/app-render-render-utils.ts
@@ -0,0 +1,31 @@
+import { InvariantError } from '../../shared/lib/invariant-error'
+
+/**
+ * This is a utility function to make scheduling sequential tasks that run back to back easier.
+ * We schedule on the same queue (setImmediate) at the same time to ensure no other events can sneak in between.
+ */
+export function scheduleInSequentialTasks<R>(
+  render: () => R | Promise<R>,
+  followup: () => void
+): Promise<R> {
+  if (process.env.NEXT_RUNTIME === 'edge') {
+    throw new InvariantError(
+      '`scheduleInSequentialTasks` should not be called in edge runtime.'
+    )
+  } else {
+    return new Promise((resolve, reject) => {
+      let pendingResult: R | Promise<R>
+      setImmediate(() => {
+        try {
+          pendingResult = render()
+        } catch (err) {
+          reject(err)
+        }
+      })
+      setImmediate(() => {
+        followup()
+        resolve(pendingResult)
+      })
+    })
+  }
+}
diff --git a/packages/next/src/server/app-render/app-render.tsx b/packages/next/src/server/app-render/app-render.tsx
index a35963b3c3351c..b7b825a21f0936 100644
--- a/packages/next/src/server/app-render/app-render.tsx
+++ b/packages/next/src/server/app-render/app-render.tsx
@@ -146,7 +146,8 @@ import {
   createReactServerPrerenderResult,
   createReactServerPrerenderResultFromRender,
   prerenderAndAbortInSequentialTasks,
-} from '../app-render/app-render-prerender-utils'
+} from './app-render-prerender-utils'
+import { scheduleInSequentialTasks } from './app-render-render-utils'
 import { waitAtLeastOneReactRenderTask } from '../../lib/scheduler'
 import {
   workUnitAsyncStorage,
@@ -1445,17 +1446,45 @@ async function renderToStream(
       ctx,
       res.statusCode === 404
     )
-    reactServerResult = new ReactServerResult(
-      workUnitAsyncStorage.run(
-        requestStore,
-        ComponentMod.renderToReadableStream,
-        RSCPayload,
-        clientReferenceManifest.clientModules,
-        {
-          onError: serverComponentsErrorHandler,
-        }
+
+    if (
+      process.env.NODE_ENV === 'development' &&
+      process.env.NEXT_RUNTIME !== 'edge' &&
+      renderOpts.experimental.dynamicIO
+    ) {
+      let environmentName = 'Prerender'
+      reactServerResult = new ReactServerResult(
+        await workUnitAsyncStorage.run(
+          requestStore,
+          scheduleInSequentialTasks,
+          () => {
+            return ComponentMod.renderToReadableStream(
+              RSCPayload,
+              clientReferenceManifest.clientModules,
+              {
+                onError: serverComponentsErrorHandler,
+                environmentName: () => environmentName,
+              }
+            )
+          },
+          () => {
+            environmentName = 'Server'
+          }
+        )
       )
-    )
+    } else {
+      reactServerResult = new ReactServerResult(
+        workUnitAsyncStorage.run(
+          requestStore,
+          ComponentMod.renderToReadableStream,
+          RSCPayload,
+          clientReferenceManifest.clientModules,
+          {
+            onError: serverComponentsErrorHandler,
+          }
+        )
+      )
+    }
 
     // React doesn't start rendering synchronously but we want the RSC render to have a chance to start
     // before we begin SSR rendering because we want to capture any available preload headers so we tick
diff --git a/packages/next/types/$$compiled.internal.d.ts b/packages/next/types/$$compiled.internal.d.ts
index 80d64c2bfdd18a..e4ff405b5c7a6b 100644
--- a/packages/next/types/$$compiled.internal.d.ts
+++ b/packages/next/types/$$compiled.internal.d.ts
@@ -48,7 +48,7 @@ declare module 'react-server-dom-webpack/server.edge' {
     },
     options?: {
       temporaryReferences?: string
-      environmentName?: string
+      environmentName?: string | (() => string)
       filterStackFrame?: (url: string, functionName: string) => boolean
       onError?: (error: unknown) => void
       onPostpone?: (reason: string) => void