Fix: Optimistic update does not get reset

acdlite · acdlite · commit e6a8f505325e · 2023-10-03T14:27:39.000-04:00
I found a bug where if an optimistic update causes a component to
rerender, and there are no other state updates during that render,
React bails out without applying the update.

Whenever a hook detects a change, we must mark the component as dirty to
prevent a bailout. We check for changes by comparing the new state to
`hook.memoizedState`. However, when implementing optimistic state
rebasing, I incorrectly reset `hook.memoizedState` to the incoming base
state, even though I only needed to reset `hook.baseState`. This was
just a mistake on my part.

This wasn't caught by the existing tests because usually when the
optimistic state changes, there's also some other state that marks the
component as dirty in the same render.

I fixed the bug and added a regression test.
diff --git a/packages/react-reconciler/src/ReactFiberHooks.js b/packages/react-reconciler/src/ReactFiberHooks.js
@@ -1817,9 +1817,9 @@ function updateOptimisticImpl<S, A>(
   // as an argument. It's called a passthrough because if there are no pending
   // updates, it will be returned as-is.
   //
-  // Reset the base state and memoized state to the passthrough. Future
-  // updates will be applied on top of this.
-  hook.baseState = hook.memoizedState = passthrough;
+  // Reset the base state to the passthrough. Future updates will be applied
+  // on top of this.
+  hook.baseState = passthrough;
 
   // If a reducer is not provided, default to the same one used by useState.
   const resolvedReducer: (S, A) => S =
@@ -1853,9 +1853,9 @@ function rerenderOptimistic<S, A>(
 
   // This is a mount. No updates to process.
 
-  // Reset the base state and memoized state to the passthrough. Future
-  // updates will be applied on top of this.
-  hook.baseState = hook.memoizedState = passthrough;
+  // Reset the base state to the passthrough. Future updates will be applied
+  // on top of this.
+  hook.baseState = passthrough;
   const dispatch = hook.queue.dispatch;
   return [passthrough, dispatch];
 }
diff --git a/packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js b/packages/react-reconciler/src/__tests__/ReactAsyncActions-test.js
@@ -1112,4 +1112,67 @@ describe('ReactAsyncActions', () => {
       </>,
     );
   });
+
+  // @gate enableAsyncActions
+  test('useOptimistic can update repeatedly in the same async action', async () => {
+    let startTransition;
+    let setLoadingProgress;
+    let setText;
+    function App() {
+      const [, _startTransition] = useTransition();
+      const [text, _setText] = useState('A');
+      const [loadingProgress, _setLoadingProgress] = useOptimistic(0);
+      startTransition = _startTransition;
+      setText = _setText;
+      setLoadingProgress = _setLoadingProgress;
+
+      return (
+        <>
+          {loadingProgress !== 0 ? (
+            <div key="progress">
+              <Text text={`Loading... (${loadingProgress})`} />
+            </div>
+          ) : null}
+          <div key="real">
+            <Text text={text} />
+          </div>
+        </>
+      );
+    }
+
+    // Initial render
+    const root = ReactNoop.createRoot();
+    await act(() => root.render(<App text="A" />));
+    assertLog(['A']);
+    expect(root).toMatchRenderedOutput(<div>A</div>);
+
+    await act(async () => {
+      startTransition(async () => {
+        setLoadingProgress('25%');
+        await getText('Wait 1');
+        setLoadingProgress('75%');
+        await getText('Wait 2');
+        startTransition(() => setText('B'));
+      });
+    });
+    assertLog(['Loading... (25%)', 'A']);
+    expect(root).toMatchRenderedOutput(
+      <>
+        <div>Loading... (25%)</div>
+        <div>A</div>
+      </>,
+    );
+
+    await act(() => resolveText('Wait 1'));
+    assertLog(['Loading... (75%)', 'A']);
+    expect(root).toMatchRenderedOutput(
+      <>
+        <div>Loading... (75%)</div>
+        <div>A</div>
+      </>,
+    );
+
+    await act(() => resolveText('Wait 2'));
+    assertLog(['B']);
+  });
 });
