[Flight] use microtask for scheduling during prerenders (#30768)

In #29491 I updated the work
scheduler for Flight to use microtasks to perform work when something
pings. This is useful but it does have some downsides in terms of our
ability to do task prioritization. Additionally the initial work is not
instantiated using a microtask which is inconsistent with how pings
work.

In this change I update the scheduling logic to use microtasks
consistently for prerenders and use regular tasks for renders both for
the initial work and pings.

packages/internal-test-utils/ReactInternalTestUtils.js

@@ -16,7 +16,7 @@ import {
   clearErrors,
   createLogAssertion,
 } from './consoleMock';
-export {act} from './internalAct';
+export {act, serverAct} from './internalAct';
 const {assertConsoleLogsCleared} = require('internal-test-utils/consoleMock');
 
 import {thrownErrors, actingUpdatesScopeDepth} from './internalAct';

packages/internal-test-utils/internalAct.js

@@ -192,3 +192,90 @@ export async function act<T>(scope: () => Thenable<T>): Thenable<T> {
     }
   }
 }
+
+export async function serverAct<T>(scope: () => Thenable<T>): Thenable<T> {
+  // We require every `act` call to assert console logs
+  // with one of the assertion helpers. Fails if not empty.
+  assertConsoleLogsCleared();
+
+  // $FlowFixMe[cannot-resolve-name]: Flow doesn't know about global Jest object
+  if (!jest.isMockFunction(setTimeout)) {
+    throw Error(
+      "This version of `act` requires Jest's timer mocks " +
+        '(i.e. jest.useFakeTimers).',
+    );
+  }
+
+  // Create the error object before doing any async work, to get a better
+  // stack trace.
+  const error = new Error();
+  Error.captureStackTrace(error, act);
+
+  // Call the provided scope function after an async gap. This is an extra
+  // precaution to ensure that our tests do not accidentally rely on the act
+  // scope adding work to the queue synchronously. We don't do this in the
+  // public version of `act`, though we maybe should in the future.
+  await waitForMicrotasks();
+
+  const errorHandlerNode = function (err: mixed) {
+    thrownErrors.push(err);
+  };
+  // We track errors that were logged globally as if they occurred in this scope and then rethrow them.
+  if (typeof process === 'object') {
+    // Node environment
+    process.on('uncaughtException', errorHandlerNode);
+  } else if (
+    typeof window === 'object' &&
+    typeof window.addEventListener === 'function'
+  ) {
+    throw new Error('serverAct is not supported in JSDOM environments');
+  }
+
+  try {
+    const result = await scope();
+
+    do {
+      // Wait until end of current task/microtask.
+      await waitForMicrotasks();
+
+      // $FlowFixMe[cannot-resolve-name]: Flow doesn't know about global Jest object
+      if (jest.isEnvironmentTornDown()) {
+        error.message =
+          'The Jest environment was torn down before `act` completed. This ' +
+          'probably means you forgot to `await` an `act` call.';
+        throw error;
+      }
+
+      // $FlowFixMe[cannot-resolve-name]: Flow doesn't know about global Jest object
+      const j = jest;
+      if (j.getTimerCount() > 0) {
+        // There's a pending timer. Flush it now. We only do this in order to
+        // force Suspense fallbacks to display; the fact that it's a timer
+        // is an implementation detail. If there are other timers scheduled,
+        // those will also fire now, too, which is not ideal. (The public
+        // version of `act` doesn't do this.) For this reason, we should try
+        // to avoid using timers in our internal tests.
+        j.runOnlyPendingTimers();
+        // If a committing a fallback triggers another update, it might not
+        // get scheduled until a microtask. So wait one more time.
+        await waitForMicrotasks();
+      } else {
+        break;
+      }
+    } while (true);
+
+    if (thrownErrors.length > 0) {
+      // Rethrow any errors logged by the global error handling.
+      const thrownError = aggregateErrors(thrownErrors);
+      thrownErrors.length = 0;
+      throw thrownError;
+    }
+
+    return result;
+  } finally {
+    if (typeof process === 'object') {
+      // Node environment
+      process.off('uncaughtException', errorHandlerNode);
+    }
+  }
+}

packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMEdge-test.js

@@ -23,10 +23,6 @@ if (typeof File === 'undefined' || typeof FormData === 'undefined') {
 // Patch for Edge environments for global scope
 global.AsyncLocalStorage = require('async_hooks').AsyncLocalStorage;
 
-const {
-  patchMessageChannel,
-} = require('../../../../scripts/jest/patchMessageChannel');
-
 let serverExports;
 let clientExports;
 let webpackMap;
@@ -39,7 +35,6 @@ let ReactServerDOMServer;
 let ReactServerDOMStaticServer;
 let ReactServerDOMClient;
 let use;
-let ReactServerScheduler;
 let reactServerAct;
 
 function normalizeCodeLocInfo(str) {
@@ -55,9 +50,7 @@ describe('ReactFlightDOMEdge', () => {
   beforeEach(() => {
     jest.resetModules();
 
-    ReactServerScheduler = require('scheduler');
-    patchMessageChannel(ReactServerScheduler);
-    reactServerAct = require('internal-test-utils').act;
+    reactServerAct = require('internal-test-utils').serverAct;
 
     // Simulate the condition resolution
     jest.mock('react', () => require('react/react.react-server'));

packages/react-server/src/ReactFlightServer.js

@@ -1794,7 +1794,11 @@ function pingTask(request: Request, task: Task): void {
   pingedTasks.push(task);
   if (pingedTasks.length === 1) {
     request.flushScheduled = request.destination !== null;
-    scheduleMicrotask(() => performWork(request));
+    if (request.type === PRERENDER) {
+      scheduleMicrotask(() => performWork(request));
+    } else {
+      scheduleWork(() => performWork(request));
+    }
   }
 }
 
@@ -4056,10 +4060,20 @@ function flushCompletedChunks(
 
 export function startWork(request: Request): void {
   request.flushScheduled = request.destination !== null;
-  if (supportsRequestStorage) {
-    scheduleWork(() => requestStorage.run(request, performWork, request));
+  if (request.type === PRERENDER) {
+    if (supportsRequestStorage) {
+      scheduleMicrotask(() => {
+        requestStorage.run(request, performWork, request);
+      });
+    } else {
+      scheduleMicrotask(() => performWork(request));
+    }
   } else {
-    scheduleWork(() => performWork(request));
+    if (supportsRequestStorage) {
+      scheduleWork(() => requestStorage.run(request, performWork, request));
+    } else {
+      scheduleWork(() => performWork(request));
+    }
   }
 }
 
@@ -4073,6 +4087,8 @@ function enqueueFlush(request: Request): void {
     request.destination !== null
   ) {
     request.flushScheduled = true;
+    // Unlike startWork and pingTask we intetionally use scheduleWork
+    // here even during prerenders to allow as much batching as possible
     scheduleWork(() => {
       request.flushScheduled = false;
       const destination = request.destination;