fix(react): defer ExternalNode creation until useEffect runs (#321)

* fix(react): defer ExternalNode creation until useEffect runs

* update snapshots

Author: bowheart

packages/react/src/hooks/useAtomInstance.ts

@@ -9,6 +9,7 @@ import {
   ParamsOf,
   Selectable,
   SelectorInstance,
+  ZeduxNode,
   zi,
 } from '@zedux/atoms'
 import { Dispatch, SetStateAction, useEffect, useState } from 'react'
@@ -22,6 +23,9 @@ import {
 import { useEcosystem } from './useEcosystem'
 import { useReactComponentId } from './useReactComponentId'
 
+const unmaterializedNodes = new Set<ZeduxNode>()
+let isQueued = false
+
 /**
  * Creates an atom instance for the passed atom template based on the passed
  * params. If an instance has already been created for the passed params, reuses
@@ -116,31 +120,13 @@ export const useAtomInstance: {
   const renderedValue = instance.v
   const isSelector = is(instance, SelectorInstance)
 
-  if (isSelector) render.i = instance as SelectorInstance
-
-  let node =
-    (ecosystem.n.get(observerId) as ExternalNode) ??
-    new ExternalNode(ecosystem, observerId, render)
-
-  const addEdge = () => {
-    node.l === zi.D && (node = new ExternalNode(ecosystem, observerId, render))
-
-    // cancel edge cleanup if the below effect cleanup ran and scheduled it but
-    // the component rerendered or the effect ran again before it happened
-    node.c?.()
-
-    if (node.i !== instance) {
-      node.u(
-        instance,
-        operation,
-        External | (subscribe ? Eventless : EventlessStatic)
-      )
+  if (isSelector) {
+    if (!render.i && !instance.o.size) {
+      unmaterializedNodes.add(instance)
     }
-  }
 
-  // Yes, subscribe during render. This operation is idempotent and we handle
-  // React's StrictMode specifically.
-  addEdge()
+    render.i = instance as SelectorInstance
+  }
 
   // Only remove the graph edge when the instance id changes or on component
   // destruction.
@@ -153,11 +139,41 @@ export const useAtomInstance: {
 
     if (isSelector) render.i = instance as SelectorInstance
 
+    let node =
+      (ecosystem.n.get(observerId) as ExternalNode) ??
+      new ExternalNode(ecosystem, observerId, render)
+
+    const addEdge = () => {
+      node.l === zi.D &&
+        (node = new ExternalNode(ecosystem, observerId, render))
+
+      // cancel edge cleanup if the below effect cleanup ran and scheduled it but
+      // the component rerendered or the effect ran again before it happened
+      node.c?.()
+
+      if (node.i !== instance) {
+        node.u(
+          instance,
+          operation,
+          External | (subscribe ? Eventless : EventlessStatic)
+        )
+      }
+    }
+
     // Try adding the edge again (will be a no-op unless React's StrictMode ran
     // this effect's cleanup unnecessarily)
     addEdge()
     render.m = true
 
+    if (!isQueued) {
+      isQueued = true
+      ecosystem.asyncScheduler.queue(() => {
+        isQueued = false
+        unmaterializedNodes.forEach(node => node.destroy())
+        unmaterializedNodes.clear()
+      })
+    }
+
     // an unmounting component's effect cleanup can update or force-destroy the
     // atom instance before this component is mounted. If that happened, trigger
     // a rerender to recreate the atom instance and/or get its new state
@@ -172,7 +188,7 @@ export const useAtomInstance: {
       node.k(instance, true)
       // don't set `render.m = false` here
     }
-  }, [instance.id])
+  }, [instance])
 
   if (suspend !== false) {
     const status = (instance as AtomInstance).promiseStatus

packages/react/test/integrations/suspense.test.tsx

@@ -3,12 +3,13 @@ import {
   atom,
   injectAtomValue,
   injectPromise,
+  injectSignal,
   StateOf,
   useAtomInstance,
   useAtomState,
   useAtomValue,
 } from '@zedux/react'
-import React, { Suspense } from 'react'
+import React, { Suspense, useState } from 'react'
 import { ErrorBoundary } from '../utils/ErrorBoundary'
 import { renderInEcosystem } from '../utils/renderInEcosystem'
 import { mockConsole } from '../utils/console'
@@ -330,4 +331,83 @@ describe('suspense', () => {
 
     expect(calls).toEqual([undefined, undefined, 1, 1])
   })
+
+  test('suspending component does not create extra ExternalNodes or register extra observers', async () => {
+    const testSuspenseAtom = atom(
+      'testSuspenseAtom',
+      () => {
+        const signal = injectSignal({ loaded: false })
+
+        const { promise } = injectPromise(async () => {
+          await new Promise(resolve => setTimeout(resolve, 500))
+          signal.set({ loaded: true })
+        }, [])
+
+        return api(signal).setPromise(promise)
+      },
+      { ttl: 0 }
+    )
+
+    const TestComponent: React.FC = () => {
+      const instance = useAtomInstance(testSuspenseAtom)
+
+      return (
+        <div data-testid="mounted">
+          <p>Atom ID: {instance.id}</p>
+        </div>
+      )
+    }
+
+    const TestWrapper: React.FC = () => {
+      const [showComponent, setShowComponent] = useState(false)
+
+      return (
+        <>
+          <button
+            data-testid="toggle"
+            onClick={() => setShowComponent(prev => !prev)}
+          >
+            {showComponent ? 'Hide Component' : 'Show Component'}
+          </button>
+
+          {showComponent && (
+            <Suspense fallback={<div data-testid="loading">Loading...</div>}>
+              <TestComponent />
+            </Suspense>
+          )}
+        </>
+      )
+    }
+
+    const { findByTestId, queryByTestId, getByTestId } = renderInEcosystem(
+      <TestWrapper />
+    )
+
+    expect(ecosystem.find(testSuspenseAtom)).toBeUndefined()
+
+    act(() => {
+      getByTestId('toggle').click()
+    })
+
+    await findByTestId('loading')
+    await findByTestId('mounted')
+
+    expect(ecosystem.find(testSuspenseAtom)?.o.size).toBe(1)
+
+    act(() => {
+      getByTestId('toggle').click()
+    })
+
+    expect(queryByTestId('mounted')).toBeNull()
+
+    // atom should be destroyed when component is unmounted
+    await act(async () => {
+      await Promise.resolve()
+    })
+
+    expect(ecosystem.find(testSuspenseAtom)).toBeUndefined()
+  })
+
+  // TODO: Would be nice if we could support this. It's really out of React's realm of practicality:
+  // test('floating suspense nodes are cleaned up when the component finally mounts', () => {})
 })

packages/react/test/units/__snapshots__/useAtomSelector.test.tsx.snap

@@ -115,7 +115,7 @@ exports[`useAtomSelector inline selector that returns a different object referen
   "1": {
     "className": "AtomInstance",
     "observers": {
-      "@selector(unknown)-1": {
+      "@selector(unknown)-2": {
         "flags": 1,
         "operation": "get",
         "p": undefined,
@@ -132,7 +132,7 @@ exports[`useAtomSelector inline selector that returns a different object referen
     "className": "ExternalNode",
     "observers": {},
     "sources": {
-      "@selector(unknown)-1": {
+      "@selector(unknown)-2": {
         "flags": 5,
         "operation": "useAtomValue",
       },
@@ -141,7 +141,7 @@ exports[`useAtomSelector inline selector that returns a different object referen
     "status": "Active",
     "weight": 3,
   },
-  "@selector(unknown)-1": {
+  "@selector(unknown)-2": {
     "className": "SelectorInstance",
     "observers": {
       "@component(Test)-:rb:": {
@@ -170,7 +170,7 @@ exports[`useAtomSelector inline selector that returns a different object referen
   "1": {
     "className": "AtomInstance",
     "observers": {
-      "@selector(unknown)-1": {
+      "@selector(unknown)-2": {
         "flags": 1,
         "operation": "get",
         "p": undefined,
@@ -187,7 +187,7 @@ exports[`useAtomSelector inline selector that returns a different object referen
     "className": "ExternalNode",
     "observers": {},
     "sources": {
-      "@selector(unknown)-1": {
+      "@selector(unknown)-2": {
         "flags": 5,
         "operation": "useAtomValue",
       },
@@ -196,7 +196,7 @@ exports[`useAtomSelector inline selector that returns a different object referen
     "status": "Active",
     "weight": 3,
   },
-  "@selector(unknown)-1": {
+  "@selector(unknown)-2": {
     "className": "SelectorInstance",
     "observers": {
       "@component(Test)-:rb:": {
@@ -222,7 +222,7 @@ exports[`useAtomSelector inline selector that returns a different object referen
 
 exports[`useAtomSelector inline selectors are swapped out in strict mode double-renders 1`] = `
 {
-  "@selector(unknown)-1": {
+  "@selector(unknown)-2": {
     "className": "SelectorInstance",
     "observers": {
       "@component(Test)-:r9:": {
@@ -246,7 +246,7 @@ exports[`useAtomSelector inline selectors are swapped out in strict mode double-
 
 exports[`useAtomSelector inline selectors are swapped out in strict mode double-renders 2`] = `
 {
-  "@selector(unknown)-1": {
+  "@selector(unknown)-2": {
     "className": "SelectorInstance",
     "observers": {
       "@component(Test)-:r9:": {