Fix regional caching of tag cache

packages/gitbook/openNext/customWorkers/default.js

@@ -2,6 +2,9 @@ import { runWithCloudflareRequestContext } from '../../.open-next/cloudflare/ini
 
 import { DurableObject } from 'cloudflare:workers';
 
+//Only needed to run locally, in prod we'll use the one from do.js
+export { DOShardedTagCache } from '../../.open-next/.build/durable-objects/sharded-tag-cache.js';
+
 // Only needed to run locally, in prod we'll use the one from do.js
 export class R2WriteBuffer extends DurableObject {
     writePromise;

packages/gitbook/openNext/customWorkers/defaultWrangler.jsonc

@@ -36,13 +36,17 @@
                     {
                         "name": "WRITE_BUFFER",
                         "class_name": "R2WriteBuffer"
+                    },
+                    {
+                        "name": "NEXT_TAG_CACHE_DO_SHARDED",
+                        "class_name": "DOShardedTagCache"
                     }
                 ]
             },
             "migrations": [
                 {
                     "tag": "v1",
-                    "new_sqlite_classes": ["R2WriteBuffer"]
+                    "new_sqlite_classes": ["R2WriteBuffer", "DOShardedTagCache"]
                 }
             ]
         },

packages/gitbook/openNext/customWorkers/middlewareWrangler.jsonc

