Data: add useSuspenseSelect hook

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

@@ -246,3 +246,131 @@ 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();
+	// 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 latestMapSelect = useRef();
+	const latestIsAsync = useRef( isAsync );
+	const latestMapOutput = useRef();
+	const latestMapOutputError = useRef();
+	const isMountedAndNotUnsubscribing = useRef();
+
+	// Keep track of the stores being selected in the `mapSelect` function,
+	// and only subscribe to those stores later.
+	const listeningStores = useRef( [] );
+	const trapSelect = useCallback(
+		( callback ) =>
+			registry.__experimentalMarkListeningStores(
+				() => callback( registry.suspendSelect, registry ),
+				listeningStores
+			),
+		[ registry ]
+	);
+
+	let mapOutput = latestMapOutput.current;
+	let mapOutputError = latestMapOutputError.current;
+
+	if ( latestMapSelect.current !== _mapSelect ) {
+		try {
+			mapOutput = trapSelect( _mapSelect );
+		} catch ( error ) {
+			mapOutputError = error;
+		}
+	}
+
+	useIsomorphicLayoutEffect( () => {
+		latestMapSelect.current = _mapSelect;
+		latestMapOutput.current = mapOutput;
+		latestMapOutputError.current = mapOutputError;
+		isMountedAndNotUnsubscribing.current = true;
+
+		// This has to run after the other ref updates
+		// to avoid using stale values in the flushed
+		// callbacks or potentially overwriting a
+		// changed `latestMapOutput.current`.
+		if ( latestIsAsync.current !== isAsync ) {
+			latestIsAsync.current = isAsync;
+			renderQueue.flush( queueContext );
+		}
+	} );
+
+	// 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 || [] );
+
+	useIsomorphicLayoutEffect( () => {
+		const onStoreChange = () => {
+			if ( ! isMountedAndNotUnsubscribing.current ) {
+				return;
+			}
+
+			try {
+				const newMapOutput = trapSelect( latestMapSelect.current );
+
+				if ( isShallowEqual( latestMapOutput.current, newMapOutput ) ) {
+					return;
+				}
+
+				latestMapOutput.current = newMapOutput;
+			} catch ( error ) {
+				latestMapOutputError.current = error;
+			}
+
+			forceRender();
+		};
+
+		const onChange = () => {
+			if ( latestIsAsync.current ) {
+				renderQueue.add( queueContext, onStoreChange );
+			} else {
+				onStoreChange();
+			}
+		};
+
+		// catch any possible state changes during mount before the subscription
+		// could be set.
+		onChange();
+
+		const unsubscribers = listeningStores.current.map( ( storeName ) =>
+			registry.__experimentalSubscribeStore( storeName, onChange )
+		);
+
+		return () => {
+			isMountedAndNotUnsubscribing.current = false;
+			// The return value of the subscribe function could be undefined if the store is a custom generic store.
+			unsubscribers.forEach( ( unsubscribe ) => unsubscribe?.() );
+			renderQueue.flush( queueContext );
+		};
+	}, [ registry, trapSelect, depsChangedFlag ] );
+
+	if ( mapOutputError ) {
+		throw mapOutputError;
+	}
+
+	return mapOutput;
+}

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.

packages/data/src/redux-store/index.js

@@ -162,10 +162,12 @@ export default function createReduxStore( key, options ) {
 			}
 
 			const resolveSelectors = mapResolveSelectors( selectors, store );
+			const suspendSelectors = mapSuspendSelectors( selectors, store );
 
 			const getSelectors = () => selectors;
 			const getActions = () => actions;
 			const getResolveSelectors = () => resolveSelectors;
+			const getSuspendSelectors = () => suspendSelectors;
 
 			// We have some modules monkey-patching the store object
 			// It's wrong to do so but until we refactor all of our effects to controls
@@ -200,6 +202,7 @@ export default function createReduxStore( key, options ) {
 				resolvers,
 				getSelectors,
 				getResolveSelectors,
+				getSuspendSelectors,
 				getActions,
 				subscribe,
 			};
@@ -375,6 +378,42 @@ function mapResolveSelectors( selectors, store ) {
 	} );
 }
 
+/**
+ * Maps selectors to functions that throw a suspense promise if not yet resolved.
+ *
+ * @param {Object} selectors Selectors to map.
+ * @param {Object} store     The redux store the selectors select from.
+ *
+ * @return {Object} Selectors mapped to their suspense functions.
+ */
+function mapSuspendSelectors( selectors, store ) {
+	return mapValues( selectors, ( selector, selectorName ) => {
+		// Selector without a resolver doesn't have any extra suspense behavior.
+		if ( ! selector.hasResolver ) {
+			return selector;
+		}
+
+		return ( ...args ) => {
+			const result = selector.apply( null, args );
+
+			if ( selectors.hasFinishedResolution( selectorName, args ) ) {
+				return result;
+			}
+
+			throw new Promise( ( resolve ) => {
+				const unsubscribe = store.subscribe( () => {
+					if (
+						selectors.hasFinishedResolution( selectorName, args )
+					) {
+						resolve();
+						unsubscribe();
+					}
+				} );
+			} );
+		};
+	} );
+}
+
 /**
  * Returns resolvers with matched selectors for a given namespace.
  * Resolvers are side effects invoked once per argument set of a given selector call,

packages/data/src/registry.js

@@ -93,14 +93,18 @@ export function createRegistry( storeConfigs = {}, parent = null ) {
 			return store.getSelectors();
 		}
 
-		return parent && parent.select( storeName );
+		return parent?.select( storeName );
 	}
 
-	function __experimentalMarkListeningStores( callback, ref ) {
+	function __experimentalMarkListeningStores( callback, listeningStoresRef ) {
 		__experimentalListeningStores.clear();
-		const result = callback.call( this );
-		ref.current = Array.from( __experimentalListeningStores );
-		return result;
+		try {
+			return callback.call( this );
+		} finally {
+			listeningStoresRef.current = Array.from(
+				__experimentalListeningStores
+			);
+		}
 	}
 
 	/**
@@ -127,6 +131,29 @@ export function createRegistry( storeConfigs = {}, parent = null ) {
 		return parent && parent.resolveSelect( storeName );
 	}
 
+	/**
+	 * 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.
+	 */
+	function suspendSelect( storeNameOrDescriptor ) {
+		const storeName = isObject( storeNameOrDescriptor )
+			? storeNameOrDescriptor.name
+			: storeNameOrDescriptor;
+		__experimentalListeningStores.add( storeName );
+		const store = stores[ storeName ];
+		if ( store ) {
+			return store.getSuspendSelectors();
+		}
+
+		return parent && parent.suspendSelect( storeName );
+	}
+
 	/**
 	 * Returns the available actions for a part of the state.
 	 *
@@ -276,6 +303,7 @@ export function createRegistry( storeConfigs = {}, parent = null ) {
 		subscribe,
 		select,
 		resolveSelect,
+		suspendSelect,
 		dispatch,
 		use,
 		register,