support-custom-key

uriva · Nov 11, 2023 · 7c0c2ea · 7c0c2ea
1 parent d9fa48c
commit 7c0c2ea
Show file tree

Hide file tree

Showing 3 changed files with 293 additions and 258 deletions.
diff --git a/client/src/index.test.ts b/client/src/index.test.ts
@@ -70,6 +70,26 @@ const testVariadic = async (cacher: any) => {
 
 Deno.test("local cache", () => testCache(false, false, { cacheId: "some-id" }));
 
+Deno.test("custom key function", async () => {
+  let nCalled = 0;
+  const f = cache({
+    cacheId: "some id",
+    customKeyFn: (x, y) => x % 2 === 0,
+  })((x: number, y: number) => {
+    nCalled++;
+    return Promise.resolve(x);
+  });
+  await f(1, 1);
+  await waitAllWrites();
+  await f(1, 1); // should not cause a call
+  assertEquals(nCalled, 1);
+  await f(3, 1); // should also not cause a call
+  assertEquals(nCalled, 1);
+  await f(2, 1);
+  assertEquals(nCalled, 2);
+  await waitAllWrites();
+});
+
 const cleanRedis = async () => {
   const { REDIS_PASSWORD, REDIS_URL, REDIS_PORT } = config();
   const redisClient = await connect({

diff --git a/client/src/index.ts b/client/src/index.ts
@@ -90,46 +90,51 @@ const abstractCache = <F extends Func>({
 
 export const waitAllWrites = () => Promise.all(writePromises);
 
+type CustomKeyFn = (..._: any[]) => any;
+
 const inputToCacheKey =
   // deno-lint-ignore no-explicit-any
-  <Args extends any[]>(secret: string) => (...x: Args): string =>
-    hash(jsonStableStringify(x) + secret);
-
-type MemParams = { ttl?: number };
-
-export const memCache = ({ ttl }: MemParams) => <F extends Func>(f: F) => {
-  const keyToValue: Record<string, Awaited<ReturnType<F>>> = {};
-  const keyToTimestamp: Record<string, number> = {};
-  return abstractCache({
-    key: inputToCacheKey<Parameters<F>>(""),
-    f,
-    // @ts-expect-error Promise<Awaited<Awaited<X>>> is just Promise<X>
-    read: (key: string) => {
-      if (!(key in keyToValue)) {
-        return Promise.reject();
-      }
-      if (ttl && Date.now() - keyToTimestamp[key] > ttl * 1000) {
-        delete keyToTimestamp[key];
-        delete keyToValue[key];
-        return Promise.reject();
-      }
-      return Promise.resolve(keyToValue[key]);
-    },
-    write: (key: string, value: Awaited<ReturnType<F>>) => {
-      keyToValue[key] = value;
-      keyToTimestamp[key] = Date.now();
-      return Promise.resolve();
-    },
-  });
-};
+  <Args extends any[]>(secret: string, customKeyFn: CustomKeyFn | undefined) =>
+  (...x: Args): string =>
+    hash(jsonStableStringify(customKeyFn ? customKeyFn(...x) : x) + secret);
+
+type MemParams = { ttl?: number; customKeyFn?: (..._: any[]) => string };
+
+export const memCache =
+  ({ ttl, customKeyFn }: MemParams) => <F extends Func>(f: F) => {
+    const keyToValue: Record<string, Awaited<ReturnType<F>>> = {};
+    const keyToTimestamp: Record<string, number> = {};
+    return abstractCache({
+      key: inputToCacheKey<Parameters<F>>("", customKeyFn),
+      f,
+      // @ts-expect-error Promise<Awaited<Awaited<X>>> is just Promise<X>
+      read: (key: string) => {
+        if (!(key in keyToValue)) {
+          return Promise.reject();
+        }
+        if (ttl && Date.now() - keyToTimestamp[key] > ttl * 1000) {
+          delete keyToTimestamp[key];
+          delete keyToValue[key];
+          return Promise.reject();
+        }
+        return Promise.resolve(keyToValue[key]);
+      },
+      write: (key: string, value: Awaited<ReturnType<F>>) => {
+        keyToValue[key] = value;
+        keyToTimestamp[key] = Date.now();
+        return Promise.resolve();
+      },
+    });
+  };
 
-const localCache = ({ cacheId }: LocalCacheParams) => <F extends Func>(f: F) =>
-  // @ts-expect-error Promise+Awaited = nothing
-  abstractCache({
-    key: inputToCacheKey<Parameters<F>>(""),
-    f,
-    ...makeLocalReadWrite<Awaited<ReturnType<F>>>(cacheId),
-  });
+const localCache =
+  ({ cacheId, customKeyFn }: LocalCacheParams) => <F extends Func>(f: F) =>
+    // @ts-expect-error Promise+Awaited = nothing
+    abstractCache({
+      key: inputToCacheKey<Parameters<F>>("", customKeyFn),
+      f,
+      ...makeLocalReadWrite<Awaited<ReturnType<F>>>(cacheId),
+    });
 
 // deno-lint-ignore no-explicit-any
 type ServerParams = any;
@@ -183,22 +188,29 @@ const getRemote =
 
 export type CacheParams = LocalCacheParams | CloudCacheParams;
 
-type LocalCacheParams = { cacheId: string };
+type LocalCacheParams = {
+  cacheId: string;
+  customKeyFn?: CustomKeyFn;
+};
 
 type CloudCacheParams = {
   cacheId: string;
   token: string;
   url: string;
   ttl?: number;
   encryptionKey?: string;
+  customKeyFn?: CustomKeyFn;
 };
 
 export const cache = (params: CacheParams) =>
   "token" in params ? cloudCache(params) : localCache(params);
 
 const cloudCache = (params: CloudCacheParams) => <F extends Func>(f: F) =>
   abstractCache({
-    key: inputToCacheKey<Parameters<F>>(params.encryptionKey || ""),
+    key: inputToCacheKey<Parameters<F>>(
+      params.encryptionKey || "",
+      params.customKeyFn,
+    ),
     f,
     read: (key) =>
       getRemote(params)(key)