From 804dced80ed104c8c9e9d44db9cfb8c760b38412 Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sebastian@calyptus.eu>
Date: Mon, 1 May 2023 18:53:45 -0400
Subject: [PATCH] Gracefully handle suspending in DOM configs

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.

There was a todo about this already but I'm not 100% sure it's always safe.
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