Garbage-collect orphaned thumbnails via reference markers

functions/webdav/copy.ts

@@ -1,7 +1,7 @@
 import pLimit from "p-limit";
 
 import { notFound } from "./utils";
-import { listAll, RequestHandlerParams, WEBDAV_ENDPOINT } from "./utils";
+import { copyObject, listAll, RequestHandlerParams, WEBDAV_ENDPOINT } from "./utils";
 
 export async function handleRequestCopy({
   bucket,
@@ -33,10 +33,7 @@ export async function handleRequestCopy({
   const destinationExists = await bucket.head(destination);
   if (dontOverwrite && destinationExists)
     return new Response("Precondition Failed", { status: 412 });
-  await bucket.put(destination, src.body, {
-    httpMetadata: src.httpMetadata,
-    customMetadata: src.customMetadata,
-  });
+  await copyObject(bucket, src, destination);
 
   const isDirectory =
     src.httpMetadata?.contentType === "application/x-directory";
@@ -51,10 +48,7 @@ export async function handleRequestCopy({
           const target = `${destination}/${object.key.slice(prefix.length)}`;
           const src = await bucket.get(object.key);
           if (src === null) return;
-          await bucket.put(target, src.body, {
-            httpMetadata: object.httpMetadata,
-            customMetadata: object.customMetadata,
-          });
+          await copyObject(bucket, src, target);
         };
         const limit = pLimit(5);
         const promises = [];

functions/webdav/delete.ts

@@ -1,5 +1,13 @@
 import { notFound } from "./utils";
-import { listAll, RequestHandlerParams } from "./utils";
+import {
+  deletePrefix,
+  gcThumbnail,
+  isValidThumbnailDigest,
+  listAll,
+  RequestHandlerParams,
+  THUMBNAIL_PREFIX,
+  thumbnailRefKey,
+} from "./utils";
 
 export async function handleRequestDelete({
   bucket,
@@ -9,13 +17,36 @@ export async function handleRequestDelete({
     const obj = await bucket.head(path);
     if (obj === null) return notFound();
     await bucket.delete(path);
-    if (obj.httpMetadata?.contentType !== "application/x-directory")
+    if (obj.httpMetadata?.contentType !== "application/x-directory") {
+      const digest = obj.customMetadata?.thumbnail;
+      if (isValidThumbnailDigest(digest)) {
+        await bucket.delete(thumbnailRefKey(digest, path));
+        await gcThumbnail(bucket, digest);
+      }
       return new Response(null, { status: 204 });
+    }
   }
 
+  // Remove one marker per deleted child first, then garbage-collect each
+  // touched thumbnail once, so duplicates sharing a thumbnail are handled.
+  const digests = new Set<string>();
   const children = listAll(bucket, path === "" ? undefined : `${path}/`);
   for await (const child of children) {
     await bucket.delete(child.key);
+    const { thumbnail } = child.customMetadata ?? {};
+    if (isValidThumbnailDigest(thumbnail)) {
+      digests.add(thumbnail);
+      await bucket.delete(thumbnailRefKey(thumbnail, child.key));
+    }
+  }
+  for (const digest of digests) {
+    await gcThumbnail(bucket, digest);
+  }
+
+  // "Delete all" removes every user object; drop the shared thumbnails with
+  // them, as listAll skips the internal `_$flaredrive$/` subtree.
+  if (path === "") {
+    await deletePrefix(bucket, THUMBNAIL_PREFIX);
   }
 
   return new Response(null, { status: 204 });

functions/webdav/post.ts

@@ -1,14 +1,25 @@
 import { notFound } from "./utils";
-import { RequestHandlerParams } from "./utils";
+import {
+  addThumbnailRef,
+  isValidThumbnailDigest,
+  RequestHandlerParams,
+} from "./utils";
 
 export async function handleRequestPostCreateMultipart({
   bucket,
   path,
   request,
 }: RequestHandlerParams) {
-  const thumbnail = request.headers.get("fd-thumbnail");
+  const rawThumbnail = request.headers.get("fd-thumbnail");
+  const thumbnail = isValidThumbnailDigest(rawThumbnail)
+    ? rawThumbnail
+    : undefined;
   const customMetadata = thumbnail ? { thumbnail } : undefined;
 
+  // Reference the thumbnail before the upload starts, so a concurrent delete
+  // of the last other file using it cannot garbage-collect it mid-upload.
+  if (thumbnail) await addThumbnailRef(bucket, thumbnail, path);
+
   const multipartUpload = await bucket.createMultipartUpload(path, {
     httpMetadata: request.headers,
     customMetadata,

functions/webdav/put.ts

@@ -1,4 +1,12 @@
-import { RequestHandlerParams, ROOT_OBJECT } from "./utils";
+import {
+  RequestHandlerParams,
+  ROOT_OBJECT,
+  addThumbnailRef,
+  isValidThumbnailDigest,
+  releaseThumbnailRef,
+  thumbnailObjectKey,
+  thumbnailRefKey,
+} from "./utils";
 
 async function handleRequestPutMultipart({
   bucket,
@@ -46,16 +54,42 @@ export async function handleRequestPut({
     if (parentDir === null) return new Response("Conflict", { status: 409 });
   }
 
-  const thumbnail = request.headers.get("fd-thumbnail");
-  const customMetadata = thumbnail ? { thumbnail } : undefined;
+  const rawThumbnail = request.headers.get("fd-thumbnail");
+  const thumbnail = isValidThumbnailDigest(rawThumbnail)
+    ? rawThumbnail
+    : undefined;
+
+  // Overwriting a file replaces its customMetadata, so remember whether the
+  // previous object referenced a thumbnail that will need releasing.
+  const prev = await bucket.head(path);
+  const prevThumbnail = prev?.customMetadata?.thumbnail;
+
+  if (thumbnail) {
+    // Reference the thumbnail before the file exists, so a concurrent delete
+    // of the last other file using it cannot garbage-collect it mid-upload.
+    await addThumbnailRef(bucket, thumbnail, path);
+    // The client uploads the thumbnail blob just before the file; if a
+    // concurrent GC removed it in between, fail so the upload can be retried.
+    if ((await bucket.head(thumbnailObjectKey(thumbnail))) === null) {
+      await bucket.delete(thumbnailRefKey(thumbnail, path));
+      return new Response("Thumbnail is missing", { status: 409 });
+    }
+  }
 
   const result = await bucket.put(path, request.body, {
     onlyIf: request.headers,
     httpMetadata: request.headers,
-    customMetadata,
+    customMetadata: thumbnail ? { thumbnail } : undefined,
   });
 
-  if (!result) return new Response("Preconditions failed", { status: 412 });
+  if (!result) {
+    if (thumbnail) await bucket.delete(thumbnailRefKey(thumbnail, path));
+    return new Response("Preconditions failed", { status: 412 });
+  }
+
+  if (prevThumbnail !== thumbnail) {
+    await releaseThumbnailRef(bucket, prevThumbnail, path);
+  }
 
   return new Response("", { status: 201 });
 }

functions/webdav/utils.ts

@@ -34,6 +34,99 @@ export function parseBucketPath(context: any): [R2Bucket, string] {
   return [env[driveid] || env["BUCKET"], path];
 }
 
+// Thumbnails are content-addressed (`_$flaredrive$/thumbnails/<digest>.png`)
+// and shared by every file whose generated thumbnail has the same digest, so a
+// file delete cannot remove one outright. Instead, each referencing file keeps
+// an immutable marker object under THUMBNAIL_REFS_PREFIX; a thumbnail is
+// deleted only once no marker remains. Markers are written before the
+// referencing object and removed after the reference is gone, so a concurrent
+// delete can at worst over-retain a thumbnail, never break one still in use.
+export const THUMBNAIL_PREFIX = "_$flaredrive$/thumbnails/";
+export const THUMBNAIL_REFS_PREFIX = `${THUMBNAIL_PREFIX}refs/`;
+
+// Digests are hex-encoded hashes produced by the client (e.g. SHA-1).
+const THUMBNAIL_DIGEST_RE = /^[a-f0-9]{16,128}$/;
+
+export function isValidThumbnailDigest(digest: unknown): digest is string {
+  return typeof digest === "string" && THUMBNAIL_DIGEST_RE.test(digest);
+}
+
+export function thumbnailObjectKey(digest: string): string {
+  return `${THUMBNAIL_PREFIX}${digest}.png`;
+}
+
+export function thumbnailRefKey(digest: string, path: string): string {
+  return `${THUMBNAIL_REFS_PREFIX}${digest}/${path}`;
+}
+
+export async function addThumbnailRef(
+  bucket: R2Bucket,
+  digest: unknown,
+  path: string
+) {
+  if (!isValidThumbnailDigest(digest)) return;
+  await bucket.put(thumbnailRefKey(digest, path), "");
+}
+
+/** Deletes a thumbnail if no reference marker is left for it. */
+export async function gcThumbnail(bucket: R2Bucket, digest: unknown) {
+  if (!isValidThumbnailDigest(digest)) return;
+  const refs = await bucket.list({
+    prefix: thumbnailRefKey(digest, ""),
+    limit: 1,
+  });
+  if (refs.objects.length === 0) {
+    await bucket.delete(thumbnailObjectKey(digest));
+  }
+}
+
+/** Removes one reference marker, then garbage-collects the thumbnail. */
+export async function releaseThumbnailRef(
+  bucket: R2Bucket,
+  digest: unknown,
+  path: string
+) {
+  if (!isValidThumbnailDigest(digest)) return;
+  await bucket.delete(thumbnailRefKey(digest, path));
+  await gcThumbnail(bucket, digest);
+}
+
+/** Deletes every object under a key prefix (used for internal subtrees). */
+export async function deletePrefix(bucket: R2Bucket, prefix: string) {
+  let cursor: string | undefined = undefined;
+  do {
+    const listed = await bucket.list({ prefix, cursor });
+    if (listed.objects.length > 0) {
+      await bucket.delete(listed.objects.map((obj) => obj.key));
+    }
+    if (!listed.truncated) return;
+    cursor = listed.cursor;
+  } while (cursor);
+}
+
+/**
+ * Copies an object while keeping thumbnail references in sync: adds a marker
+ * for the copy before writing it, and releases the destination's previous
+ * thumbnail (if any) after a successful overwrite.
+ */
+export async function copyObject(
+  bucket: R2Bucket,
+  source: R2ObjectBody,
+  destination: string
+) {
+  const prev = await bucket.head(destination);
+  const prevDigest = prev?.customMetadata?.thumbnail;
+  const digest = source.customMetadata?.thumbnail;
+  await addThumbnailRef(bucket, digest, destination);
+  await bucket.put(destination, source.body, {
+    httpMetadata: source.httpMetadata,
+    customMetadata: source.customMetadata,
+  });
+  if (prevDigest !== digest) {
+    await releaseThumbnailRef(bucket, prevDigest, destination);
+  }
+}
+
 export async function* listAll(
   bucket: R2Bucket,
   prefix?: string,