Add suspense support to the data module (#37261)

* Data: add useSuspenseSelect hook

* useSuspenseSelect: throw resolution error, add unit test suite

* Sync with recent useSelect changes

Co-authored-by: Jarda Snajdr <jsnajdr@gmail.com>

packages/data/README.md

@@ -662,6 +662,20 @@ _Parameters_
 
 -   _listener_ `Function`: Callback function.
 
+### suspendSelect
+
+Given the name of a registered store, returns an object containing the store's
+selectors pre-bound to state so that you only need to supply additional arguments,
+and modified so that they throw promises in case the selector is not resolved yet.
+
+_Parameters_
+
+-   _storeNameOrDescriptor_ `string|StoreDescriptor`: Unique namespace identifier for the store or the store descriptor.
+
+_Returns_
+
+-   `Object`: Object containing the store's suspense-wrapped selectors.
+
 ### use
 
 Extends a registry to inherit functionality provided by a given plugin. A
@@ -827,6 +841,20 @@ _Returns_
 
 -   `Function`: A custom react hook.
 
+### useSuspenseSelect
+
+A variant of the `useSelect` hook that has the same API, but will throw a
+suspense Promise if any of the called selectors is in an unresolved state.
+
+_Parameters_
+
+-   _mapSelect_ `Function`: Function called on every state change. The returned value is exposed to the component using this hook. The function receives the `registry.suspendSelect` method as the first argument and the `registry` as the second one.
+-   _deps_ `Array`: A dependency array used to memoize the `mapSelect` so that the same `mapSelect` is invoked on every state change unless the dependencies change.
+
+_Returns_
+
+-   `Object`: Data object returned by the `mapSelect` function.
+
 ### withDispatch
 
 Higher-order component used to add dispatch props using registered action

packages/data/src/components/use-select/index.js

@@ -241,3 +241,130 @@ export default function useSelect( mapSelect, deps ) {
 
 	return hasMappingFunction ? mapOutput : registry.select( mapSelect );
 }
+
+/**
+ * A variant of the `useSelect` hook that has the same API, but will throw a
+ * suspense Promise if any of the called selectors is in an unresolved state.
+ *
+ * @param {Function} mapSelect Function called on every state change. The
+ *                             returned value is exposed to the component
+ *                             using this hook. The function receives the
+ *                             `registry.suspendSelect` method as the first
+ *                             argument and the `registry` as the second one.
+ * @param {Array}    deps      A dependency array used to memoize the `mapSelect`
+ *                             so that the same `mapSelect` is invoked on every
+ *                             state change unless the dependencies change.
+ *
+ * @return {Object} Data object returned by the `mapSelect` function.
+ */
+export function useSuspenseSelect( mapSelect, deps ) {
+	const _mapSelect = useCallback( mapSelect, deps );
+
+	const registry = useRegistry();
+	const isAsync = useAsyncMode();
+
+	const latestRegistry = useRef( registry );
+	const latestMapSelect = useRef();
+	const latestIsAsync = useRef( isAsync );
+	const latestMapOutput = useRef();
+	const latestMapOutputError = useRef();
+
+	// Keep track of the stores being selected in the `mapSelect` function,
+	// and only subscribe to those stores later.
+	const listeningStores = useRef( [] );
+	const wrapSelect = useCallback(
+		( callback ) =>
+			registry.__unstableMarkListeningStores(
+				() => callback( registry.suspendSelect, registry ),
+				listeningStores
+			),
+		[ registry ]
+	);
+
+	// Generate a "flag" for used in the effect dependency array.
+	// It's different than just using `mapSelect` since deps could be undefined,
+	// in that case, we would still want to memoize it.
+	const depsChangedFlag = useMemo( () => ( {} ), deps || [] );
+
+	let mapOutput = latestMapOutput.current;
+	let mapOutputError = latestMapOutputError.current;
+
+	const hasReplacedRegistry = latestRegistry.current !== registry;
+	const hasReplacedMapSelect = latestMapSelect.current !== _mapSelect;
+	const hasLeftAsyncMode = latestIsAsync.current && ! isAsync;
+
+	if ( hasReplacedRegistry || hasReplacedMapSelect || hasLeftAsyncMode ) {
+		try {
+			mapOutput = wrapSelect( _mapSelect );
+		} catch ( error ) {
+			mapOutputError = error;
+		}
+	}
+
+	useIsomorphicLayoutEffect( () => {
+		latestRegistry.current = registry;
+		latestMapSelect.current = _mapSelect;
+		latestIsAsync.current = isAsync;
+		latestMapOutput.current = mapOutput;
+		latestMapOutputError.current = mapOutputError;
+	} );
+
+	// React can sometimes clear the `useMemo` cache.
+	// We use the cache-stable `useMemoOne` to avoid
+	// losing queues.
+	const queueContext = useMemoOne( () => ( { queue: true } ), [ registry ] );
+	const [ , forceRender ] = useReducer( ( s ) => s + 1, 0 );
+	const isMounted = useRef( false );
+
+	useIsomorphicLayoutEffect( () => {
+		const onStoreChange = () => {
+			try {
+				const newMapOutput = wrapSelect( latestMapSelect.current );
+
+				if ( isShallowEqual( latestMapOutput.current, newMapOutput ) ) {
+					return;
+				}
+				latestMapOutput.current = newMapOutput;
+			} catch ( error ) {
+				latestMapOutputError.current = error;
+			}
+
+			forceRender();
+		};
+
+		const onChange = () => {
+			if ( ! isMounted.current ) {
+				return;
+			}
+
+			if ( latestIsAsync.current ) {
+				renderQueue.add( queueContext, onStoreChange );
+			} else {
+				onStoreChange();
+			}
+		};
+
+		// catch any possible state changes during mount before the subscription
+		// could be set.
+		onStoreChange();
+
+		const unsubscribers = listeningStores.current.map( ( storeName ) =>
+			registry.__unstableSubscribeStore( storeName, onChange )
+		);
+
+		isMounted.current = true;
+
+		return () => {
+			// The return value of the subscribe function could be undefined if the store is a custom generic store.
+			unsubscribers.forEach( ( unsubscribe ) => unsubscribe?.() );
+			renderQueue.cancel( queueContext );
+			isMounted.current = false;
+		};
+	}, [ registry, wrapSelect, depsChangedFlag ] );
+
+	if ( mapOutputError ) {
+		throw mapOutputError;
+	}
+
+	return mapOutput;
+}

packages/data/src/components/use-select/test/suspense.js

@@ -0,0 +1,160 @@
+/**
+ * External dependencies
+ */
+import { render, waitFor } from '@testing-library/react';
+
+/**
+ * WordPress dependencies
+ */
+import {
+	createRegistry,
+	createReduxStore,
+	useSuspenseSelect,
+	RegistryProvider,
+} from '@wordpress/data';
+import { Component, Suspense } from '@wordpress/element';
+
+jest.useRealTimers();
+
+function createRegistryWithStore() {
+	const initialState = {
+		prefix: 'pre-',
+		token: null,
+		data: null,
+		fails: true,
+	};
+
+	const reducer = ( state = initialState, action ) => {
+		switch ( action.type ) {
+			case 'RECEIVE_TOKEN':
+				return { ...state, token: action.token };
+			case 'RECEIVE_DATA':
+				return { ...state, data: action.data };
+			default:
+				return state;
+		}
+	};
+
+	const selectors = {
+		getPrefix: ( state ) => state.prefix,
+		getToken: ( state ) => state.token,
+		getData: ( state, token ) => {
+			if ( ! token ) {
+				throw 'missing token in selector';
+			}
+			return state.data;
+		},
+		getThatFails: ( state ) => state.fails,
+	};
+
+	const sleep = ( ms ) => new Promise( ( r ) => setTimeout( () => r(), ms ) );
+
+	const resolvers = {
+		getToken: () => async ( { dispatch } ) => {
+			await sleep( 10 );
+			dispatch( { type: 'RECEIVE_TOKEN', token: 'token' } );
+		},
+		getData: ( token ) => async ( { dispatch } ) => {
+			await sleep( 10 );
+			if ( ! token ) {
+				throw 'missing token in resolver';
+			}
+			dispatch( { type: 'RECEIVE_DATA', data: 'therealdata' } );
+		},
+		getThatFails: () => async () => {
+			await sleep( 10 );
+			throw 'resolution failed';
+		},
+	};
+
+	const store = createReduxStore( 'test', {
+		reducer,
+		selectors,
+		resolvers,
+	} );
+
+	const registry = createRegistry();
+	registry.register( store );
+
+	return { registry, store };
+}
+
+describe( 'useSuspenseSelect', () => {
+	it( 'renders after suspending a few times', async () => {
+		const { registry, store } = createRegistryWithStore();
+		let attempts = 0;
+		let renders = 0;
+
+		const UI = () => {
+			attempts++;
+			const { result } = useSuspenseSelect( ( select ) => {
+				const prefix = select( store ).getPrefix();
+				const token = select( store ).getToken();
+				const data = select( store ).getData( token );
+				return { result: prefix + data };
+			}, [] );
+			renders++;
+			return <div aria-label="loaded">{ result }</div>;
+		};
+
+		const App = () => (
+			<RegistryProvider value={ registry }>
+				<Suspense fallback="loading">
+					<UI />
+				</Suspense>
+			</RegistryProvider>
+		);
+
+		const rendered = render( <App /> );
+		await waitFor( () => rendered.getByLabelText( 'loaded' ) );
+
+		// Verify there were 3 attempts to render. Suspended twice because of
+		// `getToken` and `getData` selectors not being resolved, and then finally
+		// rendered after all data got loaded.
+		expect( attempts ).toBe( 3 );
+		expect( renders ).toBe( 1 );
+	} );
+
+	it( 'shows error when resolution fails', async () => {
+		const { registry, store } = createRegistryWithStore();
+
+		const UI = () => {
+			const { token } = useSuspenseSelect( ( select ) => {
+				// Call a selector whose resolution fails. The `useSuspenseSelect`
+				// is then supposed to throw the resolution error.
+				return { token: select( store ).getThatFails() };
+			}, [] );
+			return <div aria-label="loaded">{ token }</div>;
+		};
+
+		class Error extends Component {
+			state = { error: null };
+
+			static getDerivedStateFromError( error ) {
+				return { error };
+			}
+
+			render() {
+				if ( this.state.error ) {
+					return <div aria-label="error">{ this.state.error }</div>;
+				}
+				return this.props.children;
+			}
+		}
+
+		const App = () => (
+			<RegistryProvider value={ registry }>
+				<Error>
+					<Suspense fallback="loading">
+						<UI />
+					</Suspense>
+				</Error>
+			</RegistryProvider>
+		);
+
+		const rendered = render( <App /> );
+		const label = await waitFor( () => rendered.getByLabelText( 'error' ) );
+		expect( label.textContent ).toBe( 'resolution failed' );
+		expect( console ).toHaveErrored();
+	} );
+} );

packages/data/src/index.js

@@ -19,7 +19,10 @@ export {
 	RegistryConsumer,
 	useRegistry,
 } from './components/registry-provider';
-export { default as useSelect } from './components/use-select';
+export {
+	default as useSelect,
+	useSuspenseSelect,
+} from './components/use-select';
 export { useDispatch } from './components/use-dispatch';
 export { AsyncModeProvider } from './components/async-mode-provider';
 export { createRegistry } from './registry';
@@ -115,6 +118,18 @@ export const select = defaultRegistry.select;
  */
 export const resolveSelect = defaultRegistry.resolveSelect;
 
+/**
+ * Given the name of a registered store, returns an object containing the store's
+ * selectors pre-bound to state so that you only need to supply additional arguments,
+ * and modified so that they throw promises in case the selector is not resolved yet.
+ *
+ * @param {string|StoreDescriptor} storeNameOrDescriptor Unique namespace identifier for the store
+ *                                                       or the store descriptor.
+ *
+ * @return {Object} Object containing the store's suspense-wrapped selectors.
+ */
+export const suspendSelect = defaultRegistry.suspendSelect;
+
 /**
  * Given the name of a registered store, returns an object of the store's action creators.
  * Calling an action creator will cause it to be dispatched, updating the state value accordingly.