dofs: link staged chunks during sync apply

.changeset/link-staged-chunks-on-apply.md

@@ -0,0 +1,6 @@
+---
+"@cloudflare/dofs": patch
+"@cloudflare/computer": patch
+---
+
+Cut peak memory during a sync pull. Applying a file entry now links the chunks the sender already staged instead of reading them back and joining them into one whole-file buffer, which used to hold roughly twice the file size in the isolate at once.

packages/dofs/src/fs/gc.test.ts

@@ -1,7 +1,12 @@
+import { createHash } from "node:crypto";
+
 import { describe, expect, it } from "vitest";
 
 import type { Database } from "../storage.js";
+import { applyChanges } from "../sync/apply.js";
+import { stageBlob } from "../sync/blobs.js";
 import { gc } from "./gc.js";
+import { readFile } from "./readFile.js";
 import { rm } from "./rm.js";
 import { withDB } from "./with-db.js";
 import { writeFile } from "./writeFile.js";
@@ -33,6 +38,44 @@ describe("gc", () => {
     });
   });
 
+  it("does not free blobs a sync apply linked", async () => {
+    await withDB(async (db) => {
+      // A pull stages chunks before linking them, so vfs_chunks reachability
+      // must protect blobs independently of their staging timestamp.
+      const bytes = new TextEncoder().encode("staged content");
+      const hash = new Uint8Array(createHash("sha256").update(bytes).digest());
+      stageBlob(db, hash, bytes, 1000);
+
+      await applyChanges(
+        db,
+        [
+          {
+            kind: "file",
+            rev: 1,
+            path: "/staged.txt",
+            mode: 0o644,
+            mtime: 1000,
+            size: bytes.byteLength,
+            chunks: [{ hash, size: bytes.byteLength }],
+          },
+        ],
+        new Map(),
+        { source: "upstream" },
+      );
+      expect(blobCount(db)).toBe(1);
+
+      // A zero safety window and far-future clock make every blob stale;
+      // only references from vfs_chunks protect this blob and manifest.
+      expect(gc(db, { now: () => 999_999_999, safetyWindowMs: 0 })).toEqual({
+        blobsFreed: 0,
+        manifestsFreed: 0,
+      });
+      expect(blobCount(db)).toBe(1);
+      expect(blobBytesCount(db)).toBe(1);
+      expect(await readFile(db, "/staged.txt", "utf8")).toBe("staged content");
+    });
+  });
+
   it("frees orphan blobs left behind by overwrite", async () => {
     await withDB(async (db) => {
       await writeFile(db, "/x.txt", "first", {}, () => 1000);

packages/dofs/src/fs/mount-guard.test.ts

@@ -1,6 +1,9 @@
+import { createHash } from "node:crypto";
+
 import { describe, expect, it } from "vitest";
 
 import type { Database } from "../storage.js";
+import { stageBlob } from "../sync/blobs.js";
 import { mkdir } from "./mkdir.js";
 import {
   assertNotReadOnly,
@@ -12,7 +15,7 @@ import { resolveInode } from "./resolve.js";
 import { rm } from "./rm.js";
 import { symlink } from "./symlink.js";
 import { withDB } from "./with-db.js";
-import { writeFile, writeFileSync } from "./writeFile.js";
+import { linkStagedChunksSync, writeFile, writeFileSync } from "./writeFile.js";
 
 // Stage a read-only mount the way the workspace-side indexer
 // eventually will: a row in `_vfs_mounts` plus an actual subtree
@@ -141,6 +144,29 @@ describe("writeFile under a read-only mount", () => {
     });
   });
 
+  it("rejects linkStagedChunksSync under the mount root with EROFS", async () => {
+    await withDB(async (db) => {
+      mkdir(db, "/workspace/r2", { recursive: true }, () => 0);
+      stageMount(db, "/workspace/r2", "read-only");
+
+      const bytes = new TextEncoder().encode("blocked");
+      const hash = new Uint8Array(createHash("sha256").update(bytes).digest());
+      stageBlob(db, hash, bytes, 0);
+
+      expect(() =>
+        linkStagedChunksSync(
+          db,
+          "/workspace/r2/hello.txt",
+          ["workspace", "r2", "hello.txt"],
+          [{ hash, size: bytes.byteLength }],
+          {},
+          0,
+        ),
+      ).toThrow(/EROFS|read-only/);
+      expect(resolveInode(db, "/workspace/r2/hello.txt")).toBeNull();
+    });
+  });
+
   it("allows writes under a read-write mount", async () => {
     await withDB(async (db) => {
       mkdir(db, "/workspace/rw", { recursive: true }, () => 0);

packages/dofs/src/fs/writeFile.ts

@@ -211,10 +211,47 @@ async function writeFileStreaming(
     flush(carry);
   }
 
-  // Wire up the inode against the staged blobs in one short
-  // transaction. From this point on the SQL is the same shape as the
-  // synchronous path — only the chunk-bytes step is skipped because
-  // stageBlob already landed them above.
+  linkStagedChunksSync(db, canonical, parts, chunkRefs, { ...options, mode }, mtime);
+}
+
+// Reject a chunk list that positional reads could not address.
+// readRangeSync finds the chunk covering an offset by dividing that
+// offset by CHUNK_SIZE, and takes a chunk's start offset to be its
+// index times CHUNK_SIZE, so every chunk but the last has to fill a
+// whole window and none may overflow one. Local writers chunk with
+// chunksOf and satisfy this by construction; a chunk list that
+// arrived over sync does not have to.
+export function assertChunkWindows(chunkRefs: ChunkRef[], canonical: string): void {
+  for (let idx = 0; idx < chunkRefs.length; idx++) {
+    const { size } = chunkRefs[idx];
+    const last = idx === chunkRefs.length - 1;
+    if (size === CHUNK_SIZE || (last && size < CHUNK_SIZE)) continue;
+    throw createWorkspaceError(
+      "EINVAL",
+      `chunk ${idx} of ${chunkRefs.length} holds ${size} bytes; only the last chunk may be shorter than ${CHUNK_SIZE}: ${canonical}`,
+      canonical,
+    );
+  }
+}
+
+// Link a path to chunks already staged in content-addressed storage.
+// This keeps payload bytes out of memory during sync apply.
+//
+// The chunk list comes from a caller that did its own chunking, so
+// the guards every other write path gets from writeFile have to run
+// here too: the read-only mount check, and the fixed-window layout
+// that positional reads depend on.
+export function linkStagedChunksSync(
+  db: Database,
+  canonical: string,
+  parts: string[],
+  chunkRefs: { hash: Uint8Array; size: number }[],
+  options: WriteFileOptions,
+  mtime: number,
+): void {
+  assertNotReadOnly(db, canonical);
+  assertChunkWindows(chunkRefs, canonical);
+  const mode = (options.mode ?? 0o644) & 0o7777;
   db.transactionSync(() => {
     const parentInode = resolveParent(db, parts, canonical);
     const leafName = parts[parts.length - 1];
@@ -251,6 +288,8 @@ async function writeFileStreaming(
         ref.size,
       );
     }
+    // Referenced chunks cannot be collected, so last_seen only protects
+    // blobs during the staging window before this transaction.
     const manifestHash = buildManifest(db, chunkRefs, mtime);
     const rev = incrementRev(db);
     let totalSize = 0;

packages/dofs/src/sync/apply.test.ts

@@ -1,4 +1,6 @@
-import { describe, expect, it } from "vitest";
+import { createHash } from "node:crypto";
+
+import { describe, expect, it, vi } from "vitest";
 
 import { link } from "../fs/link.js";
 import { mkdir } from "../fs/mkdir.js";
@@ -10,8 +12,9 @@ import { resolveInode } from "../fs/resolve.js";
 import { rm } from "../fs/rm.js";
 import { symlink } from "../fs/symlink.js";
 import { withDB, withTwoDBs } from "../fs/with-db.js";
-import { writeFile, writeFileSync } from "../fs/writeFile.js";
+import { CHUNK_SIZE, writeFile, writeFileSync } from "../fs/writeFile.js";
 import { applyChanges, applyChangesSync } from "./apply.js";
+import { stageBlob } from "./blobs.js";
 import type { ChangeEntry } from "./changes.js";
 import { coalesceChanges } from "./coalesce.js";
 import { fetchObjects } from "./fetch.js";
@@ -73,6 +76,171 @@ describe("applyChanges", () => {
     );
   });
 
+  it("links staged chunks without reading payload bytes", async () => {
+    const content = `${"a".repeat(CHUNK_SIZE)}b`;
+    await withTwoDBs(
+      async (a) => {
+        await writeFile(a, "/large.txt", content, {}, () => 1);
+        const entries = await drain(coalesceChanges(a, 0));
+        const entry = entries.find((candidate) => candidate.path === "/large.txt");
+        if (entry?.kind !== "file") throw new Error("missing file entry");
+        expect(entry.chunks).toHaveLength(2);
+        return { entry, objects: await collectObjects(a, [entry]) };
+      },
+      async (b, { entry, objects }) => {
+        for (const chunk of entry.chunks) {
+          const bytes = objects.get(hex(chunk.hash));
+          if (bytes === undefined) throw new Error("missing chunk bytes");
+          stageBlob(b, chunk.hash, bytes, 2);
+        }
+
+        const all = vi.spyOn(b, "all");
+        try {
+          await applyChanges(b, [entry], new Map(), { source: "upstream" });
+          const payloadReads = all.mock.calls.filter(([query]) =>
+            query.includes("SELECT bytes FROM vfs_blob_bytes"),
+          );
+          expect(all.mock.calls.length).toBeGreaterThan(0);
+          expect(payloadReads).toHaveLength(0);
+        } finally {
+          all.mockRestore();
+        }
+        expect(await readFile(b, "/large.txt", "utf8")).toBe(content);
+      },
+    );
+  });
+
+  it("links staged chunks synchronously without reading payload bytes", async () => {
+    const content = `${"a".repeat(CHUNK_SIZE)}b`;
+    await withTwoDBs(
+      async (a) => {
+        await writeFile(a, "/large.txt", content, {}, () => 1);
+        const entries = await drain(coalesceChanges(a, 0));
+        const entry = entries.find((candidate) => candidate.path === "/large.txt");
+        if (entry?.kind !== "file") throw new Error("missing file entry");
+        expect(entry.chunks).toHaveLength(2);
+        return { entry, objects: await collectObjects(a, [entry]) };
+      },
+      async (b, { entry, objects }) => {
+        for (const chunk of entry.chunks) {
+          const bytes = objects.get(hex(chunk.hash));
+          if (bytes === undefined) throw new Error("missing chunk bytes");
+          stageBlob(b, chunk.hash, bytes, 2);
+        }
+
+        const all = vi.spyOn(b, "all");
+        try {
+          applyChangesSync(b, [entry], new Map(), { source: "upstream" });
+          const payloadReads = all.mock.calls.filter(([query]) =>
+            query.includes("SELECT bytes FROM vfs_blob_bytes"),
+          );
+          expect(all.mock.calls.length).toBeGreaterThan(0);
+          expect(payloadReads).toHaveLength(0);
+        } finally {
+          all.mockRestore();
+        }
+        expect(await readFile(b, "/large.txt", "utf8")).toBe(content);
+      },
+    );
+  });
+
+  it("rejects a file entry whose interior chunks are not chunk-aligned", async () => {
+    // Positional reads locate a chunk by dividing the offset by
+    // CHUNK_SIZE, so only the final chunk may be short. Linking a
+    // sender's chunk list verbatim has to enforce that.
+    await withDB(async (db) => {
+      const first = new TextEncoder().encode("first");
+      const second = new TextEncoder().encode("second");
+      const chunks = [first, second].map((bytes) => {
+        const hash = new Uint8Array(createHash("sha256").update(bytes).digest());
+        stageBlob(db, hash, bytes, 1000);
+        return { hash, size: bytes.byteLength };
+      });
+
+      await expect(
+        applyChanges(
+          db,
+          [
+            {
+              kind: "file",
+              rev: 1,
+              path: "/ragged.txt",
+              mode: 0o644,
+              mtime: 1000,
+              size: first.byteLength + second.byteLength,
+              chunks,
+            },
+          ],
+          new Map(),
+          { source: "upstream" },
+        ),
+      ).rejects.toThrow(/chunk/);
+      expect(resolveInode(db, "/ragged.txt")).toBeNull();
+    });
+  });
+
+  it("leaves the existing file in place when it rejects a ragged entry", async () => {
+    await withDB(async (db) => {
+      await writeFile(db, "/keep.txt", "original", {}, () => 1000);
+
+      const first = new TextEncoder().encode("first");
+      const second = new TextEncoder().encode("second");
+      const chunks = [first, second].map((bytes) => {
+        const hash = new Uint8Array(createHash("sha256").update(bytes).digest());
+        stageBlob(db, hash, bytes, 1000);
+        return { hash, size: bytes.byteLength };
+      });
+
+      await expect(
+        applyChanges(
+          db,
+          [
+            {
+              kind: "file",
+              rev: 2,
+              path: "/keep.txt",
+              mode: 0o644,
+              mtime: 2000,
+              size: first.byteLength + second.byteLength,
+              chunks,
+            },
+          ],
+          new Map(),
+          { source: "upstream" },
+        ),
+      ).rejects.toThrow(/chunk/);
+      expect(await readFile(db, "/keep.txt", "utf8")).toBe("original");
+    });
+  });
+
+  it("rejects a file entry with a chunk larger than the chunk size", async () => {
+    await withDB(async (db) => {
+      const bytes = new Uint8Array(CHUNK_SIZE + 1);
+      const hash = new Uint8Array(createHash("sha256").update(bytes).digest());
+      stageBlob(db, hash, bytes, 1000);
+
+      await expect(
+        applyChanges(
+          db,
+          [
+            {
+              kind: "file",
+              rev: 1,
+              path: "/oversized.bin",
+              mode: 0o644,
+              mtime: 1000,
+              size: bytes.byteLength,
+              chunks: [{ hash, size: bytes.byteLength }],
+            },
+          ],
+          new Map(),
+          { source: "upstream" },
+        ),
+      ).rejects.toThrow(/chunk/);
+      expect(resolveInode(db, "/oversized.bin")).toBeNull();
+    });
+  });
+
   it("commits in batches capped by byte budget", async () => {
     // Force many small files; with a tiny byte budget the apply
     // path should still converge, just across more batches. We