From ec9a2f84c2db3a751791a3379aba1312323fdfa8 Mon Sep 17 00:00:00 2001
From: Josh Story <story@hey.com>
Date: Fri, 13 Sep 2024 15:57:45 -0700
Subject: [PATCH] [Flight] Start initial work immediately

In a past update we made render and prerender have different work scheduling behavior because these methods are meant to be used in differeent environments with different performance tradeoffs in mind. For instance to prioritize streaming we want to allow as much IO to complete before triggering a round of work because we want to flush as few intermediate UI states. With Prerendering there will never be any intermediate UI states so we can more aggressively render tasks as they complete.

One thing we've found is that even during render we should ideally kick off work immediately. This update normalizes the intitial work for render and prerender to start in a microtask. Choosing microtask over sync is somewhat arbitrary but there really isn't a reason to make them different between render/prerender so for now we'll unify them and keep it as a microtask for now.

This change also updates pinging behavior. If the request is still in the initial task that spawned it then pings will schedule on the microtask queue. This allows immediately available async APIs to resolve right away. The concern with doing this for normal pings is that it might crowd out IO events but since this is the initial task there would be IO to already be scheduled.
---
 .../react-server/src/ReactFlightServer.js     | 49 +++++++++----------
 1 file changed, 24 insertions(+), 25 deletions(-)

diff --git a/packages/react-server/src/ReactFlightServer.js b/packages/react-server/src/ReactFlightServer.js
index b40c0628b2c28..f0e632c5499b0 100644
--- a/packages/react-server/src/ReactFlightServer.js
+++ b/packages/react-server/src/ReactFlightServer.js
@@ -352,8 +352,17 @@ type Task = {
 
 interface Reference {}
 
+const OPENING = 10;
+const OPEN = 11;
+const ABORTING = 12;
+const CLOSING = 13;
+const CLOSED = 14;
+
+const RENDER = 20;
+const PRERENDER = 21;
+
 export type Request = {
-  status: 10 | 11 | 12 | 13,
+  status: 10 | 11 | 12 | 13 | 14,
   type: 20 | 21,
   flushScheduled: boolean,
   fatalError: mixed,
@@ -426,14 +435,6 @@ function defaultPostponeHandler(reason: string) {
   // Noop
 }
 
-const OPEN = 10;
-const ABORTING = 11;
-const CLOSING = 12;
-const CLOSED = 13;
-
-const RENDER = 20;
-const PRERENDER = 21;
-
 function RequestInstance(
   this: $FlowFixMe,
   type: 20 | 21,
@@ -472,7 +473,7 @@ function RequestInstance(
   }
   const hints = createHints();
   this.type = type;
-  this.status = OPEN;
+  this.status = OPENING;
   this.flushScheduled = false;
   this.fatalError = null;
   this.destination = null;
@@ -1794,7 +1795,7 @@ function pingTask(request: Request, task: Task): void {
   pingedTasks.push(task);
   if (pingedTasks.length === 1) {
     request.flushScheduled = request.destination !== null;
-    if (request.type === PRERENDER) {
+    if (request.type === PRERENDER || request.status === OPENING) {
       scheduleMicrotask(() => performWork(request));
     } else {
       scheduleWork(() => performWork(request));
@@ -4062,21 +4063,18 @@ function flushCompletedChunks(
 
 export function startWork(request: Request): void {
   request.flushScheduled = request.destination !== null;
-  if (request.type === PRERENDER) {
-    if (supportsRequestStorage) {
-      scheduleMicrotask(() => {
-        requestStorage.run(request, performWork, request);
-      });
-    } else {
-      scheduleMicrotask(() => performWork(request));
-    }
+  if (supportsRequestStorage) {
+    scheduleMicrotask(() => {
+      requestStorage.run(request, performWork, request);
+    });
   } else {
-    if (supportsRequestStorage) {
-      scheduleWork(() => requestStorage.run(request, performWork, request));
-    } else {
-      scheduleWork(() => performWork(request));
-    }
+    scheduleMicrotask(() => performWork(request));
   }
+  scheduleWork(() => {
+    if (request.status === OPENING) {
+      request.status = OPEN;
+    }
+  });
 }
 
 function enqueueFlush(request: Request): void {
@@ -4129,7 +4127,8 @@ export function stopFlowing(request: Request): void {
 
 export function abort(request: Request, reason: mixed): void {
   try {
-    if (request.status === OPEN) {
+    // We define any status below OPEN as OPEN equivalent
+    if (request.status <= OPEN) {
       request.status = ABORTING;
     }
     const abortableTasks = request.abortableTasks;