Share memory cache between middleware and request handler (#2347)

Author: jpreynat

package.json

@@ -11,7 +11,7 @@
         "format": "prettier ./ --ignore-unknown --write",
         "format:check": "prettier ./ --ignore-unknown --list-different",
         "typecheck": "tsc --noEmit",
-        "unit": "bun test {src,packages}/**/*.test.ts",
+        "unit": "bun test {src,packages}",
         "e2e": "playwright test",
         "postinstall": "rm -rf ./public/~gitbook/static/[email protected] && mkdir -p ./public/~gitbook/static/ && cp -R node_modules/mathjax/es5 ./public/~gitbook/static/[email protected]"
     },

src/lib/cache/cache.test.ts

@@ -7,21 +7,24 @@ describe('cache', () => {
 
     let fn: CacheFunction<[string], string>;
     let testId = 0;
+    let getTtl: () => number;
 
     beforeEach(() => {
         impl.mockClear();
 
         testId += 1;
+        getTtl = () => 1000;
 
         fn = cache(`cache-${testId}`, async (arg: string, options: CacheFunctionOptions) => {
             await new Promise((resolve) => setTimeout(resolve, 20));
             return {
                 data: impl(arg),
+                ttl: getTtl(),
             };
         });
     });
 
-    it('should only execute once for same argument', async () => {
+    it('should only execute once for same argument if there is a ttl', async () => {
         const result = await Promise.all([fn('a'), fn('a')]);
 
         expect(result).toEqual(['test-a', 'test-a']);
@@ -34,6 +37,21 @@ describe('cache', () => {
         expect(impl).toHaveBeenCalledTimes(1);
     });
 
+    it('should execute multiple times for same argument if ttl is 0 or undefined', async () => {
+        getTtl = () => 0;
+
+        const result = await Promise.all([fn('a'), fn('a')]);
+
+        expect(result).toEqual(['test-a', 'test-a']);
+
+        expect(impl).toHaveBeenCalled();
+        expect(impl).toHaveBeenCalledTimes(1);
+
+        expect(await fn('a')).toEqual('test-a');
+
+        expect(impl).toHaveBeenCalledTimes(2);
+    });
+
     it('should execute for different arguments', async () => {
         const result = await Promise.all([fn('a'), fn('b')]);
 

src/lib/cache/memory.ts

@@ -1,11 +1,12 @@
 import { CacheBackend, CacheEntry } from './types';
 import { NON_IMMUTABLE_LOCAL_CACHE_MAX_AGE_SECONDS, isCacheEntryImmutable } from './utils';
+import { getGlobalContext } from '../waitUntil';
 
 export const memoryCache: CacheBackend = {
     name: 'memory',
     replication: 'local',
     async get(key) {
-        const memoryCache = getMemoryCache();
+        const memoryCache = await getMemoryCache();
         const memoryEntry = memoryCache.get(key);
 
         if (!memoryEntry) {
@@ -21,7 +22,7 @@ export const memoryCache: CacheBackend = {
         return null;
     },
     async set(key, entry) {
-        const memoryCache = getMemoryCache();
+        const memoryCache = await getMemoryCache();
         // When the entry is immutable, we can cache it for the entire duration.
         // Else we cache it for a very short time.
         const expiresAt =
@@ -41,11 +42,11 @@ export const memoryCache: CacheBackend = {
         memoryCache.set(key, { ...entry, meta });
     },
     async del(keys) {
-        const memoryCache = getMemoryCache();
+        const memoryCache = await getMemoryCache();
         keys.forEach((key) => memoryCache.delete(key));
     },
     async revalidateTags(tags) {
-        const memoryCache = getMemoryCache();
+        const memoryCache = await getMemoryCache();
         const keys: string[] = [];
 
         memoryCache.forEach((entry, key) => {
@@ -66,13 +67,12 @@ export const memoryCache: CacheBackend = {
  * With next-on-pages, the code seems to be isolated between the middleware and the handler.
  * To share the cache between the two, we use a global variable.
  */
-function getMemoryCache(): Map<string, CacheEntry> {
-    // @ts-ignore
-    if (!globalThis.gitbookMemoryCache) {
-        // @ts-ignore
-        globalThis.gitbookMemoryCache = new Map();
+async function getMemoryCache(): Promise<Map<string, CacheEntry>> {
+    let globalThisForMemoryCache: any = await getGlobalContext();
+
+    if (!globalThisForMemoryCache.gitbookMemoryCache) {
+        globalThisForMemoryCache.gitbookMemoryCache = new Map();
     }
 
-    // @ts-ignore
-    return globalThis.gitbookMemoryCache;
+    return globalThisForMemoryCache.gitbookMemoryCache;
 }

src/lib/waitUntil.ts

@@ -12,7 +12,7 @@ export async function getGlobalContext(): Promise<object> {
 
     // We lazy-load the next-on-pages package to avoid errors when running tests because of 'server-only'.
     const { getOptionalRequestContext } = await import('@cloudflare/next-on-pages');
-    return getOptionalRequestContext()?.cf ?? globalThis;
+    return getOptionalRequestContext()?.ctx ?? globalThis;
 }
 
 /**