Data: Fix issue with in-stack unsubscribe

data/index.js

@@ -28,7 +28,7 @@ let listeners = [];
  * Global listener called for each store's update.
  */
 export function globalListener() {
-	listeners.forEach( listener => listener() );
+	listeners.forEach( ( listener ) => listener() );
 }
 
 /**
@@ -163,13 +163,23 @@ export const withSelect = ( mapStateToProps ) => ( WrappedComponent ) => {
 
 		componentWillUnmount() {
 			this.unsubscribe();
+
+			// While above unsubscribe avoids future listener calls, callbacks
+			// are snapshotted before being invoked, so if unmounting occurs
+			// during a previous callback, we need to explicitly track and
+			// avoid the `runSelection` that is scheduled to occur.
+			this.isUnmounting = true;
 		}
 
 		subscribe() {
 			this.unsubscribe = subscribe( this.runSelection );
 		}
 
 		runSelection( props = this.props ) {
+			if ( this.isUnmounting ) {
+				return;
+			}
+
 			const newState = mapStateToProps( select, props );
 			if ( ! isEqualShallow( newState, this.state ) ) {
 				this.setState( newState );

data/test/index.js

@@ -66,6 +66,20 @@ describe( 'select', () => {
 } );
 
 describe( 'withSelect', () => {
+	const unsubscribes = [];
+	afterEach( () => {
+		let unsubscribe;
+		while ( ( unsubscribe = unsubscribes.shift() ) ) {
+			unsubscribe();
+		}
+	} );
+
+	function subscribeWithUnsubscribe( ...args ) {
+		const unsubscribe = subscribe( ...args );
+		unsubscribes.push( unsubscribe );
+		return unsubscribe;
+	}
+
 	it( 'passes the relevant data to the component', () => {
 		registerReducer( 'reactReducer', () => ( { reactKey: 'reactState' } ) );
 		registerSelectors( 'reactReducer', {
@@ -162,6 +176,36 @@ describe( 'withSelect', () => {
 
 		wrapper.unmount();
 	} );
+
+	it( 'ensures component is still mounted before setting state', () => {
+		// This test verifies that even though unsubscribe doesn't take effect
+		// until after the current listener stack is called, we don't attempt
+		// to setState on an unmounting `withSelect` component. It will fail if
+		// an attempt is made to `setState` on an unmounted component.
+		const store = registerReducer( 'counter', ( state = 0, action ) => {
+			if ( action.type === 'increment' ) {
+				return state + 1;
+			}
+
+			return state;
+		} );
+
+		registerSelectors( 'counter', {
+			getCount: ( state, offset ) => state + offset,
+		} );
+
+		subscribeWithUnsubscribe( () => {
+			wrapper.unmount();
+		} );
+
+		const Component = withSelect( ( _select, ownProps ) => ( {
+			count: _select( 'counter' ).getCount( ownProps.offset ),
+		} ) )( ( props ) => <div>{ props.count }</div> );
+
+		const wrapper = mount( <Component offset={ 0 } /> );
+
+		store.dispatch( { type: 'increment' } );
+	} );
 } );
 
 describe( 'withDispatch', () => {
@@ -209,6 +253,20 @@ describe( 'withDispatch', () => {
 } );
 
 describe( 'subscribe', () => {
+	const unsubscribes = [];
+	afterEach( () => {
+		let unsubscribe;
+		while ( ( unsubscribe = unsubscribes.shift() ) ) {
+			unsubscribe();
+		}
+	} );
+
+	function subscribeWithUnsubscribe( ...args ) {
+		const unsubscribe = subscribe( ...args );
+		unsubscribes.push( unsubscribe );
+		return unsubscribe;
+	}
+
 	it( 'registers multiple selectors to the public API', () => {
 		let incrementedValue = null;
 		const store = registerReducer( 'myAwesomeReducer', ( state = 0 ) => state + 1 );
@@ -233,6 +291,35 @@ describe( 'subscribe', () => {
 
 		expect( incrementedValue ).toBe( 3 );
 	} );
+
+	it( 'snapshots listeners on change, avoiding a later listener if subscribed during earlier callback', () => {
+		const store = registerReducer( 'myAwesomeReducer', ( state = 0 ) => state + 1 );
+		const secondListener = jest.fn();
+		const firstListener = jest.fn( () => {
+			subscribeWithUnsubscribe( secondListener );
+		} );
+
+		subscribeWithUnsubscribe( firstListener );
+
+		store.dispatch( { type: 'dummy' } );
+
+		expect( secondListener ).not.toHaveBeenCalled();
+	} );
+
+	it( 'snapshots listeners on change, calling a later listener even if unsubscribed during earlier callback', () => {
+		const store = registerReducer( 'myAwesomeReducer', ( state = 0 ) => state + 1 );
+		const firstListener = jest.fn( () => {
+			secondUnsubscribe();
+		} );
+		const secondListener = jest.fn();
+
+		subscribeWithUnsubscribe( firstListener );
+		const secondUnsubscribe = subscribeWithUnsubscribe( secondListener );
+
+		store.dispatch( { type: 'dummy' } );
+
+		expect( secondListener ).toHaveBeenCalled();
+	} );
 } );
 
 describe( 'dispatch', () => {