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/mathjax@3.2.2 && mkdir -p ./public/~gitbook/static/ && cp -R node_modules/mathjax/es5 ./public/~gitbook/static/mathjax@3.2.2"
  },

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;
 }
 
 /**