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 {