From 46637d0e6a70bd944572580fce34c973e63a4a98 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Tue, 17 Dec 2024 15:26:57 -0800
Subject: [PATCH] process all batch functions even on error

---
 src/vanilla/store.ts         | 14 +++----
 tests/vanilla/store.test.tsx | 75 ++++++++++++------------------------
 2 files changed, 30 insertions(+), 59 deletions(-)

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index 5cd5bb719d..beef6e0fdb 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -192,7 +192,7 @@ const registerBatchAtom = (
   if (!batch.D.has(atom)) {
     batch.D.set(atom, new Set())
     addBatchFuncMedium(batch, () => {
-      atomState.m?.l.forEach((listener) => listener())
+      atomState.m?.l.forEach((listener) => addBatchFuncMedium(batch, listener))
     })
   }
 }
@@ -211,12 +211,6 @@ const addBatchAtomDependent = (
 const getBatchAtomDependents = (batch: Batch, atom: AnyAtom) =>
   batch.D.get(atom)
 
-const copySetAndClear = <T>(origSet: Set<T>): Set<T> => {
-  const newSet = new Set(origSet)
-  origSet.clear()
-  return newSet
-}
-
 const flushBatch = (batch: Batch) => {
   let error: AnyError
   let hasError = false
@@ -232,8 +226,10 @@ const flushBatch = (batch: Batch) => {
   }
   while (batch.M.size || batch.L.size) {
     batch.D.clear()
-    copySetAndClear(batch.M).forEach(call)
-    copySetAndClear(batch.L).forEach(call)
+    batch.M.forEach(call)
+    batch.M.clear()
+    batch.L.forEach(call)
+    batch.L.clear()
   }
   if (hasError) {
     throw error
diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx
index 4746bf568c..977a9a0e2b 100644
--- a/tests/vanilla/store.test.tsx
+++ b/tests/vanilla/store.test.tsx
@@ -2,10 +2,6 @@ import { waitFor } from '@testing-library/react'
 import { assert, describe, expect, it, vi } from 'vitest'
 import { atom, createStore } from 'jotai/vanilla'
 import type { Atom, Getter, PrimitiveAtom } from 'jotai/vanilla'
-import type {
-  INTERNAL_DevStoreRev4,
-  INTERNAL_PrdStore,
-} from 'jotai/vanilla/store'
 
 it('should not fire on subscribe', async () => {
   const store = createStore()
@@ -316,52 +312,6 @@ it('should update derived atoms during write (#2107)', async () => {
   expect(store.get(countAtom)).toBe(2)
 })
 
-it.only('mounts dependencies in async edge case', async () => {
-  const store = createStore().unstable_derive((getAtomState, ...rest) => [
-    (a) => Object.assign(getAtomState(a), { label: a.debugLabel }),
-    ...rest,
-  ]) as INTERNAL_DevStoreRev4 & INTERNAL_PrdStore
-  const getAtomState = store.dev4_get_internal_weak_map().get
-
-  const a = atom(0)
-  a.debugLabel = 'a'
-  const resolve: (() => void)[] = []
-  const b = atom((get) => {
-    get(a)
-    return new Promise<void>((r) => {
-      resolve.push(() => {
-        r()
-      })
-    })
-  })
-  b.debugLabel = 'b'
-  const c = atom(async (get) => {
-    await Promise.resolve()
-    await get(b)
-  })
-  c.debugLabel = 'c'
-
-  store.sub(c, () => {})
-
-  await Promise.resolve()
-  expect(resolve.length).toBe(1)
-  resolve[0]!()
-  // --- Need to wait two microtasks to make it work ---
-  await Promise.resolve()
-  await Promise.resolve()
-  const aState = getAtomState(a)
-  const bState = getAtomState(b)
-  const cState = getAtomState(c)
-  console.log('aState', aState)
-  console.log('bState', bState)
-  console.log('cState', cState)
-
-  store.set(a, 20)
-  store.set(a, 30)
-  await Promise.resolve()
-  expect(resolve.length).toBe(3)
-})
-
 it('resolves dependencies reliably after a delay (#2192)', async () => {
   expect.assertions(1)
   const countAtom = atom(0)
@@ -1016,3 +966,28 @@ it('processes deep atom a graph beyond maxDepth', () => {
   expect(() => store.set(baseAtom, 1)).not.toThrow()
   // store.set(lastAtom) // FIXME: This is causing a stack overflow
 })
+
+it('should process all atom listeners even if some of them throw errors', () => {
+  const store = createStore()
+  const a = atom(0)
+  const listenerA = vi.fn()
+  const listenerB = vi.fn(() => {
+    throw new Error('error')
+  })
+  const listenerC = vi.fn()
+  const listenerD = vi.fn()
+
+  store.sub(a, listenerA)
+  store.sub(a, listenerB)
+  store.sub(a, listenerC)
+  store.sub(a, listenerD)
+  try {
+    store.set(a, 1)
+  } catch {
+    // expect empty
+  }
+  expect(listenerA).toHaveBeenCalledTimes(1)
+  expect(listenerB).toHaveBeenCalledTimes(1)
+  expect(listenerC).toHaveBeenCalledTimes(1)
+  expect(listenerD).toHaveBeenCalledTimes(1)
+})