From 041b39e1b791fd261ebffdd81ba914c2d94ecc28 Mon Sep 17 00:00:00 2001 From: Zoltan Kochan Date: Fri, 17 Apr 2026 01:20:36 +0200 Subject: [PATCH] feat(cafs-writer): streaming class + darwin codesign MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds `CafsStreamWriter` to the crate — a class-based NAPI interface that accepts gunzipped chunks via push() and dispatches each completed file to rayon's global pool, letting disk writes overlap with the remainder of the download. The buffered writeFiles() remains the primary entry point because on real-install workloads the streaming path performed within noise of it (lock churn per chunk offset the overlap benefit). Also: - copy-artifact.js now ad-hoc codesigns the .node on macOS; unsigned dylibs get SIGKILLed silently on load, which was painful to debug. - crossbeam-channel added for the streaming writer's result pipe. --- agent/cafs-writer/Cargo.toml | 1 + agent/cafs-writer/index.d.ts | 12 ++ agent/cafs-writer/scripts/copy-artifact.js | 9 + agent/cafs-writer/src/lib.rs | 197 +++++++++++++++++++++ worker/src/start.ts | 4 + 5 files changed, 223 insertions(+) diff --git a/agent/cafs-writer/Cargo.toml b/agent/cafs-writer/Cargo.toml index e3fa8b845f..67b929ecf9 100644 --- a/agent/cafs-writer/Cargo.toml +++ b/agent/cafs-writer/Cargo.toml @@ -10,6 +10,7 @@ crate-type = ["cdylib"] napi = { version = "2", default-features = false, features = ["napi4"] } napi-derive = "2" rayon = "1" +crossbeam-channel = "0.5" [target.'cfg(target_os = "linux")'.dependencies] tokio-uring = "0.5" diff --git a/agent/cafs-writer/index.d.ts b/agent/cafs-writer/index.d.ts index 4c70ff041e..78c4835896 100644 --- a/agent/cafs-writer/index.d.ts +++ b/agent/cafs-writer/index.d.ts @@ -3,3 +3,15 @@ * Returns the number of files newly written (EEXIST is not counted). */ export function writeFiles (storeDir: string, payload: Buffer): number + +/** + * Streaming writer: push chunks as they arrive off the wire; the parser + * dispatches each complete file to a thread pool so writes overlap with + * the rest of the download. Call `finish` once to block until all + * in-flight writes are done and receive the total count. + */ +export class CafsStreamWriter { + constructor (storeDir: string) + push (chunk: Buffer): void + finish (): number +} diff --git a/agent/cafs-writer/scripts/copy-artifact.js b/agent/cafs-writer/scripts/copy-artifact.js index fab8a3f9cd..66644ca53d 100644 --- a/agent/cafs-writer/scripts/copy-artifact.js +++ b/agent/cafs-writer/scripts/copy-artifact.js @@ -33,4 +33,13 @@ if (!existsSync(source)) { const triple = `${process.platform}-${process.arch}` const dest = join(root, `pnpm-cafs-writer.${triple}.node`) copyFileSync(source, dest) + +// macOS silently SIGKILLs dylibs loaded by node if they aren't signed; +// rebuild produces an unsigned artifact, so ad-hoc sign it after copy. +if (process.platform === 'darwin') { + const { spawnSync } = require('node:child_process') + const r = spawnSync('codesign', ['--sign', '-', dest], { stdio: 'inherit' }) + if (r.status !== 0) process.exit(r.status ?? 1) +} + console.log(`Wrote ${dest}`) diff --git a/agent/cafs-writer/src/lib.rs b/agent/cafs-writer/src/lib.rs index 930f165eaf..b659a494de 100644 --- a/agent/cafs-writer/src/lib.rs +++ b/agent/cafs-writer/src/lib.rs @@ -247,3 +247,200 @@ fn hex_encode(bytes: &[u8]) -> String { } out } + +// ======================================================================= +// Streaming writer +// ======================================================================= +// +// Incremental variant of `write_files`: the caller pushes bytes as they +// arrive off the wire, and the Rust side parses and writes each file +// the moment its full content is in the buffer. This overlaps the HTTP +// download with disk writes — the buffered `write_files` call has to +// wait for the full response to gunzip before any file can be written. +// +// The Linux io_uring backend remains batch-only for now; on Linux this +// class falls back to std::fs::write on a rayon-backed task pool, same +// as non-Linux platforms. + +/// A streaming CAFS writer. Push chunks as they arrive; call `finish` +/// once to wait for all in-flight writes and get the total file count. +/// +/// Not thread-safe; a single writer serves one /v1/files response. +#[napi] +pub struct CafsStreamWriter { + state: std::sync::Mutex, +} + +struct StreamState { + files_dir: PathBuf, + buffer: Vec, + header_skipped: bool, + parent_dirs_created: std::collections::HashSet, + // tx is dropped during finish() so the collector can observe EOF. + // Each dispatched write task clones tx; when all tasks complete and + // the owning tx is dropped, rx.iter() terminates. + tx: Option>>, + rx: crossbeam_channel::Receiver>, +} + +#[napi] +impl CafsStreamWriter { + #[napi(constructor)] + pub fn new(store_dir: String) -> Result { + let files_dir = PathBuf::from(store_dir).join("files"); + let (tx, rx) = crossbeam_channel::unbounded(); + Ok(CafsStreamWriter { + state: std::sync::Mutex::new(StreamState { + files_dir, + buffer: Vec::with_capacity(64 * 1024), + header_skipped: false, + parent_dirs_created: std::collections::HashSet::new(), + tx: Some(tx), + rx, + }), + }) + } + + /// Append a chunk of the gunzipped payload and dispatch any now-complete + /// file entries to the write pool. + #[napi] + pub fn push(&self, chunk: Buffer) -> Result<()> { + let mut s = self.state.lock().map_err(poisoned)?; + s.buffer.extend_from_slice(&chunk); + drain_buffer(&mut s).map_err(|e| Error::from_reason(format!("cafs-writer parse error: {e}"))) + } + + /// Signal end-of-stream; blocks until all dispatched writes have + /// completed and returns the total number of files newly written. + #[napi] + pub fn finish(&self) -> Result { + // Drop our Sender clone so rx.iter() terminates once all task clones + // also go out of scope. + let rx = { + let mut s = self.state.lock().map_err(poisoned)?; + // Buffer should be empty or just contain the end marker at this point. + drain_buffer(&mut s) + .map_err(|e| Error::from_reason(format!("cafs-writer parse error: {e}")))?; + s.tx.take(); + s.rx.clone() + }; + + let mut total = 0u32; + for result in rx.iter() { + total += result.map_err(|e| Error::from_reason(format!("cafs-writer write error: {e}")))?; + } + Ok(total) + } +} + +fn poisoned(_: std::sync::PoisonError) -> Error { + Error::from_reason("cafs-writer: state mutex poisoned (worker panic?)") +} + +// Parse as many complete entries as fit in the buffer, dispatching +// each to the write pool. Leaves any trailing incomplete bytes in the +// buffer for the next push. +fn drain_buffer(s: &mut StreamState) -> std::result::Result<(), String> { + let mut pos = 0usize; + + // Skip the one-time JSON header on the very first bytes we see. + if !s.header_skipped { + if s.buffer.len() < 4 { + return Ok(()); + } + let json_len = u32::from_be_bytes([s.buffer[0], s.buffer[1], s.buffer[2], s.buffer[3]]) + as usize; + if s.buffer.len() < 4 + json_len { + return Ok(()); + } + pos = 4 + json_len; + s.header_skipped = true; + } + + loop { + if s.buffer.len() - pos < 64 { + break; + } + // End marker: 64 zero bytes. Everything beyond is ignored. + if s.buffer[pos..pos + 64].iter().all(|&b| b == 0) { + pos += 64; + break; + } + if s.buffer.len() - pos < 69 { + break; // need digest + 4-byte size + 1-byte mode + } + let size = u32::from_be_bytes([ + s.buffer[pos + 64], + s.buffer[pos + 65], + s.buffer[pos + 66], + s.buffer[pos + 67], + ]) as usize; + let entry_len = 69 + size; + if s.buffer.len() - pos < entry_len { + break; // content not yet fully arrived + } + let digest_hex = hex_encode(&s.buffer[pos..pos + 64]); + let executable = (s.buffer[pos + 68] & 0x01) != 0; + let content = s.buffer[pos + 69..pos + entry_len].to_vec(); + pos += entry_len; + + // Create the files/XX/ prefix once per unique prefix. + let prefix = digest_hex[..2].to_string(); + if s.parent_dirs_created.insert(prefix.clone()) { + std::fs::create_dir_all(s.files_dir.join(&prefix)) + .map_err(|e| format!("mkdir {prefix}: {e}"))?; + } + + let files_dir = s.files_dir.clone(); + let tx = s.tx.as_ref().ok_or("push() after finish()")?.clone(); + // Use rayon's global thread pool — a fresh per-writer pool oversubscribes + // the machine when multiple worker threads run batches in parallel. + rayon::spawn(move || { + let res = write_entry_owned(&files_dir, &digest_hex, executable, &content); + let _ = tx.send(res); + }); + } + + // Drop the consumed prefix of the buffer in one shot. + if pos > 0 { + s.buffer.drain(..pos); + } + Ok(()) +} + +fn write_entry_owned( + files_dir: &std::path::Path, + digest_hex: &str, + executable: bool, + content: &[u8], +) -> std::result::Result { + let (prefix, suffix) = digest_hex.split_at(2); + let filename = if executable { + format!("{suffix}-exec") + } else { + suffix.to_string() + }; + let path = files_dir.join(prefix).join(filename); + let mode = if executable { 0o755 } else { 0o644 }; + + let mut opts = std::fs::OpenOptions::new(); + opts.write(true).create_new(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + opts.mode(mode); + } + #[cfg(not(unix))] + let _ = mode; + + match opts.open(&path) { + Ok(mut f) => { + use std::io::Write; + f.write_all(content) + .map_err(|e| format!("write {}: {e}", path.display()))?; + Ok(1) + } + Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => Ok(0), + Err(e) => Err(format!("open {}: {e}", path.display())), + } +} diff --git a/worker/src/start.ts b/worker/src/start.ts index 91916506fb..9fa025bc92 100644 --- a/worker/src/start.ts +++ b/worker/src/start.ts @@ -503,6 +503,10 @@ function writeCafsFiles (message: WriteCafsFilesMessage): { status: string, file // Try to load the native CAFS writer. Falls back to the pure-JS streaming // parser below when the addon hasn't been built for this platform. +// +// The native module also exports a CafsStreamWriter class that parses +// chunk-by-chunk; empirically on real-install workloads it performs the +// same as the buffered path, so we stay on the simpler writeFiles call. let nativeWriteFiles: ((storeDir: string, payload: Buffer) => number) | undefined if (!process.env.PNPM_CAFS_WRITER_DISABLE_NATIVE) { try {