use DO tag cache in the server

Fix NextRequest body being incorrect

Author: Nicolas Dorseuil

packages/gitbook-v2/open-next.config.ts

@@ -8,7 +8,7 @@ export default {
             proxyExternalRequest: 'fetch',
             queue: () => import('./openNext/queue/server').then((m) => m.default),
             incrementalCache: () => import('./openNext/incrementalCache').then((m) => m.default),
-            tagCache: () => import('./openNext/tagCache/server').then((m) => m.default),
+            tagCache: () => import('./openNext/tagCache/middleware').then((m) => m.default),
         },
     },
     middleware: {

packages/gitbook-v2/openNext/customWorkers/default.js

@@ -1,10 +1,13 @@
 import { runWithCloudflareRequestContext } from '../../.open-next/cloudflare/init.js';
 
-import { handler } from '../../.open-next/server-functions/default/handler.mjs';
-
 export default {
     async fetch(request, env, ctx) {
         return runWithCloudflareRequestContext(request, env, ctx, async () => {
+            // We can't move the handler import to the top level, otherwise the runtime will not be properly initialized
+            const { handler } = await import(
+                '../../.open-next/server-functions/default/handler.mjs'
+            );
+
             // - `Request`s are handled by the Next server
             return handler(request, env, ctx);
         });

packages/gitbook-v2/openNext/customWorkers/defaultWrangler.jsonc

@@ -28,16 +28,14 @@
                 {
                     "binding": "WORKER_SELF_REFERENCE",
                     "service": "gitbook-open-v2-server-dev"
-                },
-                {
-                    "binding": "MIDDLEWARE_REFERENCE",
-                    "service": "gitbook-open-v2-dev"
                 }
             ]
         },
         "preview": {
             "vars": {
-                "STAGE": "preview"
+                "STAGE": "preview",
+                // Just as a test for the preview environment to check that everything works
+                "NEXT_PRIVATE_DEBUG_CACHE": "true"
             },
             "r2_buckets": [
                 {
@@ -49,10 +47,6 @@
                 {
                     "binding": "WORKER_SELF_REFERENCE",
                     "service": "gitbook-open-v2-server-preview"
-                },
-                {
-                    "binding": "MIDDLEWARE_REFERENCE",
-                    "service": "gitbook-open-v2-preview"
                 }
             ]
             // No durable objects on preview, as they block the generation of preview URLs
@@ -69,10 +63,21 @@
                 {
                     "binding": "WORKER_SELF_REFERENCE",
                     "service": "gitbook-open-v2-server-staging"
-                },
+                }
+            ],
+            "durable_objects": {
+                "bindings": [
+                    {
+                        "name": "NEXT_TAG_CACHE_DO_SHARDED",
+                        "class_name": "DOShardedTagCache",
+                        "script_name": "gitbook-open-v2-staging"
+                    }
+                ]
+            },
+            "migrations": [
                 {
-                    "binding": "MIDDLEWARE_REFERENCE",
-                    "service": "gitbook-open-v2-staging"
+                    "tag": "v1",
+                    "new_sqlite_classes": ["DOShardedTagCache"]
                 }
             ],
             "tail_consumers": [
@@ -92,10 +97,21 @@
                 {
                     "binding": "WORKER_SELF_REFERENCE",
                     "service": "gitbook-open-v2-server-production"
-                },
+                }
+            ],
+            "durable_objects": {
+                "bindings": [
+                    {
+                        "name": "NEXT_TAG_CACHE_DO_SHARDED",
+                        "class_name": "DOShardedTagCache",
+                        "script_name": "gitbook-open-v2-production"
+                    }
+                ]
+            },
+            "migrations": [
                 {
-                    "binding": "MIDDLEWARE_REFERENCE",
-                    "service": "gitbook-open-v2-production"
+                    "tag": "v1",
+                    "new_sqlite_classes": ["DOShardedTagCache"]
                 }
             ],
             "tail_consumers": [

packages/gitbook-v2/openNext/customWorkers/middleware.js

@@ -1,8 +1,6 @@
 import { WorkerEntrypoint } from 'cloudflare:workers';
 import { runWithCloudflareRequestContext } from '../../.open-next/cloudflare/init.js';
 
-import onConfig from '../../.open-next/middleware/open-next.config.mjs';
-
 import { handler as middlewareHandler } from '../../.open-next/middleware/handler.mjs';
 
 export { DOQueueHandler } from '../../.open-next/.build/durable-objects/queue.js';
@@ -19,10 +17,11 @@ export default class extends WorkerEntrypoint {
             }
 
             if (this.env.STAGE !== 'preview') {
+                reqOrResp.headers.set(
+                    'Cloudflare-Workers-Version-Overrides',
+                    `gitbook-open-v2-${this.env.STAGE}="${this.env.WORKER_VERSION_ID}"`
+                );
                 return this.env.DEFAULT_WORKER?.fetch(reqOrResp, {
-                    headers: {
-                        'Cloudflare-Workers-Version-Overrides': `gitbook-open-v2-${this.env.STAGE}="${this.env.WORKER_VERSION_ID}"`,
-                    },
                     cf: {
                         cacheEverything: false,
                     },
@@ -39,37 +38,4 @@ export default class extends WorkerEntrypoint {
             });
         });
     }
