useSelect: incrementally subscribe to stores when first selected from

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

@@ -42,56 +42,103 @@ function Store( registry, suspense ) {
 	let lastMapResult;
 	let lastMapResultValid = false;
 	let lastIsAsync;
-	let subscribe;
-
-	const createSubscriber = ( stores ) => ( listener ) => {
-		// Invalidate the value right after subscription was created. React will
-		// call `getValue` after subscribing, to detect store updates that happened
-		// in the interval between the `getValue` call during render and creating
-		// the subscription, which is slightly delayed. We need to ensure that this
-		// second `getValue` call will compute a fresh value.
-		lastMapResultValid = false;
-
-		const onStoreChange = () => {
-			// Invalidate the value on store update, so that a fresh value is computed.
+	let subscriber;
+
+	const createSubscriber = ( stores ) => {
+		// The set of stores the `subscribe` function is supposed to subscribe to. Here it is
+		// initialized, and then the `updateStores` function can add new stores to it.
+		const activeStores = [ ...stores ];
+
+		// The `subscribe` function, which is passed to the `useSyncExternalStore` hook, could
+		// be called multiple times to establish multiple subscriptions. That's why we need to
+		// keep a set of active subscriptions;
+		const activeSubscriptions = new Set();
+
+		function subscribe( listener ) {
+			// Invalidate the value right after subscription was created. React will
+			// call `getValue` after subscribing, to detect store updates that happened
+			// in the interval between the `getValue` call during render and creating
+			// the subscription, which is slightly delayed. We need to ensure that this
+			// second `getValue` call will compute a fresh value.
 			lastMapResultValid = false;
-			listener();
-		};
 
-		const onChange = () => {
-			if ( lastIsAsync ) {
-				renderQueue.add( queueContext, onStoreChange );
-			} else {
-				onStoreChange();
+			const onStoreChange = () => {
+				// Invalidate the value on store update, so that a fresh value is computed.
+				lastMapResultValid = false;
+				listener();
+			};
+
+			const onChange = () => {
+				if ( lastIsAsync ) {
+					renderQueue.add( queueContext, onStoreChange );
+				} else {
+					onStoreChange();
+				}
+			};
+
+			const unsubs = [];
+			function subscribeStore( storeName ) {
+				unsubs.push( registry.subscribe( onChange, storeName ) );
 			}
-		};
 
-		const unsubs = stores.map( ( storeName ) => {
-			return registry.subscribe( onChange, storeName );
-		} );
+			for ( const storeName of activeStores ) {
+				subscribeStore( storeName );
+			}
+
+			activeSubscriptions.add( subscribeStore );
+
+			return () => {
+				activeSubscriptions.delete( subscribeStore );
+
+				for ( const unsub of unsubs.values() ) {
+					// The return value of the subscribe function could be undefined if the store is a custom generic store.
+					unsub?.();
+				}
+				// Cancel existing store updates that were already scheduled.
+				renderQueue.cancel( queueContext );
+			};
+		}
+
+		// Check if `newStores` contains some stores we're not subscribed to yet, and add them.
+		function updateStores( newStores ) {
+			for ( const newStore of newStores ) {
+				if ( activeStores.includes( newStore ) ) {
+					continue;
+				}
+
+				// New `subscribe` calls will subscribe to `newStore`, too.
+				activeStores.push( newStore );
 
-		return () => {
-			// The return value of the subscribe function could be undefined if the store is a custom generic store.
-			for ( const unsub of unsubs ) {
-				unsub?.();
+				// Add `newStore` to existing subscriptions.
+				for ( const subscription of activeSubscriptions ) {
+					subscription( newStore );
+				}
 			}
-			// Cancel existing store updates that were already scheduled.
-			renderQueue.cancel( queueContext );
-		};
-	};
+		}
 
-	return ( mapSelect, resubscribe, isAsync ) => {
-		const selectValue = () => mapSelect( select, registry );
+		return { subscribe, updateStores };
+	};
 
-		function updateValue( selectFromStore ) {
+	return ( mapSelect, isAsync ) => {
+		function updateValue() {
 			// If the last value is valid, and the `mapSelect` callback hasn't changed,
 			// then we can safely return the cached value. The value can change only on
 			// store update, and in that case value will be invalidated by the listener.
 			if ( lastMapResultValid && mapSelect === lastMapSelect ) {
 				return lastMapResult;
 			}
 
-			const mapResult = selectFromStore();
+			const listeningStores = { current: null };
+			const mapResult = registry.__unstableMarkListeningStores(
+				() => mapSelect( select, registry ),
+				listeningStores
+			);
+
+			if ( ! subscriber ) {
+				subscriber = createSubscriber( listeningStores.current );
+			} else {
+				subscriber.updateStores( listeningStores.current );
+			}
 
 			// If the new value is shallow-equal to the old one, keep the old one so
 			// that we don't trigger unwanted updates that do a `===` check.
@@ -103,7 +150,7 @@ function Store( registry, suspense ) {
 
 		function getValue() {
 			// Update the value in case it's been invalidated or `mapSelect` has changed.
-			updateValue( selectValue );
+			updateValue();
 			return lastMapResult;
 		}
 
@@ -115,30 +162,13 @@ function Store( registry, suspense ) {
 			renderQueue.cancel( queueContext );
 		}
 
-		// Either initialize the `subscribe` function, or create a new one if `mapSelect`
-		// changed and has dependencies.
-		// Usage without dependencies, `useSelect( ( s ) => { ... } )`, will subscribe
-		// only once, at mount, and won't resubscibe even if `mapSelect` changes.
-		if ( ! subscribe || ( resubscribe && mapSelect !== lastMapSelect ) ) {
-			// Find out what stores the `mapSelect` callback is selecting from and
-			// use that list to create subscriptions to specific stores.
-			const listeningStores = { current: null };
-			updateValue( () =>
-				registry.__unstableMarkListeningStores(
-					selectValue,
-					listeningStores
-				)
-			);
-			subscribe = createSubscriber( listeningStores.current );
-		} else {
-			updateValue( selectValue );
-		}
+		updateValue();
 
 		lastIsAsync = isAsync;
 		lastMapSelect = mapSelect;
 
 		// Return a pair of functions that can be passed to `useSyncExternalStore`.
-		return { subscribe, getValue };
+		return { subscribe: subscriber.subscribe, getValue };
 	};
 }
 
@@ -151,7 +181,7 @@ function useMappingSelect( suspense, mapSelect, deps ) {
 	const isAsync = useAsyncMode();
 	const store = useMemo( () => Store( registry, suspense ), [ registry ] );
 	const selector = useCallback( mapSelect, deps );
-	const { subscribe, getValue } = store( selector, !! deps, isAsync );
+	const { subscribe, getValue } = store( selector, isAsync );
 	const result = useSyncExternalStore( subscribe, getValue, getValue );
 	useDebugValue( result );
 	return result;

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

@@ -21,10 +21,10 @@ import useSelect from '..';
 
 jest.useRealTimers();
 
-function counterStore( initialCount = 0 ) {
+function counterStore( initialCount = 0, step = 1 ) {
 	return {
 		reducer: ( state = initialCount, action ) =>
-			action.type === 'INC' ? state + 1 : state,
+			action.type === 'INC' ? state + step : state,
 		actions: {
 			inc: () => ( { type: 'INC' } ),
 		},
@@ -124,7 +124,7 @@ describe( 'useSelect', () => {
 		);
 
 		expect( selectSpyFoo ).toHaveBeenCalledTimes( 2 );
-		expect( selectSpyBar ).toHaveBeenCalledTimes( 2 );
+		expect( selectSpyBar ).toHaveBeenCalledTimes( 1 );
 		expect( TestComponent ).toHaveBeenCalledTimes( 3 );
 
 		// Ensure expected state was rendered.
@@ -195,6 +195,81 @@ describe( 'useSelect', () => {
 		expect( Child ).toHaveBeenCalledTimes( 1 );
 	} );
 
+	it( 'incrementally subscribes to newly selected stores', () => {
+		registry.registerStore( 'store-main', counterStore() );
+		registry.registerStore( 'store-even', counterStore( 0, 2 ) );
+		registry.registerStore( 'store-odd', counterStore( 1, 2 ) );
+
+		const mapSelect = jest.fn( ( select ) => {
+			const first = select( 'store-main' ).get();
+			// select from other stores depending on whether main value is even or odd
+			const secondStore = first % 2 === 1 ? 'store-odd' : 'store-even';
+			const second = select( secondStore ).get();
+			return first + ':' + second;
+		} );
+
+		const TestComponent = jest.fn( () => {
+			const data = useSelect( mapSelect, [] );
+			return <div role="status">{ data }</div>;
+		} );
+
+		render(
+			<RegistryProvider value={ registry }>
+				<TestComponent />
+			</RegistryProvider>
+		);
+
+		expect( mapSelect ).toHaveBeenCalledTimes( 2 );
+		expect( TestComponent ).toHaveBeenCalledTimes( 1 );
+		expect( screen.getByRole( 'status' ) ).toHaveTextContent( '0:0' );
+
+		// check that increment in store-even triggers a render
+		act( () => {
+			registry.dispatch( 'store-even' ).inc();
+		} );
+
+		expect( mapSelect ).toHaveBeenCalledTimes( 3 );
+		expect( TestComponent ).toHaveBeenCalledTimes( 2 );
+		expect( screen.getByRole( 'status' ) ).toHaveTextContent( '0:2' );
+
+		// check that increment in store-odd doesn't trigger a render (not listening yet)
+		act( () => {
+			registry.dispatch( 'store-odd' ).inc();
+		} );
+
+		expect( mapSelect ).toHaveBeenCalledTimes( 3 );
+		expect( TestComponent ).toHaveBeenCalledTimes( 2 );
+		expect( screen.getByRole( 'status' ) ).toHaveTextContent( '0:2' );
+
+		// check that increment in main store switches to store-odd
+		act( () => {
+			registry.dispatch( 'store-main' ).inc();
+		} );
+
+		expect( mapSelect ).toHaveBeenCalledTimes( 4 );
+		expect( TestComponent ).toHaveBeenCalledTimes( 3 );
+		expect( screen.getByRole( 'status' ) ).toHaveTextContent( '1:3' );
+
+		// check that increment in store-odd triggers a render
+		act( () => {
+			registry.dispatch( 'store-odd' ).inc();
+		} );
+
+		expect( mapSelect ).toHaveBeenCalledTimes( 5 );
+		expect( TestComponent ).toHaveBeenCalledTimes( 4 );
+		expect( screen.getByRole( 'status' ) ).toHaveTextContent( '1:5' );
+
+		// check that increment in store-even triggers a mapSelect call (still listening)
+		// but not a render (not used for selected value which doesn't change)
+		act( () => {
+			registry.dispatch( 'store-even' ).inc();
+		} );
+
+		expect( mapSelect ).toHaveBeenCalledTimes( 6 );
+		expect( TestComponent ).toHaveBeenCalledTimes( 4 );
+		expect( screen.getByRole( 'status' ) ).toHaveTextContent( '1:5' );
+	} );
+
 	describe( 'rerenders as expected with various mapSelect return types', () => {
 		const getComponent = ( mapSelectSpy ) => () => {
 			const data = useSelect( mapSelectSpy, [] );
@@ -411,7 +486,7 @@ describe( 'useSelect', () => {
 				setDep( 1 );
 			} );
 
-			expect( selectCount1AndDep ).toHaveBeenCalledTimes( 4 );
+			expect( selectCount1AndDep ).toHaveBeenCalledTimes( 3 );
 			expect( screen.getByRole( 'status' ) ).toHaveTextContent(
 				'count:0,dep:1'
 			);
@@ -420,7 +495,7 @@ describe( 'useSelect', () => {
 				registry.dispatch( 'store-1' ).inc();
 			} );
 
-			expect( selectCount1AndDep ).toHaveBeenCalledTimes( 5 );
+			expect( selectCount1AndDep ).toHaveBeenCalledTimes( 4 );
 			expect( screen.getByRole( 'status' ) ).toHaveTextContent(
 				'count:1,dep:1'
 			);
@@ -550,13 +625,11 @@ describe( 'useSelect', () => {
 					( select ) => {
 						if ( shouldSelectCount1 ) {
 							selectCount1();
-							select( 'store-1' ).get();
-							return 'count1';
+							return 'count1:' + select( 'store-1' ).get();
 						}
 
 						selectCount2();
-						select( 'store-2' ).get();
-						return 'count2';
+						return 'count2:' + select( 'store-2' ).get();
 					},
 					[ shouldSelectCount1 ]
 				);
@@ -578,15 +651,26 @@ describe( 'useSelect', () => {
 			expect( selectCount1 ).toHaveBeenCalledTimes( 0 );
 			expect( selectCount2 ).toHaveBeenCalledTimes( 2 );
 			expect( screen.getByRole( 'status' ) ).toHaveTextContent(
-				'count2'
+				'count2:0'
 			);
 
 			act( () => screen.getByText( 'Toggle' ).click() );
 
+			expect( selectCount1 ).toHaveBeenCalledTimes( 1 );
+			expect( selectCount2 ).toHaveBeenCalledTimes( 2 );
+			expect( screen.getByRole( 'status' ) ).toHaveTextContent(
+				'count1:0'
+			);
+
+			// Verify that the component subscribed to store-1 after selected from
+			act( () => {
+				registry.dispatch( 'store-1' ).inc();
+			} );
+
 			expect( selectCount1 ).toHaveBeenCalledTimes( 2 );
 			expect( selectCount2 ).toHaveBeenCalledTimes( 2 );
 			expect( screen.getByRole( 'status' ) ).toHaveTextContent(
-				'count1'
+				'count1:1'
 			);
 		} );
 
@@ -1097,11 +1181,11 @@ describe( 'useSelect', () => {
 			rerender( <App store="counter-2" /> );
 			expect( screen.getByRole( 'status' ) ).toHaveTextContent( '10' );
 
-			// update from counter-2 is ignored because component is subcribed only to counter-1
+			// update from counter-2 is processed because component has subscribed to counter-2
 			act( () => {
 				registry.dispatch( 'counter-2' ).inc();
 			} );
-			expect( screen.getByRole( 'status' ) ).toHaveTextContent( '10' );
+			expect( screen.getByRole( 'status' ) ).toHaveTextContent( '11' );
 		} );
 	} );