Add suspense support to the data module

packages/block-library/src/post-template/edit.js

@@ -7,7 +7,7 @@ import classnames from 'classnames';
  * WordPress dependencies
  */
 import { useState, useMemo } from '@wordpress/element';
-import { useSelect } from '@wordpress/data';
+import { useSuspenseSelect } from '@wordpress/data';
 import { __ } from '@wordpress/i18n';
 import {
 	BlockContextProvider,
@@ -55,7 +55,7 @@ export default function PostTemplateEdit( {
 	const [ { page } ] = queryContext;
 	const [ activeBlockContext, setActiveBlockContext ] = useState();
 
-	const { posts, blocks } = useSelect(
+	const { posts, blocks } = useSuspenseSelect(
 		( select ) => {
 			const { getEntityRecords } = select( coreStore );
 			const { getBlocks } = select( blockEditorStore );

packages/block-library/src/site-logo/edit.js

@@ -32,7 +32,7 @@ import {
 	__experimentalImageEditor as ImageEditor,
 	__experimentalImageEditingProvider as ImageEditingProvider,
 } from '@wordpress/block-editor';
-import { useSelect, useDispatch } from '@wordpress/data';
+import { useSuspenseSelect, useDispatch } from '@wordpress/data';
 import { store as coreStore } from '@wordpress/core-data';
 import { crop, upload } from '@wordpress/icons';
 import { SVG, Path } from '@wordpress/primitives';
@@ -72,7 +72,7 @@ const SiteLogo = ( {
 	const classes = classnames( 'custom-logo-link', {
 		'is-transient': isBlobURL( logoUrl ),
 	} );
-	const { imageEditing, maxWidth, title } = useSelect( ( select ) => {
+	const { imageEditing, maxWidth, title } = useSuspenseSelect( ( select ) => {
 		const { getSettings } = select( blockEditorStore );
 		const siteEntities = select( coreStore ).getEditedEntityRecord(
 			'root',
@@ -318,7 +318,7 @@ export default function LogoEdit( {
 		url,
 		mediaItemData,
 		isRequestingMediaItem,
-	} = useSelect( ( select ) => {
+	} = useSuspenseSelect( ( select ) => {
 		const { canUser, getEntityRecord, getEditedEntityRecord } = select(
 			coreStore
 		);

packages/block-library/src/template-part/edit/index.js

@@ -1,7 +1,7 @@
 /**
  * WordPress dependencies
  */
-import { useSelect } from '@wordpress/data';
+import { useSuspenseSelect } from '@wordpress/data';
 import {
 	BlockControls,
 	useBlockProps,
@@ -50,7 +50,7 @@ export default function TemplatePartEdit( {
 		area,
 		enableSelection,
 		hasResolvedReplacements,
-	} = useSelect(
+	} = useSuspenseSelect(
 		( select ) => {
 			const {
 				getEditedEntityRecord,

packages/data/README.md

@@ -827,6 +827,70 @@ _Returns_
 
 -   `Function`: A custom react hook.
 
+### useSuspenseSelect
+
+Custom react hook for retrieving props from registered selectors.
+
+In general, this custom React hook follows the
+[rules of hooks](https://reactjs.org/docs/hooks-rules.html).
+
+_Usage_
+
+```js
+import { useSelect } from '@wordpress/data';
+
+function HammerPriceDisplay( { currency } ) {
+	const price = useSelect(
+		( select ) => {
+			return select( 'my-shop' ).getPrice( 'hammer', currency );
+		},
+		[ currency ]
+	);
+	return new Intl.NumberFormat( 'en-US', {
+		style: 'currency',
+		currency,
+	} ).format( price );
+}
+
+// Rendered in the application:
+// <HammerPriceDisplay currency="USD" />
+```
+
+In the above example, when `HammerPriceDisplay` is rendered into an
+application, the price will be retrieved from the store state using the
+`mapSelect` callback on `useSelect`. If the currency prop changes then
+any price in the state for that currency is retrieved. If the currency prop
+doesn't change and other props are passed in that do change, the price will
+not change because the dependency is just the currency.
+
+When data is only used in an event callback, the data should not be retrieved
+on render, so it may be useful to get the selectors function instead.
+
+**Don't use `useSelect` this way when calling the selectors in the render
+function because your component won't re-render on a data change.**
+
+```js
+import { useSelect } from '@wordpress/data';
+
+function Paste( { children } ) {
+	const { getSettings } = useSelect( 'my-shop' );
+	function onPaste() {
+		// Do something with the settings.
+		const settings = getSettings();
+	}
+	return <div onPaste={ onPaste }>{ children }</div>;
+}
+```
+
+_Parameters_
+
+-   _mapSelect_ `Function|StoreDescriptor|string`: Function called on every state change. The returned value is exposed to the component implementing this hook. The function receives the `registry.select` method on the first argument and the `registry` on the second argument. When a store key is passed, all selectors for the store will be returned. This is only meant for usage of these selectors in event callbacks, not for data needed to create the element tree.
+-   _deps_ `Array`: If provided, this memoizes the mapSelect so the same `mapSelect` is invoked on every state change unless the dependencies change.
+
+_Returns_
+
+-   `Function`: A custom react hook.
+
 ### withDispatch
 
 Higher-order component used to add dispatch props using registered action

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

@@ -246,3 +246,238 @@ export default function useSelect( mapSelect, deps ) {
 
 	return hasMappingFunction ? mapOutput : registry.select( mapSelect );
 }
+
+/**
+ * Custom react hook for retrieving props from registered selectors.
+ *
+ * In general, this custom React hook follows the
+ * [rules of hooks](https://reactjs.org/docs/hooks-rules.html).
+ *
+ * @param {Function|StoreDescriptor|string} mapSelect Function called on every state change. The
+ *                                                    returned value is exposed to the component
+ *                                                    implementing this hook. The function receives
+ *                                                    the `registry.select` method on the first
+ *                                                    argument and the `registry` on the second
+ *                                                    argument.
+ *                                                    When a store key is passed, all selectors for
+ *                                                    the store will be returned. This is only meant
+ *                                                    for usage of these selectors in event
+ *                                                    callbacks, not for data needed to create the
+ *                                                    element tree.
+ * @param {Array}                           deps      If provided, this memoizes the mapSelect so the
+ *                                                    same `mapSelect` is invoked on every state
+ *                                                    change unless the dependencies change.
+ *
+ * @example
+ * ```js
+ * import { useSelect } from '@wordpress/data';
+ *
+ * function HammerPriceDisplay( { currency } ) {
+ *   const price = useSelect( ( select ) => {
+ *     return select( 'my-shop' ).getPrice( 'hammer', currency )
+ *   }, [ currency ] );
+ *   return new Intl.NumberFormat( 'en-US', {
+ *     style: 'currency',
+ *     currency,
+ *   } ).format( price );
+ * }
+ *
+ * // Rendered in the application:
+ * // <HammerPriceDisplay currency="USD" />
+ * ```
+ *
+ * In the above example, when `HammerPriceDisplay` is rendered into an
+ * application, the price will be retrieved from the store state using the
+ * `mapSelect` callback on `useSelect`. If the currency prop changes then
+ * any price in the state for that currency is retrieved. If the currency prop
+ * doesn't change and other props are passed in that do change, the price will
+ * not change because the dependency is just the currency.
+ *
+ * When data is only used in an event callback, the data should not be retrieved
+ * on render, so it may be useful to get the selectors function instead.
+ *
+ * **Don't use `useSelect` this way when calling the selectors in the render
+ * function because your component won't re-render on a data change.**
+ *
+ * ```js
+ * import { useSelect } from '@wordpress/data';
+ *
+ * function Paste( { children } ) {
+ *   const { getSettings } = useSelect( 'my-shop' );
+ *   function onPaste() {
+ *     // Do something with the settings.
+ *     const settings = getSettings();
+ *   }
+ *   return <div onPaste={ onPaste }>{ children }</div>;
+ * }
+ * ```
+ *
+ * @return {Function}  A custom react hook.
+ */
+export function useSuspenseSelect( mapSelect, deps ) {
+	const hasMappingFunction = 'function' === typeof mapSelect;
+
+	// If we're recalling a store by its name or by
+	// its descriptor then we won't be caching the
+	// calls to `mapSelect` because we won't be calling it.
+	if ( ! hasMappingFunction ) {
+		deps = [];
+	}
+
+	// Because of the "rule of hooks" we have to call `useCallback`
+	// on every invocation whether or not we have a real function
+	// for `mapSelect`. we'll create this intermediate variable to
+	// fulfill that need and then reference it with our "real"
+	// `_mapSelect` if we can.
+	const callbackMapper = useCallback(
+		hasMappingFunction ? mapSelect : noop,
+		deps
+	);
+	const _mapSelect = hasMappingFunction ? callbackMapper : null;
+
+	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();
+	const shouldBeSuspended = useRef( false );
+
+	// 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,
+				listeningStores,
+				shouldBeSuspended
+			),
+		[ 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;
+
+	if ( _mapSelect ) {
+		mapOutput = latestMapOutput.current;
+		const hasReplacedMapSelect = latestMapSelect.current !== _mapSelect;
+		const lastMapSelectFailed = !! latestMapOutputError.current;
+
+		if ( hasReplacedMapSelect || lastMapSelectFailed ) {
+			try {
+				mapOutput = trapSelect( () =>
+					_mapSelect( registry.select, registry )
+				);
+			} catch ( error ) {
+				let errorMessage = `An error occurred while running 'mapSelect': ${ error.message }`;
+
+				if ( latestMapOutputError.current ) {
+					errorMessage += `\nThe error may be correlated with this previous error:\n`;
+					errorMessage += `${ latestMapOutputError.current.stack }\n\n`;
+					errorMessage += 'Original stack trace:';
+				}
+
+				// eslint-disable-next-line no-console
+				console.error( errorMessage );
+			}
+
+			if ( shouldBeSuspended.current ) {
+				throw shouldBeSuspended.current;
+			}
+		}
+	}
+
+	useIsomorphicLayoutEffect( () => {
+		if ( ! hasMappingFunction ) {
+			return;
+		}
+
+		latestMapSelect.current = _mapSelect;
+		latestMapOutput.current = mapOutput;
+		latestMapOutputError.current = undefined;
+		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 );
+		}
+	} );
+
+	useIsomorphicLayoutEffect( () => {
+		if ( ! hasMappingFunction ) {
+			return;
+		}
+
+		const onStoreChange = () => {
+			if ( isMountedAndNotUnsubscribing.current ) {
+				try {
+					const newMapOutput = trapSelect( () =>
+						latestMapSelect.current( registry.select, registry )
+					);
+
+					if (
+						isShallowEqual( latestMapOutput.current, newMapOutput )
+					) {
+						return;
+					}
+					latestMapOutput.current = newMapOutput;
+				} catch ( error ) {
+					latestMapOutputError.current = error;
+				}
+
+				if ( shouldBeSuspended.current ) {
+					throw shouldBeSuspended.current;
+				}
+				forceRender();
+			}
+		};
+
+		// catch any possible state changes during mount before the subscription
+		// could be set.
+		if ( latestIsAsync.current ) {
+			renderQueue.add( queueContext, onStoreChange );
+		} else {
+			onStoreChange();
+		}
+
+		const onChange = () => {
+			if ( latestIsAsync.current ) {
+				renderQueue.add( queueContext, onStoreChange );
+			} else {
+				onStoreChange();
+			}
+		};
+
+		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 );
+		};
+		// If you're tempted to eliminate the spread dependencies below don't do it!
+		// We're passing these in from the calling function and want to make sure we're
+		// examining every individual value inside the `deps` array.
+	}, [ registry, trapSelect, hasMappingFunction, depsChangedFlag ] );
+
+	return hasMappingFunction ? mapOutput : registry.select( mapSelect );
+}

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';