[Cache Components] Allow sync IO inside console methods

console methods may be patched to do things like add a timestamp to the log row. While we can't determine the meaning of readin the current time in arbitrary contexts we can know with reasonable certainty that any sync IO inside a console method is not being used in the logical output of the program. To account for this we can make a special affordance for sync IO inside of these methods that allows them to run without interrupting the prerender early.

Originall I wanted to make it so that even if you patch these methods in userland we can still exit the context but it turns out this is very complex and difficult because we have many patches layered together and they need to run in the correct order. So essentially this just permits sync IO in the preexisting implementations of the console methods. This is probably sufficient since typically these methods aren't patched in userspace but in the runtime host itself. If we recieve feedback that this is insufficient we can consider alternative later.

The original implementation exited the workUnitAsyncStorage but it turns out that some console methods like assert call others from their internal implementation. This means that you lose the context at the top of the patch stack for things like dimming console logs if you do this. Instead we simply set a temporary flag to allow sync IO during the a portion of the prerender. It is important to guard this with a try/finally and always synchronously otherwise you can accidentally and racily allow sync IO in other parts of the program that you did not intend

Author: gnoff

packages/next/src/server/app-render/app-render.tsx

@@ -174,6 +174,7 @@ import {
   workUnitAsyncStorage,
   type PrerenderStore,
 } from './work-unit-async-storage.external'
+import { devLogsAsyncStorage } from './dev-logs-async-storage.external'
 import { CacheSignal } from './cache-signal'
 import { getTracedMetadata } from '../lib/trace/utils'
 import { InvariantError } from '../../shared/lib/invariant-error'
@@ -2302,7 +2303,9 @@ async function renderToStream(
         }
       )
 
-      spawnDynamicValidationInDev(
+      devLogsAsyncStorage.run(
+        { dim: true },
+        spawnDynamicValidationInDev,
         resolveValidation,
         tree,
         ctx,

packages/next/src/server/app-render/dev-logs-async-storage-instance.ts

@@ -0,0 +1,5 @@
+import { createAsyncLocalStorage } from './async-local-storage'
+import type { DevLogsAsyncStorage } from './dev-logs-async-storage.external'
+
+export const devLogsAsyncStorageInstance: DevLogsAsyncStorage =
+  createAsyncLocalStorage()

packages/next/src/server/app-render/dev-logs-async-storage.external.ts

@@ -0,0 +1,12 @@
+import type { AsyncLocalStorage } from 'async_hooks'
+
+// Share the instance module in the next-shared layer
+import { devLogsAsyncStorageInstance } from './dev-logs-async-storage-instance' with { 'turbopack-transition': 'next-shared' }
+
+export interface DevLogsStore {
+  dim: boolean
+}
+
+export type DevLogsAsyncStorage = AsyncLocalStorage<DevLogsStore>
+
+export { devLogsAsyncStorageInstance as devLogsAsyncStorage }

packages/next/src/server/node-environment-extensions/console-dev.tsx

@@ -1,5 +1,5 @@
 import { dim } from '../../lib/picocolors'
-import { workUnitAsyncStorage } from '../app-render/work-unit-async-storage.external'
+import { devLogsAsyncStorage } from '../app-render/dev-logs-async-storage.external'
 
 type InterceptableConsoleMethod =
   | 'error'
@@ -163,25 +163,13 @@ function patchConsoleMethodDEV(methodName: InterceptableConsoleMethod): void {
     const originalMethod = descriptor.value
     const originalName = Object.getOwnPropertyDescriptor(originalMethod, 'name')
     const wrapperMethod = function (this: typeof console, ...args: any[]) {
-      const workUnitStore = workUnitAsyncStorage.getStore()
+      const devLogsStore = devLogsAsyncStorage.getStore()
 
-      switch (workUnitStore?.type) {
-        case 'prerender':
-        case 'prerender-client':
-        case 'prerender-runtime':
-          originalMethod.apply(this, dimConsoleCall(methodName, args))
-          break
-        case 'prerender-ppr':
-        case 'prerender-legacy':
-        case 'request':
-        case 'cache':
-        case 'private-cache':
-        case 'unstable-cache':
-        case undefined:
-          originalMethod.apply(this, args)
-          break
+      switch (devLogsStore?.dim) {
+        case true:
+          return originalMethod.apply(this, dimConsoleCall(methodName, args))
         default:
-          workUnitStore satisfies never
+          return originalMethod.apply(this, args)
       }
     }
     if (originalName) {