From fc5ef50da8e975a569622d477f1fed54cb8b193d Mon Sep 17 00:00:00 2001
From: Josh Story <story@hey.com>
Date: Sat, 14 Sep 2024 09:26:01 -0700
Subject: [PATCH] [Flight] Start initial work immediately (#30961)

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 b40c0628b2..f0e632c549 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;