From cb81942caf3d6dad0302704e45575e5975973ab5 Mon Sep 17 00:00:00 2001
From: XLor <yjl9903@vip.qq.com>
Date: Thu, 7 Mar 2024 23:08:59 +0800
Subject: [PATCH] fix: missing set removing state

---
 src/external.ts   | 56 +++++++++++++++++++++++++++--------------------
 test/memo.test.ts |  2 +-
 2 files changed, 33 insertions(+), 25 deletions(-)

diff --git a/src/external.ts b/src/external.ts
index 5c07d58..71a6fb8 100644
--- a/src/external.ts
+++ b/src/external.ts
@@ -28,7 +28,7 @@ export function memoExternal<F extends AsyncFn>(
     } else {
       try {
         // Waiting or Removing
-        if (cur.state === State.Removing) {
+        while (cur.state === State.Removing) {
           await new Promise<void>((res) => {
             if (!cur.removingCallbacks) {
               cur.removingCallbacks = new Set();
@@ -90,33 +90,41 @@ export function memoExternal<F extends AsyncFn>(
     const cur = walkOrBreak<F, any[]>(root, path);
 
     if (cur) {
-      if (cur.state === State.Waiting) {
-        await new Promise((res) => {
-          if (!cur.callbacks) {
-            cur.callbacks = new Set();
-          }
-          // Ignore error
-          cur.callbacks!.add({ res, rej: () => {} });
-        });
-      } else if (cur.state === State.Removing) {
-        await new Promise<void>((res) => {
-          if (!cur.removingCallbacks) {
-            cur.removingCallbacks = new Set();
-          }
-          // Ignore error
-          cur.removingCallbacks!.add({ res, rej: () => {} });
-        });
+      while (cur.state === State.Waiting || cur.state === State.Removing) {
+        if (cur.state === State.Waiting) {
+          await new Promise((res) => {
+            if (!cur.callbacks) {
+              cur.callbacks = new Set();
+            }
+            // Ignore error
+            cur.callbacks!.add({ res, rej: () => {} });
+          });
+        } else if (cur.state === State.Removing) {
+          await new Promise<void>((res) => {
+            if (!cur.removingCallbacks) {
+              cur.removingCallbacks = new Set();
+            }
+            // Ignore error
+            cur.removingCallbacks!.add({ res, rej: () => {} });
+          });
+        }
       }
 
-      await options.external.remove
-        .bind(memoFunc)(args as Parameters<F>)
-        .catch(options.external?.error ?? (() => undefined));
+      try {
+        cur.state = State.Removing;
+
+        await options.external.remove
+          .bind(memoFunc)(args as Parameters<F>)
+          .catch(options.external?.error ?? (() => undefined));
 
-      // Resolve other waiting callbacks
-      for (const callback of cur.removingCallbacks ?? []) {
-        callback.res();
+        // Resolve other waiting callbacks
+        for (const callback of cur.removingCallbacks ?? []) {
+          callback.res();
+        }
+        cur.removingCallbacks = undefined;
+      } finally {
+        cur.state = State.Empty;
       }
-      cur.removingCallbacks = undefined;
     }
   };
 
diff --git a/test/memo.test.ts b/test/memo.test.ts
index f7d7fb3..86ee420 100644
--- a/test/memo.test.ts
+++ b/test/memo.test.ts
@@ -236,7 +236,7 @@ describe('memo external', () => {
       func(),
       func()
     ]);
-    expect(tasks2).toStrictEqual([undefined, 1, 1, 1, undefined, 1, 1]);
+    expect(tasks2).toStrictEqual([undefined, 1, 2, 3, undefined, 4, 5]);
   });
 });