Bun.color: reuse one mimalloc heap per VM instead of creating one per call

src/css_jsc/color_js.rs

@@ -210,6 +210,38 @@ fn zero_if_none(component: f32) -> f32 {
     if component.is_nan() { 0.0 } else { component }
 }
 
+/// The VM's scratch arena for one call; `Drop` hands it back on every exit path.
+struct ScratchArena<'a> {
+    global: &'a JSGlobalObject,
+    arena: Option<Arena>,
+}
+
+impl<'a> ScratchArena<'a> {
+    fn take(global: &'a JSGlobalObject) -> Self {
+        let arena = global.bun_vm().as_mut().rare_data().take_scratch_arena();
+        Self {
+            global,
+            arena: Some(arena),
+        }
+    }
+
+    fn get(&self) -> &Arena {
+        self.arena.as_ref().expect("held until drop")
+    }
+}
+
+impl Drop for ScratchArena<'_> {
+    fn drop(&mut self) {
+        if let Some(arena) = self.arena.take() {
+            self.global
+                .bun_vm()
+                .as_mut()
+                .rare_data()
+                .put_back_scratch_arena(arena);
+        }
+    }
+}
+
 pub fn js_function_color(global: &JSGlobalObject, frame: &CallFrame) -> JsResult<JSValue> {
     use bun_ast::symbol::Map as SymbolMap;
     use bun_core::Utf8Bytes;
@@ -241,6 +273,8 @@ pub fn js_function_color(global: &JSGlobalObject, frame: &CallFrame) -> JsResult
         break 'brk OutputColorFormat::Css;
     };
     let input: Utf8Bytes;
