Fix render ratom watches being lost for class components with StrictMode

Same as #640

Author: Deraen

src/reagent/impl/component.cljs

@@ -116,6 +116,33 @@
   (binding [*current-component* c]
     (wrap-render c compiler)))
 
+;; Ratom cleanup
+;; On StrictMode React calls mount->unmount->mount to
+;; ensure there all start is correctly setup there.
+;; This means we need to use tricks to keep the render ratom
+;; non-disposed when unmount is called on the component
+;; initialization. It we let the dispose to be called,
+;; we lose the ratom watches from derefs inside reagent-render,
+;; and there is no good way to restore them.
+
+(defn cancel-cleanup [^clj component-state]
+  (when (.-cleanup-queued component-state)
+    (set! (.-cleanup-cancelled component-state) true)
+    (set! (.-cleanup-queued component-state) false)))
+
+;; Delay ratom dispose call so it can be cancelled if StrictMode
+;; calls setup the second time.
+(defn queue-cleanup [^clj component-state]
+  (set! (.-cleanup-cancelled component-state) false)
+  (set! (.-cleanup-queued component-state) true)
+  ;; Promise.resolve creates a microtask, vs setTimeout regular task.
+  ;; A scheduled microtask runs before the next event loop begings (where a regular task would run).
+  (.then (.resolve js/Promise nil)
+         (fn []
+           (when (false? (.-cleanup-cancelled component-state))
+             (set! (.-cleanup-queued component-state) false)
+             (batch/mark-rendered component-state)
+             (some-> (gobj/get component-state "cljsRatom") ratom/dispose!)))))
 
 ;;; Method wrapping
 
@@ -193,13 +220,15 @@
 
     :componentDidMount
     (fn componentDidMount []
-      (this-as c (.call f c c)))
+      (this-as c
+        (cancel-cleanup c)
+        (when-not (nil? f)
+          (.call f c c))))
 
     :componentWillUnmount
     (fn componentWillUnmount []
       (this-as c
-               (some-> (gobj/get c "cljsRatom") ratom/dispose!)
-               (batch/mark-rendered c)
+               (queue-cleanup c)
                (when-not (nil? f)
                  (.call f c c))))
 
@@ -218,6 +247,7 @@
 ;; Though the value is nil here, the wrapper function will be
 ;; added to class to manage Reagent ratom lifecycle.
 (def obligatory {:shouldComponentUpdate nil
+                 :componentDidMount nil
                  :componentWillUnmount nil})
 
 (def dash-to-method-name (util/memoize-1 util/dash-to-method-name))
@@ -427,26 +457,11 @@
           ;; FIXME: Access cljsRatom using interop forms
           rat ^ratom/Reaction (gobj/get reagent-state "cljsRatom")]
 
-      ;; Delay ratom dispose call so it can be cancelled if StrictMode
-      ;; calls setup the second time.
       (react/useEffect
         (fn mount []
-          (when (.-cleanup-queued reagent-state)
-            ;; (js/console.log "cancel cleanup")
-            (set! (.-cleanup-cancelled reagent-state) true)
-            (set! (.-cleanup-queued reagent-state) false))
+          (cancel-cleanup reagent-state)
           (fn unmount []
-            ;; (js/console.log "queue cleanup")
-            (set! (.-cleanup-cancelled reagent-state) false)
-            (set! (.-cleanup-queued reagent-state) true)
-            ;; Promise.resolve creates a microtask, vs setTimeout regular task.
-            ;; A scheduled microtask runs before the next event loop begings (where a regular task would run).
-            (.then (.resolve js/Promise nil)
-                   (fn []
-                     (when (false? (.-cleanup-cancelled reagent-state))
-                       ;; (js/console.log "run cleanup")
-                       (set! (.-cleanup-queued reagent-state) false)
-                       (some-> (gobj/get reagent-state "cljsRatom") ratom/dispose!))))))
+            (queue-cleanup reagent-state)))
         ;; Ignore props - only run effect once on mount and unmount
         ;; (which means always twice under the React strict mode).
         #js [])

test/reagenttest/testreagent.cljs

@@ -1448,6 +1448,36 @@
         ;; TODO: Test that component RAtom is disposed
         (is (= "Count 6" (.-innerText div)))))))
 
+(deftest ^:dom class-strict
+  ;; Ensure mount->unmount->mount done by StrictMode
+  ;; doesn't prematurely dispose the render ratom
+  (let [numbers (r/atom [0 1])
+        num-c (fn [i]
+                (let [x (get @numbers i)]
+                  [:span x]))
+        c (fn []
+            (let [xs @numbers]
+              [:span
+               "counting "
+               (for [i (range (count xs))]
+                 ^{:key i}
+                 [num-c i])]))
+        runs (rv/running)]
+    (u/async
+      (u/with-render [div [c]]
+        {:compiler u/class-compiler
+         :strict? true}
+
+        (is (= "counting 01" (.-innerText div)))
+
+        (u/act (reset! numbers [2 4 6]))
+
+        (is (= "counting 246" (.-innerText div))))
+
+      (testing "after unmount ratom watches are cleaned"
+        (is (= {} (.-watches ^clj numbers)))
+        (is (= runs (rv/running)))))))
+
 (deftest ^:dom functional-strict
   (let [numbers (r/atom [0 1])
         num-c (fn [i]
@@ -1459,7 +1489,7 @@
                   ref (react/useRef nil)]
               [:span
                "counting "
-               (for [i xs]
+               (for [i (range (count xs))]
                  ^{:key i}
                  [num-c i])]))
         runs (rv/running)]
@@ -1480,8 +1510,8 @@
         (is (= runs (rv/running)))))))
 
 (deftest ^:dom functional-strict-transitive
-  (let [count (r/atom 2)
-        numbers (rv/reaction (vec (range @count)))
+  (let [cnt (r/atom 2)
+        numbers (rv/reaction (vec (range @cnt)))
         num-c (fn [i]
                 (let [x (get @numbers i)]
                   [:span x]))
@@ -1491,7 +1521,7 @@
                   ref (react/useRef nil)]
               [:span
                "counting "
-               (for [i xs]
+               (for [i (range (count xs))]
                  ^{:key i}
                  [num-c i])]))
 
@@ -1502,11 +1532,11 @@
          :capture-errors true
          :strict? true}
 
-        (u/act (reset! count 3))
+        (u/act (reset! cnt 3))
         (is (= "counting 012" (.-innerText div)))
         (is (= [0 1 2] @numbers))
 
-        (u/act (reset! count 2))
+        (u/act (reset! cnt 2))
         (is (= "counting 01" (.-innerText div)))
 
         (when (dev?)

test/reagenttest/utils.cljs

@@ -115,6 +115,7 @@
                          (resolve))
                        RENDER-WAIT)
         (catch :default e
+          ;; NOTE: errors from act body aren't logged currently
           (reject e))))))
 
 (def ^:dynamic *render-error* nil)
@@ -137,7 +138,10 @@
                                   (init-capture))
          ;; Magic setup to make exception from render available to the
          ;; with-render body.
-         render-error (atom nil)]
+         render-error (atom nil)
+         comp (if (:strict-mode? options)
+                [:> react/StrictMode comp]
+                comp)]
      (try
        (if compiler
          (rdom/render comp div {:compiler compiler