Fix hook state loss when Suspense re-suspends

Author: JoviDeCroock

compat/src/suspense.js

@@ -38,14 +38,19 @@ function initSuspenseHooks() {
 	};
 }
 
-function detachedClone(vnode, detachedParent, parentDom) {
+function detachedClone(vnode, detachedParent, parentDom, preserveHooks) {
 	if (vnode) {
 		if (vnode._component && vnode._component.__hooks) {
-			vnode._component.__hooks._list.forEach(effect => {
-				if (typeof effect._cleanup == 'function') effect._cleanup();
+			const hooks = vnode._component.__hooks;
+			hooks._list.forEach(effect => {
+				if (!preserveHooks || effect._passive != null) {
+					const cleanup = effect._cleanup;
+					effect._cleanup = effect._args = UNDEFINED;
+					if (typeof cleanup == 'function') cleanup();
+				}
 			});
-
-			vnode._component.__hooks = null;
+			if (preserveHooks) hooks._pendingEffects = [];
+			else vnode._component.__hooks = null;
 		}
 
 		vnode = assign({ constructor: UNDEFINED }, vnode);
@@ -62,7 +67,7 @@ function detachedClone(vnode, detachedParent, parentDom) {
 		vnode._children =
 			vnode._children &&
 			vnode._children.map(child =>
-				detachedClone(child, detachedParent, parentDom)
+				detachedClone(child, detachedParent, parentDom, preserveHooks)
 			);
 	}
 
@@ -186,6 +191,10 @@ function createSuspense() {
 	Suspense.prototype.componentWillUnmount = function () {
 		this._suspenders = [];
 	};
+	Suspense.prototype.componentDidMount =
+		Suspense.prototype.componentDidUpdate = function () {
+			if (!this._pendingSuspensionCount) this._unmounted = false;
+		};
 
 	/**
 	 * @this {import('./internal').SuspenseComponent}
@@ -203,7 +212,8 @@ function createSuspense() {
 				this._vnode._children[0] = detachedClone(
 					this._detachOnNextRender,
 					detachedParent,
-					(detachedComponent._originalParentDom = detachedComponent._parentDom)
+					(detachedComponent._originalParentDom = detachedComponent._parentDom),
+					this._unmounted == false
 				);
 			}
 

compat/test/browser/suspense-hydration.test.jsx

@@ -2,8 +2,10 @@ import { setupRerender } from 'preact/test-utils';
 import React, {
 	createElement,
 	hydrate,
+	render,
 	Fragment,
 	Suspense,
+	use,
 	memo,
 	useState,
 	useSyncExternalStore
@@ -874,6 +876,75 @@ describe('suspense hydration', () => {
 			});
 	});
 
+	it('should preserve component state when re-suspending after streaming-style hydration', async () => {
+		scratch.innerHTML =
+			'<!--$s:2--><div><p>Hello</p><button>Count: 0</button></div><!--/$s:2-->';
+
+		let promise = Promise.resolve('Hello');
+		let increment;
+		function App() {
+			const message = use(promise);
+			const [count, setCount] = useState(0);
+			increment = () => setCount(value => value + 1);
+			return (
+				<div>
+					<p>{message}</p>
+					<button>Count: {count}</button>
+				</div>
+			);
+		}
+
+		hydrate(
+			<Suspense fallback={<div>Fallback</div>}>
+				<App />
+			</Suspense>,
+			scratch
+		);
+		await promise;
+		rerender();
+		rerender();
+
+		increment();
+		rerender();
+		expect(scratch.querySelector('button').textContent).to.equal('Count: 1');
+
+		promise = Promise.resolve('Hello');
+		increment();
+		rerender();
+		await promise;
+		rerender();
+		rerender();
+
+		expect(scratch.querySelector('button').textContent).to.equal('Count: 2');
+	});
+
+	it('should preserve component state when re-suspending after client render', async () => {
+		let promise = Promise.resolve('Hello');
+		let increment;
+		function App() {
+			use(promise);
+			const [count, setCount] = useState(0);
+			increment = () => setCount(value => value + 1);
+			return <button>Count: {count}</button>;
+		}
+
+		render(<Suspense fallback="Fallback"><App /></Suspense>, scratch);
+		await promise;
+		rerender();
+		rerender();
+		increment();
+		rerender();
+		expect(scratch.textContent).to.equal('Count: 1');
+
+		promise = Promise.resolve('Hello');
+		increment();
+		rerender();
+		await promise;
+		rerender();
+		rerender();
+		expect(scratch.textContent).to.equal('Count: 2');
+	});
+
 	it('should correctly hydrate and rerender a memoized lazy data loader', () => {
 		const originalHtml = '<p>Count: 5</p>';
 		scratch.innerHTML = originalHtml;

compat/test/browser/suspense.test.jsx

@@ -8,6 +8,7 @@ import React, {
 	Fragment,
 	createContext,
 	useState,
+	useRef,
 	useEffect,
 	useLayoutEffect,
 	memo
@@ -202,9 +203,43 @@ describe('suspense', () => {
 		resolve().then(assert).catch(assert);
 	});
 
-	it('should reset hooks of components', () => {
+	it('should reset hooks when a component subtree suspends during mount', async () => {
+		let resolve;
+		let resolved = false;
+		let initializations = 0;
+		const promise = new Promise(r => {
+			resolve = () => {
+				resolved = true;
+				r();
+				return promise;
+			};
+		});
+
+		function App() {
+			const [value] = useState(() => ++initializations);
+			if (!resolved) throw promise;
+			return <p>{value}</p>;
+		}
+
+		render(
+			<Suspense fallback="loading">
+				<App />
+			</Suspense>,
+			scratch
+		);
+		rerender();
+		expect(scratch.textContent).to.equal('loading');
+
+		await resolve();
+		rerender();
+		expect(scratch.innerHTML).to.equal('<p>2</p>');
+		expect(initializations).to.equal(2);
+	});
+
+	it('should preserve hooks of mounted components', () => {
 		/** @type {(v) => void} */
 		let set;
+		let initialRef;
 		const LazyComp = ({ name }) => <div>Hello from {name}</div>;
 
 		/** @type {() => Promise<void>} */
@@ -222,7 +257,10 @@ describe('suspense', () => {
 
 		const Parent = ({ children }) => {
 			const [state, setState] = useState(false);
+			const ref = useRef({});
 			set = setState;
+			if (!initialRef) initialRef = ref;
+			else expect(ref).to.equal(initialRef);
 
 			return (
 				<div>
@@ -249,15 +287,21 @@ describe('suspense', () => {
 
 		return resolve().then(() => {
 			rerender();
-			expect(scratch.innerHTML).to.eql(`<div><p>hi</p></div>`);
+			expect(scratch.innerHTML).to.eql(
+				`<div><p>hi</p><div>Hello from LazyComp</div></div>`
+			);
 		});
 	});
 
-	it('should call effect cleanups', () => {
+	it('should call effect cleanups and setups when hiding and revealing', async () => {
 		/** @type {(v) => void} */
 		let set;
+		const effectSetupSpy = vi.fn();
 		const effectSpy = vi.fn();
+		const effectWithoutCleanupSpy = vi.fn();
+		const layoutEffectSetupSpy = vi.fn();
 		const layoutEffectSpy = vi.fn();
+		const layoutEffectWithoutCleanupSpy = vi.fn();
 		const LazyComp = ({ name }) => <div>Hello from {name}</div>;
 
 		/** @type {() => Promise<void>} */
@@ -277,16 +321,20 @@ describe('suspense', () => {
 			const [state, setState] = useState(false);
 			set = setState;
 			useEffect(() => {
+				effectSetupSpy();
 				return () => {
 					effectSpy();
 				};
-			}, []);
+			}, [state]);
+			useEffect(effectWithoutCleanupSpy, [state]);
 
 			useLayoutEffect(() => {
+				layoutEffectSetupSpy();
 				return () => {
 					layoutEffectSpy();
 				};
 			}, []);
+			useLayoutEffect(layoutEffectWithoutCleanupSpy, []);
 
 			return state ? (
 				<div>{children}</div>
@@ -305,19 +353,29 @@ describe('suspense', () => {
 			</Suspense>,
 			scratch
 		);
+		expect(layoutEffectSetupSpy).toHaveBeenCalledOnce();
+		expect(layoutEffectWithoutCleanupSpy).toHaveBeenCalledOnce();
 
 		set(true);
 		rerender();
 		expect(scratch.innerHTML).to.eql('<div>Suspended...</div>');
+
+		expect(effectSetupSpy).toHaveBeenCalledOnce();
+		expect(effectWithoutCleanupSpy).toHaveBeenCalledOnce();
 		expect(effectSpy).toHaveBeenCalledOnce();
 		expect(layoutEffectSpy).toHaveBeenCalledOnce();
 
-		return resolve().then(() => {
-			rerender();
-			expect(effectSpy).toHaveBeenCalledOnce();
-			expect(layoutEffectSpy).toHaveBeenCalledOnce();
-			expect(scratch.innerHTML).to.eql(`<div><p>hi</p></div>`);
-		});
+		await resolve();
+		await act(() => rerender());
+		expect(effectSpy).toHaveBeenCalledOnce();
+		expect(layoutEffectSpy).toHaveBeenCalledOnce();
+		expect(effectSetupSpy).toHaveBeenCalledTimes(2);
+		expect(effectWithoutCleanupSpy).toHaveBeenCalledTimes(2);
+		expect(layoutEffectSetupSpy).toHaveBeenCalledTimes(2);
+		expect(layoutEffectWithoutCleanupSpy).toHaveBeenCalledTimes(2);
+		expect(scratch.innerHTML).to.eql(
+			`<div><div>Hello from LazyComp</div></div>`
+		);
 	});
 
 	it('should support a call to setState before rendering the fallback', () => {

hooks/src/index.js

@@ -299,6 +299,7 @@ export function useLayoutEffect(callback, args) {
 	/** @type {import('./internal').EffectHookState} */
 	const state = getHookState(currentIndex++, 4);
 	if (!options._skipEffects && argsChanged(state._args, args)) {
+		state._passive = false;
 		state._value = callback;
 		state._pendingArgs = args;
 

mangle.json

@@ -32,6 +32,7 @@
       "$_hydrationMismatch": "__m",
       "$_list": "__",
       "$_pendingEffects": "__h",
+      "$_passive": "__P",
       "$_value": "__",
       "$_nextValue": "__N",
       "$_original": "__v",