+    // Shared by the parse and the print of one call.
+    let mut scratch: Option<ScratchArena> = None;
 
     let parsed_color: css::CssColorParseResult = 'brk: {
         if args[0].is_number() {
@@ -324,10 +358,10 @@ pub fn js_function_color(global: &JSGlobalObject, frame: &CallFrame) -> JsResult
 
         input = args[0].to_utf8(global)?;
 
-        // MimallocArena::new() calls mi_heap_new(), so defer creation to the
-        // paths that actually allocate.
-        let arena = Arena::new();
-        let mut parser_input = css::ParserInput::new(input.slice(), &arena);
+        let arena = scratch
+            .get_or_insert_with(|| ScratchArena::take(global))
+            .get();
+        let mut parser_input = css::ParserInput::new(input.slice(), arena);
         let mut parser = css::Parser::new(
             &mut parser_input,
             None,
@@ -599,13 +633,15 @@ pub fn js_function_color(global: &JSGlobalObject, frame: &CallFrame) -> JsResult
             }
 
             // Fallback to CSS string output
-            let arena = Arena::new();
+            let arena = scratch
+                .get_or_insert_with(|| ScratchArena::take(global))
+                .get();
             let mut dest: Vec<u8> = Vec::new();
 
             let symbols = SymbolMap::init_list(Default::default());
             let mut printer = css::Printer::new(
-                &arena,
-                bun_alloc::ArenaVec::<u8>::new_in(&arena),
+                arena,
+                bun_alloc::ArenaVec::<u8>::new_in(arena),
                 &mut dest,
                 &css::PrinterOptions::default(),
                 None,

src/jsc/rare_data.rs

@@ -273,6 +273,8 @@ pub struct RareData {
     websocket_inflate_scratch: Option<Vec<u8>>,
     /// One libdeflate handle for every JS-thread one-shot inflate; see [`Self::libdeflate_decompressor`].
     libdeflate_decompressor: Option<libdeflate::OwnedDecompressor>,
+    /// Parked mimalloc heap for short synchronous parses; see [`Self::take_scratch_arena`].
+    scratch_arena: Option<bun_alloc::Arena>,
 
     // There is intentionally no `aws_signature_cache` field — storage lives in
     // `bun_s3_signing::credentials::AWS_SIGNATURE_CACHE` (process static; it
@@ -334,6 +336,7 @@ impl Default for RareData {
             compression_scratch: None,
             websocket_inflate_scratch: None,
             libdeflate_decompressor: None,
+            scratch_arena: None,
             s3_default_client: Strong::empty(),
             node_quic_callbacks: Strong::empty(),
             default_csrf_secret: Box::default(),
@@ -678,6 +681,21 @@ impl RareData {
         }
     }
 
+    /// Per VM, not a thread-local: mimalloc keeps a heap whose thread exits, and
+    /// `VirtualMachine::destroy` frees this one on the Worker's own thread.
+    pub fn take_scratch_arena(&mut self) -> bun_alloc::Arena {
+        self.scratch_arena.take().unwrap_or_default()
+    }
+
+    /// Hand a taken arena back; the slot keeps the first one returned and recycles it past `KEEP`.
+    pub fn put_back_scratch_arena(&mut self, mut arena: bun_alloc::Arena) {
+        const KEEP: usize = 64 * 1024;
+        if self.scratch_arena.is_none() {
+            arena.reset_retain_with_limit(KEEP);
+            self.scratch_arena = Some(arena);
+        }
+    }
+
     /// libdeflate keeps no state between calls, so one handle serves every one-shot inflate on
     /// this thread. `None` when libdeflate cannot allocate it; callers fall back to zlib.
     pub fn libdeflate_decompressor(&mut self) -> Option<&mut libdeflate::Decompressor> {

test/js/bun/css/color.test.ts

@@ -1,6 +1,6 @@
 import { color } from "bun";
 import { describe, expect, test } from "bun:test";
-import { bunEnv, bunExe, isASAN, isDebug, withoutAggressiveGC } from "harness";
+import { bunEnv, bunExe, isASAN, isDebug, tempDir, withoutAggressiveGC } from "harness";
 
 const namedColors = ["red", "green", "blue", "yellow", "purple", "orange", "pink", "brown", "gray"];
 
@@ -572,6 +572,45 @@ describe("input forms", () => {
     expect(color("#gg0000", "hex")).toBeNull();
   });
 
+  // A CSS escape or a NUL byte makes the tokenizer copy the token into the
+  // arena, which is reused between calls.
+  const escaped: [string, string | null, string | null, number | null][] = [
+    ["\\72 ed", "red", "#ff0000", 0xff0000],
+    ["\\72\\65\\64", "red", "#ff0000", 0xff0000],
+    ["r\\65 d", "red", "#ff0000", 0xff0000],
+    ["#ff\\38 800", "#f80", "#ff8800", 0xff8800],
+    ["re\0d", null, null, null],
+    ["\0", null, null, null],
+    ["\\", null, null, null],
+    ["r\\", null, null, null],
+  ];
+
+  test.each(escaped)("escaped input %j", (input, css, hex, number) => {
+    expect({ css: color(input, "css"), hex: color(input, "hex"), number: color(input, "number") }).toEqual({
+      css,
+      hex,
+      number,
+    });
+  });
+
+  test("escaped and plain inputs give the same answers when interleaved", () => {
+    const plain: [string, string | null][] = [
+      ["red", "red"],
+      ["#ff8800", "#f80"],
+      ["hsl(120, 100%, 50%)", "#0f0"],
+      ["rgb(255, 0, 0)", "red"],
+      ["notacolor", null],
+    ];
+    const inputs = [...plain, ...escaped.map(([input, css]): [string, string | null] => [input, css])];
+    const wrong: string[] = [];
+    for (let i = 0; i < 2000; i++) {
+      const [input, expected] = inputs[(i * 7) % inputs.length];
+      const actual = color(input, "css");
+      if (actual !== expected) wrong.push(`${i}: color(${JSON.stringify(input)}) = ${JSON.stringify(actual)}`);
+    }
+    expect(wrong).toEqual([]);
+  });
+
   test("alpha survives the object and array forms", () => {
     expect(color("#f00", "{rgba}")).toEqual({ r: 255, g: 0, b: 0, a: 1 });
     expect(color("#f00", "[rgba]")).toEqual([255, 0, 0, 255]);
@@ -764,3 +803,76 @@ describe("conversions between color spaces", () => {
     );
   });
 });
+
+// The CSS parser and printer take an arena. Bun.color reuses one mimalloc heap
+// per VM for it; creating and destroying a heap costs more than the conversion.
+describe.concurrent("mimalloc heaps", () => {
+  test("a string to css conversion does not create a heap per call", async () => {
+    await using proc = Bun.spawn({
+      cmd: [
+        bunExe(),
+        "-e",
+        `
+          import { heapStats } from "bun:jsc";
+          const seqs = () => heapStats({ dump: true }).mimallocDump.heaps.map(h => h.seq);
+          // Heaps are numbered in creation order, and a live Bun.Transpiler owns one,
+          // so a new Transpiler shows how many heaps the process has created so far.
+          const keep = [];
+          const newestHeap = () => {
+            const before = new Set(seqs());
+            const transpiler = new Bun.Transpiler();
+            transpiler.transformSync("1");
+            keep.push(transpiler);
+            return Math.max(...seqs().filter(seq => !before.has(seq)));
+          };
+          Bun.color("red", "css");
+          const a = newestHeap();
+          // Nothing runs between these two, so the difference is what the probe itself creates.
+          const b = newestHeap();
+          for (let i = 0; i < 1000; i++) Bun.color("#ff8800", "css");
+          const c = newestHeap();
+          console.log(JSON.stringify({ observed: b > a, heapsPerThousandCalls: c - b - (b - a) }));
+        `,
+      ],
+      env: bunEnv,
+      stderr: "pipe",
+    });
+    const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]);
+    expect(stderr).toBe("");
+    expect(JSON.parse(stdout)).toEqual({ observed: true, heapsPerThousandCalls: 0 });
+    expect(exitCode).toBe(0);
+  });
+
+  test("a Worker that called Bun.color leaves no heap behind when it exits", async () => {
+    using dir = tempDir("color-worker-heap", {
+      "color-worker-heap-fixture.js": `
+        import { heapStats } from "bun:jsc";
+        import { Worker, isMainThread } from "node:worker_threads";
+        if (!isMainThread) {
+          for (const input of ["red", "#ff8800", "hsl(120, 50%, 50%)", "\\\\72 ed"]) Bun.color(input, "css");
+        } else {
+          const liveHeaps = () => heapStats({ dump: true }).mimallocDump.heaps.length;
+          const runWorker = () =>
+            new Promise((resolve, reject) => {
+              const worker = new Worker(import.meta.filename);
+              worker.on("error", reject);
+              worker.on("exit", code => (code === 0 ? resolve() : reject(new Error("worker exited with " + code))));
+            });
+          const before = liveHeaps();
+          await runWorker();
+          console.log(JSON.stringify({ leaked: liveHeaps() - before }));
+        }
+      `,
+    });
+    await using proc = Bun.spawn({
+      cmd: [bunExe(), "color-worker-heap-fixture.js"],
+      env: bunEnv,
+      cwd: String(dir),
+      stderr: "pipe",
+    });
+    const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]);
+    expect(stderr).toBe("");
+    expect(JSON.parse(stdout)).toEqual({ leaked: 0 });
+    expect(exitCode).toBe(0);
+  });
+});