useObserver registry refactor

.changeset/forty-books-dance.md

@@ -0,0 +1,5 @@
+---
+"mobx-react-lite": patch
+---
+
+refactor reaction tracking

packages/mobx-react-lite/__tests__/strictAndConcurrentModeUsingFinalizationRegistry.test.tsx

@@ -1,21 +1,20 @@
 import { cleanup, render } from "@testing-library/react"
 import * as mobx from "mobx"
 import * as React from "react"
-
 import { useObserver } from "../src/useObserver"
 import { sleep } from "./utils"
-import { FinalizationRegistry } from "../src/utils/FinalizationRegistryWrapper"
-
-// @ts-ignore
 import gc from "expose-gc/function"
+import { observerFinalizationRegistry } from "../src/utils/observerFinalizationRegistry"
+
+if (typeof globalThis.FinalizationRegistry !== "function") {
+    throw new Error("This test must run with node >= 14")
+}
+
+expect(observerFinalizationRegistry).toBeInstanceOf(globalThis.FinalizationRegistry)
 
 afterEach(cleanup)
 
 test("uncommitted components should not leak observations", async () => {
-    if (!FinalizationRegistry) {
-        throw new Error("This test must run with node >= 14")
-    }
-
     const store = mobx.observable({ count1: 0, count2: 0 })
 
     // Track whether counts are observed

packages/mobx-react-lite/__tests__/strictAndConcurrentModeUsingTimers.test.tsx

@@ -2,21 +2,22 @@ import "./utils/killFinalizationRegistry"
 import { act, cleanup, render } from "@testing-library/react"
 import * as mobx from "mobx"
 import * as React from "react"
-
 import { useObserver } from "../src/useObserver"
 import {
-    forceCleanupTimerToRunNowForTests,
-    resetCleanupScheduleForTests
-} from "../src/utils/reactionCleanupTracking"
-import {
-    CLEANUP_LEAKED_REACTIONS_AFTER_MILLIS,
-    CLEANUP_TIMER_LOOP_MILLIS
-} from "../src/utils/reactionCleanupTrackingCommon"
+    REGISTRY_FINALIZE_AFTER,
+    REGISTRY_SWEEP_INTERVAL
+} from "../src/utils/UniversalFinalizationRegistry"
+import { observerFinalizationRegistry } from "../src/utils/observerFinalizationRegistry"
+import { TimerBasedFinalizationRegistry } from "../src/utils/UniversalFinalizationRegistry"
+
+expect(observerFinalizationRegistry).toBeInstanceOf(TimerBasedFinalizationRegistry)
+
+const registry = observerFinalizationRegistry as TimerBasedFinalizationRegistry<unknown>
 
 afterEach(cleanup)
 
 test("uncommitted components should not leak observations", async () => {
-    resetCleanupScheduleForTests()
+    registry.finalizeAllImmediately()
 
     // Unfortunately, Jest fake timers don't mock out Date.now, so we fake
     // that out in parallel to Jest useFakeTimers
@@ -51,7 +52,7 @@ test("uncommitted components should not leak observations", async () => {
     )
 
     // Allow any reaction-disposal cleanup timers to run
-    const skip = Math.max(CLEANUP_LEAKED_REACTIONS_AFTER_MILLIS, CLEANUP_TIMER_LOOP_MILLIS)
+    const skip = Math.max(REGISTRY_FINALIZE_AFTER, REGISTRY_SWEEP_INTERVAL)
     fakeNow += skip
     jest.advanceTimersByTime(skip)
 
@@ -72,7 +73,7 @@ test("cleanup timer should not clean up recently-pended reactions", () => {
     // 5. The commit phase runs for component A, but reaction R2 has already been disposed. Game over.
 
     // This unit test attempts to replicate that scenario:
-    resetCleanupScheduleForTests()
+    registry.finalizeAllImmediately()
 
     // Unfortunately, Jest fake timers don't mock out Date.now, so we fake
     // that out in parallel to Jest useFakeTimers
@@ -106,7 +107,7 @@ test("cleanup timer should not clean up recently-pended reactions", () => {
     // We force our cleanup loop to run even though enough time hasn't _really_
     // elapsed.  In theory, it won't do anything because not enough time has
     // elapsed since the reactions were queued, and so they won't be disposed.
-    forceCleanupTimerToRunNowForTests()
+    registry.sweep()
 
     // Advance time enough to allow any timer-queued effects to run
     jest.advanceTimersByTime(500)
@@ -137,7 +138,7 @@ test.skip("component should recreate reaction if necessary", () => {
 
     // This unit test attempts to replicate that scenario:
 
-    resetCleanupScheduleForTests()
+    registry.finalizeAllImmediately()
 
     // Unfortunately, Jest fake timers don't mock out Date.now, so we fake
     // that out in parallel to Jest useFakeTimers
@@ -166,9 +167,9 @@ test.skip("component should recreate reaction if necessary", () => {
     // and _then_ the component commits.
 
     // Force everything to be disposed.
-    const skip = Math.max(CLEANUP_LEAKED_REACTIONS_AFTER_MILLIS, CLEANUP_TIMER_LOOP_MILLIS)
+    const skip = Math.max(REGISTRY_FINALIZE_AFTER, REGISTRY_SWEEP_INTERVAL)
     fakeNow += skip
-    forceCleanupTimerToRunNowForTests()
+    registry.sweep()
 
     // The reaction should have been cleaned up.
     expect(countIsObserved).toBeFalsy()

packages/mobx-react-lite/src/index.ts

@@ -5,6 +5,7 @@ import { observerBatching } from "./utils/observerBatching"
 import { useDeprecated } from "./utils/utils"
 import { useObserver as useObserverOriginal } from "./useObserver"
 import { enableStaticRendering } from "./staticRendering"
+import { observerFinalizationRegistry } from "./utils/observerFinalizationRegistry"
 
 observerBatching(batch)
 
@@ -14,7 +15,8 @@ export { Observer } from "./ObserverComponent"
 export { useLocalObservable } from "./useLocalObservable"
 export { useLocalStore } from "./useLocalStore"
 export { useAsObservableSource } from "./useAsObservableSource"
-export { resetCleanupScheduleForTests as clearTimers } from "./utils/reactionCleanupTracking"
+
+export const clearTimes = observerFinalizationRegistry["finalizeAllImmediately"] ?? (() => {})
 
 export function useObserver<T>(fn: () => T, baseComponentName: string = "observed"): T {
     if ("production" !== process.env.NODE_ENV) {

packages/mobx-react-lite/src/useObserver.ts

@@ -1,17 +1,30 @@
 import { Reaction } from "mobx"
 import React from "react"
 import { printDebugValue } from "./utils/printDebugValue"
-import {
-    addReactionToTrack,
-    IReactionTracking,
-    recordReactionAsCommitted
-} from "./utils/reactionCleanupTracking"
+import { observerFinalizationRegistry } from "./utils/observerFinalizationRegistry"
 import { isUsingStaticRendering } from "./staticRendering"
 
 function observerComponentNameFor(baseComponentName: string) {
     return `observer${baseComponentName}`
 }
 
+type ObserverAdministration = {
+    /** The Reaction created during first render, which may be leaked */
+    reaction: Reaction | null
+
+    /**
+     * Whether the component has yet completed mounting (for us, whether
+     * its useEffect has run)
+     */
+    mounted: boolean
+
+    /**
+     * Whether the observables that the component is tracking changed between
+     * the first render and the first useEffect.
+     */
+    changedBeforeMount: boolean
+}
+
 /**
  * We use class to make it easier to detect in heap snapshots by name
  */
@@ -34,50 +47,51 @@ export function useObserver<T>(fn: () => T, baseComponentName: string = "observe
     // StrictMode/ConcurrentMode/Suspense may mean that our component is
     // rendered and abandoned multiple times, so we need to track leaked
     // Reactions.
-    const reactionTrackingRef = React.useRef<IReactionTracking | null>(null)
+    const admRef = React.useRef<ObserverAdministration | null>(null)
+
+    if (!admRef.current) {
+        // First render
+        admRef.current = {
+            reaction: null,
+            mounted: false,
+            changedBeforeMount: false
+        }
+    }
 
-    if (!reactionTrackingRef.current) {
-        // First render for this component (or first time since a previous
-        // reaction from an abandoned render was disposed).
+    const adm = admRef.current!
 
-        const newReaction = new Reaction(observerComponentNameFor(baseComponentName), () => {
+    if (!adm.reaction) {
+        // First render or component was not committed and reaction was disposed by registry
+        adm.reaction = new Reaction(observerComponentNameFor(baseComponentName), () => {
             // Observable has changed, meaning we want to re-render
             // BUT if we're a component that hasn't yet got to the useEffect()
             // stage, we might be a component that _started_ to render, but
             // got dropped, and we don't want to make state changes then.
             // (It triggers warnings in StrictMode, for a start.)
-            if (trackingData.mounted) {
+            if (adm.mounted) {
                 // We have reached useEffect(), so we're mounted, and can trigger an update
                 forceUpdate()
             } else {
                 // We haven't yet reached useEffect(), so we'll need to trigger a re-render
                 // when (and if) useEffect() arrives.
-                trackingData.changedBeforeMount = true
+                adm.changedBeforeMount = true
             }
         })
 
-        const trackingData = addReactionToTrack(
-            reactionTrackingRef,
-            newReaction,
-            objectRetainedByReact
-        )
+        observerFinalizationRegistry.register(objectRetainedByReact, adm, adm)
     }
 
-    const { reaction } = reactionTrackingRef.current!
-    React.useDebugValue(reaction, printDebugValue)
+    React.useDebugValue(adm.reaction, printDebugValue)
 
     React.useEffect(() => {
-        // Called on first mount only
-        recordReactionAsCommitted(reactionTrackingRef)
-
-        if (reactionTrackingRef.current) {
-            // Great. We've already got our reaction from our render;
-            // all we need to do is to record that it's now mounted,
-            // to allow future observable changes to trigger re-renders
-            reactionTrackingRef.current.mounted = true
-            // Got a change before first mount, force an update
-            if (reactionTrackingRef.current.changedBeforeMount) {
-                reactionTrackingRef.current.changedBeforeMount = false
+        observerFinalizationRegistry.unregister(adm)
+
+        adm.mounted = true
+
+        if (adm.reaction) {
+            if (adm.changedBeforeMount) {
+                // Got a change before mount, force an update
+                adm.changedBeforeMount = false
                 forceUpdate()
             }
         } else {
@@ -87,21 +101,18 @@ export function useObserver<T>(fn: () => T, baseComponentName: string = "observe
             // reaction got cleaned up
 
             // Re-create the reaction
-            reactionTrackingRef.current = {
-                reaction: new Reaction(observerComponentNameFor(baseComponentName), () => {
-                    // We've definitely already been mounted at this point
-                    forceUpdate()
-                }),
-                mounted: true,
-                changedBeforeMount: false,
-                cleanAt: Infinity
-            }
+            adm.reaction = new Reaction(observerComponentNameFor(baseComponentName), () => {
+                // We've definitely already been mounted at this point
+                forceUpdate()
+            })
             forceUpdate()
         }
 
         return () => {
-            reactionTrackingRef.current!.reaction.dispose()
-            reactionTrackingRef.current = null
+            adm.reaction!.dispose()
+            adm.reaction = null
+            adm.mounted = false
+            adm.changedBeforeMount = false
         }
     }, [])
 
@@ -110,7 +121,7 @@ export function useObserver<T>(fn: () => T, baseComponentName: string = "observe
     // can be invalidated (see above) once a dependency changes
     let rendering!: T
     let exception
-    reaction.track(() => {
+    adm.reaction.track(() => {
         try {
             rendering = fn()
         } catch (e) {

packages/mobx-react-lite/src/utils/UniversalFinalizationRegistry.ts

@@ -0,0 +1,62 @@
+export declare class FinalizationRegistryType<T> {
+    constructor(finalize: (value: T) => void)
+    register(target: object, value: T, token?: object): void
+    unregister(token: object): void
+}
+
+declare const FinalizationRegistry: typeof FinalizationRegistryType | undefined
+
+export const REGISTRY_FINALIZE_AFTER = 10_000
+export const REGISTRY_SWEEP_INTERVAL = 10_000
+
+export class TimerBasedFinalizationRegistry<T> implements FinalizationRegistryType<T> {
+    private registrations: Map<unknown, { value: T; registeredAt: number }> = new Map()
+    private sweepTimeout: ReturnType<typeof setTimeout> | undefined
+
+    constructor(private readonly finalize: (value: T) => void) {}
+
+    // Token is actually required with this impl
+    register(target: object, value: T, token?: object) {
+        this.registrations.set(token, {
+            value,
+            registeredAt: Date.now()
+        })
+        this.scheduleSweep()
+    }
+
+    unregister(token: unknown) {
+        this.registrations.delete(token)
+    }
+
+    // Bound so it can be used directly as setTimeout callback.
+    sweep = (maxAge = REGISTRY_FINALIZE_AFTER) => {
+        // cancel timeout so we can force sweep anytime
+        clearTimeout(this.sweepTimeout)
+        this.sweepTimeout = undefined
+
+        const now = Date.now()
+        this.registrations.forEach((registration, token) => {
+            if (now - registration.registeredAt >= maxAge) {
+                this.finalize(registration.value)
+                this.registrations.delete(token)
+            }
+        })
+
+        if (this.registrations.size > 0) {
+            this.scheduleSweep()
+        }
+    }
+
+    // Bound so it can be exported directly as clearTimers test utility.
+    finalizeAllImmediately = () => {
+        this.sweep(0)
+    }
+
+    private scheduleSweep() {
+        if (this.sweepTimeout === undefined) {
+            this.sweepTimeout = setTimeout(this.sweep, REGISTRY_SWEEP_INTERVAL)
+        }
+    }
+}
+
+export const UniversalFinalizationRegistry = FinalizationRegistry ?? TimerBasedFinalizationRegistry