Repository: vercel/next.js
Pull request: Improve task pipelining

Stacked on https://github.com/vercel/next.js/pull/89972

The pipelining function has some awkward structure. This change cleans up the structure a bit

Notably there is a potentially meaningful change here where the final user provided task in runInSequentialTasks gets the same Fast Immediate treatment as the other tasks do. There is an implicit final task which does the assertion that there are no remaining immediates unflushed and then resolves the original tasks value.

Changed files:
filename: packages/next/errors.json
patch: @@ -1051,5 +1051,6 @@
   "1050": "Turbopack is not supported on this platform (%s/%s) because native bindings are not available. Only WebAssembly (WASM) bindings were loaded, and Turbopack requires native bindings.\\n\\nTo build on this platform, use Webpack instead:\\n  next build --webpack\\n\\nFor more information, see: https://nextjs.org/docs/app/api-reference/turbopack#supported-platforms",
   "1051": "Turbopack trace server is not supported on this platform (%s/%s) because native bindings are not available. Only WebAssembly (WASM) bindings were loaded, and Turbopack requires native bindings.",
   "1052": "Turbopack is not supported on this platform (%s/%s) because native bindings are not available. Only WebAssembly (WASM) bindings were loaded, and Turbopack requires native bindings. Use the --webpack flag instead.",
-  "1053": "Export encountered errors on %s %s:\\n\\t%s"
+  "1053": "Export encountered errors on %s %s:\\n\\t%s",
+  "1054": "`runInSequentialTasks` should not be called in edge runtime."
 }

