node compile cache: unmap an accepted entry's stored source, keep only the bytecode mapped

src/jsc/NodeCompileCache.rs

@@ -62,7 +62,7 @@ struct Entry {
 /// 128-byte-aligned blob. JSC's bytecode decoder reads the blob in place and
 /// requires the same alignment the standalone graph provides (see
 /// StandaloneModuleGraph.rs "Bytecode alignment" note). Either a heap buffer
-/// or a span inside a whole-file mapping.
+/// or a span inside a mapping of the cache file.
 struct AlignedBlob {
     ptr: core::ptr::NonNull<u8>,
     len: usize,
@@ -71,11 +71,12 @@ struct AlignedBlob {
 
 enum Backing {
     Heap,
-    /// `PROT_READ`/`MAP_PRIVATE` mapping of the whole cache file; `ptr` points
-    /// at [`blob_file_offset`] inside it, 128-aligned because the mapping base
-    /// is page-aligned. Safe against entry rewrites: writers go through
-    /// tmpfile + rename, so a replaced file's old inode stays live under the
-    /// mapping.
+    /// `PROT_READ`/`MAP_PRIVATE` mapping of the cache file from the page the
+    /// blob starts on (so `ptr` stays 128-aligned) to its end; the header and
+    /// stored code before that page are unmapped once validated, or still
+    /// included here if that unmap failed. Safe against entry rewrites: writers
+    /// go through tmpfile + rename, so a replaced file's old inode stays live
+    /// under the mapping.
     Map {
         base: core::ptr::NonNull<u8>,
         map_len: usize,
@@ -110,6 +111,33 @@ impl AlignedBlob {
         // SAFETY: `ptr` is valid for `len` bytes for the lifetime of `self`.
         unsafe { core::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len) }
     }
+
+    /// Takes over a validated whole-file mapping, unmapping everything before
+    /// the blob's page (see [`Backing::Map`]).
+    ///
+    /// # Safety
+    /// `(base, map_len)` must be a live mapping nothing else reads any more,
+    /// with `blob_off + blob_len <= map_len` and `base + blob_off` 128-aligned.
+    unsafe fn from_mapping(
+        base: core::ptr::NonNull<u8>,
+        map_len: usize,
+        blob_off: usize,
+        blob_len: usize,
+    ) -> Self {
+        let prefix_len = blob_off - blob_off % bun_alloc::page_size();
+        // SAFETY: `prefix_len <= blob_off <= map_len` per the contract.
+        let (ptr, tail) = unsafe { (base.add(blob_off), base.add(prefix_len)) };
+        let (base, map_len) = if prefix_len != 0 && sys::munmap(base.as_ptr(), prefix_len).is_ok() {
+            (tail, map_len - prefix_len)
+        } else {
+            (base, map_len)
+        };
+        Self {
+            ptr,
+            len: blob_len,
+            backing: Backing::Map { base, map_len },
+        }
+    }
 }
 
 impl Drop for AlignedBlob {
@@ -831,14 +859,11 @@ fn read_cache_file(state: &CacheState, key: u64, entry: &mut Entry, code: Option
     });
     let blob = if map_is_aligned {
         let (base, map_len) = map_guard.take().expect("checked above");
-        // SAFETY: `blob_off + cache_size == map_len` was just validated.
-        let ptr =
-            unsafe { core::ptr::NonNull::new_unchecked(base.as_ptr().add(blob_off as usize)) };
-        AlignedBlob {
-            ptr,
-            len: cache_size as usize,
-            backing: Backing::Map { base, map_len },
-        }
+        // SAFETY: `blob_off + cache_size == map_len` and the blob's alignment
+        // were just validated, and `bytes`/`blob_bytes` (views of the whole
+        // mapping, part of which is unmapped here) are not used past this
+        // point.
+        unsafe { AlignedBlob::from_mapping(base, map_len, blob_off as usize, cache_size as usize) }
     } else {
         // No mapping, or the blob would be misaligned in it: copy to an
         // aligned heap buffer instead (map_guard unmaps on return).

test/js/node/module/node-module-module.test.js

@@ -1,6 +1,6 @@
 import { describe, expect, test } from "bun:test";
 import fs from "fs";
-import { bunEnv, bunExe, isWindows, ospath, tempDir } from "harness";
+import { bunEnv, bunExe, isLinux, isWindows, ospath, tempDir } from "harness";
 import Module, { _nodeModulePaths, builtinModules, createRequire, isBuiltin, wrap } from "module";
 import path from "path";
 
@@ -216,6 +216,79 @@ console.log("survived", require("./late.js"));`,
     expect(exitCode).toBe(0);
   });
 
+  test.skipIf(!isLinux)("compile cache keeps only the bytecode pages of an accepted entry mapped", async () => {
+    // An accepted entry file is mmapped so JSC reads the bytecode in place. The
+    // stored source in front of the bytecode is only read by the byte compare
+    // (the module's live source holds the same text), so its pages are unmapped
+    // afterwards instead of staying resident for the rest of the process.
+    const sourceSize = 256 * 1024;
+    using dir = tempDir("compile-cache-unmap", {
+      // Spans several pages even with 64 KiB pages.
+      "big.js": `module.exports = "${Buffer.alloc(sourceSize, "x").toString()}";`,
+      // Prints the mappings (file offset + mapped length) of every file whose
+      // path contains argv[2].
+      "main.js": `
+        const length = require("./big.js").length;
+        const mapped = require("fs")
+          .readFileSync("/proc/self/maps", "utf8")
+          .split("\\n")
+          .filter(line => line.includes(process.argv[2]))
+          .map(line => {
+            const [range, , offset] = line.split(/\\s+/);
+            const [start, end] = range.split("-");
+            return { offset: parseInt(offset, 16), length: parseInt(end, 16) - parseInt(start, 16) };
+          });
+        console.log(JSON.stringify({ length, mapped }));
+      `,
+    });
+    const cacheDir = path.join(String(dir), "cc");
+    const env = { ...bunEnv, NODE_COMPILE_CACHE: cacheDir };
+    async function run(filter) {
+      await using proc = Bun.spawn({ cmd: [bunExe(), "main.js", filter], env, cwd: String(dir), stderr: "pipe" });
+      const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]);
+      expect(stderr).toBe("");
+      expect(exitCode).toBe(0);
+      return JSON.parse(stdout);
+    }
+
+    // Misses populate the cache and leave nothing mapped.
+    expect(await run(cacheDir)).toEqual({ length: sourceSize, mapped: [] });
+
+    // big.js has the large entry. An entry (NodeCompileCache.rs) starts with
+    // `magic u32 | stored source size u32 | bytecode size u32`, holds the stored
+    // source after the header, and ends with the bytecode.
+    const entries = [...new Bun.Glob("**/*").scanSync({ cwd: cacheDir, onlyFiles: true })]
+      .map(f => path.join(cacheDir, f))
+      .sort((a, b) => fs.statSync(b).size - fs.statSync(a).size);
+    expect(entries).toHaveLength(2);
+    const entry = fs.readFileSync(entries[0]);
+    const bytecodeOffset = entry.length - entry.readUInt32LE(8);
+    expect(entry.readUInt32LE(4)).toBeGreaterThanOrEqual(sourceSize);
+    expect(bytecodeOffset).toBeGreaterThan(sourceSize);
+
+    // /proc/self/auxv is a list of (type, value) u64 pairs; AT_PAGESZ is type 6.
+    const auxv = fs.readFileSync("/proc/self/auxv");
+    let pageSize = 0;
+    for (let i = 0; i + 16 <= auxv.length; i += 16) {
+      if (auxv.readBigUInt64LE(i) === 6n) pageSize = Number(auxv.readBigUInt64LE(i + 8));
+    }
+    expect(pageSize).toBeGreaterThan(0);
+    const roundDown = n => n - (n % pageSize);
+    const roundUp = n => roundDown(n + pageSize - 1);
+
+    // Only the pages from the one the bytecode starts on through the end of
+    // the file stay mapped.
+    expect(await run(path.basename(entries[0]))).toEqual({
+      length: sourceSize,
+      mapped: [
+        {
+          offset: roundDown(bytecodeOffset),
+          length: roundUp(entry.length) - roundDown(bytecodeOffset),
+        },
+      ],
+    });
+  });
+
   const compileCacheEnv = { ...bunEnv };
   delete compileCacheEnv.NODE_COMPILE_CACHE;
   delete compileCacheEnv.NODE_COMPILE_CACHE_PORTABLE;