Fixes subscription bug when a new store is passed as a prop. (reduxjs#589)

* Fixes subscription bug when a new store is passed as a prop.

Reported here: davidkpiano/react-redux-form#592

* Refactors and adds more comments to connectAdvanced

Author: jimbolla

src/components/connectAdvanced.js

@@ -6,6 +6,30 @@ import Subscription from '../utils/Subscription'
 import storeShape from '../utils/storeShape'
 
 let hotReloadingVersion = 0
+const dummyState = {}
+function noop() {}
+
+function makeSelectorStateful(sourceSelector, store) {
+  // wrap the selector in an object that tracks its results between runs.
+  const selector = {
+    run: function runComponentSelector(props) {
+      try {
+        const nextProps = sourceSelector(store.getState(), props)
+        if (nextProps !== selector.props || selector.error) {
+          selector.shouldComponentUpdate = true
+          selector.props = nextProps
+          selector.error = null
+        }
+      } catch (error) {
+        selector.shouldComponentUpdate = true
+        selector.error = error
+      }
+    }
+  }
+
+  return selector
+}
+
 export default function connectAdvanced(
   /*
     selectorFactory is a func that is responsible for returning the selector function used to
@@ -96,28 +120,27 @@ export default function connectAdvanced(
         this.version = version
         this.state = {}
         this.renderCount = 0
-        this.store = this.props[storeKey] || this.context[storeKey]
-        this.parentSub = props[subscriptionKey] || context[subscriptionKey]
-
+        this.store = props[storeKey] || context[storeKey]
+        this.propsMode = Boolean(props[storeKey])
         this.setWrappedInstance = this.setWrappedInstance.bind(this)
 
         invariant(this.store,
-          `Could not find "${storeKey}" in either the context or ` +
-          `props of "${displayName}". ` +
-          `Either wrap the root component in a <Provider>, ` +
+          `Could not find "${storeKey}" in either the context or props of ` +
+          `"${displayName}". Either wrap the root component in a <Provider>, ` +
           `or explicitly pass "${storeKey}" as a prop to "${displayName}".`
         )
 
-        // make sure `getState` is properly bound in order to avoid breaking
-        // custom store implementations that rely on the store's context
-        this.getState = this.store.getState.bind(this.store);
-
         this.initSelector()
         this.initSubscription()
       }
 
       getChildContext() {
-        return { [subscriptionKey]: this.subscription || this.parentSub }
+        // If this component received store from props, its subscription should be transparent
+        // to any descendants receiving store+subscription from context; it passes along
+        // subscription passed to it. Otherwise, it shadows the parent subscription, which allows
+        // Connect to control ordering of notifications to flow top-down.
+        const subscription = this.propsMode ? null : this.subscription
+        return { [subscriptionKey]: subscription || this.context[subscriptionKey] }
       }
 
       componentDidMount() {
@@ -144,12 +167,11 @@ export default function connectAdvanced(
 
       componentWillUnmount() {
         if (this.subscription) this.subscription.tryUnsubscribe()
-        // these are just to guard against extra memory leakage if a parent element doesn't
-        // dereference this instance properly, such as an async callback that never finishes
         this.subscription = null
+        this.notifyNestedSubs = noop
         this.store = null
-        this.parentSub = null
-        this.selector.run = () => {}
+        this.selector.run = noop
+        this.selector.shouldComponentUpdate = false
       }
 
       getWrappedInstance() {
@@ -165,65 +187,63 @@ export default function connectAdvanced(
       }
 
       initSelector() {
-        const { dispatch } = this.store
-        const { getState } = this;
-        const sourceSelector = selectorFactory(dispatch, selectorFactoryOptions)
-
-        // wrap the selector in an object that tracks its results between runs
-        const selector = this.selector = {
-          shouldComponentUpdate: true,
-          props: sourceSelector(getState(), this.props),
-          run: function runComponentSelector(props) {
-            try {
-              const nextProps = sourceSelector(getState(), props)
-              if (selector.error || nextProps !== selector.props) {
-                selector.shouldComponentUpdate = true
-                selector.props = nextProps
-                selector.error = null
-              }
-            } catch (error) {
-              selector.shouldComponentUpdate = true
-              selector.error = error
-            }
-          }
-        }
+        const sourceSelector = selectorFactory(this.store.dispatch, selectorFactoryOptions)
+        this.selector = makeSelectorStateful(sourceSelector, this.store)
+        this.selector.run(this.props)
       }
 
       initSubscription() {
-        if (shouldHandleStateChanges) {
-          const subscription = this.subscription = new Subscription(this.store, this.parentSub)
-          const dummyState = {}
-
-          subscription.onStateChange = function onStateChange() {
-            this.selector.run(this.props)
-
-            if (!this.selector.shouldComponentUpdate) {
-              subscription.notifyNestedSubs()
-            } else {
-              this.componentDidUpdate = function componentDidUpdate() {
-                this.componentDidUpdate = undefined
-                subscription.notifyNestedSubs()
-              }
-
-              this.setState(dummyState)
-            }
-          }.bind(this)
+        if (!shouldHandleStateChanges) return
+        
+        // parentSub's source should match where store came from: props vs. context. A component
+        // connected to the store via props shouldn't use subscription from context, or vice versa.
+        const parentSub = (this.propsMode ? this.props : this.context)[subscriptionKey]
+        this.subscription = new Subscription(this.store, parentSub, this.onStateChange.bind(this))
+        
+        // `notifyNestedSubs` is duplicated to handle the case where the component is  unmounted in
+        // the middle of the notification loop, where `this.subscription` will then be null. An
+        // extra null check every change can be avoided by copying the method onto `this` and then
+        // replacing it with a no-op on unmount. This can probably be avoided if Subscription's 
+        // listeners logic is changed to not call listeners that have been unsubscribed in the
+        // middle of the notification loop.
+        this.notifyNestedSubs = this.subscription.notifyNestedSubs.bind(this.subscription)
+      }
+
+      onStateChange() {
+        this.selector.run(this.props)
+
+        if (!this.selector.shouldComponentUpdate) {
+          this.notifyNestedSubs()
+        } else {
+          this.componentDidUpdate = this.notifyNestedSubsOnComponentDidUpdate
+          this.setState(dummyState)
         }
+      }      
+
+      notifyNestedSubsOnComponentDidUpdate() {
+        // `componentDidUpdate` is conditionally implemented when `onStateChange` determines it
+        // needs to notify nested subs. Once called, it unimplements itself until further state
+        // changes occur. Doing it this way vs having a permanent `componentDidMount` that does
+        // a boolean check every time avoids an extra method call most of the time, resulting
+        // in some perf boost.
+        this.componentDidUpdate = undefined
+        this.notifyNestedSubs()
       }
 
       isSubscribed() {
         return Boolean(this.subscription) && this.subscription.isSubscribed()
       }
 
       addExtraProps(props) {
-        if (!withRef && !renderCountProp) return props
+        if (!withRef && !renderCountProp && !(this.propsMode && this.subscription)) return props
         // make a shallow copy so that fields added don't leak to the original selector.
         // this is especially important for 'ref' since that's a reference back to the component
         // instance. a singleton memoized selector would then be holding a reference to the
         // instance, preventing the instance from being garbage collected, and that would be bad
         const withExtras = { ...props }
         if (withRef) withExtras.ref = this.setWrappedInstance
         if (renderCountProp) withExtras[renderCountProp] = this.renderCount++
+        if (this.propsMode && this.subscription) withExtras[subscriptionKey] = this.subscription
         return withExtras
       }
 

src/utils/Subscription.js

@@ -41,9 +41,10 @@ function createListenerCollection() {
 }
 
 export default class Subscription {
-  constructor(store, parentSub) {
+  constructor(store, parentSub, onStateChange) {
     this.store = store
     this.parentSub = parentSub
+    this.onStateChange = onStateChange
     this.unsubscribe = null
     this.listeners = nullListeners
   }
@@ -63,7 +64,6 @@ export default class Subscription {
 
   trySubscribe() {
     if (!this.unsubscribe) {
-      // this.onStateChange is set by connectAdvanced.initSubscription()
       this.unsubscribe = this.parentSub
         ? this.parentSub.addNestedSub(this.onStateChange)
         : this.store.subscribe(this.onStateChange)

test/components/connect.spec.js

@@ -2207,5 +2207,47 @@ describe('React', () => {
       store.dispatch({ type: 'INC' })
     })
 
+    it('should subscribe properly when a new store is provided via props', () => {
+      const store1 = createStore((state = 0, action) => (action.type === 'INC' ? state + 1 : state))
+      const store2 = createStore((state = 0, action) => (action.type === 'INC' ? state + 1 : state))
+
+      @connect(state => ({ count: state }))
+      class A extends Component {
+        render() { return <B store={store2} /> }
+      }
+
+      const mapStateToPropsB = expect.createSpy().andCall(state => ({ count: state }))
+      @connect(mapStateToPropsB)
+      class B extends Component {
+        render() { return <C {...this.props} /> }
+      }
+
+      const mapStateToPropsC = expect.createSpy().andCall(state => ({ count: state }))
+      @connect(mapStateToPropsC)
+      class C extends Component {
+        render() { return <D /> }
+      }
+
+      const mapStateToPropsD = expect.createSpy().andCall(state => ({ count: state }))
+      @connect(mapStateToPropsD)
+      class D extends Component {
+        render() { return <div>{this.props.count}</div> }
+      }
+
+      TestUtils.renderIntoDocument(<ProviderMock store={store1}><A /></ProviderMock>)
+      expect(mapStateToPropsB.calls.length).toBe(1)
+      expect(mapStateToPropsC.calls.length).toBe(1)
+      expect(mapStateToPropsD.calls.length).toBe(1)
+
+      store1.dispatch({ type: 'INC' })
+      expect(mapStateToPropsB.calls.length).toBe(1)
+      expect(mapStateToPropsC.calls.length).toBe(1)
+      expect(mapStateToPropsD.calls.length).toBe(2)
+
+      store2.dispatch({ type: 'INC' })
+      expect(mapStateToPropsB.calls.length).toBe(2)
+      expect(mapStateToPropsC.calls.length).toBe(2)
+      expect(mapStateToPropsD.calls.length).toBe(2)
+    })
   })
 })