From c10010a6a00911fe99452bc561dd47c3e15f4eb8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Markb=C3=A5ge?= <sebastian@calyptus.eu>
Date: Wed, 3 May 2023 18:13:47 -0400
Subject: [PATCH] [Fizz] Gracefully handle suspending in DOM configs (#26768)

E.g. if we suspend (throw a promise) in pushStartInstance today we might
have already pushed some chunks (or even child segments potentially). We
should revert back to where we were.

This doesn't usually happen because when we suspend in a component it
doesn't write anything itself, it'll always defer to som host instance
to do the writing.

There was a todo about this already but I'm not 100% sure it's always
safe when suspending. It should be safe when suspending just regularly
because it's just a noop. We might not even want "throwing a promise" in
this mechanism to be supported longer term but for now that's how a
suspend in internals.
---
 packages/react-server/src/ReactFizzServer.js | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
index 489c440ee78d5..fc56c8d7c8583 100644
--- a/packages/react-server/src/ReactFizzServer.js
+++ b/packages/react-server/src/ReactFizzServer.js
@@ -1603,8 +1603,11 @@ function spawnNewSuspendedTask(
 // This is a non-destructive form of rendering a node. If it suspends it spawns
 // a new task and restores the context of this task to what it was before.
 function renderNode(request: Request, task: Task, node: ReactNodeList): void {
-  // TODO: Store segment.children.length here and reset it in case something
+  // Store how much we've pushed at this point so we can reset it in case something
   // suspended partially through writing something.
+  const segment = task.blockedSegment;
+  const childrenLength = segment.children.length;
+  const chunkLength = segment.chunks.length;
 
   // Snapshot the current context in case something throws to interrupt the
   // process.
@@ -1620,6 +1623,10 @@ function renderNode(request: Request, task: Task, node: ReactNodeList): void {
   } catch (thrownValue) {
     resetHooksState();
 
+    // Reset the write pointers to where we started.
+    segment.children.length = childrenLength;
+    segment.chunks.length = chunkLength;
+
     const x =
       thrownValue === SuspenseException
         ? // This is a special type of exception used for Suspense. For historical
@@ -1895,6 +1902,9 @@ function retryTask(request: Request, task: Task): void {
     prevTaskInDEV = currentTaskInDEV;
     currentTaskInDEV = task;
   }
+
+  const childrenLength = segment.children.length;
+  const chunkLength = segment.chunks.length;
   try {
     // We call the destructive form that mutates this task. That way if something
     // suspends again, we can reuse the same task instead of spawning a new one.
@@ -1919,6 +1929,10 @@ function retryTask(request: Request, task: Task): void {
   } catch (thrownValue) {
     resetHooksState();
 
+    // Reset the write pointers to where we started.
+    segment.children.length = childrenLength;
+    segment.chunks.length = chunkLength;
+
     const x =
       thrownValue === SuspenseException
         ? // This is a special type of exception used for Suspense. For historical