feat[r2]: add gc routine

Author: SharzyL

src/index.ts

@@ -4,11 +4,17 @@ import { handleOptions, corsWrapResponse } from "./handlers/handleCors.js"
 import { handlePostOrPut } from "./handlers/handleWrite.js"
 import { handleGet } from "./handlers/handleRead.js"
 import { handleDelete } from "./handlers/handleDelete.js"
+import { cleanExpiredInR2 } from "./storage/storage.js"
 
 export default {
   async fetch(request: Request, env: Env, ctx: ExecutionContext): Promise<Response> {
     return await handleRequest(request, env, ctx)
   },
+
+  // eslint-disable-next-line @typescript-eslint/require-await
+  async scheduled(controller: ScheduledController, env, ctx) {
+    ctx.waitUntil(cleanExpiredInR2(env, controller))
+  },
 } satisfies ExportedHandler<Env>
 
 async function handleRequest(request: Request, env: Env, ctx: ExecutionContext): Promise<Response> {

src/storage/storage.ts

@@ -3,6 +3,15 @@ import { parseSize } from "../shared.js"
 
 export type PasteLocation = "KV" | "R2"
 
+// since CF does not allow expiration shorter than 60s, extend the expiration to 70s
+const PASTE_EXPIRE_SPECIFIED_MIN = 70
+
+/* Since we need the metadata stored in KV to perform R2 cleanup,
+ the paste in KV should not be deleted until it is cleaned in R2.
+ We extend the lifetime by 2 days to avoid it being cleaned in VK too early
+ */
+const PASTE_EXPIRE_EXTENSION_FOR_R2 = 2 * 24 * 60 * 60
+
 // TODO: allow admin to upload permanent paste
 // TODO: add filename length check
 export type PasteMetadata = {
@@ -102,7 +111,7 @@ export async function getPaste(env: Env, short: string, ctx: ExecutionContext):
     if (metadata.location === "R2") {
       const object = await env.R2.get(short)
       if (object === null) {
-        throw new WorkerError(404, `cannot find R2 bucket of name '${short}'`)
+        return null
       }
       return { paste: object.body, metadata }
     } else {
@@ -143,10 +152,12 @@ export async function updatePaste(
   },
 ) {
   const expirationUnix = dateToUnix(options.now) + options.expirationSeconds
-  // since CF does not allow expiration shorter than 60s, extend the expiration to 70s
-  const expirationUnixSpecified = dateToUnix(options.now) + Math.max(options.expirationSeconds, 70)
+  let expirationUnixSpecified =
+    dateToUnix(options.now) + Math.max(options.expirationSeconds, PASTE_EXPIRE_SPECIFIED_MIN)
 
   if (originalMetadata.location === "R2") {
+    expirationUnixSpecified = expirationUnixSpecified + PASTE_EXPIRE_EXTENSION_FOR_R2
+
     await env.R2.put(pasteName, content)
   }
   const metadata: PasteMetadata = {
@@ -182,11 +193,13 @@ export async function createPaste(
 ) {
   const expirationUnix = dateToUnix(options.now) + options.expirationSeconds
 
-  // since CF does not allow expiration shorter than 60s, extend the expiration to 70s
-  const expirationUnixSpecified = dateToUnix(options.now) + Math.max(options.expirationSeconds, 70)
+  let expirationUnixSpecified =
+    dateToUnix(options.now) + Math.max(options.expirationSeconds, PASTE_EXPIRE_SPECIFIED_MIN)
 
   const location = options.contentLength > parseSize(env.R2_THRESHOLD)! ? "R2" : "KV"
   if (location === "R2") {
+    expirationUnixSpecified = expirationUnixSpecified + PASTE_EXPIRE_EXTENSION_FOR_R2
+
     await env.R2.put(pasteName, content)
   }
 
@@ -226,3 +239,49 @@ export async function deletePaste(env: Env, pasteName: string, originalMetadata:
     await env.R2.delete(pasteName)
   }
 }
+
+export async function cleanExpiredInR2(env: Env, controller: ScheduledController) {
+  // types generated by wrangler somehow not working, so cast manually
+  type Listed = {
+    list_complete: false
+    keys: KVNamespaceListKey<PasteMetadataInStorage, string>[]
+    cursor: string
+    cacheStatus: string | null
+  }
+
+  const nowUnix = controller.scheduledTime / 1000
+
+  let numCleaned = 0
+  const r2NamesToClean: string[] = []
+
+  async function clean() {
+    await env.R2.delete(r2NamesToClean)
+    numCleaned += r2NamesToClean.length
+    r2NamesToClean.length = 0
+  }
+
+  let cursor: string | null = null
+  while (true) {
+    const listed = (await env.PB.list<PasteMetadataInStorage>({ cursor })) as Listed
+
+    cursor = listed.cursor
+
+    for (const key of listed.keys) {
+      if (key.metadata !== undefined) {
+        const metadata = migratePasteMetadata(key.metadata)
+        if (metadata.location === "R2" && metadata.willExpireAtUnix < nowUnix) {
+          r2NamesToClean.push(key.name)
+
+          if (r2NamesToClean.length === 1000) {
+            await clean()
+          }
+        }
+      }
+    }
+
+    if (listed.list_complete) break
+  }
+  await clean()
+
+  console.log(`${numCleaned} buckets cleaned`)
+}

test/r2.spec.ts

@@ -1,7 +1,16 @@
-import { expect, test } from "vitest"
+import { expect, test, vi, beforeEach, afterEach } from "vitest"
 import { areBlobsEqual, genRandomBlob, upload, workerFetch } from "./testUtils"
 import { parseSize } from "../src/shared"
-import { createExecutionContext, env } from "cloudflare:test"
+import { createExecutionContext, createScheduledController, env, waitOnExecutionContext } from "cloudflare:test"
+import worker from "../src/index.js"
+
+beforeEach(() => {
+  vi.useFakeTimers()
+})
+
+afterEach(() => {
+  vi.useRealTimers()
+})
 
 test("basic", async () => {
   const blob1 = genRandomBlob(parseSize(env.R2_THRESHOLD)! * 2)
@@ -26,3 +35,34 @@ test("basic", async () => {
   expect(revisitResp.status).toStrictEqual(200)
   expect(areBlobsEqual(await revisitResp.blob(), blob2)).toBeTruthy()
 })
+
+test("schedule", async () => {
+  const ctx = createExecutionContext()
+
+  // upload
+  vi.setSystemTime(new Date(2035, 0, 0))
+  const blob1 = genRandomBlob(parseSize(env.R2_THRESHOLD)! * 2)
+  const uploadResponse = await upload(ctx, { c: blob1, e: "7d" })
+  const url = uploadResponse.url
+
+  // test get
+  const getResp = await workerFetch(ctx, url)
+  expect(getResp.status).toStrictEqual(200)
+  await getResp.blob() // we must consume body to prevent breaking isolated storage
+
+  // go to past, nothing will be cleaned
+  await worker.scheduled(createScheduledController({ scheduledTime: new Date(2000, 0, 0) }), env, ctx)
+  await waitOnExecutionContext(ctx)
+
+  // test get after cleanup
+  const getResp1 = await workerFetch(ctx, url)
+  expect(getResp.status).toStrictEqual(200)
+  await getResp1.blob()
+
+  // jump to 1 year later, now all pastes are expired
+  await worker.scheduled(createScheduledController({ scheduledTime: new Date(2040, 0, 0) }), env, ctx)
+  await waitOnExecutionContext(ctx)
+
+  // test get after cleanup
+  expect((await workerFetch(ctx, url)).status).toStrictEqual(404)
+})

wrangler.toml

@@ -14,6 +14,10 @@ directory = "dist/frontend"
 run_worker_first = true
 binding = "ASSETS"
 
+[triggers]
+# clean r2 garbage every day
+crons = ["0 0 * * *"]
+
 #----------------------------------------
 # lines below are what you should modify
 #----------------------------------------