Fix deletion cron so that it can handle large bucket sizes (#45)

Author: AdiRishi

src/crons/deleteOldCache.test.ts

@@ -1,7 +1,7 @@
 import { MockedFunction, afterEach, beforeEach, expect, test, vi } from 'vitest';
 import { Env } from '../';
 import { isDateOlderThan } from '../utils/date';
-import { deleteOldCache } from './deleteOldCache';
+import { CURSOR_SIZE, deleteOldCache } from './deleteOldCache';
 
 vi.mock('../utils/date', async (importActual) => {
   const actual = await importActual<typeof import('../utils/date')>();
@@ -48,7 +48,21 @@ describe('R2 delete cron', () => {
 
   test('should delete all artifacts when there are more than CURSOR_SIZE', async () => {
     isDateOlderThanMock.mockReturnValue(true);
-    for (let i = 0; i < 1000; i++) {
+    for (let i = 0; i < Math.round(CURSOR_SIZE * 1.5); i++) {
+      await workerEnv.R2_STORE.put(`${teamId}/${artifactId}-${i}`, artifactContent, {
+        customMetadata: { artifactTag },
+      });
+    }
+
+    await deleteOldCache(workerEnv);
+
+    const artifacts = await workerEnv.R2_STORE.list();
+    expect(artifacts.objects.length).toEqual(0);
+  });
+
+  test('should delete all artifacts when there are over 1000 items ready for deletion', async () => {
+    isDateOlderThanMock.mockReturnValue(true);
+    for (let i = 0; i < 1001; i++) {
       await workerEnv.R2_STORE.put(`${teamId}/${artifactId}-${i}`, artifactContent, {
         customMetadata: { artifactTag },
       });

src/crons/deleteOldCache.ts

@@ -1,31 +1,46 @@
 import { Env } from '..';
 import { isDateOlderThan } from '../utils/date';
 
-const CURSOR_SIZE = 500;
+// Cursor size should be kept below 1000 to avoid limits on bulk operations
+export const CURSOR_SIZE = 500;
 
-export async function deleteOldCache(env: Env) {
+class R2KeysForDeletion {
+  keys: string[] = [];
+  add(key: string) {
+    this.keys.push(key);
+  }
+}
+
+export async function deleteOldCache(env: Env): Promise<void> {
   const BUCKET_CUTOFF_HOURS = env.BUCKET_OBJECT_EXPIRATION_HOURS;
   let truncated = false;
   let cursor: string | undefined;
   let list: R2Objects;
-  const keysMarkedForDeletion: string[] = [];
+  const keysMarkedForDeletion: R2KeysForDeletion[] = [];
 
   do {
     list = await env.R2_STORE.list({ limit: CURSOR_SIZE, cursor });
     truncated = list.truncated;
     cursor = list.truncated ? list.cursor : undefined;
 
+    /**
+     * Deleting keys while iterating over the list can sometimes cause the list to be truncated.
+     * So we mark the keys for deletion and delete after at least one additional iteration.
+     */
+    const keysAvailableForDeletion = keysMarkedForDeletion.shift();
+    if (keysAvailableForDeletion) {
+      await env.R2_STORE.delete(keysAvailableForDeletion.keys);
+    }
+
+    const keysForDeletion = new R2KeysForDeletion();
     for (const object of list.objects) {
       if (isDateOlderThan(object.uploaded, BUCKET_CUTOFF_HOURS)) {
-        keysMarkedForDeletion.push(object.key);
+        keysForDeletion.add(object.key);
       }
     }
+    keysMarkedForDeletion.push(keysForDeletion);
   } while (truncated);
-
-  if (keysMarkedForDeletion.length > 0) {
-    // if (env.ENVIRONMENT === 'development') {
-    //   console.log(`Deleting ${keysMarkedForDeletion.length} keys`, keysMarkedForDeletion);
-    // }
-    await env.R2_STORE.delete(keysMarkedForDeletion);
+  for (const keysForDeletion of keysMarkedForDeletion) {
+    await env.R2_STORE.delete(keysForDeletion.keys);
   }
 }