fix: should update dependents with the value of the unwrapped atom when the promise resolves

src/vanilla/store.ts

@@ -342,11 +342,6 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     // Compute a new state for this atom.
     atomState.d.clear()
     let isSync = true
-    const mountDependenciesIfAsync = () => {
-      if (atomState.m) {
-        runWithTransaction(() => mountDependencies(atom, atomState))
-      }
-    }
     const getter: Getter = <V>(a: Atom<V>) => {
       if (isSelfAtom(atom, a)) {
         const aState = ensureAtomState(a)
@@ -366,8 +361,8 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
         return returnAtomValue(aState)
       } finally {
         addDependency(atom, atomState, a, aState)
-        if (!isSync) {
-          mountDependenciesIfAsync()
+        if (!isSync && atomState.m) {
+          runWithTransaction(() => mountDependencies(atom, atomState))
         }
       }
     }
@@ -405,7 +400,6 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
       setAtomStateValueOrPromise(atom, atomState, valueOrPromise)
       if (isPromiseLike(valueOrPromise)) {
         valueOrPromise.onCancel?.(() => controller?.abort())
-        valueOrPromise.then(mountDependenciesIfAsync, mountDependenciesIfAsync)
       }
       return atomState
     } catch (error) {
@@ -565,7 +559,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
   ): Result => runWithTransaction(() => writeAtomState(atom, ...args))
 
   const mountDependencies = (atom: AnyAtom, atomState: AtomState) => {
-    if (atomState.m && !isPendingPromise(atomState.v)) {
+    if (atomState.m) {
       for (const a of atomState.d.keys()) {
         if (!atomState.m.d.has(a)) {
           const aMounted = mountAtom(a, ensureAtomState(a))

tests/vanilla/dependency.test.tsx

@@ -94,7 +94,7 @@ it('correctly handles the same promise being returned twice from an atom getter
   await expect(store.get(derivedAtom)).resolves.toBe('Asynchronous Data')
 })
 
-it('keeps atoms mounted between recalculations', async () => {
+it('keeps sync atoms mounted between recalculations', async () => {
   const metrics1 = {
     mounted: 0,
     unmounted: 0,
@@ -148,9 +148,11 @@ it('keeps atoms mounted between recalculations', async () => {
     mounted: 1,
     unmounted: 0,
   })
+  await Promise.resolve()
+  // async atom has been re-mounted
   expect(metrics2).toEqual({
-    mounted: 1,
-    unmounted: 0,
+    mounted: 2,
+    unmounted: 1,
   })
 })
 
@@ -425,3 +427,96 @@ it('batches sync writes', () => {
   expect(fetch).toBeCalledWith(1)
   expect(store.get(a)).toBe(1)
 })
+
+it('mounts and unmounts sync and async dependencies correctly', async () => {
+  const mounted = [0, 0, 0, 0, 0] as [number, number, number, number, number]
+  const a = atom(0)
+  a.debugLabel = 'a'
+  a.onMount = () => {
+    ++mounted[0]
+    return () => {
+      --mounted[0]
+    }
+  }
+
+  const b = atom(0)
+  b.debugLabel = 'b'
+  b.onMount = () => {
+    ++mounted[1]
+    return () => {
+      --mounted[1]
+    }
+  }
+
+  const c = atom(0)
+  c.debugLabel = 'c'
+  c.onMount = () => {
+    ++mounted[2]
+    return () => {
+      --mounted[2]
+    }
+  }
+
+  const d = atom(0)
+  d.debugLabel = 'd'
+  d.onMount = () => {
+    ++mounted[3]
+    return () => {
+      --mounted[3]
+    }
+  }
+
+  const e = atom(0)
+  e.debugLabel = 'e'
+  e.onMount = () => {
+    ++mounted[4]
+    return () => {
+      --mounted[4]
+    }
+  }
+
+  let resolve: (() => Promise<void>) | undefined
+  const f = atom((get) => {
+    if (!get(a)) {
+      get(b)
+    } else {
+      get(c)
+    }
+    const promise = new Promise<void>((r) => {
+      resolve = () => {
+        r()
+        return promise
+      }
+    }).then(() => {
+      if (!get(a)) {
+        get(d)
+      } else {
+        get(e)
+      }
+    })
+    return promise
+  })
+  f.debugLabel = 'f'
+
+  const store = createStore()
+  // mount a, b synchronously
+  const unsub = store.sub(f, () => {})
+  expect(mounted).toEqual([1, 1, 0, 0, 0])
+
+  // mount d asynchronously
+  await resolve!()
+  expect(mounted).toEqual([1, 1, 0, 1, 0])
+
+  // unmount b, mount c synchronously
+  // d is also unmounted
+  store.set(a, 1)
+  expect(mounted).toEqual([1, 0, 1, 0, 0])
+
+  // mount e asynchronously
+  await resolve!()
+  expect(mounted).toEqual([1, 0, 1, 0, 1])
+
+  // unmount a, b, d, e synchronously
+  unsub()
+  expect(mounted).toEqual([0, 0, 0, 0, 0])
+})

tests/vanilla/memoryleaks.test.ts

@@ -3,7 +3,7 @@ import { describe, expect, it } from 'vitest'
 import { atom, createStore } from 'jotai/vanilla'
 import type { Atom } from 'jotai/vanilla'
 
-describe('test memory leaks (get & set only)', () => {
+describe('memory leaks (get & set only)', () => {
   it('one atom', async () => {
     const store = createStore()
     let objAtom: Atom<object> | undefined = atom({})
@@ -48,7 +48,7 @@ describe('test memory leaks (get & set only)', () => {
   })
 })
 
-describe('test memory leaks (with subscribe)', () => {
+describe('memory leaks (with subscribe)', () => {
   it('one atom', async () => {
     const store = createStore()
     let objAtom: Atom<object> | undefined = atom({})
@@ -91,3 +91,63 @@ describe('test memory leaks (with subscribe)', () => {
     expect(await detector.isLeaking()).toBe(false)
   })
 })
+
+describe('memory leaks (with dependencies)', () => {
+  it('sync dependency', async () => {
+    const store = createStore()
+    let objAtom: Atom<object> | undefined = atom({})
+    const detector = new LeakDetector(store.get(objAtom))
+    const atom1 = atom(0)
+    const atom2 = atom((get) => get(atom1) || (objAtom && get(objAtom)))
+    store.sub(atom2, () => {})
+    store.set(atom1, 1)
+    objAtom = undefined
+    expect(await detector.isLeaking()).toBe(false)
+  })
+
+  it('async dependency', async () => {
+    const store = createStore()
+    let objAtom: Atom<object> | undefined = atom({})
+    const detector = new LeakDetector(store.get(objAtom))
+    const atom1 = atom(0)
+    const atom2 = atom(async (get) => get(atom1) || (objAtom && get(objAtom)))
+    store.sub(atom2, () => {})
+    store.set(atom1, 1)
+    objAtom = undefined
+    await Promise.resolve()
+    expect(await detector.isLeaking()).toBe(false)
+  })
+
+  it('async await dependency', async () => {
+    const store = createStore()
+    let objAtom: Atom<object> | undefined = atom({})
+    const detector = new LeakDetector(store.get(objAtom))
+    const atom1 = atom(0)
+    const atom2 = atom(async (get) => {
+      await Promise.resolve()
+      return get(atom1) || (objAtom && get(objAtom))
+    })
+    store.sub(atom2, () => {})
+    store.set(atom1, 1)
+    objAtom = undefined
+    expect(await detector.isLeaking()).toBe(false)
+  })
+
+  it('infinite async dependency', async () => {
+    const store = createStore()
+    let objAtom: Atom<object> | undefined = atom({})
+    const detector = new LeakDetector(store.get(objAtom))
+    const atom1 = atom(0)
+    const atom2 = atom(async (get) => {
+      if (get(atom1)) {
+        return new Promise(() => {})
+      }
+      return objAtom && get(objAtom)
+    })
+    store.sub(atom2, () => {})
+    store.set(atom1, 1)
+    objAtom = undefined
+    await Promise.resolve()
+    expect(await detector.isLeaking()).toBe(false)
+  })
+})

tests/vanilla/utils/unwrap.test.ts

@@ -150,3 +150,17 @@ describe('unwrap', () => {
     expect(store.get(syncAtom)).toEqual('concrete')
   })
 })
+
+it('should update dependents with the value of the unwrapped atom when the promise resolves', async () => {
+  const store = createStore()
+  const asyncTarget = atom(() => Promise.resolve('value'))
+  const target = unwrap(asyncTarget)
+  const results: string[] = []
+  const derived = atom(async (get) => {
+    await Promise.resolve()
+    results.push('effect ' + get(target))
+  })
+  store.sub(derived, () => {})
+  await new Promise((r) => setTimeout(r))
+  expect(results).toEqual(['effect undefined', 'effect value'])
+})