Respond with all tags of all cache used for page generation (#2973)

Author: SamyPesse

packages/gitbook/src/cloudflare-entrypoint.ts

@@ -1,6 +1,7 @@
 // @ts-ignore
 import nextOnPagesHandler from '@cloudflare/next-on-pages/fetch-handler';
 
+import { withResponseCacheTags } from './lib/cache/response';
 import { withMiddlewareHeadersStorage } from './lib/middleware';
 
 /**
@@ -9,8 +10,8 @@ import { withMiddlewareHeadersStorage } from './lib/middleware';
  */
 export default {
     async fetch(request, env, ctx) {
-        const response = await withMiddlewareHeadersStorage(() =>
-            nextOnPagesHandler.fetch(request, env, ctx)
+        const response = await withResponseCacheTags(() =>
+            withMiddlewareHeadersStorage(() => nextOnPagesHandler.fetch(request, env, ctx))
         );
 
         return response;

packages/gitbook/src/lib/api.ts

@@ -178,6 +178,7 @@ export const getUserById = cache({
             tag: 'user',
             user: userId,
         }),
+    tagImmutable: true,
     get: async (userId: string, options: CacheFunctionOptions) => {
         try {
             const apiCtx = await api();
@@ -379,6 +380,7 @@ export const getRevision = cache({
     name: 'api.getRevision.v2',
     tag: (spaceId, revisionId) =>
         getCacheTag({ tag: 'revision', space: spaceId, revision: revisionId }),
+    tagImmutable: true,
     getKeySuffix: getAPIContextId,
     get: async (
         spaceId: string,
@@ -411,6 +413,7 @@ export const getRevisionPages = cache({
     name: 'api.getRevisionPages.v4',
     tag: (spaceId, revisionId) =>
         getCacheTag({ tag: 'revision', space: spaceId, revision: revisionId }),
+    tagImmutable: true,
     getKeySuffix: getAPIContextId,
     get: async (
         spaceId: string,
@@ -446,6 +449,7 @@ export const getRevisionPageByPath = cache({
     name: 'api.getRevisionPageByPath.v3',
     tag: (spaceId, revisionId) =>
         getCacheTag({ tag: 'revision', space: spaceId, revision: revisionId }),
+    tagImmutable: true,
     getKeySuffix: getAPIContextId,
     get: async (
         spaceId: string,
@@ -492,6 +496,7 @@ const getRevisionFileById = cache({
     name: 'api.getRevisionFile.v3',
     tag: (spaceId, revisionId) =>
         getCacheTag({ tag: 'revision', space: spaceId, revision: revisionId }),
+    tagImmutable: true,
     get: async (
         spaceId: string,
         revisionId: string,
@@ -528,6 +533,7 @@ const getRevisionReusableContentById = cache({
     name: 'api.getRevisionReusableContentById.v1',
     tag: (spaceId, revisionId) =>
         getCacheTag({ tag: 'revision', space: spaceId, revision: revisionId }),
+    tagImmutable: true,
     getKeySuffix: getAPIContextId,
     get: async (
         spaceId: string,
@@ -569,6 +575,7 @@ const getRevisionAllFiles = cache({
     name: 'api.getRevisionAllFiles.v2',
     tag: (spaceId, revisionId) =>
         getCacheTag({ tag: 'revision', space: spaceId, revision: revisionId }),
+    tagImmutable: true,
     get: async (spaceId: string, revisionId: string, options: CacheFunctionOptions) => {
         const response = await getAll(
             async (params) => {
@@ -684,6 +691,7 @@ export const getDocument = cache({
     name: 'api.getDocument.v2',
     tag: (spaceId, documentId) =>
         getCacheTag({ tag: 'document', space: spaceId, document: documentId }),
+    tagImmutable: true,
     getKeySuffix: getAPIContextId,
     get: async (spaceId: string, documentId: string, options: CacheFunctionOptions) => {
         const apiCtx = await api();

packages/gitbook/src/lib/cache/cache.ts

@@ -7,6 +7,7 @@ import { assertIsNotV2 } from '../v2';
 import { waitUntil } from '../waitUntil';
 import { cacheBackends } from './backends';
 import { memoryCache } from './memory';
+import { addResponseCacheTag } from './response';
 import type { CacheBackend, CacheEntry } from './types';
 
 export type CacheFunctionOptions = {
@@ -54,6 +55,9 @@ export interface CacheDefinition<Args extends any[], Result> {
     /** Tag to associate to the entry */
     tag?: (...args: Args) => string;
 
+    /** If true, the tag will not be sent to the HTTP response, as we consider it immutable */
+    tagImmutable?: boolean;
+
     /** Filter the arguments that should be taken into consideration for the cache key */
     getKeyArgs?: (args: Args) => any[];
 
@@ -134,6 +138,11 @@ export function cache<Args extends any[], Result>(
             let result: readonly [CacheEntry, string] | null = null;
             const tag = cacheDef.tag?.(...args);
 
+            // Add the cache tag to the HTTP response
+            if (tag && !cacheDef.tagImmutable) {
+                addResponseCacheTag(tag);
+            }
+
             // Try the memory backend, independently of the other backends as it doesn't have a network cost
             const memoryEntry = await memoryCache.get({ key, tag });
             if (memoryEntry) {

packages/gitbook/src/lib/cache/response.ts

@@ -0,0 +1,44 @@
+import { AsyncLocalStorage } from 'node:async_hooks';
+
+const symbolResponseCacheTagsLocalStorage = Symbol.for('__gitbook-response-cache-tags__');
+
+/**
+ * Wrap a response handler with a storage to store cache tags and write them as header to the response.
+ */
+export async function withResponseCacheTags(handler: () => Promise<Response>): Promise<Response> {
+    const responseCacheTagsLocalStorage =
+        // @ts-ignore
+        (globalThis[symbolResponseCacheTagsLocalStorage] as AsyncLocalStorage<Set<string>>) ??
+        new AsyncLocalStorage<Set<string>>();
+    // @ts-ignore
+    globalThis[symbolResponseCacheTagsLocalStorage] = responseCacheTagsLocalStorage;
+
+    const responseCacheTags = new Set<string>();
+    const response = await responseCacheTagsLocalStorage.run(responseCacheTags, handler);
+
+    if (responseCacheTags.size > 0) {
+        const headerCacheTag = Array.from(responseCacheTags).join(',');
+        response.headers.set('cache-tag', headerCacheTag);
+        response.headers.set('x-gitbook-cache-tag', headerCacheTag);
+    }
+
+    return response;
+}
+
+/**
+ * Add a cache tag to the HTTP response.
+ * This is a hack for v1, for v2, we use ISR and do not cache the whole response.
+ */
+export function addResponseCacheTag(...tags: string[]) {
+    const responseCacheTagsLocalStorage =
+        // @ts-ignore
+        (globalThis[symbolResponseCacheTagsLocalStorage] as AsyncLocalStorage<Set<string>>) ??
+        new AsyncLocalStorage<Set<string>>();
+    // @ts-ignore
+    globalThis[symbolResponseCacheTagsLocalStorage] = responseCacheTagsLocalStorage;
+
+    const store = responseCacheTagsLocalStorage.getStore();
+    tags.forEach((tag) => {
+        store?.add(tag);
+    });
+}

packages/gitbook/src/lib/middleware.ts

@@ -1,5 +1,9 @@
 import { AsyncLocalStorage } from 'node:async_hooks';
 
+/*
+ * This code is ONLY for v1.
+ */
+
 /**
  * Set a header on the middleware response.
  * We do this because of https://github.com/opennextjs/opennextjs-cloudflare/issues/92

packages/gitbook/src/middleware.ts

@@ -33,6 +33,7 @@ import {
 import { joinPath, normalizePathname } from '@/lib/paths';
 import { getProxyModeBasePath } from '@/lib/proxy';
 import { MiddlewareHeaders } from '@v2/lib/middleware';
+import { addResponseCacheTag } from './lib/cache/response';
 
 export const config = {
     matcher:
@@ -287,10 +288,8 @@ export async function middleware(request: NextRequest) {
         }
     }
 
-    if (resolved.cacheTags && resolved.cacheTags.length > 0) {
-        const headerCacheTag = resolved.cacheTags.join(',');
-        setMiddlewareHeader(response, 'cache-tag', headerCacheTag);
-        setMiddlewareHeader(response, 'x-gitbook-cache-tag', headerCacheTag);
+    if (resolved.cacheTags) {
+        addResponseCacheTag(...resolved.cacheTags);
     }
 
     return response;