From 810e3a83a7a4c0a7301d9e29e8f96b72a9f19e92 Mon Sep 17 00:00:00 2001
From: Karolis Narkevicius <karolis.n@gmail.com>
Date: Wed, 29 Dec 2021 22:04:32 +0000
Subject: [PATCH] Fix effect cleanup

---
 src/index.tsx                  | 11 ++++++++++-
 test/useEffectReducer.test.tsx |  4 ++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/src/index.tsx b/src/index.tsx
index ea049ac..d8735b3 100644
--- a/src/index.tsx
+++ b/src/index.tsx
@@ -191,12 +191,21 @@ export function useEffectReducer<
     >,
     event: TEvent | FlushEvent
   ): AggregatedEffectsState<TState, TEvent> => {
+    // Remove stopped entities if any, but otherwise keep them around
+    // if they haven't been stopped yet since this means our cleanup
+    // useEffect hasn't been executed yet
+    entitiesToStop = entitiesToStop.some(
+      entity => entity.status === EntityStatus.Stopped
+    )
+      ? entitiesToStop.filter(entity => entity.status !== EntityStatus.Stopped)
+      : entitiesToStop;
+
     const nextEffectEntities: Array<EffectEntity<TState, TEvent>> = [];
     const nextEntitiesToStop: Array<EffectEntity<TState, TEvent>> = [];
 
     if (event.type === flushEffectsSymbol) {
       // Record that effects have already been executed
-      return [state, stateEffectTuples.slice(event.count), nextEntitiesToStop];
+      return [state, stateEffectTuples.slice(event.count), entitiesToStop];
     }
 
     const exec = (
diff --git a/test/useEffectReducer.test.tsx b/test/useEffectReducer.test.tsx
index 1ffb900..808543d 100644
--- a/test/useEffectReducer.test.tsx
+++ b/test/useEffectReducer.test.tsx
@@ -507,10 +507,14 @@ describe('useEffectReducer', () => {
     const helloButton = getByTestId('send-hello');
     const goodbyeButton = getByTestId('send-goodbye');
 
+    // Click each button twice to make sure no effect cleanup
+    // is skipped due to batching
+    fireEvent.click(helloButton);
     fireEvent.click(helloButton);
 
     setTimeout(() => {
       fireEvent.click(goodbyeButton);
+      fireEvent.click(goodbyeButton);
     }, 30);
 
     await waitFor(() => {