From dc2d50e455ba4da3205b19e2f1210358ee5cc23d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Samy=20Pess=C3=A9?= <samypesse@gmail.com>
Date: Wed, 20 Mar 2024 17:48:17 +0000
Subject: [PATCH] Iterate over all keys in KV when purging cache instead of
 just first page (#312)

* Iterate over all keys in KV when purging cache instead of just first key

* Call getTagPrefix once
---
 src/lib/cache/cloudflare-kv.ts | 49 +++++++++++++++++++---------------
 src/lib/cache/types.ts         |  2 +-
 2 files changed, 28 insertions(+), 23 deletions(-)

diff --git a/src/lib/cache/cloudflare-kv.ts b/src/lib/cache/cloudflare-kv.ts
index 8a04f6753..8845b6352 100644
--- a/src/lib/cache/cloudflare-kv.ts
+++ b/src/lib/cache/cloudflare-kv.ts
@@ -7,7 +7,6 @@ import { trace } from '../tracing';
 const cacheVersion = 1;
 
 interface KVTagMetadata {
-    key: string;
     meta: CacheEntryMeta;
 }
 
@@ -69,7 +68,6 @@ export const cloudflareKVCache: CacheBackend = {
 
                 if (entry.meta.tags.length > 0) {
                     const metadata: KVTagMetadata = {
-                        key, // TODO: Remove this key from the metadata in 1 day
                         meta: entry.meta,
                     };
                     const jsonMetadata = JSON.stringify(metadata);
@@ -115,28 +113,35 @@ export const cloudflareKVCache: CacheBackend = {
 
         const pendingDeletions: Array<Promise<unknown>> = [];
 
+        const iterateKVPage = async (prefix: string, cursor: string | null, max: number = 3) => {
+            const entries = await kv.list<KVTagMetadata>({
+                prefix,
+                cursor,
+                limit: 100,
+            });
+
+            for (const entry of entries.keys) {
+                if (entry.metadata) {
+                    const metadata = entry.metadata;
+                    const key = metadata.meta.key;
+
+                    result.metas.push(metadata.meta);
+                    result.keys.push(key);
+
+                    // Delete the tag key and the value key
+                    pendingDeletions.push(kv.delete(getValueKey(key)));
+                    pendingDeletions.push(kv.delete(entry.name));
+                }
+            }
+
+            if (!entries.list_complete && max > 0) {
+                await iterateKVPage(prefix, entries.cursor, max - 1);
+            }
+        };
+
         await Promise.all(
             tags.map(async (tag) => {
-                const entries = await kv.list({
-                    prefix: getTagPrefix(tag),
-                    // We don't paginate at the moment and only get the first 100 keys.
-                    limit: 100,
-                });
-
-                for (const entry of entries.keys) {
-                    if (entry.metadata) {
-                        const metadata = entry.metadata as KVTagMetadata;
-
-                        const key = metadata.meta.key ?? metadata.key;
-
-                        result.metas.push(metadata.meta);
-                        result.keys.push(key);
-
-                        // Delete the tag key and the value key
-                        pendingDeletions.push(kv.delete(getValueKey(key)));
-                        pendingDeletions.push(kv.delete(entry.name));
-                    }
-                }
+                await iterateKVPage(getTagPrefix(tag), null);
             }),
         );
 
diff --git a/src/lib/cache/types.ts b/src/lib/cache/types.ts
index 097ea95cd..ec12b9b46 100644
--- a/src/lib/cache/types.ts
+++ b/src/lib/cache/types.ts
@@ -1,5 +1,5 @@
 export interface CacheEntryMeta {
-    key?: string;
+    key: string;
 
     /**
      * Timestamp when the entry was created.