From db45c6d77348bb7fa9017656ecc0a7aac1542e3f Mon Sep 17 00:00:00 2001
From: Johan Preynat <johan.preynat@gmail.com>
Date: Tue, 25 Jun 2024 14:30:38 +0200
Subject: [PATCH] Rework context to actually share memory cache between
 middleware and handlers for all requests (#2358)

---
 src/lib/async.ts        |  8 ++++----
 src/lib/cache/memory.ts | 24 +++++++++++++++++-------
 src/lib/waitUntil.ts    | 29 ++++++++++++++++++++++-------
 3 files changed, 43 insertions(+), 18 deletions(-)

diff --git a/src/lib/async.ts b/src/lib/async.ts
index e99cb26f7..3e99ddff2 100644
--- a/src/lib/async.ts
+++ b/src/lib/async.ts
@@ -1,6 +1,6 @@
 import { MaybePromise } from 'p-map';
 
-import { waitUntil, getGlobalContext } from './waitUntil';
+import { waitUntil, getRequestContext } from './waitUntil';
 
 /**
  * Execute a function for each input in parallel and return the first result.
@@ -238,7 +238,7 @@ export function singleton<R>(execute: () => Promise<R>): () => Promise<R> {
         }
 
         // Promises are not shared between requests in Cloudflare Workers
-        const ctx = await getGlobalContext();
+        const ctx = await getRequestContext();
         const current = states.get(ctx);
         if (current) {
             return current;
@@ -269,7 +269,7 @@ export function singletonMap<Key extends string, Args extends any[], Result>(
     const states = new WeakMap<object, Map<string, Promise<Result>>>();
 
     const fn: SingletonFunction<Key, Args, Result> = async (key, ...args) => {
-        const ctx = await getGlobalContext();
+        const ctx = await getRequestContext();
         let current = states.get(ctx);
         if (current) {
             const existing = current.get(key);
@@ -292,7 +292,7 @@ export function singletonMap<Key extends string, Args extends any[], Result>(
     };
 
     fn.isRunning = async (key: string) => {
-        const ctx = await getGlobalContext();
+        const ctx = await getRequestContext();
         const current = states.get(ctx);
         return current?.has(key) ?? false;
     };
diff --git a/src/lib/cache/memory.ts b/src/lib/cache/memory.ts
index 15654aadc..70cc8e1f0 100644
--- a/src/lib/cache/memory.ts
+++ b/src/lib/cache/memory.ts
@@ -1,6 +1,6 @@
-import { CacheBackend } from './types';
+import { CacheBackend, CacheEntry } from './types';
 import { NON_IMMUTABLE_LOCAL_CACHE_MAX_AGE_SECONDS, isCacheEntryImmutable } from './utils';
-import { singleton } from '../async';
+import { getGlobalContext } from '../waitUntil';
 
 export const memoryCache: CacheBackend = {
     name: 'memory',
@@ -64,9 +64,19 @@ 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.
- * By using a singleton, we ensure that the cache is only created once and stored in the
- * current request context.
+ * In memory cache shared globally.
  */
-const getMemoryCache = singleton(async () => new Map());
+async function getMemoryCache(): Promise<Map<string, CacheEntry>> {
+    const ctx: Awaited<ReturnType<typeof getGlobalContext>> & {
+        gitbookMemoryCache?: Map<string, CacheEntry>;
+    } = await getGlobalContext();
+
+    if (ctx.gitbookMemoryCache) {
+        return ctx.gitbookMemoryCache;
+    }
+
+    const gitbookMemoryCache = new Map<string, CacheEntry>();
+    ctx.gitbookMemoryCache = gitbookMemoryCache;
+
+    return gitbookMemoryCache;
+}
diff --git a/src/lib/waitUntil.ts b/src/lib/waitUntil.ts
index 43ea13aee..6e2940f17 100644
--- a/src/lib/waitUntil.ts
+++ b/src/lib/waitUntil.ts
@@ -1,10 +1,28 @@
+import type { ExecutionContext, IncomingRequestCfProperties } from '@cloudflare/workers-types';
+
 let pendings: Array<Promise<unknown>> = [];
 
+/**
+ * Get a global context object for the current execution.
+ * This object can be used to store data that should be shared between the middleware and the handler
+ * and re-used for all requests.
+ */
+export async function getGlobalContext(): Promise<ExecutionContext | object> {
+    if (process.env.NODE_ENV === 'test') {
+        // Do not try loading the next-on-pages package in tests as it'll fail
+        return globalThis;
+    }
+
+    // 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()?.ctx ?? globalThis;
+}
+
 /**
  * Get a global context object for the current request.
  * This object can be used as a key to store request-specific data in a WeakMap.
  */
-export async function getGlobalContext(): Promise<object> {
+export async function getRequestContext(): Promise<IncomingRequestCfProperties | object> {
     if (process.env.NODE_ENV === 'test') {
         // Do not try loading the next-on-pages package in tests as it'll fail
         return globalThis;
@@ -29,12 +47,9 @@ export async function waitUntil(promise: Promise<unknown>) {
         return;
     }
 
-    // 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');
-
-    const cloudflare = getOptionalRequestContext();
-    if (cloudflare) {
-        cloudflare.ctx.waitUntil(promise);
+    const cloudflareContext = await getGlobalContext();
+    if ('waitUntil' in cloudflareContext) {
+        cloudflareContext.waitUntil(promise);
     } else {
         await promise;
     }