From b66d94c4c9db30246cb301cda175fc90695b2fe5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mathis=20M=C3=B8ller?= <thisen-dk@hotmail.com>
Date: Fri, 26 Mar 2021 11:18:01 +0100
Subject: [PATCH] Fix #375 - Atom state is not flushed when using sync invoked
 async setter (#380)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Add failing test

* Fix flushing

* Proper test name

* fix with pending promises status

Co-authored-by: Mathis Møller <mmm@cryosinternational.com>
Co-authored-by: daishi <daishi@axlight.com>
---
 .vscode/launch.json  | 20 ++++++++++++++++
 src/core/vanilla.ts  | 56 +++++++++++++++++++++++++++-----------------
 tests/async.test.tsx | 44 ++++++++++++++++++++++++++++++++++
 3 files changed, 98 insertions(+), 22 deletions(-)
 create mode 100644 .vscode/launch.json

diff --git a/.vscode/launch.json b/.vscode/launch.json
new file mode 100644
index 0000000000..50623af35f
--- /dev/null
+++ b/.vscode/launch.json
@@ -0,0 +1,20 @@
+{
+  "version": "0.2.0",
+  "configurations": [
+    {
+      "name": "Debug Jest Tests",
+      "type": "node",
+      "request": "launch",
+      "runtimeArgs": [
+        "--inspect-brk",
+        "${workspaceRoot}/node_modules/.bin/jest",
+        "--runInBand",
+        "--watch",
+        "--no-coverage"
+      ],
+      "console": "integratedTerminal",
+      "internalConsoleOptions": "neverOpen",
+      "port": 9229
+    }
+  ]
+}
diff --git a/src/core/vanilla.ts b/src/core/vanilla.ts
index e79849bad2..31fad600d2 100644
--- a/src/core/vanilla.ts
+++ b/src/core/vanilla.ts
@@ -359,18 +359,21 @@ const writeAtomState = <Value, Update>(
   state: State,
   atom: WritableAtom<Value, Update>,
   update: Update,
-  pendingPromises?: Promise<void>[]
+  pendingPromises: Promise<void>[]
 ): void => {
+  const isPendingPromisesExpired = !pendingPromises.length
   const atomState = getAtomState(state, atom)
   if (
     atomState &&
     atomState.w // write promise
   ) {
     const promise = atomState.w.then(() => {
-      writeAtomState(state, atom, update)
-      flushPending(state)
+      writeAtomState(state, atom, update, pendingPromises)
+      if (isPendingPromisesExpired) {
+        flushPending(state)
+      }
     })
-    if (pendingPromises) {
+    if (!isPendingPromisesExpired) {
       pendingPromises.push(promise)
     }
     return
@@ -413,33 +416,39 @@ const writeAtomState = <Value, Update>(
           setAtomValue(state, a, v)
           invalidateDependents(state, a)
         } else {
-          writeAtomState(state, a, v)
+          const isPendingPromisesExpired = !pendingPromises.length
+          writeAtomState(state, a, v, pendingPromises)
+          if (isPendingPromisesExpired) {
+            flushPending(state)
+          }
         }
       }) as Setter,
       update
     )
     if (promiseOrVoid instanceof Promise) {
-      if (pendingPromises) {
-        pendingPromises.push(promiseOrVoid)
-      }
-      setAtomWritePromise(
-        state,
-        atom,
-        promiseOrVoid.then(() => {
-          setAtomWritePromise(state, atom)
+      const promise = promiseOrVoid.then(() => {
+        setAtomWritePromise(state, atom)
+        if (isPendingPromisesExpired) {
           flushPending(state)
-        })
-      )
+        }
+      })
+      if (!isPendingPromisesExpired) {
+        pendingPromises.push(promise)
+      }
+      setAtomWritePromise(state, atom, promise)
     }
   } catch (e) {
-    if (pendingPromises && pendingPromises.length) {
+    if (pendingPromises.length === 1) {
+      // still in sync, throw it right away
+      throw e
+    } else if (!isPendingPromisesExpired) {
       pendingPromises.push(
         new Promise((_resolve, reject) => {
           reject(e)
         })
       )
     } else {
-      throw e
+      console.error('Uncaught exception: Use promise to catch error', e)
     }
   }
 }
@@ -449,21 +458,24 @@ export const writeAtom = <Value, Update>(
   writingAtom: WritableAtom<Value, Update>,
   update: Update
 ): void | Promise<void> => {
-  const pendingPromises: Promise<void>[] = []
+  const pendingPromises: Promise<void>[] = [Promise.resolve()]
 
   writeAtomState(state, writingAtom, update, pendingPromises)
   flushPending(state)
 
-  if (pendingPromises.length) {
+  if (pendingPromises.length <= 1) {
+    pendingPromises.splice(0)
+  } else {
     return new Promise<void>((resolve, reject) => {
       const loop = () => {
-        const len = pendingPromises.length
-        if (len === 0) {
+        if (pendingPromises.length <= 1) {
+          pendingPromises.splice(0)
           resolve()
         } else {
           Promise.all(pendingPromises)
             .then(() => {
-              pendingPromises.splice(0, len)
+              pendingPromises.splice(1)
+              flushPending(state)
               loop()
             })
             .catch(reject)
diff --git a/tests/async.test.tsx b/tests/async.test.tsx
index 1ef4aa4d1f..8bfb9cfd13 100644
--- a/tests/async.test.tsx
+++ b/tests/async.test.tsx
@@ -581,3 +581,47 @@ it('a derived atom from a newly created async atom (#351)', async () => {
   await findByText('loading')
   await findByText('derived: 13, commits: 3')
 })
+
+it('Handles synchronously invoked async set (#375)', async () => {
+  const loadingAtom = atom(false)
+  const documentAtom = atom<string | undefined>(undefined)
+  const loadDocumentAtom = atom(null, (_get, set) => {
+    const fetch = async () => {
+      set(loadingAtom, true)
+      const response = await new Promise<string>((resolve) =>
+        setTimeout(() => resolve('great document'), 10)
+      )
+      set(documentAtom, response)
+      set(loadingAtom, false)
+    }
+    fetch()
+  })
+
+  const ListDocuments: React.FC = () => {
+    const [loading] = useAtom(loadingAtom)
+    const [document] = useAtom(documentAtom)
+    const [, loadDocument] = useAtom(loadDocumentAtom)
+
+    useEffect(() => {
+      loadDocument()
+    }, [loadDocument])
+
+    return (
+      <>
+        {loading && <div>loading</div>}
+        {!loading && <div>{document}</div>}
+      </>
+    )
+  }
+
+  const { getByText, findByText } = render(
+    <StrictMode>
+      <Provider>
+        <ListDocuments />
+      </Provider>
+    </StrictMode>
+  )
+
+  await findByText('loading')
+  await findByText('great document')
+})