From aec74d68fdae95e950f7d17fc4a9d9d2b581aa21 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Samy=20Pess=C3=A9?= <samypesse@gmail.com>
Date: Tue, 23 Jan 2024 23:19:23 +0100
Subject: [PATCH] Stop using shared promises (#114)

* Try cf.cacheTags

* Remove shared promises

* Use a pendings set per context
---
 src/lib/cache/cache.ts      | 13 +++++++++++--
 src/lib/cache/cloudflare.ts |  6 +++++-
 src/lib/cache/waitUntil.ts  |  8 ++++++++
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/src/lib/cache/cache.ts b/src/lib/cache/cache.ts
index 2a992737e..0fc56f9c9 100644
--- a/src/lib/cache/cache.ts
+++ b/src/lib/cache/cache.ts
@@ -2,7 +2,7 @@ import hash from 'object-hash';
 
 import { cacheBackends } from './backends';
 import { CacheEntry } from './types';
-import { waitUntil } from './waitUntil';
+import { getGlobalContext, waitUntil } from './waitUntil';
 
 export type CacheFunction<Args extends any[], Result> = ((...args: Args) => Promise<Result>) & {
     /**
@@ -106,12 +106,21 @@ export function cache<Args extends any[], Result>(
     // During development, for now it fetches data twice between the middleware and the handler.
     // TODO: find a way to share the cache between the two.
 
-    const pendings = new Map<string, Promise<any>>();
+    // On Cloudflare Workers, we can't share promises between requests:
+    // > Cannot perform I/O on behalf of a different request. I/O objects (such as streams, request/response bodies, and others)
+    // > created in the context of one request handler cannot be accessed from a different request's handler.
+    //
+    // To avoid this limitation and still avoid concurrent requests, we use a WeakMap to store the pending requests.
+    const contextPendings = new WeakMap<object, Map<string, Promise<any>>>();
 
     const cacheFn = async (...args: Args) => {
         const cacheArgs = options.extractArgs ? options.extractArgs(args) : args;
         const key = getCacheKey(cacheName, cacheArgs);
 
+        const context = getGlobalContext();
+        const pendings = contextPendings.get(context) ?? new Map<string, Promise<any>>();
+        contextPendings.set(context, pendings);
+
         // If a pending request exists, wait for it
         if (pendings.has(key)) {
             return await pendings.get(key);
diff --git a/src/lib/cache/cloudflare.ts b/src/lib/cache/cloudflare.ts
index 016218e91..151cbda9b 100644
--- a/src/lib/cache/cloudflare.ts
+++ b/src/lib/cache/cloudflare.ts
@@ -79,17 +79,21 @@ async function serializeKey(key: string): Promise<string> {
 function serializeEntry(entry: CacheEntry): WorkerResponse {
     const headers = new Headers();
     headers.set('Content-Type', 'application/json');
+    const cacheTags = ['gitbook-open', ...entry.meta.tags];
 
     // Limit the cloudflare cache to 5 minutes
     headers.set(
         'Cache-Control',
         `public, max-age=${Math.min((entry.meta.expiresAt - Date.now()) / 1000, 5 * 60)}`,
     );
-    headers.set('Cache-Tag', ['gitbook-open', ...entry.meta.tags].join(','));
+    headers.set('Cache-Tag', cacheTags.join(','));
 
     // @ts-ignore
     return new Response(JSON.stringify(entry), {
         headers,
+        cf: {
+            cacheTags,
+        },
     });
 }
 
diff --git a/src/lib/cache/waitUntil.ts b/src/lib/cache/waitUntil.ts
index 8ec6ba4f0..e7c7daa14 100644
--- a/src/lib/cache/waitUntil.ts
+++ b/src/lib/cache/waitUntil.ts
@@ -1,5 +1,13 @@
 import { getRequestCloudflareContext } from '@cloudflare/next-on-pages/helpers';
 
+/**
+ * Get the global context object for the current request.
+ */
+export function getGlobalContext(): object {
+    const cloudflare = getRequestCloudflareContext();
+    return cloudflare ?? globalThis;
+}
+
 /**
  * Extend the lifetime of the event handler until the promise is resolved.
  * https://developers.cloudflare.com/workers/runtime-apis/handlers/fetch/#contextwaituntil