[Fizz] model fb bundle's external work scheduling explicitly (react#36576)

There are parts of Fizz that need to schedule work regardless of whether
the primary rendering pathway is drive externally through performWork.
Historically scheduleWork and later s cheduleMicrotask were noops in
this bundle but it makes it hard to reason about the code because you
cannot be assured that calling scheduleWork will actually result in the
function
ever executing. Now we model this explicitly through config. For builds
that drive work through external calls to performaWork we simply omit
any work scheduling in startWork or pi
ngTask. Now that this is modeled explicitly we can implement
scheduleMicrotask and scheduleWork to actually provide a guarantee that
the callbacks will get invoked.

For now I've implemented these two for the fb build as synchronous
however it is likely that queueMicrotask and setTimeout or similar are
preferred.

Author: gnoff

packages/react-server/src/ReactFizzServer.js

@@ -88,6 +88,7 @@ import {
   hoistHoistables,
   createHoistableState,
   createPreambleState,
+  isWorkLoopExternallyDriven,
   supportsRequestStorage,
   requestStorage,
   pushFormStateMarkerIsMatching,
@@ -527,7 +528,7 @@ function RequestInstance(
     progressiveChunkSize === undefined
       ? DEFAULT_PROGRESSIVE_CHUNK_SIZE
       : progressiveChunkSize;
-  this.status = OPENING;
+  this.status = isWorkLoopExternallyDriven ? OPEN : OPENING;
   this.fatalError = null;
   this.nextSegmentId = 0;
   this.allPendingTasks = 0;
@@ -790,12 +791,16 @@ export function resolveRequest(): null | Request {
 function pingTask(request: Request, task: Task): void {
   const pingedTasks = request.pingedTasks;
   pingedTasks.push(task);
-  if (request.pingedTasks.length === 1) {
-    request.flushScheduled = request.destination !== null;
-    if (request.trackedPostpones !== null || request.status === OPENING) {
-      scheduleMicrotask(() => performWork(request));
-    } else {
-      scheduleWork(() => performWork(request));
+  if (isWorkLoopExternallyDriven) {
+    return;
+  } else {
+    if (request.pingedTasks.length === 1) {
+      request.flushScheduled = request.destination !== null;
+      if (request.trackedPostpones !== null || request.status === OPENING) {
+        scheduleMicrotask(() => performWork(request));
+      } else {
+        scheduleWork(() => performWork(request));
+      }
     }
   }
 }
@@ -6030,39 +6035,45 @@ function flushCompletedQueues(
 }
 
 export function startWork(request: Request): void {
-  request.flushScheduled = request.destination !== null;
-  // When prerendering we use microtasks for pinging work
-  if (supportsRequestStorage) {
-    scheduleMicrotask(() => requestStorage.run(request, performWork, request));
+  if (isWorkLoopExternallyDriven) {
+    return;
   } else {
-    scheduleMicrotask(() => performWork(request));
-  }
-  scheduleWork(() => {
-    if (request.status === OPENING) {
-      request.status = OPEN;
-    }
-
-    if (request.trackedPostpones === null) {
-      // this is either a regular render or a resume. For regular render we want
-      // to call emitEarlyPreloads after the first performWork because we want
-      // are responding to a live request and need to balance sending something early
-      // (i.e. don't want for the shell to finish) but we need something to send.
-      // The only implementation of this is for DOM at the moment and during resumes nothing
-      // actually emits but the code paths here are the same.
-      // During a prerender we don't want to be too aggressive in emitting early preloads
-      // because we aren't responding to a live request and we can wait for the prerender to
-      // postpone before we emit anything.
-      if (supportsRequestStorage) {
-        requestStorage.run(
-          request,
-          enqueueEarlyPreloadsAfterInitialWork,
-          request,
-        );
-      } else {
-        enqueueEarlyPreloadsAfterInitialWork(request);
-      }
+    request.flushScheduled = request.destination !== null;
+    // When prerendering we use microtasks for pinging work
+    if (supportsRequestStorage) {
+      scheduleMicrotask(() =>
+        requestStorage.run(request, performWork, request),
+      );
+    } else {
+      scheduleMicrotask(() => performWork(request));
     }
-  });
+    scheduleWork(() => {
+      if (request.status === OPENING) {
+        request.status = OPEN;
+      }
+
+      if (request.trackedPostpones === null) {
+        // this is either a regular render or a resume. For regular render we want
+        // to call emitEarlyPreloads after the first performWork because we want
+        // are responding to a live request and need to balance sending something early
+        // (i.e. don't want for the shell to finish) but we need something to send.
+        // The only implementation of this is for DOM at the moment and during resumes nothing
+        // actually emits but the code paths here are the same.
+        // During a prerender we don't want to be too aggressive in emitting early preloads
+        // because we aren't responding to a live request and we can wait for the prerender to
+        // postpone before we emit anything.
+        if (supportsRequestStorage) {
+          requestStorage.run(
+            request,
+            enqueueEarlyPreloadsAfterInitialWork,
+            request,
+          );
+        } else {
+          enqueueEarlyPreloadsAfterInitialWork(request);
+        }
+      }
+    });
+  }
 }
 
 function enqueueEarlyPreloadsAfterInitialWork(request: Request) {

packages/react-server/src/ReactServerStreamConfigFB.js

@@ -18,6 +18,18 @@ export opaque type PrecomputedChunk = string;
 export opaque type Chunk = string;
 export opaque type BinaryChunk = string;
 
+export function scheduleMicrotask(callback: () => void) {
+  // TODO: Consider unifying this with the FB Flight stream config and
+  // adopting its microtask scheduling implementation.
+  callback();
+}
+
+export function scheduleWork(callback: () => void) {
+  // TODO: Consider unifying this with the FB Flight stream config and
+  // adopting its work scheduling implementation.
+  callback();
+}
+
 export function flushBuffered(destination: Destination) {}
 
 export const supportsRequestStorage = false;

packages/react-server/src/forks/ReactFizzConfig.custom.js

@@ -40,6 +40,8 @@ export const isPrimaryRenderer = false;
 
 export const supportsClientAPIs = true;
 
+export const isWorkLoopExternallyDriven =
+  $$$config.isWorkLoopExternallyDriven === true;
 export const supportsRequestStorage = false;
 export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);
 

packages/react-server/src/forks/ReactFizzConfig.dom-edge.js

@@ -12,6 +12,8 @@ export * from 'react-dom-bindings/src/server/ReactFizzConfigDOM';
 
 export * from 'react-client/src/ReactClientConsoleConfigServer';
 
+export const isWorkLoopExternallyDriven = false;
+
 // For now, we get this from the global scope, but this will likely move to a module.
 export const supportsRequestStorage = typeof AsyncLocalStorage === 'function';
 export const requestStorage: AsyncLocalStorage<Request | void> =

packages/react-server/src/forks/ReactFizzConfig.dom-fb.js

@@ -0,0 +1,20 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+import type {Request} from 'react-server/src/ReactFizzServer';
+
+export * from 'react-dom-bindings/src/server/ReactFizzConfigDOM';
+
+export * from 'react-client/src/ReactClientConsoleConfigBrowser';
+
+// This renderer is pulled from the outside through renderNextChunk. Promises
+// can ping tasks, but the outer caller decides when to process them.
+export const isWorkLoopExternallyDriven = true;
+
+export const supportsRequestStorage = false;
+export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);

packages/react-server/src/forks/ReactFizzConfig.dom-legacy.js

@@ -12,5 +12,6 @@ export * from 'react-dom-bindings/src/server/ReactFizzConfigDOMLegacy';
 
 export * from 'react-client/src/ReactClientConsoleConfigPlain';
 
+export const isWorkLoopExternallyDriven = false;
 export const supportsRequestStorage = false;
 export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);

packages/react-server/src/forks/ReactFizzConfig.dom-node.js

@@ -15,6 +15,7 @@ export * from 'react-dom-bindings/src/server/ReactFizzConfigDOM';
 
 export * from 'react-client/src/ReactClientConsoleConfigServer';
 
+export const isWorkLoopExternallyDriven = false;
 export const supportsRequestStorage = true;
 export const requestStorage: AsyncLocalStorage<Request | void> =
   new AsyncLocalStorage();

packages/react-server/src/forks/ReactFizzConfig.dom.js

@@ -12,5 +12,6 @@ export * from 'react-dom-bindings/src/server/ReactFizzConfigDOM';
 
 export * from 'react-client/src/ReactClientConsoleConfigBrowser';
 
+export const isWorkLoopExternallyDriven = false;
 export const supportsRequestStorage = false;
 export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);

packages/react-server/src/forks/ReactFizzConfig.markup.js

@@ -12,5 +12,6 @@ export * from 'react-markup/src/ReactFizzConfigMarkup.js';
 
 export * from 'react-client/src/ReactClientConsoleConfigPlain';
 
+export const isWorkLoopExternallyDriven = false;
 export const supportsRequestStorage = false;
 export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);

packages/react-server/src/forks/ReactFizzConfig.noop.js

@@ -40,6 +40,8 @@ export const isPrimaryRenderer = false;
 
 export const supportsClientAPIs = true;
 
+export const isWorkLoopExternallyDriven =
+  $$$config.isWorkLoopExternallyDriven === true;
 export const supportsRequestStorage = false;
 export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);