Rework context to actually share memory cache between middleware and handlers for all requests (#2358)

Author: jpreynat

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

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;
+}

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