feat: update and not remove

yjl9903 · Mar 7, 2024 · 79388da · 79388da
1 parent 2743a4a
commit 79388da
Show file tree

Hide file tree

Showing 4 changed files with 140 additions and 56 deletions.
diff --git a/src/external.ts b/src/external.ts
@@ -25,19 +25,19 @@ export function memoExternal<F extends AsyncFn>(
         }
         cur.callbacks!.add({ res, rej });
       });
+    } else if (cur.state === State.Updating) {
+      let updatedValue: ReturnType<F>;
+      while (cur.state === State.Updating) {
+        updatedValue = await new Promise((res, rej) => {
+          if (!cur.updatingCallbacks) {
+            cur.updatingCallbacks = new Set();
+          }
+          cur.updatingCallbacks!.add({ res, rej });
+        });
+      }
+      return updatedValue!;
     } else {
       try {
-        // Waiting or Removing
-        while (cur.state === State.Removing) {
-          await new Promise<void>((res) => {
-            if (!cur.removingCallbacks) {
-              cur.removingCallbacks = new Set();
-            }
-            // Ignore error
-            cur.removingCallbacks!.add({ res, rej: () => {} });
-          });
-        }
-
         cur.state = State.Waiting;
 
         const externalOnError = options.external.error ?? (() => undefined);
@@ -69,12 +69,16 @@ export function memoExternal<F extends AsyncFn>(
         cur.state = State.Empty;
         cur.error = error;
 
-        // Reject other waiting callbacks
-        for (const callback of cur.callbacks ?? []) {
-          callback.rej(error);
+        try {
+          // Reject other waiting callbacks
+          for (const callback of cur.callbacks ?? []) {
+            callback.rej(error);
+          }
+          // Release callbacks
+          cur.callbacks = undefined;
+        } catch {
+          // Ignore errors
         }
-        // Release callbacks
-        cur.callbacks = undefined;
 
         throw error;
       }
@@ -94,45 +98,72 @@ export function memoExternal<F extends AsyncFn>(
     const cur = walkOrBreak<F, any[]>(root, path);
 
     if (cur) {
-      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));
+    }
+  };
+
+  memoFunc.update = async (...args) => {
+    const path = options.serialize ? options.serialize.bind(memoFunc)(...args) : args;
+    const cur = walkAndCreate<F, any[]>(root, path);
+
+    while (cur.state === State.Waiting || cur.state === State.Updating) {
+      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.Updating) {
+        await new Promise((res) => {
+          if (!cur.updatingCallbacks) {
+            cur.updatingCallbacks = new Set();
+          }
+          // Ignore error
+          cur.updatingCallbacks!.add({ res, rej: () => {} });
+        });
+      } else {
+        break;
       }
+    }
 
-      try {
-        cur.state = State.Removing;
+    try {
+      cur.state = State.Updating;
 
-        await options.external.remove
-          .bind(memoFunc)(args as Parameters<F>)
-          .catch(options.external?.error ?? (() => undefined));
-      } finally {
-        cur.state = State.Empty;
+      const externalOnError = options.external.error ?? (() => undefined);
+      const value = await fn(...args);
+      await options.external.set.bind(memoFunc)(args, value).catch(externalOnError);
+
+      cur.state = State.Empty;
+
+      try {
+        // Resolve other waiting callbacks
+        for (const callback of cur.updatingCallbacks ?? []) {
+          callback.res(value);
+        }
+        cur.updatingCallbacks = undefined;
+      } catch {
+        // Should not have errors here
       }
 
+      return value;
+    } catch (error) {
+      cur.state = State.Empty;
+
       try {
         // Resolve other waiting callbacks
-        for (const callback of cur.removingCallbacks ?? []) {
-          callback.res();
+        for (const callback of cur.updatingCallbacks ?? []) {
+          callback.rej(error);
         }
-        cur.removingCallbacks = undefined;
+        cur.updatingCallbacks = undefined;
       } catch {
         // Should not have errors here
       }
+
+      throw error;
     }
   };
 

diff --git a/src/trie.ts b/src/trie.ts
@@ -5,7 +5,7 @@ export const enum State {
   Empty,
   Ok,
   Waiting,
-  Removing, // Used for external cache
+  Updating,
   Error
 }
 
@@ -16,7 +16,7 @@ export interface Node<T extends Fn> {
   primitive: Map<any, Node<T>>;
   reference: WeakMap<any, Node<T>>;
   callbacks?: Set<{ res: (value: ReturnType<T>) => void; rej: (error: unknown) => void }>;
-  removingCallbacks?: Set<{ res: () => void; rej: (error: unknown) => void }>;
+  updatingCallbacks?: Set<{ res: (value: ReturnType<T>) => void; rej: (error: unknown) => void }>;
 }
 
 export function makeNode<T extends Fn>(): Node<T> {

diff --git a/src/types.ts b/src/types.ts
@@ -112,6 +112,9 @@ export interface MemoExternalFunc<F extends Fn> {
   // Remove cache
   remove(...args: Parameters<F>): Promise<void>;
 
+  // Update cache
+  update(...args: Parameters<F>): Promise<void>;
+
   // External cache
   external?: MemoExternalOptions<F>['external'];
 }
diff --git a/test/memo.test.ts b/test/memo.test.ts
@@ -198,45 +198,95 @@ describe('memo external', () => {
     expect(tasks2).toStrictEqual([1, 1, 1, 1, 1]);
   });
 
-  it('should not interleave get and remove external cache', async () => {
+  it('should not interleave get and update external cache', async () => {
     let cnt = 0;
     const func = memoExternal(async () => 0, {
       external: {
         async get() {
           await sleep(200);
           return ++cnt;
         },
-        async set() {},
-        async remove() {
+        async set() {
           await sleep(100);
           cnt = 0;
-        }
+        },
+        async remove() {}
       }
     });
 
     const tasks = await Promise.all([
       func(),
-      func.remove(),
+      func.update(),
       func(),
       func(),
       func(),
-      func.remove(),
+      func.update(),
       func()
     ]);
-    expect(tasks).toStrictEqual([1, undefined, 1, 1, 1, undefined, 1]);
+    expect(tasks).toStrictEqual([1, 0, 1, 1, 1, 0, 1]);
     expect(cnt).toBe(0);
 
     cnt = 10;
     const tasks2 = await Promise.all([
-      func.remove(),
+      func.update(),
       func(),
       func(),
       func(),
-      func.remove(),
+      func.update(),
       func(),
       func()
     ]);
-    expect(tasks2).toStrictEqual([undefined, 1, 2, 3, undefined, 4, 5]);
+    expect(tasks2).toStrictEqual([0, 0, 0, 0, 0, 0, 0]);
+  });
+
+  it('should get and update external cache', async () => {
+    let cnt = 0;
+    const func = memoExternal(async () => 0, {
+      external: {
+        async get() {
+          await sleep(200);
+          return ++cnt;
+        },
+        async set() {
+          await sleep(100);
+          cnt = 0;
+        },
+        async remove() {}
+      }
+    });
+
+    expect(await func.get()).toBe(1);
+    expect(await func.get()).toBe(2);
+    expect(await func.get()).toBe(3);
+    expect(await func.update()).toBe(0);
+    expect(await func.get()).toBe(1);
+    expect(await func.get()).toBe(2);
+    expect(await func.get()).toBe(3);
+  });
+
+  it('should get and update external cache', async () => {
+    let cnt = 1;
+    const func = memoExternal(async () => 0, {
+      external: {
+        async get() {
+          await sleep(200);
+          return ++cnt;
+        },
+        async set(p, v) {
+          await sleep(100);
+          cnt = v;
+        },
+        async remove() {}
+      }
+    });
+
+    expect(await func.get()).toBe(2);
+    expect(await func.get()).toBe(3);
+    expect(await func.get()).toBe(4);
+    expect(await func.update()).toBe(0);
+    expect(await func.get()).toBe(1);
+    expect(await func.get()).toBe(2);
+    expect(await func.get()).toBe(3);
   });
 });