@@ -22,6 +22,9 @@
             "vars": {
                 "STAGE": "dev",
                 "NEXT_PRIVATE_DEBUG_CACHE": "true",
+                // When deployed locally, we don't have access to the tag cache here,
+                // we should just bypass the cache to go to the server directly
+                "SHOULD_BYPASS_CACHE": "true",
                 "OPEN_NEXT_REQUEST_ID_HEADER": "true"
             },
             "r2_buckets": [

packages/gitbook/openNext/incrementalCache.ts

@@ -36,6 +36,11 @@ class GitbookIncrementalCache implements IncrementalCache {
         const r2 = getCloudflareContext().env[BINDING_NAME];
         const localCache = await this.getCacheInstance();
         if (!r2) throw new Error('No R2 bucket');
+        if (process.env.SHOULD_BYPASS_CACHE === 'true') {
+            // We are in a local middleware environment, we should bypass the cache
+            // and go directly to the server.
+            return null;
+        }
         try {
             // Check local cache first if available
             const localCacheEntry = await localCache.match(this.getCacheUrlKey(cacheKey));
@@ -134,19 +139,27 @@ class GitbookIncrementalCache implements IncrementalCache {
     }
 
     async writeToR2(key: string, value: string): Promise<void> {
-        const env = getCloudflareContext().env as {
-            WRITE_BUFFER: DurableObjectNamespace<
-                Rpc.DurableObjectBranded & {
-                    write: (key: string, value: string) => Promise<void>;
-                }
-            >;
-        };
-        const id = env.WRITE_BUFFER.idFromName(key);
-
-        // A stub is a client used to invoke methods on the Durable Object
-        const stub = env.WRITE_BUFFER.get(id);
-
-        await stub.write(key, value);
+        try {
+            const env = getCloudflareContext().env as {
+                WRITE_BUFFER: DurableObjectNamespace<
+                    Rpc.DurableObjectBranded & {
+                        write: (key: string, value: string) => Promise<void>;
+                    }
+                >;
+            };
+            const id = env.WRITE_BUFFER.idFromName(key);
+
+            // A stub is a client used to invoke methods on the Durable Object
+            const stub = env.WRITE_BUFFER.get(id);
+
+            await stub.write(key, value);
+        } catch {
+            // We fallback to writing directly to R2
+            // it can fail locally because the limit is 1Mb per args
+            // It is 32Mb in production, so we should be fine
+            const r2 = getCloudflareContext().env[BINDING_NAME];
+            r2?.put(key, value);
+        }
     }
 
     async getCacheInstance(): Promise<Cache> {

packages/gitbook/openNext/tagCache/middleware.ts

@@ -1,13 +1,15 @@
-import { getLogger } from '@/lib/logger';
+import { createLogger, getLogger } from '@/lib/logger';
+import { filterOutNullable } from '@/lib/typescript';
 import type { NextModeTagCache } from '@opennextjs/aws/types/overrides.js';
+import { getCloudflareContext } from '@opennextjs/cloudflare';
 import doShardedTagCache from '@opennextjs/cloudflare/overrides/tag-cache/do-sharded-tag-cache';
 import { softTagFilter } from '@opennextjs/cloudflare/overrides/tag-cache/tag-cache-filter';
 
 const originalTagCache = doShardedTagCache({
     baseShardSize: 12,
-    regionalCache: true,
-    // We can probably increase this value even further
-    regionalCacheTtlSec: 60,
+    // It is broken right now, need to be fixed in OpenNext - We might still not use it depending on how
+    // it will be implemented there.
+    regionalCache: false,
     shardReplication: {
         numberOfSoftReplicas: 2,
         numberOfHardReplicas: 1,
@@ -17,24 +19,124 @@ const originalTagCache = doShardedTagCache({
     },
 });
 
+function deduplicateTags(tags: string[]): string[] {
+    return Array.from(new Set(tags));
+}
+
+function getCacheKey(tag: string) {
+    return `http://regional.cache/${tag}`;
+}
+
+async function getFromRegionalCache(tags: string[]): Promise<(readonly [string, number])[]> {
+    try {
+        const cache = await caches.open('tag');
+
+        const responses = await Promise.all(
+            tags.map(async (tag) => {
+                const resp = await cache.match(getCacheKey(tag));
+                if (!resp) {
+                    return null;
+                }
+                return { tag, resp };
+            })
+        );
+
+        const result = responses
+            .filter(filterOutNullable)
+            .map(
+                async (response) => [response.tag, (await response.resp.json()) as number] as const
+            );
+
+        return Promise.all(result);
+    } catch {
+        //If we have an error here, we just fallback to the original tag cache
+        return [];
+    }
+}
+
+async function updateRegionalCache(tags: string[]) {
+    const regionalCache = await caches.open('tag');
+    for (const tag of tags) {
+        const cacheKey = getCacheKey(tag);
+        const lastRevalidated = (await originalTagCache.getLastRevalidated([tag])) || 0;
+
+        await regionalCache.put(
+            cacheKey,
+            new Response(JSON.stringify(lastRevalidated), {
+                headers: {
+                    'Content-Type': 'application/json',
+                    // We should be safe to cache this for a while.
+                    'Cache-Control': 'public, max-age=300',
+                    'Cache-Tag': tag,
+                },
+            })
+        );
+    }
+}
+
+async function deleteFromRegionalCache(tags: string[]) {
+    const regionalCache = await caches.open('tag');
+    await Promise.all(
+        tags.map(async (tag) => {
+            const cacheKey = getCacheKey(tag);
+            await regionalCache.delete(cacheKey);
+        })
+    );
+}
+
 export default {
     name: 'GitbookTagCache',
     mode: 'nextMode',
+    // We don't really use this one, as of now it is only used for soft tags
     getLastRevalidated: async (tags: string[]) => {
         const tagsToCheck = tags.filter(softTagFilter);
         if (tagsToCheck.length === 0) {
             return 0; // If no tags to check, return 0
         }
+        const deduplicatedTags = deduplicateTags(tagsToCheck);
 
-        return await originalTagCache.getLastRevalidated(tagsToCheck);
+        return await originalTagCache.getLastRevalidated(deduplicatedTags);
     },
     hasBeenRevalidated: async (tags: string[], lastModified?: number) => {
-        const tagsToCheck = tags.filter(softTagFilter);
-        if (tagsToCheck.length === 0) {
-            return false; // If no tags to check, return false
-        }
+        try {
+            const tagsToCheck = tags.filter(softTagFilter);
+            if (tagsToCheck.length === 0) {
+                return false; // If no tags to check, return false
+            }
+
+            const deduplicatedTags = deduplicateTags(tagsToCheck);
+            const regionalCacheResult = await getFromRegionalCache(deduplicatedTags);
+            if (regionalCacheResult.length > 0) {
+                // If we have results from the regional cache, check if any of them are newer than lastModified
+                const cacheRevalidated = regionalCacheResult.some(
+                    ([_, timestamp]) => timestamp >= (lastModified ?? 0)
+                );
+                if (cacheRevalidated) {
+                    // If any tag is revalidated, we can return true
+                    return true;
+                }
+            }
 
-        return await originalTagCache.hasBeenRevalidated(tagsToCheck, lastModified);
+            const remainingTags = deduplicatedTags.filter(
+                (tag) => !regionalCacheResult.some(([cachedTag]) => cachedTag === tag)
+            );
+            if (remainingTags.length > 0) {
+                // If there are remaining tags, check their status in the original cache
+                const result = await originalTagCache.hasBeenRevalidated(
+                    remainingTags,
+                    lastModified
+                );
+                getCloudflareContext().ctx.waitUntil(updateRegionalCache(remainingTags));
+                return result;
+            }
+            return false; // If no tags were found in the regional cache and no remaining tags, return false
+        } catch (e) {
+            createLogger('gitbookTagCache', {}).error(
+                `hasBeenRevalidated - Error checking tags ${tags.join(', ')}`,
+                e
+            );
+            return false; // In case of error, return false
+        }
     },
     writeTags: async (tags: string[]) => {
         const tagsToWrite = tags.filter(softTagFilter);
@@ -43,7 +145,13 @@ export default {
             logger.warn('writeTags - No valid tags to write');
             return; // If no tags to write, exit early
         }
+
+        const deduplicatedTags = deduplicateTags(tagsToWrite);
+
         // Write only the filtered tags
-        await originalTagCache.writeTags(tagsToWrite);
+        await originalTagCache.writeTags(deduplicatedTags);
+
+        // delete from regional cache
+        await deleteFromRegionalCache(deduplicatedTags);
     },
 } satisfies NextModeTagCache;