fix(vanilla): deal with promise resolving race condition

src/vanilla/store.ts

@@ -259,48 +259,44 @@ export const createStore = () => {
   ): AtomState<Value> => {
     if (isPromiseLike(valueOrPromise)) {
       let continuePromise: (next: Promise<Awaited<Value>>) => void
+      const updatePromiseDependencies = () => {
+        const prevAtomState = getAtomState(atom)
+        if (
+          !prevAtomState ||
+          !hasPromiseAtomValue(prevAtomState) ||
+          prevAtomState.v !== promise
+        ) {
+          // not the latest promise
+          return
+        }
+        // update dependencies, that could have changed
+        const nextAtomState = setAtomValue(
+          atom,
+          promise as Value,
+          nextDependencies
+        )
+        if (mountedMap.has(atom) && prevAtomState.d !== nextAtomState.d) {
+          mountDependencies(atom, nextAtomState, prevAtomState.d)
+        }
+      }
       const promise: Promise<Awaited<Value>> & PromiseMeta<Awaited<Value>> =
         new Promise((resolve, reject) => {
           let settled = false
           valueOrPromise.then(
             (v) => {
               if (!settled) {
                 settled = true
-                const prevAtomState = getAtomState(atom)
-                // update dependencies, that could have changed
-                const nextAtomState = setAtomValue(
-                  atom,
-                  promise as Value,
-                  nextDependencies
-                )
                 resolvePromise(promise, v)
                 resolve(v as Awaited<Value>)
-                if (
-                  mountedMap.has(atom) &&
-                  prevAtomState?.d !== nextAtomState.d
-                ) {
-                  mountDependencies(atom, nextAtomState, prevAtomState?.d)
-                }
+                updatePromiseDependencies()
               }
             },
             (e) => {
               if (!settled) {
                 settled = true
-                const prevAtomState = getAtomState(atom)
-                // update dependencies, that could have changed
-                const nextAtomState = setAtomValue(
-                  atom,
-                  promise as Value,
-                  nextDependencies
-                )
                 rejectPromise(promise, e)
                 reject(e)
-                if (
-                  mountedMap.has(atom) &&
-                  prevAtomState?.d !== nextAtomState.d
-                ) {
-                  mountDependencies(atom, nextAtomState, prevAtomState?.d)
-                }
+                updatePromiseDependencies()
               }
             }
           )

tests/react/async2.test.tsx

@@ -1,7 +1,7 @@
 import { StrictMode, Suspense } from 'react'
-import { fireEvent, render } from '@testing-library/react'
-import { describe, it } from 'vitest'
-import { useAtomValue, useSetAtom } from 'jotai/react'
+import { fireEvent, render, waitFor } from '@testing-library/react'
+import { assert, describe, expect, it } from 'vitest'
+import { useAtom, useAtomValue, useSetAtom } from 'jotai/react'
 import { atom } from 'jotai/vanilla'
 
 describe('useAtom delay option test', () => {
@@ -141,3 +141,77 @@ describe('atom read function setSelf option test', () => {
     await findByText('text: hello1')
   })
 })
+
+describe('timing issue with setSelf', () => {
+  it('resolves dependencies reliably after a delay (#2192)', async () => {
+    expect.assertions(1)
+    const countAtom = atom(0)
+
+    let result: number | null = null
+    const resolve: (() => void)[] = []
+    const asyncAtom = atom(async (get) => {
+      const count = get(countAtom)
+      await new Promise<void>((r) => resolve.push(r))
+      return count
+    })
+
+    const derivedAtom = atom(
+      async (get, { setSelf }) => {
+        get(countAtom)
+        await Promise.resolve()
+        const resultCount = await get(asyncAtom)
+        result = resultCount
+        if (resultCount === 2) setSelf() // <-- necessary
+      },
+      () => {}
+    )
+
+    const derivedSyncAtom = atom((get) => {
+      get(derivedAtom)
+    })
+
+    const increment = (c: number) => c + 1
+    function TestComponent() {
+      useAtom(derivedSyncAtom)
+      const [count, setCount] = useAtom(countAtom)
+      const onClick = () => {
+        setCount(increment)
+        setCount(increment)
+      }
+      return (
+        <>
+          count: {count}
+          <button onClick={onClick}>button</button>
+        </>
+      )
+    }
+
+    const { getByText, findByText } = render(
+      <StrictMode>
+        <TestComponent />
+      </StrictMode>
+    )
+
+    await waitFor(() => assert(resolve.length === 1))
+    resolve[0]!()
+
+    // The use of fireEvent is required to reproduce the issue
+    fireEvent.click(getByText('button'))
+
+    await waitFor(() => assert(resolve.length === 3))
+    resolve[1]!()
+    resolve[2]!()
+
+    await waitFor(() => assert(result === 2))
+
+    // The use of fireEvent is required to reproduce the issue
+    fireEvent.click(getByText('button'))
+
+    await waitFor(() => assert(resolve.length === 5))
+    resolve[3]!()
+    resolve[4]!()
+
+    await findByText('count: 4')
+    expect(result).toBe(4) // 3
+  })
+})

tests/vanilla/store.test.tsx

@@ -1,4 +1,5 @@
-import { describe, expect, it, vi } from 'vitest'
+import { waitFor } from '@testing-library/dom'
+import { assert, describe, expect, it, vi } from 'vitest'
 import { atom, createStore } from 'jotai/vanilla'
 import type { Getter } from 'jotai/vanilla'
 
@@ -424,6 +425,58 @@ it('should update derived atoms during write (#2107)', async () => {
   expect(store.get(countAtom)).toBe(2)
 })
 
+it('resolves dependencies reliably after a delay (#2192)', async () => {
+  expect.assertions(1)
+  const countAtom = atom(0)
+  let result: number | null = null
+
+  const resolve: (() => void)[] = []
+  const asyncAtom = atom(async (get) => {
+    const count = get(countAtom)
+    await new Promise<void>((r) => resolve.push(r))
+    return count
+  })
+
+  const derivedAtom = atom(
+    async (get, { setSelf }) => {
+      get(countAtom)
+      await Promise.resolve()
+      result = await get(asyncAtom)
+      if (result === 2) setSelf() // <-- necessary
+    },
+    () => {}
+  )
+
+  const store = createStore()
+  store.sub(derivedAtom, () => {})
+
+  await waitFor(() => assert(resolve.length === 1))
+
+  resolve[0]!()
+  const increment = (c: number) => c + 1
+  store.set(countAtom, increment)
+  store.set(countAtom, increment)
+
+  await waitFor(() => assert(resolve.length === 3))
+
+  resolve[1]!()
+  resolve[2]!()
+  await waitFor(() => assert(result === 2))
+
+  store.set(countAtom, increment)
+  store.set(countAtom, increment)
+
+  await waitFor(() => assert(resolve.length === 5))
+
+  resolve[3]!()
+  resolve[4]!()
+
+  await new Promise(setImmediate)
+  await waitFor(() => assert(store.get(countAtom) === 4))
+
+  expect(result).toBe(4) // 3
+})
+
 it('should not recompute a derived atom value if unchanged (#2168)', async () => {
   const store = createStore()
   const countAtom = atom(1)