filename: packages/next/src/server/app-render/app-render-render-utils.ts
patch: @@ -47,65 +47,69 @@ export function scheduleInSequentialTasks<R>(
 /**
  * This is a utility function to make scheduling sequential tasks that run back to back easier.
  * We schedule on the same queue (setTimeout) at the same time to ensure no other events can sneak in between.
- * The function that runs in the second task gets access to the first tasks's result.
+ *
+ * The first function runs in the first task. Each subsequent function runs in its own task.
+ * The returned promise resolves after the last task completes.
  */
-export function pipelineInSequentialTasks<A, B, C>(
-  one: () => A,
-  two: (a: A) => B,
-  three: (b: B) => C
-): Promise<C> {
+export function runInSequentialTasks<R>(
+  first: () => R,
+  ...rest: Array<() => void>
+): Promise<R> {
   if (process.env.NEXT_RUNTIME === 'edge') {
     throw new InvariantError(
-      '`pipelineInSequentialTasks` should not be called in edge runtime.'
+      '`runInSequentialTasks` should not be called in edge runtime.'
     )
   } else {
     return new Promise((resolve, reject) => {
       const scheduleTimeout = createAtomicTimerGroup()
+      const ids: ReturnType<typeof scheduleTimeout>[] = []
 
-      let oneResult: A
-      scheduleTimeout(() => {
-        try {
-          DANGEROUSLY_runPendingImmediatesAfterCurrentTask()
-          oneResult = one()
-        } catch (err) {
-          clearTimeout(twoId)
-          clearTimeout(threeId)
-          clearTimeout(fourId)
-          reject(err)
-        }
-      })
+      let result: R
+      ids.push(
+        scheduleTimeout(() => {
+          try {
+            DANGEROUSLY_runPendingImmediatesAfterCurrentTask()
+            result = first()
+          } catch (err) {
+            for (let i = 1; i < ids.length; i++) {
+              clearTimeout(ids[i])
+            }
+            reject(err)
+          }
+        })
+      )
 
-      let twoResult: B
-      const twoId = scheduleTimeout(() => {
-        // if `one` threw, then this timeout would've been cleared,
-        // so if we got here, we're guaranteed to have a value.
-        try {
-          DANGEROUSLY_runPendingImmediatesAfterCurrentTask()
-          twoResult = two(oneResult!)
-        } catch (err) {
-          clearTimeout(threeId)
-          clearTimeout(fourId)
-          reject(err)
-        }
-      })
+      for (let i = 0; i < rest.length; i++) {
+        const fn = rest[i]
+        let index = ids.length
 
-      let threeResult: C
-      const threeId = scheduleTimeout(() => {
-        // if `two` threw, then this timeout would've been cleared,
-        // so if we got here, we're guaranteed to have a value.
-        try {
-          expectNoPendingImmediates()
-          threeResult = three(twoResult!)
-        } catch (err) {
-          clearTimeout(fourId)
-          reject(err)
-        }
-      })
+        ids.push(
+          scheduleTimeout(() => {
+            try {
+              DANGEROUSLY_runPendingImmediatesAfterCurrentTask()
+              fn()
+            } catch (err) {
+              // clear remaining timeouts
+              while (++index < ids.length) {
+                clearTimeout(ids[index])
+              }
+              reject(err)
+            }
+          })
+        )
+      }
 
-      // We wait a task before resolving/rejecting
-      const fourId = scheduleTimeout(() => {
-        resolve(threeResult)
-      })
+      // We wait a task before resolving
+      ids.push(
+        scheduleTimeout(() => {
+          try {
+            expectNoPendingImmediates()
+            resolve(result)
+          } catch (err) {
+            reject(err)
+          }
+        })
+      )
     })
   }
 }

filename: packages/next/src/server/app-render/app-render.tsx
patch: @@ -174,7 +174,7 @@ import {
   printDebugThrownValueForProspectiveRender,
 } from './prospective-render-utils'
 import {
-  pipelineInSequentialTasks,
+  runInSequentialTasks,
   scheduleInSequentialTasks,
 } from './app-render-render-utils'
 import { waitAtLeastOneReactRenderTask } from '../../lib/scheduler'
@@ -3323,99 +3323,75 @@ async function renderWithRestartOnCacheMissInDev(
   // where sync IO does not cause aborts, so it's okay if it happens before render.
   const initialRscPayload = await getPayload(requestStore)
 
-  const maybeInitialStreamResult = await workUnitAsyncStorage.run(
+  const initialStreamResult = await workUnitAsyncStorage.run(
     requestStore,
-    () =>
-      pipelineInSequentialTasks(
-        () => {
-          // Static stage
-          initialStageController.advanceStage(RenderStage.Static)
-          startTime = performance.now() + performance.timeOrigin
-
-          const stream = ComponentMod.renderToReadableStream(
-            initialRscPayload,
-            clientModules,
-            {
-              onError,
-              environmentName,
-              startTime,
-              filterStackFrame,
-              debugChannel: debugChannel?.serverSide,
-              signal: initialReactController.signal,
-            }
-          )
-          // If we abort the render, we want to reject the stage-dependent promises as well.
-          // Note that we want to install this listener after the render is started
-          // so that it runs after react is finished running its abort code.
-          initialReactController.signal.addEventListener('abort', () => {
-            initialDataController.abort(initialReactController.signal.reason)
-          })
-
-          const [continuationStream, accumulatingStream] = stream.tee()
-          const accumulatedChunksPromise = accumulateStreamChunks(
-            accumulatingStream,
-            initialStageController,
-            initialDataController.signal
-          )
-          return { stream: continuationStream, accumulatedChunksPromise }
-        },
-        ({ stream, accumulatedChunksPromise }) => {
-          // Runtime stage
+    runInSequentialTasks,
+    () => {
+      initialStageController.advanceStage(RenderStage.Static)
+      startTime = performance.now() + performance.timeOrigin
 
-          if (initialStageController.currentStage === RenderStage.Abandoned) {
-            // If we abandoned the render in the static stage, we won't proceed further.
-            return null
-          }
+      const streamPair = ComponentMod.renderToReadableStream(
+        initialRscPayload,
+        clientModules,
+        {
+          onError,
+          environmentName,
+          startTime,
+          filterStackFrame,
+          debugChannel: debugChannel?.serverSide,
+          signal: initialReactController.signal,
+        }
+      ).tee()
 
-          // If we had a cache miss in the static stage, we'll have to discard this stream
-          // and render again once the caches are warm.
-          // If we already advanced stages we similarly had sync IO that might be from module loading
-          // and need to render again once the caches are warm.
-          if (cacheSignal.hasPendingReads()) {
-            // Regardless of whether we are going to abandon this
-            // render we need the unblock runtime b/c it's essential
-            // filling caches.
-            initialAbandonController.abort()
-            return null
-          }
+      // If we abort the render, we want to reject the stage-dependent promises as well.
+      // Note that we want to install this listener after the render is started
+      // so that it runs after react is finished running its abort code.
+      initialReactController.signal.addEventListener('abort', () => {
+        initialDataController.abort(initialReactController.signal.reason)
+      })
 
-          initialStageController.advanceStage(RenderStage.Runtime)
-          return { stream, accumulatedChunksPromise }
-        },
-        (result) => {
-          // Dynamic stage
-          if (
-            result === null ||
-            initialStageController.currentStage === RenderStage.Abandoned
-          ) {
-            // If we abandoned the render in the static or runtime stage, we won't proceed further.
-            return null
-          }
+      const stream = streamPair[0]
+      const accumulatedChunksPromise = accumulateStreamChunks(
+        streamPair[1],
+        initialStageController,
+        initialDataController.signal
+      )
 
-          // If we had cache misses in either of the previous stages,
-          // then we'll only use this render for filling caches.
-          // We won't advance the stage, and thus leave dynamic APIs hanging,
-          // because they won't be cached anyway, so it'd be wasted work.
-          if (cacheSignal.hasPendingReads()) {
-            initialAbandonController.abort()
-            return null
-          }
+      initialDataController.signal.addEventListener('abort', () => {
+        accumulatedChunksPromise.catch(() => {})
+        stream.cancel()
+      })
 
-          // Regardless of whether we are going to abandon this
-          // render we need the unblock runtime b/c it's essential
-          // filling caches.
-          initialStageController.advanceStage(RenderStage.Dynamic)
-          return result
-        }
-      )
+      return {
+        stream,
+        accumulatedChunksPromise,
+      }
+    },
+    () => {
+      if (initialAbandonController.signal.aborted === true) {
+        return
+      } else if (cacheSignal.hasPendingReads()) {
+        initialAbandonController.abort()
+      } else {
+        initialStageController.advanceStage(RenderStage.Runtime)
+      }
+    },
+    () => {
+      if (initialAbandonController.signal.aborted === true) {
+        return
+      } else if (cacheSignal.hasPendingReads()) {
+        initialAbandonController.abort()
+      } else {
+        initialStageController.advanceStage(RenderStage.Dynamic)
+      }
+    }
   )
 
-  if (maybeInitialStreamResult !== null) {
+  if (initialStageController.currentStage !== RenderStage.Abandoned) {
     // No cache misses. We can use the result as-is.
     return {
-      stream: maybeInitialStreamResult.stream,
-      accumulatedChunksPromise:
-        maybeInitialStreamResult.accumulatedChunksPromise,
+      stream: initialStreamResult.stream,
+      accumulatedChunksPromise: initialStreamResult.accumulatedChunksPromise,
       staticInterruptReason: initialStageController.getStaticInterruptReason(),
       runtimeInterruptReason:
         initialStageController.getRuntimeInterruptReason(),
@@ -3481,44 +3457,42 @@ async function renderWithRestartOnCacheMissInDev(
   // where sync IO does not cause aborts, so it's okay if it happens before render.
   const finalRscPayload = await getPayload(requestStore)
 
-  const finalStreamResult = await workUnitAsyncStorage.run(requestStore, () =>
-    pipelineInSequentialTasks(
-      () => {
-        // Static stage
-        finalStageController.advanceStage(RenderStage.Static)
-        startTime = performance.now() + performance.timeOrigin
+  const finalStreamResult = await workUnitAsyncStorage.run(
+    requestStore,
+    runInSequentialTasks,
+    () => {
+      finalStageController.advanceStage(RenderStage.Static)
+      startTime = performance.now() + performance.timeOrigin
 
-        const stream = ComponentMod.renderToReadableStream(
-          finalRscPayload,
-          clientModules,
-          {
-            onError,
-            environmentName,
-            startTime,
-            filterStackFrame,
-            debugChannel: debugChannel?.serverSide,
-          }
-        )
+      const streamPair = ComponentMod.renderToReadableStream(
+        finalRscPayload,
+        clientModules,
+        {
+          onError,
+          environmentName,
+          startTime,
+          filterStackFrame,
+          debugChannel: debugChannel?.serverSide,
+        }
+      ).tee()
 
-        const [continuationStream, accumulatingStream] = stream.tee()
-        const accumulatedChunksPromise = accumulateStreamChunks(
-          accumulatingStream,
+      return {
+        stream: streamPair[0],
+        accumulatedChunksPromise: accumulateStreamChunks(
+          streamPair[1],
           finalStageController,
           null
-        )
-        return { stream: continuationStream, accumulatedChunksPromise }
-      },
-      (result) => {
-        // Runtime stage
-        finalStageController.advanceStage(RenderStage.Runtime)
-        return result
-      },
-      (result) => {
-        // Dynamic stage
-        finalStageController.advanceStage(RenderStage.Dynamic)
-        return result
+        ),
       }
-    )
+    },
+    () => {
+      // Runtime stage
+      finalStageController.advanceStage(RenderStage.Runtime)
+    },
+    () => {
+      // Dynamic stage
+      finalStageController.advanceStage(RenderStage.Dynamic)
+    }
   )
 
   if (process.env.__NEXT_DEV_SERVER && setCacheStatus) {

filename: packages/next/src/server/app-render/instant-validation/instant-validation.tsx
patch: @@ -10,7 +10,7 @@ import { InvariantError } from '../../../shared/lib/invariant-error'
 import { RenderStage } from '../staged-rendering'
 import { getServerModuleMap } from '../manifests-singleton'
 import {
-  pipelineInSequentialTasks,
+  runInSequentialTasks,
   scheduleInSequentialTasks,
 } from '../app-render-render-utils'
 import { workAsyncStorage } from '../work-async-storage.external'
@@ -441,7 +441,7 @@ export async function collectStagedSegmentData(
     [RenderStage.Runtime]: -1,
   }
 
-  await pipelineInSequentialTasks(
+  await runInSequentialTasks(
     () => {
       for (const [segmentPath, segmentData] of segments) {
         const segmentCacheItem: SegmentCacheItem = {