From 76c9ab35b1a266e65c84abb1770fd180ab0336c1 Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sebastian@calyptus.eu>
Date: Tue, 19 Mar 2024 14:08:20 -0400
Subject: [PATCH] Fix ref counting bug and retry of thrown promises

---
 .../src/ReactFlightReplyClient.js             | 41 +++++++++----------
 .../src/__tests__/ReactFlightDOMReply-test.js | 38 ++++++++++++++++-
 2 files changed, 56 insertions(+), 23 deletions(-)

diff --git a/packages/react-client/src/ReactFlightReplyClient.js b/packages/react-client/src/ReactFlightReplyClient.js
index 61253ab19ef5f..5f4bd00deb305 100644
--- a/packages/react-client/src/ReactFlightReplyClient.js
+++ b/packages/react-client/src/ReactFlightReplyClient.js
@@ -238,18 +238,17 @@ export function processReply(
             // Upgrade to use FormData to allow us to stream this value.
             formData = new FormData();
           }
+          pendingParts++;
           try {
             const resolvedModel = init(payload);
             // We always outline this as a separate part even though we could inline it
             // because it ensures a more deterministic encoding.
-            pendingParts++;
             const lazyId = nextPartId++;
             const partJSON = JSON.stringify(resolvedModel, resolveToJSON);
             // $FlowFixMe[incompatible-type] We know it's not null because we assigned it above.
             const data: FormData = formData;
             // eslint-disable-next-line react-internal/safe-string-coercion
             data.append(formFieldPrefix + lazyId, partJSON);
-            pendingParts--;
             return serializeByValueID(lazyId);
           } catch (x) {
             if (
@@ -261,28 +260,24 @@ export function processReply(
               pendingParts++;
               const lazyId = nextPartId++;
               const thenable: Thenable<any> = (x: any);
-              thenable.then(
-                partValue => {
-                  try {
-                    const partJSON = JSON.stringify(partValue, resolveToJSON);
-                    // $FlowFixMe[incompatible-type] We know it's not null because we assigned it above.
-                    const data: FormData = formData;
-                    // eslint-disable-next-line react-internal/safe-string-coercion
-                    data.append(formFieldPrefix + lazyId, partJSON);
-                    pendingParts--;
-                    if (pendingParts === 0) {
-                      resolve(data);
-                    }
-                  } catch (reason) {
-                    reject(reason);
+              const retry = function () {
+                // While the first promise resolved, its value isn't necessarily what we'll
+                // resolve into because we might suspend again.
+                try {
+                  const partJSON = JSON.stringify(value, resolveToJSON);
+                  // $FlowFixMe[incompatible-type] We know it's not null because we assigned it above.
+                  const data: FormData = formData;
+                  // eslint-disable-next-line react-internal/safe-string-coercion
+                  data.append(formFieldPrefix + lazyId, partJSON);
+                  pendingParts--;
+                  if (pendingParts === 0) {
+                    resolve(data);
                   }
-                },
-                reason => {
-                  // In the future we could consider serializing this as an error
-                  // that throws on the server instead.
+                } catch (reason) {
                   reject(reason);
-                },
-              );
+                }
+              };
+              thenable.then(retry, retry);
               return serializeByValueID(lazyId);
             } else {
               // In the future we could consider serializing this as an error
@@ -290,6 +285,8 @@ export function processReply(
               reject(x);
               return null;
             }
+          } finally {
+            pendingParts--;
           }
         }
       }
diff --git a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReply-test.js b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReply-test.js
index 99b1e1f1e8e01..938937dba2afb 100644
--- a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReply-test.js
+++ b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReply-test.js
@@ -257,11 +257,47 @@ describe('ReactFlightDOMReply', () => {
     let resolve;
     const lazy = React.lazy(() => new Promise(r => (resolve = r)));
     const bodyPromise = ReactServerDOMClient.encodeReply({lazy: lazy});
-    resolve('Hi');
+    resolve({default: 'Hi'});
     const result = await ReactServerDOMServer.decodeReply(await bodyPromise);
     expect(result.lazy).toBe('Hi');
   });
 
+  it('resolves a proxy throwing a promise inside React.lazy', async () => {
+    let resolve1;
+    let resolve2;
+    const lazy = React.lazy(() => new Promise(r => (resolve1 = r)));
+    const promise = new Promise(r => (resolve2 = r));
+    const bodyPromise1 = ReactServerDOMClient.encodeReply({lazy: lazy});
+    const target = {value: ''};
+    let loaded = false;
+    const proxy = new Proxy(target, {
+      get(targetObj, prop, receiver) {
+        if (prop === 'value') {
+          if (!loaded) {
+            throw promise;
+          }
+          return 'Hello';
+        }
+        return targetObj[prop];
+      },
+    });
+    await resolve1({default: proxy});
+
+    // Encode it again so that we have an already initialized lazy
+    // This is now already resolved but the proxy inside isn't. This ensures
+    // we trigger the retry code path.
+    const bodyPromise2 = ReactServerDOMClient.encodeReply({lazy: lazy});
+
+    // Then resolve the inner thrown promise.
+    loaded = true;
+    await resolve2('Hello');
+
+    const result1 = await ReactServerDOMServer.decodeReply(await bodyPromise1);
+    expect(await result1.lazy.value).toBe('Hello');
+    const result2 = await ReactServerDOMServer.decodeReply(await bodyPromise2);
+    expect(await result2.lazy.value).toBe('Hello');
+  });
+
   it('errors when called with JSX by default', async () => {
     let error;
     try {