-
-    /**
-     * All the functions below are used to interact with the tag cache.
-     * They are needed for the server who wouldn't be able to access the tag cache directly.
-     */
-
-    async getLastRevalidated() {
-        return runWithCloudflareRequestContext(
-            new Request('http://local'),
-            this.env,
-            this.ctx,
-            async () => {
-                const tagCache = await onConfig.middleware.override.tagCache();
-                if (tagCache) {
-                    return tagCache.getLastRevalidated();
-                }
-            }
-        );
-    }
-
-    async writeTags(tags) {
-        return runWithCloudflareRequestContext(
-            new Request('http://local'),
-            this.env,
-            this.ctx,
-            async () => {
-                const tagCache = await onConfig.middleware.override.tagCache();
-                if (tagCache) {
-                    return tagCache.writeTags(tags);
-                }
-            }
-        );
-    }
 }

packages/gitbook-v2/openNext/customWorkers/middlewareWrangler.jsonc

@@ -20,7 +20,8 @@
     "env": {
         "dev": {
             "vars": {
-                "STAGE": "dev"
+                "STAGE": "dev",
+                "NEXT_PRIVATE_DEBUG_CACHE": "true"
             },
             "r2_buckets": [
                 {

packages/gitbook-v2/openNext/queue/middleware.ts

@@ -1,3 +1,4 @@
+import { trace } from '@/lib/tracing';
 import type { Queue } from '@opennextjs/aws/types/overrides.js';
 import { getCloudflareContext } from '@opennextjs/cloudflare';
 import doQueue from '@opennextjs/cloudflare/overrides/queue/do-queue';
@@ -10,8 +11,11 @@ interface Env {
 export default {
     name: 'GitbookISRQueue',
     send: async (msg) => {
-        const { ctx, env } = getCloudflareContext();
-        const hasDurableObject = (env as Env).STAGE !== 'dev' && (env as Env).STAGE !== 'preview';
-        ctx.waitUntil(hasDurableObject ? memoryQueue.send(msg) : doQueue.send(msg));
+        return trace({ operation: 'gitbookISRQueueSend', name: msg.MessageBody.url }, async () => {
+            const { ctx, env } = getCloudflareContext();
+            const hasDurableObject =
+                (env as Env).STAGE !== 'dev' && (env as Env).STAGE !== 'preview';
+            ctx.waitUntil(hasDurableObject ? memoryQueue.send(msg) : doQueue.send(msg));
+        });
     },
 } satisfies Queue;

packages/gitbook-v2/openNext/tagCache/middleware.ts

@@ -1,47 +1,76 @@
+import { trace } from '@/lib/tracing';
 import type { NextModeTagCache } from '@opennextjs/aws/types/overrides.js';
 import doShardedTagCache from '@opennextjs/cloudflare/overrides/tag-cache/do-sharded-tag-cache';
-import {
-    softTagFilter,
-    withFilter,
-} from '@opennextjs/cloudflare/overrides/tag-cache/tag-cache-filter';
+import { softTagFilter } from '@opennextjs/cloudflare/overrides/tag-cache/tag-cache-filter';
 
-const filteredTagCache = withFilter({
-    tagCache: doShardedTagCache({
-        baseShardSize: 12,
-        regionalCache: true,
-        shardReplication: {
-            numberOfSoftReplicas: 2,
-            numberOfHardReplicas: 1,
-        },
-    }),
-    // We don't use `revalidatePath`, so we filter out soft tags
-    filterFn: softTagFilter,
+const originalTagCache = doShardedTagCache({
+    baseShardSize: 12,
+    regionalCache: true,
+    shardReplication: {
+        numberOfSoftReplicas: 2,
+        numberOfHardReplicas: 1,
+    },
 });
 
 export default {
     name: 'GitbookTagCache',
     mode: 'nextMode',
     getLastRevalidated: async (tags: string[]) => {
-        try {
-            return await filteredTagCache.getLastRevalidated(tags);
-        } catch (error) {
-            console.error('Error getting last revalidated tags:', error);
-            return 0; // Fallback to 0 if there's an error
+        const tagsToCheck = tags.filter(softTagFilter);
+        if (tagsToCheck.length === 0) {
+            // If we reach here, it probably means that there is an issue that we'll need to address.
+            console.warn(
+                'getLastRevalidated - No valid tags to check for last revalidation, original tags:',
+                tags
+            );
+            return 0; // If no tags to check, return 0
         }
+        return trace(
+            {
+                operation: 'gitbookTagCacheGetLastRevalidated',
+                name: tagsToCheck.join(', '),
+            },
+            async () => {
+                return await originalTagCache.getLastRevalidated(tagsToCheck);
+            }
+        );
     },
     hasBeenRevalidated: async (tags: string[]) => {
-        try {
-            return await filteredTagCache.hasBeenRevalidated(tags);
-        } catch (error) {
-            console.error('Error checking if tags have been revalidated:', error);
-            return false; // Fallback to false if there's an error
+        const tagsToCheck = tags.filter(softTagFilter);
+        if (tagsToCheck.length === 0) {
+            // If we reach here, it probably means that there is an issue that we'll need to address.
+            console.warn(
+                'hasBeenRevalidated - No valid tags to check for revalidation, original tags:',
+                tags
+            );
+            return false; // If no tags to check, return false
         }
+        return trace(
+            {
+                operation: 'gitbookTagCacheHasBeenRevalidated',
+                name: tagsToCheck.join(', '),
+            },
+            async () => {
+                const result = await originalTagCache.hasBeenRevalidated(tagsToCheck);
+                return result;
+            }
+        );
     },
     writeTags: async (tags: string[]) => {
-        try {
-            await filteredTagCache.writeTags(tags);
-        } catch (error) {
-            console.error('Error writing tags:', error);
-        }
+        return trace(
+            {
+                operation: 'gitbookTagCacheWriteTags',
+                name: tags.join(', '),
+            },
+            async () => {
+                const tagsToWrite = tags.filter(softTagFilter);
+                if (tagsToWrite.length === 0) {
+                    console.warn('writeTags - No valid tags to write');
+                    return; // If no tags to write, exit early
+                }
+                // Write only the filtered tags
+                await originalTagCache.writeTags(tagsToWrite);
+            }
+        );
     },
 } satisfies NextModeTagCache;