feat: add cache expiration

yjl9903 · Mar 18, 2024 · 75592cc · 75592cc
1 parent efec36a
commit 75592cc
Show file tree

Hide file tree

Showing 5 changed files with 99 additions and 2 deletions.
diff --git a/src/async.ts b/src/async.ts
@@ -1,6 +1,6 @@
 import type { AsyncFn, MemoAsyncFunc, MemoAsyncOptions } from './types';
 
-import { State, clearNode, makeNode, walkAndCreate, walkOrBreak } from './trie';
+import { State, clearNode, clearNodeCache, makeNode, walkAndCreate, walkOrBreak } from './trie';
 
 export function memoAsync<F extends AsyncFn>(
   fn: F,
@@ -13,6 +13,13 @@ export function memoAsync<F extends AsyncFn>(
     const path = options.serialize ? options.serialize.bind(memoFunc)(...args) : args;
     const cur = walkAndCreate<F, any[]>(root, path);
 
+    if ((cur.state === State.Ok || cur.state === State.Error) && cur.expiration !== undefined) {
+      // Cache expire
+      if (new Date().getTime() > cur.expiration) {
+        clearNodeCache(cur);
+      }
+    }
+
     if (cur.state === State.Ok) {
       return cur.value;
     } else if (cur.state === State.Error) {
@@ -38,6 +45,11 @@ export function memoAsync<F extends AsyncFn>(
         cur.state = State.Ok;
         cur.value = value;
 
+        if (memoFunc.expirationTtl !== undefined && memoFunc.expirationTtl !== null) {
+          const now = new Date();
+          cur.expiration = now.getTime() + memoFunc.expirationTtl;
+        }
+
         if (!hasExternalCache && options.external) {
           await options.external.set.bind(memoFunc)(args, value).catch(externalOnError);
         }
@@ -66,6 +78,8 @@ export function memoAsync<F extends AsyncFn>(
     }
   } as MemoAsyncFunc<F>;
 
+  memoFunc.expirationTtl = options.expirationTtl;
+
   memoFunc.get = (...args) => {
     return memoFunc(...args);
   };

diff --git a/src/sync.ts b/src/sync.ts
@@ -1,6 +1,6 @@
 import type { Fn, MemoFunc, MemoOptions } from './types';
 
-import { State, clearNode, makeNode, walkAndCreate, walkOrBreak } from './trie';
+import { State, clearNode, clearNodeCache, makeNode, walkAndCreate, walkOrBreak } from './trie';
 
 export function memo<F extends Fn>(fn: F, options: MemoOptions<F> = {}): MemoFunc<F> {
   const root = makeNode<F>();
@@ -10,6 +10,13 @@ export function memo<F extends Fn>(fn: F, options: MemoOptions<F> = {}): MemoFun
     const path = options.serialize ? options.serialize.bind(memoFunc)(...args) : args;
     const cur = walkAndCreate<F, any[]>(root, path);
 
+    if (cur.expiration !== undefined) {
+      // Cache expire
+      if (new Date().getTime() > cur.expiration) {
+        clearNodeCache(cur);
+      }
+    }
+
     if (cur.state === State.Ok) {
       return cur.value;
     } else if (cur.state === State.Error) {
@@ -19,6 +26,12 @@ export function memo<F extends Fn>(fn: F, options: MemoOptions<F> = {}): MemoFun
         const value = fn(...args);
         cur.state = State.Ok;
         cur.value = value;
+
+        if (memoFunc.expirationTtl !== undefined && memoFunc.expirationTtl !== null) {
+          const now = new Date();
+          cur.expiration = now.getTime() + memoFunc.expirationTtl;
+        }
+
         return value;
       } catch (error) {
         cur.state = State.Error;
@@ -28,6 +41,8 @@ export function memo<F extends Fn>(fn: F, options: MemoOptions<F> = {}): MemoFun
     }
   } as MemoFunc<F>;
 
+  memoFunc.expirationTtl = options.expirationTtl;
+
   memoFunc.get = (...args) => {
     return memoFunc(...args);
   };

diff --git a/src/trie.ts b/src/trie.ts
@@ -13,8 +13,16 @@ export interface Node<T extends Fn> {
   state: State;
   value: ReturnType<T> | undefined;
   error: unknown;
+
+  // Metadata
+  expiration: number | undefined;
+  meta: {} | undefined;
+
+  // Children
   primitive: Map<any, Node<T>>;
   reference: WeakMap<any, Node<T>>;
+
+  // Callbacks
   callbacks?: Set<{ res: (value: ReturnType<T>) => void; rej: (error: unknown) => void }>;
   updatingCallbacks?: Set<{ res: (value: ReturnType<T>) => void; rej: (error: unknown) => void }>;
 }
@@ -24,6 +32,8 @@ export function makeNode<T extends Fn>(): Node<T> {
     state: State.Empty,
     value: undefined,
     error: undefined,
+    expiration: undefined,
+    meta: undefined,
     primitive: new Map(),
     reference: new WeakMap()
   };
@@ -34,11 +44,20 @@ export function clearNode<T extends Fn>(node: Node<T> | undefined) {
     node.state = State.Empty;
     node.value = undefined;
     node.error = undefined;
+    node.expiration = undefined;
+    node.meta = undefined;
     node.primitive = new Map();
     node.reference = new WeakMap();
   }
 }
 
+export function clearNodeCache<T extends Fn>(node: Node<T>) {
+  node.state = State.Empty;
+  node.value = undefined;
+  node.error = undefined;
+  node.expiration = undefined;
+}
+
 function walkBase<T extends Fn, P extends any[] = Parameters<T>>(
   node: Node<T>,
   args: P,

diff --git a/src/types.ts b/src/types.ts
@@ -4,6 +4,11 @@ export interface MemoOptions<F extends Fn, S extends unknown[] = unknown[]> {
    * This is used to identify cache key
    */
   serialize?: (this: MemoFunc<F>, ...args: Parameters<F>) => S;
+
+  /**
+   * Default expiration time duration (in milliseconds)
+   */
+  expirationTtl?: number;
 }
 
 export interface MemoAsyncOptions<F extends Fn, S extends unknown[] = unknown[]> {
@@ -13,6 +18,11 @@ export interface MemoAsyncOptions<F extends Fn, S extends unknown[] = unknown[]>
    */
   serialize?: (this: MemoAsyncFunc<F>, ...args: Parameters<F>) => S;
 
+  /**
+   * Default expiration time duration (in milliseconds)
+   */
+  expirationTtl?: number;
+
   external?: {
     get: (
       this: MemoAsyncFunc<F>,
@@ -77,6 +87,11 @@ export interface MemoFunc<F extends Fn> {
 
   // Clear all the cache
   clear(): void;
+
+  /**
+   * Default expiration time duration (in seconds)
+   */
+  expirationTtl?: number;
 }
 
 export interface MemoAsyncFunc<F extends Fn> {
@@ -95,6 +110,11 @@ export interface MemoAsyncFunc<F extends Fn> {
   // Clear all the cache
   clear(): Promise<void>;
 
+  /**
+   * Default expiration time duration (in seconds)
+   */
+  expirationTtl?: number;
+
   // External cache
   external?: MemoAsyncOptions<F>['external'];
 }

diff --git a/test/memo.test.ts b/test/memo.test.ts
@@ -55,6 +55,35 @@ describe('memo sync', () => {
     expect(add(2, 1)).toBe(3);
     expect(add(2, 2)).toBe(3);
   });
+
+  it('should expire cache', async () => {
+    let count = 0;
+    const add = memo(
+      (a: number, b: number) => {
+        count++;
+        return a + b;
+      },
+      {
+        expirationTtl: 10
+      }
+    );
+
+    expect(count).toBe(0);
+    add(1, 2);
+    add(1, 2);
+    add(1, 2);
+    add(1, 2);
+    add(1, 2);
+    expect(count).toBe(1);
+
+    await sleep(100);
+    add(1, 2);
+    add(1, 2);
+    add(1, 2);
+    add(1, 2);
+    add(1, 2);
+    expect(count).toBe(2);
+  });
 });
 
 describe('memo async', () => {