업로드 및 썸네일 등록의 경쟁 조건을 수정하고 캐시 무효화와 스트리밍 캐시의 메모리 사용량을 개선

src/lib/features/filesystem/cache.svelte.ts

@@ -11,6 +11,7 @@ interface FilesystemCacheOptions<K, V> {
 interface Entry<V> {
   value?: V;
   promise?: Promise<V>;
+  isInvalidated: boolean;
 }
 
 // Caches filesystem entries as reactive objects which are updated in place, so that every view
@@ -24,7 +25,7 @@ export class FilesystemCache<K, V extends object> {
   private getOrCreateEntry(key: K) {
     let entry = this.map.get(key);
     if (!entry) {
-      const newEntry: Entry<V> = $state({});
+      const newEntry: Entry<V> = $state({ isInvalidated: false });
       entry = newEntry;
       this.map.set(key, entry);
     }
@@ -40,6 +41,15 @@ export class FilesystemCache<K, V extends object> {
     return entry.value!;
   }
 
+  private finish(key: K, entry: Entry<V>) {
+    entry.promise = undefined;
+    // Mutations may invalidate the entry while it is being revalidated, in which case it is queried
+    // again unless it was evicted or replaced in the meantime
+    if (entry.isInvalidated && this.map.get(key) === entry) {
+      void Promise.resolve(this.get(key)).catch(() => {});
+    }
+  }
+
   // Returns the cached value synchronously if available, and revalidates it in the background
   get(key: K, options?: { fetchFromServer?: (cachedValue: V | undefined) => Promise<V> }) {
     return untrack(() => {
@@ -49,6 +59,7 @@ export class FilesystemCache<K, V extends object> {
       const { promise, resolve, reject } = Promise.withResolvers<V>();
       promise.catch(() => {}); // Rejections are handled by callers waiting for the promise
       entry.promise = promise;
+      entry.isInvalidated = false;
 
       (async () => {
         let cachedValue = entry.value && ($state.snapshot(entry.value) as V);
@@ -67,9 +78,7 @@ export class FilesystemCache<K, V extends object> {
           console.error(e);
           reject(e);
         })
-        .finally(() => {
-          entry.promise = undefined;
-        });
+        .finally(() => this.finish(key, entry));
 
       return entry.value ?? promise;
     });
@@ -77,20 +86,23 @@ export class FilesystemCache<K, V extends object> {
 
   bulkGet(keys: Set<K>) {
     return untrack(() => {
+      const entries = new Map<K, Entry<V>>();
       const pending = new Map<K, PromiseWithResolvers<V>>();
       for (const key of keys) {
         const entry = this.getOrCreateEntry(key);
+        entries.set(key, entry);
         if (entry.promise) continue;
 
         const resolvers = Promise.withResolvers<V>();
         resolvers.promise.catch(() => {});
         entry.promise = resolvers.promise;
+        entry.isInvalidated = false;
         pending.set(key, resolvers);
       }
 
       const settleAll = (values: Map<K, V>) => {
         values.forEach((value, key) => {
-          pending.get(key)?.resolve(this.settle(this.map.get(key)!, value));
+          pending.get(key)?.resolve(this.settle(entries.get(key)!, value));
         });
       };
 
@@ -98,12 +110,12 @@ export class FilesystemCache<K, V extends object> {
         const pendingKeys = [...pending.keys()];
         settleAll(
           await this.options.bulkFetchFromIndexedDB!(
-            new Set(pendingKeys.filter((key) => !this.map.get(key)!.value)),
+            new Set(pendingKeys.filter((key) => !entries.get(key)!.value)),
           ),
         );
         settleAll(
           await this.options.bulkFetchFromServer!(
-            new Map(pendingKeys.map((key) => [key, { cachedValue: this.map.get(key)!.value }])),
+            new Map(pendingKeys.map((key) => [key, { cachedValue: entries.get(key)!.value }])),
           ),
         );
       })()
@@ -112,14 +124,11 @@ export class FilesystemCache<K, V extends object> {
           pending.forEach(({ reject }) => reject(e));
         })
         .finally(() => {
-          pending.forEach((_, key) => {
-            this.map.get(key)!.promise = undefined;
-          });
+          pending.forEach((_, key) => this.finish(key, entries.get(key)!));
         });
 
-      const makeResult = () => new Map([...keys].map((key) => [key, this.map.get(key)!.value!]));
-      const waiting = [...keys]
-        .map((key) => this.map.get(key)!)
+      const makeResult = () => new Map([...entries].map(([key, entry]) => [key, entry.value!]));
+      const waiting = [...entries.values()]
         .filter((entry) => !entry.value)
         .map((entry) => entry.promise!);
       return waiting.length > 0 ? Promise.all(waiting).then(makeResult) : makeResult();
@@ -128,7 +137,9 @@ export class FilesystemCache<K, V extends object> {
 
   // Revalidates the entry if it has been loaded
   invalidate(key: K) {
-    if (this.map.has(key)) {
+    const entry = this.map.get(key);
+    if (entry) {
+      entry.isInvalidated = true;
       void Promise.resolve(this.get(key)).catch(() => {});
     }
   }

src/lib/server/db/repositories/file.ts

@@ -99,17 +99,21 @@ export const lockFile = async (trx: Tx, userId: number, fileId: number) => {
   return !!file;
 };
 
-// Locks the file in share mode so that its DEK isn't changed until the transaction ends. Returns
-// its DEK version if it exists.
-export const lockFileDekVersion = async (trx: Tx, userId: number, fileId: number) => {
-  const file = await trx
+// Locks the file in the given mode so that its DEK isn't changed until the transaction ends.
+// Returns its DEK version if it exists.
+export const lockFileDekVersion = async (
+  trx: Tx,
+  userId: number,
+  fileId: number,
+  mode: "share" | "update" = "share",
+) => {
+  const query = trx
     .selectFrom("file")
     .select("data_encryption_key_version")
     .where("id", "=", fileId)
     .where("user_id", "=", userId)
-    .limit(1)
-    .forShare()
-    .executeTakeFirst();
+    .limit(1);
+  const file = await (mode === "update" ? query.forUpdate() : query.forShare()).executeTakeFirst();
   return file?.data_encryption_key_version ?? null;
 };
 

src/lib/server/services/common.ts

@@ -26,8 +26,9 @@ export const ensureFileDekVersion = async (
   userId: number,
   fileId: number,
   dekVersion: Date,
+  mode: "share" | "update" = "share",
 ) => {
-  const currentDekVersion = await FileRepo.lockFileDekVersion(trx, userId, fileId);
+  const currentDekVersion = await FileRepo.lockFileDekVersion(trx, userId, fileId, mode);
   if (!currentDekVersion) {
     throw new AppError("NOT_FOUND", "File not found");
   } else if (currentDekVersion.getTime() !== dekVersion.getTime()) {

src/lib/server/services/upload.ts

@@ -15,6 +15,7 @@ import {
   deleteFile,
   hashFile,
   reserveFile,
+  withTemporaryChunk,
   writeChunkAtOffset,
 } from "#lib/server/storage/index.ts";
 import {
@@ -126,11 +127,22 @@ export const uploadChunk = async (
     throw new AppError("CONFLICT", "Chunk already uploaded");
   }
 
-  // Concurrent uploads of the same chunk may corrupt the file, which will be detected by the
-  // hash verification on completion
   const isLastChunk = chunkIndex === session.totalChunks;
-  await writeChunkAtOffset(session.path, encChunkStream, chunkIndex, isLastChunk);
-  await UploadRepo.markChunkAsUploaded(db, sessionId, chunkIndex);
+  await withTemporaryChunk(encChunkStream, isLastChunk, (readChunk) =>
+    transaction(async (trx) => {
+      // The chunk is received before locking the session, which is checked again so that late
+      // requests can't overwrite completed chunks or files even if they passed the checks above
+      const lockedSession = await UploadRepo.lockUploadSession(trx, sessionId, userId);
+      if (!lockedSession) {
+        throw new AppError("NOT_FOUND", "Invalid upload id");
+      } else if (isChunkUploaded(lockedSession.bitmap, chunkIndex)) {
+        throw new AppError("CONFLICT", "Chunk already uploaded");
+      }
+
+      await writeChunkAtOffset(lockedSession.path, readChunk(), chunkIndex, isLastChunk);
+      await UploadRepo.markChunkAsUploaded(trx, sessionId, chunkIndex);
+    }),
+  );
 };
 
 const getCompletedUploadSession = async <T extends "file" | "thumbnail">(
@@ -212,7 +224,8 @@ export const completeFileThumbnailUpload = async (userId: number, uploadId: stri
       throw new AppError("NOT_FOUND", "Invalid upload id");
     }
 
-    await ensureFileDekVersion(trx, userId, session.fileId, session.dekVersion);
+    // Locks the file exclusively, as there may be no thumbnail row to lock yet
+    await ensureFileDekVersion(trx, userId, session.fileId, session.dekVersion, "update");
     const oldPath = await MediaRepo.lockFileThumbnailPath(trx, session.fileId);
     await MediaRepo.upsertFileThumbnail(trx, session.fileId, session.path, null);
     await UploadRepo.deleteUploadSession(trx, uploadId);

src/lib/server/storage/chunks.ts

@@ -1,10 +1,37 @@
+import { randomUUID } from "crypto";
+import { createReadStream } from "fs";
 import { open } from "fs/promises";
+import { tmpdir } from "os";
+import { join } from "path";
 import type { Readable } from "stream";
 import { AppError } from "#lib/server/errors.ts";
 import { ENCRYPTION_OVERHEAD, ENCRYPTED_CHUNK_SIZE } from "#lib/shared/format/index.ts";
+import { deleteFile } from "./files.ts";
 
 const invalidChunkSize = () => new AppError("BAD_REQUEST", "Invalid chunk size");
 
+// Receives and validates a chunk before it is copied into the upload's file. Temporary chunks
+// aren't placed in the persistent storage, and are removed even if receiving or applying fails.
+export const withTemporaryChunk = async <T>(
+  encChunkStream: Readable,
+  isLastChunk: boolean,
+  callback: (readChunk: () => Readable) => Promise<T>,
+) => {
+  const path = join(tmpdir(), `arkvault-chunk-${randomUUID()}`);
+  // Created exclusively with restricted permissions, as `writeChunkAtOffset` writes into existing
+  // files only
+  const file = await open(path, "wx", 0o600);
+  try {
+    await file.close();
+    await writeChunkAtOffset(path, encChunkStream, 1, isLastChunk);
+    // A new stream is created for each call, as the chunk may be read several times (e.g. when
+    // retrying a transaction)
+    return await callback(() => createReadStream(path));
+  } finally {
+    await deleteFile(path).catch(console.error);
+  }
+};
+
 export const writeChunkAtOffset = async (
   path: string,
   encChunkStream: Readable,

src/lib/shared/format/stream.ts

@@ -40,7 +40,7 @@ const createByteQueue = () => {
 
 // Decrypts the encrypted chunks covering the plaintext `range` (as returned by
 // `getEncryptedRange`) into the plaintext of the range. `onChunk` receives each decrypted chunk
-// as a whole, along with its index in the file.
+// as a whole, along with its index in the file. Async callbacks are awaited to apply backpressure.
 export const createDecryptionStream = ({
   range,
   firstChunkIndex = 0,
@@ -54,7 +54,7 @@ export const createDecryptionStream = ({
   chunkCount: number;
   dataKey: CryptoKey;
   isLegacy: boolean;
-  onChunk?: (chunkIndex: number, chunk: Uint8Array) => void;
+  onChunk?: (chunkIndex: number, chunk: Uint8Array) => void | Promise<void>;
 }) => {
   const queue = createByteQueue();
   let chunkIndex = 0;
@@ -64,7 +64,7 @@ export const createDecryptionStream = ({
     encryptedSize: number,
   ) => {
     const decrypted = new Uint8Array(await decryptChunk(queue.take(encryptedSize).buffer, dataKey));
-    onChunk?.(firstChunkIndex + chunkIndex, decrypted);
+    await onChunk?.(firstChunkIndex + chunkIndex, decrypted);
 
     const isFirst = chunkIndex === 0;
     const isLast = chunkIndex === chunkCount - 1;

src/service-worker/handlers/decrypt-file.ts

@@ -168,10 +168,10 @@ async function* fetchSegment(
       chunkCount: segment.lastChunkIndex - segment.firstChunkIndex + 1,
       dataKey: metadata.dataKey,
       isLegacy: metadata.isLegacy,
-      // Chunks are written atomically, so the worker being terminated doesn't corrupt the cache
-      onChunk: (chunkIndex, chunk) => {
-        storeCachedChunk(request.fileId, chunkIndex, chunk, layout).catch(console.error);
-      },
+      // The atomic write is awaited so that slow storage doesn't queue up the whole plaintext,
+      // and its failure doesn't fail the download
+      onChunk: (chunkIndex, chunk) =>
+        storeCachedChunk(request.fileId, chunkIndex, chunk, layout).catch(console.error),
     }),
   );
 }

tests/server/upload.test.ts

@@ -4,6 +4,7 @@ import { beforeEach, describe, expect, it } from "vitest";
 import { UploadRepo, db } from "#lib/server/db/index.ts";
 import { FileService, UploadService } from "#lib/server/services/index.ts";
 import { isChunkUploaded } from "#lib/server/services/upload.ts";
+import { hashFile } from "#lib/server/storage/index.ts";
 import { ciphertext, createUser, newEntryParams, resetDatabase } from "../helpers/database.ts";
 
 beforeEach(resetDatabase);
@@ -38,6 +39,61 @@ describe("upload session bitmap", () => {
 });
 
 describe("upload completion", () => {
+  it.each([false, true])(
+    "rejects a late duplicate without changing the file (completed: %s)",
+    async (completeBeforeDuplicate) => {
+      const { userId } = await createUser();
+      const { uploadId } = await UploadService.startFileUpload({
+        ...newEntryParams(userId, "root"),
+        totalChunks: 1,
+        hskVersion: null,
+        contentType: "image/png",
+        encCreatedAt: null,
+        encLastModifiedAt: ciphertext(),
+      });
+      const session = (await UploadRepo.getUploadSession(db, uploadId, userId))!;
+      const started = Promise.withResolvers<void>();
+      const release = Promise.withResolvers<void>();
+      const duplicate = UploadService.uploadChunk(
+        userId,
+        uploadId,
+        1,
+        Readable.from(
+          (async function* () {
+            started.resolve();
+            await release.promise;
+            yield Buffer.alloc(100, 2);
+          })(),
+        ),
+      );
+      await started.promise;
+
+      try {
+        const chunk = Buffer.alloc(100, 1);
+        const hash = createHash("sha256").update(chunk).digest("base64");
+        await UploadService.uploadChunk(userId, uploadId, 1, Readable.from([chunk]));
+        if (completeBeforeDuplicate) {
+          await UploadService.completeFileUpload(userId, uploadId, null, hash);
+        }
+
+        const rejected = expect(duplicate).rejects.toMatchObject({
+          code: completeBeforeDuplicate ? "NOT_FOUND" : "CONFLICT",
+        });
+        release.resolve();
+        await rejected;
+        expect(await hashFile(session.path)).toBe(hash);
+
+        if (!completeBeforeDuplicate) {
+          await UploadService.completeFileUpload(userId, uploadId, null, hash);
+        }
+        expect(await FileService.getAllFileIds(userId)).toHaveLength(1);
+      } finally {
+        release.resolve();
+        await duplicate.catch(() => {});
+      }
+    },
+  );
+
   it("registers the file only once when completed concurrently", async () => {
     const { userId } = await createUser();
     const { uploadId } = await UploadService.startFileUpload({

tests/unit/filesystem-cache.test.ts

@@ -30,6 +30,92 @@ const createCache = (options?: {
 const flush = () => new Promise((resolve) => setTimeout(resolve, 0));
 
 describe("FilesystemCache", () => {
+  it("revalidates again after invalidation during a pending request, coalescing repeated calls", async () => {
+    const pending = Promise.withResolvers<Entry>();
+    const fetchFromServer = vi
+      .fn<(id: number) => Promise<Entry>>()
+      .mockImplementationOnce(() => pending.promise)
+      .mockResolvedValue({ id: 1, name: "new" });
+    const cache = new FilesystemCache<number, Entry>({
+      fetchFromIndexedDB: async () => undefined,
+      fetchFromServer,
+    });
+    const loading = cache.get(1);
+    await vi.waitFor(() => expect(fetchFromServer).toHaveBeenCalledTimes(1));
+
+    cache.invalidate(1);
+    cache.invalidate(1);
+    cache.invalidate(1);
+    pending.resolve({ id: 1, name: "old" });
+    const entry = await loading;
+    await vi.waitFor(() => expect(entry.name).toBe("new"));
+    expect(fetchFromServer).toHaveBeenCalledTimes(2);
+  });
+
+  it("revalidates only invalidated entries after a pending bulk request", async () => {
+    const pending = Promise.withResolvers<Map<number, Entry>>();
+    const fetchFromServer = vi.fn(async (id: number) => ({ id, name: "new" }));
+    const bulkFetchFromServer = vi.fn(() => pending.promise);
+    const cache = new FilesystemCache<number, Entry>({
+      fetchFromIndexedDB: async () => undefined,
+      fetchFromServer,
+      bulkFetchFromIndexedDB: async () => new Map(),
+      bulkFetchFromServer,
+    });
+    const loading = cache.bulkGet(new Set([1, 2]));
+    await vi.waitFor(() => expect(bulkFetchFromServer).toHaveBeenCalledTimes(1));
+    cache.invalidate(1);
+    pending.resolve(new Map([1, 2].map((id) => [id, { id, name: "old" }])));
+
+    const entries = await loading;
+    await vi.waitFor(() => expect(entries.get(1)?.name).toBe("new"));
+    expect(entries.get(2)?.name).toBe("old");
+    expect(fetchFromServer).toHaveBeenCalledTimes(1);
+  });
+
+  it("preserves invalidation when a pending request fails", async () => {
+    vi.spyOn(console, "error").mockImplementation(() => {});
+    const pending = Promise.withResolvers<Entry>();
+    const fetchFromServer = vi
+      .fn<() => Promise<Entry>>()
+      .mockImplementationOnce(() => pending.promise)
+      .mockResolvedValue({ id: 1, name: "new" });
+    const cache = new FilesystemCache<number, Entry>({
+      fetchFromIndexedDB: async () => ({ id: 1, name: "local" }),
+      fetchFromServer,
+    });
+    const entry = await cache.get(1);
+    await vi.waitFor(() => expect(fetchFromServer).toHaveBeenCalledTimes(1));
+    cache.invalidate(1);
+    pending.reject(new Error("Failed"));
+
+    await vi.waitFor(() => expect(entry.name).toBe("new"));
+    expect(fetchFromServer).toHaveBeenCalledTimes(2);
+  });
+
+  it("doesn't revalidate evicted entries after a pending bulk request", async () => {
+    const pending = Promise.withResolvers<Map<number, Entry>>();
+    const fetchFromServer = vi.fn(async (id: number) => ({ id, name: "new" }));
+    const bulkFetchFromServer = vi.fn(() => pending.promise);
+    const cache = new FilesystemCache<number, Entry>({
+      fetchFromIndexedDB: async () => undefined,
+      fetchFromServer,
+      bulkFetchFromIndexedDB: async () => new Map(),
+      bulkFetchFromServer,
+    });
+    const loading = cache.bulkGet(new Set([1]));
+    await vi.waitFor(() => expect(bulkFetchFromServer).toHaveBeenCalledTimes(1));
+    cache.invalidate(1);
+    cache.evict(1);
+    const replacement = await cache.get(1);
+    pending.resolve(new Map([[1, { id: 1, name: "old" }]]));
+    await loading;
+    await flush();
+
+    expect(replacement.name).toBe("new");
+    expect(fetchFromServer).toHaveBeenCalledTimes(1);
+  });
+
   it("resolves with the local value first and updates it in place", async () => {
     const { cache } = createCache({ local: { 1: "local" }, server: { 1: "server" } });