Async atoms dependent on sync atoms sometimes behave as if mounted when pending

[a10nik]

Bug Description

I believe this issue is connected to getMountedOrPendingDependents(), the one with a // TODO(daishi): revisit this implementation remark.

It's used in recomputeInvalidatedAtoms and here is how it works now on a totally unmounted atom set: it searches for parents of changed atoms that are pending (none are mounted) and triggers their updates.

I would expect it to behave without any regard to pending statuses at all, just as a sync version would: to only trigger a dependant if it's mounted.

Not sure if it was intentional, but ignoring the pending dependants in recomputeInvalidatedAtoms indeed breaks some tests such as dependency.test.tsx > refreshes deps for each async read. But frankly, the broken tests' logic seems dubious itself:

• it relies on the idea that an async atom would behave eagerly while pending, as if mounted
• the tests seem fragile and tend to break anyways if a sync dependency is wrapped into a harmless-looking atom(get => get(syncDep))

Reproduction Link

https://stackblitz.com/edit/node-3bhhnxpm?file=index.test.ts&terminal=test

[dmaskasky]

The behavior does seem inconsistent. I think it should not matter if a sync atom is wrapped. If we want to treat pending dependents as mounted, then the first test should also fail, getMountedOrPendingDependents should look up ancestors and not stop at the parent.

[a10nik]

If we want to treat pending dependents as mounted, then the first test should also fail, getMountedOrPendingDependents should look up ancestors and not stop at the parent.

True. Though I'd prefer it to be more consistent with a sync version instead. It feels weird that the extra call of an async atom getter -- a fetch request usually -- depends on whether a previous request has ended (results in no extra requests) or not (triggers extra request without anyone listening to the result)

[dmaskasky]

There is signal available for aborting previous runs.

const asyncAtom = atom((get, { signal }) => { ... })

[dmaskasky]

Update: I had missed the unsub when I gave my previous response.

I agree with the behavior described in these tests - and your conclusion. I'd like to get @dai-shi's thoughts here on whether the pending promise keeps an atom from being unmounted.

import { atom, createStore } from 'jotai';
import { expect, it, vi } from 'vitest';

it('sub to asyncAtom -> syncAtomWrapper -> syncAtom++', async () => {
  const store = createStore();
  const callback = vi.fn();
  const syncAtom = atom(0);
  const syncAtomWrapper = atom((get) => get(syncAtom));
  const asyncAtom = atom((get) => {
    callback();
    get(syncAtomWrapper);
    return new Promise((res) => setTimeout(res));
  });

  const unsub = store.sub(asyncAtom, () => {});
  unsub();
  expect(callback).toHaveBeenCalledTimes(1);

  store.set(syncAtom, 1);
  expect(callback).toHaveBeenCalledTimes(1);
});

it('sub to syncAtom1 -> syncAtom2++', async () => {
  const store = createStore();
  const callback = vi.fn();
  const syncAtom1 = atom(0);
  const syncAtom2 = atom((get) => {
    get(syncAtom1);
    callback();
  });

  const unsub = store.sub(syncAtom2, () => {});
  unsub();
  expect(callback).toHaveBeenCalledTimes(1);

  store.set(syncAtom1, 1);
  expect(callback).toHaveBeenCalledTimes(1);
});

it('sub to asyncAtom -> syncAtom++', async () => {
  const store = createStore();
  const callback = vi.fn();
  const syncAtom = atom(0);
  const asyncAtom = atom((get) => {
    get(syncAtom);
    callback();
    return new Promise((res) => setTimeout(res));
  });

  const unsub = store.sub(asyncAtom, () => {});
  unsub();
  expect(callback).toHaveBeenCalledTimes(1);

  store.set(syncAtom, 1);
  expect(callback).toHaveBeenCalledTimes(1);
});

[a10nik]

There is signal available for aborting previous runs.

There is, but I don't see how it's gonna help with preventing a second call of async getter. I understand that can make the first call promise stop when the signal is aborted with some Promise.race([mainPromise, new Promise((res) => signal.addEventListener("abort", res))]). However that won't happen until the next call has already started. I'll be able to cancel the first request however I see no way to stop the second.

UPD: Saw you newer reply, the explanation above seems unnecessary now :) But I'll leave the rest for the sake of ideas mentioned

Abort is currently called during an atom read after the new promise is started and only if a previous value is in the promiseStateMap and its promise is still pending. It's not called when atom in unsubscribed if that's what you mean. And even if it was the case -- or we were to emulate it somehow via an onMount cleanup -- it would't be until the next microtask that the changes should be reflected on promiseStateMap anyway.

The only 2 solutions I came up with so far were:

• to wrap a sync atom -- apparently, that is not gonna stay for long, since nobody really likes the behaviour so far -- you and me included :)
• define an isMountedAtom, set it in onMount and throw or return something no-op once it's false. But the code seems unnecessary and introduces an extra NOT_MOUNTED symbol in the return signature on top of that

const isMountedAtom = atom(false)
const asyncAtom = atom(async (get) => {
    if (!get(isMountedAtom)) {
        return NOT_MOUNTED;
    }
    return await something()
}, (get, set) => {
   set(isMountedAtom, true);
   return () => {
       set(isMountedAtom, false);
   }
});
asyncAtom.onMount = (setSelf) => setSelf();

[a10nik]

whether the pending promise keeps an atom from being unmounted

