[Flight] use microtask for scheduling during prerenders

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;