fix: async atom in the middle of dependency chain (pmndrs#149)

* add failing test * fix async atom in the middld of dep chain
breeffy · Oct 22, 2020 · 9a30e63 · 9a30e63
1 parent 5773cd3
commit 9a30e63
Show file tree

Hide file tree

Showing 3 changed files with 95 additions and 33 deletions.
diff --git a/.size-snapshot.json b/.size-snapshot.json
@@ -1,8 +1,8 @@
 {
   "index.js": {
-    "bundled": 18870,
-    "minified": 8470,
-    "gzipped": 3104,
+    "bundled": 19093,
+    "minified": 8544,
+    "gzipped": 3130,
     "treeshaked": {
       "rollup": {
         "code": 318,
@@ -14,14 +14,14 @@
     }
   },
   "index.cjs.js": {
-    "bundled": 22468,
-    "minified": 10473,
-    "gzipped": 3612
+    "bundled": 22735,
+    "minified": 10562,
+    "gzipped": 3638
   },
   "index.iife.js": {
-    "bundled": 23889,
-    "minified": 8630,
-    "gzipped": 3303
+    "bundled": 24162,
+    "minified": 8716,
+    "gzipped": 3343
   },
   "utils.js": {
     "bundled": 2687,

diff --git a/src/core/Provider.ts b/src/core/Provider.ts
@@ -168,9 +168,10 @@ const readAtomState = <Value>(
   atom: Atom<Value>,
   prevState: State,
   setState: Dispatch<(prev: State) => State>,
-  atomStateCache?: AtomStateCache // read state then cache
+  atomStateCache: AtomStateCache,
+  force?: boolean
 ) => {
-  if (atomStateCache) {
+  if (!force) {
     let atomState = mGet(prevState, atom) as AtomState<Value> | undefined
     if (atomState) {
       return [atomState, prevState] as const
@@ -260,7 +261,12 @@ const readAtomState = <Value>(
     }
   } catch (errorOrPromise) {
     if (errorOrPromise instanceof Promise) {
-      promise = errorOrPromise
+      promise = errorOrPromise.then(() => {
+        setState(
+          (prev) =>
+            readAtomState(atom, mDel(prev, atom), setState, atomStateCache)[1]
+        )
+      })
     } else if (errorOrPromise instanceof Error) {
       error = errorOrPromise
     } else {
@@ -291,7 +297,8 @@ const readAtomState = <Value>(
 const updateDependentsState = <Value>(
   atom: Atom<Value>,
   prevState: State,
-  setState: Dispatch<(prev: State) => State>
+  setState: Dispatch<(prev: State) => State>,
+  atomStateCache: AtomStateCache
 ) => {
   const atomState = mGet(prevState, atom)
   if (!atomState) {
@@ -312,16 +319,25 @@ const updateDependentsState = <Value>(
     const [dependentState, nextNextState] = readAtomState(
       dependent,
       nextState,
-      setState
+      setState,
+      atomStateCache,
+      true
     )
     const promise = dependentState.readP
     if (promise) {
       promise.then(() => {
-        setState((prev) => updateDependentsState(dependent, prev, setState))
+        setState((prev) =>
+          updateDependentsState(dependent, prev, setState, atomStateCache)
+        )
       })
       nextState = nextNextState
     } else {
-      nextState = updateDependentsState(dependent, nextNextState, setState)
+      nextState = updateDependentsState(
+        dependent,
+        nextNextState,
+        setState,
+        atomStateCache
+      )
     }
   })
   return nextState
@@ -332,7 +348,7 @@ const readAtom = <Value>(
   readingAtom: Atom<Value>,
   setState: Dispatch<(prev: State) => State>,
   readPendingMap: ReadPendingMap,
-  atomStateCache?: AtomStateCache
+  atomStateCache: AtomStateCache
 ) => {
   const prevState = readPendingMap.get(state) || state
   const [atomState, nextState] = readAtomState(
@@ -351,6 +367,7 @@ const writeAtom = <Value, Update>(
   writingAtom: WritableAtom<Value, Update>,
   update: Update,
   setState: Dispatch<(prev: State) => State>,
+  atomStateCache: AtomStateCache,
   addWriteThunk: (thunk: WriteThunk) => void
 ) => {
   const pendingPromises: Promise<void>[] = []
@@ -400,14 +417,16 @@ const writeAtom = <Value, Update>(
               nextState = updateDependentsState(
                 a,
                 updateAtomState(a, nextState, partialAtomState),
-                setState
+                setState,
+                atomStateCache
               )
             } else {
               setState((prev) =>
                 updateDependentsState(
                   a,
                   updateAtomState(a, prev, partialAtomState),
-                  setState
+                  setState,
+                  atomStateCache
                 )
               )
             }
@@ -631,21 +650,27 @@ export const Provider: React.FC = ({ children }) => {
         atom: WritableAtom<Value, Update>,
         update: Update
       ) =>
-        writeAtom(atom, update, setState, (thunk: WriteThunk) => {
-          writeThunkQueueRef.current.push(thunk)
-          if (lastStateRef.current) {
-            runWriteThunk(
-              lastStateRef,
-              pendingStateRef,
-              setState,
-              contextUpdateRef.current as ContextUpdate,
-              writeThunkQueueRef.current
-            )
-          } else {
-            // force update (FIXME this is a workaround for now)
-            setState((prev) => mMerge(prev, mCreate()))
+        writeAtom(
+          atom,
+          update,
+          setState,
+          atomStateCache,
+          (thunk: WriteThunk) => {
+            writeThunkQueueRef.current.push(thunk)
+            if (lastStateRef.current) {
+              runWriteThunk(
+                lastStateRef,
+                pendingStateRef,
+                setState,
+                contextUpdateRef.current as ContextUpdate,
+                writeThunkQueueRef.current
+              )
+            } else {
+              // force update (FIXME this is a workaround for now)
+              setState((prev) => mMerge(prev, mCreate()))
+            }
           }
-        }),
+        ),
     }),
     [readPendingMap, atomStateCache, setState]
   )

diff --git a/tests/async.test.tsx b/tests/async.test.tsx
@@ -160,3 +160,40 @@ it('uses multiple async atoms at once', async () => {
   await findByText('loading')
   await findByText('ready ready2')
 })
+
+it('uses async atom in the middle of dependency chain', async () => {
+  const countAtom = atom(0)
+  const asyncCountAtom = atom(async (get) => {
+    await new Promise((r) => setTimeout(r, 10))
+    return get(countAtom)
+  })
+  const delayedCountAtom = atom((get) => get(asyncCountAtom))
+
+  const Counter: React.FC = () => {
+    const [count, setCount] = useAtom(countAtom)
+    const [delayedCount] = useAtom(delayedCountAtom)
+    return (
+      <>
+        <div>
+          count: {count}, delayed: {delayedCount}
+        </div>
+        <button onClick={() => setCount((c) => c + 1)}>button</button>
+      </>
+    )
+  }
+
+  const { getByText, findByText } = render(
+    <Provider>
+      <React.Suspense fallback="loading">
+        <Counter />
+      </React.Suspense>
+    </Provider>
+  )
+
+  await findByText('loading')
+  await findByText('count: 0, delayed: 0')
+
+  fireEvent.click(getByText('button'))
+  // no loading
+  await findByText('count: 1, delayed: 1')
+})