Not dai-shi myself of course :) but I have checked the scenarios in my tests in debugger and the mountedMap stays consistently empty, so it doesn't seem to be about not unmounting per se. More like recomputeInvalidatedAtoms treating mounted and pending atoms equally -- as getMountedOrPendingDependents suggests :)

I have actually tried to disable the logic of getMountedOrPendingDependents in recomputeInvalidatedAtoms by something ugly like that:

  const recomputeInvalidatedAtoms =
...
-        for (const d of getMountedOrPendingDependents(a, aState, mountedMap)) {
+        for (const d of getMountedOrPendingDependents(a, aState, mountedMap, true)) { // <<< temporary flag to test the idea
...

const getMountedOrPendingDependents = (
  atom: AnyAtom,
  atomState: AtomState,
  mountedMap: MountedMap,
+  recomputeInvalidatedAtoms?: true // <<< temporary flag to test the idea
): Set<AnyAtom> => {
  const dependents = new Set<AnyAtom>()
  for (const a of mountedMap.get(atom)?.t || []) {
    if (mountedMap.has(a)) {
      dependents.add(a)
    }
  }
  for (const atomWithPendingPromise of atomState.p) {
+      if (!recomputeInvalidatedAtoms) // <<< added a condition here
          dependents.add(atomWithPendingPromise) 
  }
  return dependents
}

That made my 3 tests pass alright, although several of existing jotai tests broke.

[dmaskasky]

Maybe this would resolve the concern: #3125

[a10nik]

That does indeed fix the unsub() case I was after. Still it leaves out some discrepancies that come with the current recomputeInvalidatedAtoms+getMountedOrPendingDependents behaviour. Can actually open another discussion for that if you'd like.

Here're another 3 similar tests on the PR's jotai version.

First 2 are about getting an atom without subbing-unsubbing at all. That's where wrapping an atom still matters during a pending stage.

The 3rd test is about another discrepance introduced between these cases:

• an async atom is gotten, then while it's still pending it's mounted and immediately unmounted (it turns off its eager reactivity the moment it's unmounted)
• an async atom is gotten and not mounted at all (it stays eagerly reactive until the end of a promise)

[dmaskasky]

I will be out in the wilderness for a few days. I can take another look next week.

[dai-shi]

Sorry, late to the party.

Sounds like we are hitting a tricky situation. I'm fine to drop some existing tests if we agree that the fix is more appropriate.
My general design principle is:

• Logic should be simple and so is the implementation. (even though the behavior is a little unintuitive.)
• Inconsistency should be avoided.
• I would hope Jotai core shouldn't care much about promises. But, we can't stop recomputing pending promises at all, can we?
• If the change is big (affects many users) but we should make it, it will be v3.

I'm not sure if it helps though.

---

I'd like to get @dai-shi's thoughts here on whether the pending promise keeps an atom from being unmounted.

The reason we don't mount a new promise while pending is to avoid memory leaks.
But, I'm not sure it still holds with our current implementation. Now, mountedMap is a WeakMap. And, we have memory leak tests to make it sure.
So, based on my principle, if we can mount all promises, it would feel better.

[a10nik]

But, we can't stop recomputing pending promises at all, can we?

Could you pls clarify why not? I mean when an atom is not mounted (or is unmounted) but still has a pending promise

Is it about Suspense quirks I've described in #3125 (comment)?

[dai-shi]

I'm not sure if we are on the same page, probably because I'm not following the discussion in details.

What I mean is that in our code, currently we have a function getMountedOrPendingDependents. But, ideally it should be getMountedDependents. I thought there was a reason to support some use cases. (And if we can drop it, it'll simplify.)

[dai-shi]

This is what I mean. Does it fix your problem?

diff --git a/src/vanilla/internals.ts b/src/vanilla/internals.ts
index 3923229..9abb919 100644
--- a/src/vanilla/internals.ts
+++ b/src/vanilla/internals.ts
@@ -243,7 +243,7 @@ const setAtomStateValueOrPromise = (
   const prevValue = atomState.v
   if (isPromiseLike(valueOrPromise)) {
     for (const a of atomState.d.keys()) {
-      addPendingPromiseToDependency(atom, valueOrPromise, ensureAtomState(a))
+      //addPendingPromiseToDependency(atom, valueOrPromise, ensureAtomState(a))
     }
   }
   atomState.v = valueOrPromise
@@ -601,7 +601,7 @@ const buildStore = (
         } finally {
           atomState.d.set(a, aState.n)
           if (isPendingPromise(atomState.v)) {
-            addPendingPromiseToDependency(atom, atomState.v, aState)
+            //addPendingPromiseToDependency(atom, atomState.v, aState)
           }
           mountedMap.get(a)?.t.add(atom)
           if (!isSync) {

[dmaskasky]

Do we want to support this use-case?

const fn = vi.fn()
const atomA = atom(0)
const asyncAtom = atom(async (get) => {
  fn()
  await delay(2000)
  get(atomA)
})
store.sub(asyncAtom, () => {})
await delay(3000)
fn.mockClear()
store.set(atomA, 1)
expect(fn).toHaveBeenCalled()

[dai-shi]

I this setting atomA to 1 while asyncAtom is pending?
Or, is this about not reading asyncAtom?

I think at least this should work:

const fn = vi.fn()
const atomA = atom(0)
const asyncAtom = atom(async (get) => {
  fn()
  await delay(2000)
  get(atomA)
})
store.sub(asyncAtom, () => {})
store.get(asyncAtom)
await delay(3000)
toHaveBeenCalledTimes(1)
store.set(atomA, 1)
toHaveBeenCalledTimes(2)