initial implementation of a sharded tag cache

Author: conico974

examples/e2e/app-router/open-next.config.ts

@@ -2,9 +2,10 @@ import { defineCloudflareConfig } from "@opennextjs/cloudflare";
 import d1TagCache from "@opennextjs/cloudflare/d1-tag-cache";
 import kvIncrementalCache from "@opennextjs/cloudflare/kv-cache";
 import doQueue from "@opennextjs/cloudflare/durable-queue";
+import shardedTagCache from "@opennextjs/cloudflare/do-sharded-tag-cache";
 
 export default defineCloudflareConfig({
   incrementalCache: kvIncrementalCache,
-  tagCache: d1TagCache,
+  tagCache: shardedTagCache,
   queue: doQueue,
 });

examples/e2e/app-router/wrangler.jsonc

@@ -13,13 +13,17 @@
       {
         "name": "NEXT_CACHE_REVALIDATION_DURABLE_OBJECT",
         "class_name": "DurableObjectQueueHandler"
+      },
+      {
+        "name": "NEXT_CACHE_D1_SHARDED",
+        "class_name": "DOShardedTagCache"
       }
     ]
   },
   "migrations": [
     {
       "tag": "v1",
-      "new_sqlite_classes": ["DurableObjectQueueHandler"]
+      "new_sqlite_classes": ["DurableObjectQueueHandler", "DOShardedTagCache"]
     }
   ],
   "kv_namespaces": [

packages/cloudflare/src/api/cloudflare-context.ts

@@ -1,6 +1,7 @@
 import type { Context, RunningCodeOptions } from "node:vm";
 
 import type { DurableObjectQueueHandler } from "./durable-objects/queue";
+import { DOShardedTagCache } from "./durable-objects/sharded-tag-cache";
 
 declare global {
   interface CloudflareEnv {
@@ -16,6 +17,9 @@ declare global {
     NEXT_CACHE_REVALIDATION_WORKER?: Service;
     // Durable Object namespace to use for the durable object queue handler
     NEXT_CACHE_REVALIDATION_DURABLE_OBJECT?: DurableObjectNamespace<DurableObjectQueueHandler>;
+    // Durables object namespace to use for the sharded tag cache
+    NEXT_CACHE_D1_SHARDED?: DurableObjectNamespace<DOShardedTagCache>;
+
     // Asset binding
     ASSETS?: Fetcher;
 

packages/cloudflare/src/api/do-sharded-tag-cache.ts

@@ -0,0 +1,84 @@
+import { debug } from "@opennextjs/aws/adapters/logger.js";
+import { generateMessageGroupId } from "@opennextjs/aws/core/routing/queue.js";
+import type { OpenNextConfig } from "@opennextjs/aws/types/open-next";
+import type { NextModeTagCache } from "@opennextjs/aws/types/overrides.js";
+import { IgnorableError } from "@opennextjs/aws/utils/error.js";
+
+import { getCloudflareContext } from "./cloudflare-context";
+
+class ShardedD1TagCache implements NextModeTagCache {
+  mode = "nextMode" as const;
+  public readonly name = "sharded-d1-tag-cache";
+
+  private getDurableObjectStub(shardId: string) {
+    const durableObject = getCloudflareContext().env.NEXT_CACHE_D1_SHARDED;
+    if (!durableObject) throw new IgnorableError("No durable object binding for cache revalidation");
+
+    const id = durableObject.idFromName(shardId);
+    return durableObject.get(id);
+  }
+
+  private generateShards(tags: string[]) {
+    // For each tag, we generate a message group id
+    const messageGroupIds = tags.map((tag) => ({ shardId: generateMessageGroupId(tag), tag }));
+    // We group the tags by shard
+    const shards = new Map<string, string[]>();
+    for (const { shardId, tag } of messageGroupIds) {
+      const tags = shards.get(shardId) ?? [];
+      tags.push(tag);
+      shards.set(shardId, tags);
+    }
+    return shards;
+  }
+
+  private getConfig() {
+    const cfEnv = getCloudflareContext().env;
+    const db = cfEnv.NEXT_CACHE_D1_SHARDED;
+
+    if (!db) debug("No Durable object found");
+
+    const isDisabled = !!(globalThis as unknown as { openNextConfig: OpenNextConfig }).openNextConfig
+      .dangerous?.disableTagCache;
+
+    if (!db || isDisabled) {
+      return { isDisabled: true as const };
+    }
+
+    return {
+      isDisabled: false as const,
+      db,
+    };
+  }
+
+  async hasBeenRevalidated(tags: string[], lastModified?: number): Promise<boolean> {
+    const { isDisabled } = this.getConfig();
+    if (isDisabled) return false;
+    const shards = this.generateShards(tags);
+    let hasBeenRevalidated = false;
+    console.log("Checking revalidation for tags", tags, shards);
+    // We then create a new durable object for each shard
+    await Promise.all(
+      Array.from(shards.entries()).map(async ([shardId, shardedTags]) => {
+        const stub = this.getDurableObjectStub(shardId);
+        hasBeenRevalidated = hasBeenRevalidated || (await stub.hasBeenRevalidated(shardedTags, lastModified));
+      })
+    );
+    return hasBeenRevalidated;
+  }
+
+  async writeTags(tags: string[]): Promise<void> {
+    const { isDisabled } = this.getConfig();
+    if (isDisabled) return;
+    const shards = this.generateShards(tags);
+    console.log("Writing tags", tags, shards);
+    // We then create a new durable object for each shard
+    await Promise.all(
+      Array.from(shards.entries()).map(async ([shardId, shardedTags]) => {
+        const stub = this.getDurableObjectStub(shardId);
+        await stub.writeTags(shardedTags);
+      })
+    );
+  }
+}
+
+export default new ShardedD1TagCache();

packages/cloudflare/src/api/durable-objects/sharded-tag-cache.ts

@@ -0,0 +1,36 @@
+import { DurableObject } from "cloudflare:workers";
+
+export class DOShardedTagCache extends DurableObject<CloudflareEnv> {
+  sql: SqlStorage;
+
+  constructor(state: DurableObjectState, env: CloudflareEnv) {
+    super(state, env);
+    this.sql = state.storage.sql;
+    state.blockConcurrencyWhile(async () => {
+      this.sql.exec(`CREATE TABLE IF NOT EXISTS revalidations (tag TEXT PRIMARY KEY, revalidatedAt INTEGER)`);
+    });
+  }
+
+  async hasBeenRevalidated(tags: string[], lastModified?: number): Promise<boolean> {
+    const result = this.sql
+      .exec<{
+        cnt: number;
+      }>(
+        `SELECT COUNT(*) as cnt FROM revalidations WHERE tag IN (${tags.map(() => "?").join(", ")}) AND revalidatedAt > ?`,
+        ...tags,
+        lastModified ?? Date.now()
+      )
+      .one();
+    return result.cnt > 0;
+  }
+
+  async writeTags(tags: string[]): Promise<void> {
+    tags.forEach((tag) => {
+      this.sql.exec(
+        `INSERT OR REPLACE INTO revalidations (tag, revalidatedAt) VALUES (?, ?)`,
+        tag,
+        Date.now()
+      );
+    });
+  }
+}

packages/cloudflare/src/cli/build/open-next/compileDurableObjects.ts

@@ -6,7 +6,10 @@ import { type BuildOptions, esbuildSync, getPackagePath } from "@opennextjs/aws/
 
 export function compileDurableObjects(buildOpts: BuildOptions) {
   const _require = createRequire(import.meta.url);
-  const entryPoints = [_require.resolve("@opennextjs/cloudflare/durable-objects/queue")];
+  const entryPoints = [
+    _require.resolve("@opennextjs/cloudflare/durable-objects/queue"),
+    _require.resolve("@opennextjs/cloudflare/durable-objects/sharded-tag-cache"),
+  ];
 
   const { outputDir } = buildOpts;
 

packages/cloudflare/src/cli/templates/worker.ts

@@ -19,6 +19,8 @@ Object.defineProperty(globalThis, Symbol.for("__cloudflare-context__"), {
 
 //@ts-expect-error: Will be resolved by wrangler build
 export { DurableObjectQueueHandler } from "./.build/durable-objects/queue.js";
+//@ts-expect-error: Will be resolved by wrangler build
+export { DOShardedTagCache } from "./.build/durable-objects/sharded-tag-cache.js";
 
 // Populate process.env on the first request
 let processEnvPopulated = false;