diff --git a/programs/Makefile b/programs/Makefile index eb76c1dc1..ac91210cc 100644 --- a/programs/Makefile +++ b/programs/Makefile @@ -37,6 +37,7 @@ RUST_CLI_SOURCES := $(RUST_CLI_MANIFEST) $(RUST_CLI_DIR)/Cargo.lock \ $(RUST_DIR)/src/fileio_prefs.rs \ $(RUST_DIR)/src/dibio.rs \ $(RUST_DIR)/src/util.rs \ + $(RUST_DIR)/src/zstdcli_trace.rs \ $(RUST_DIR)/src/timefn.rs $(RUST_DIR)/src/benchfn.rs \ $(RUST_DIR)/src/datagen.rs $(RUST_DIR)/src/lorem.rs @@ -232,7 +233,7 @@ ZSTDLIB_LOCAL_SRC = $(notdir $(ZSTDLIB_FULL_SRC)) ZSTDLIB_LOCAL_OBJ0 := $(ZSTDLIB_LOCAL_SRC:.c=.o) ZSTDLIB_LOCAL_OBJ := $(ZSTDLIB_LOCAL_OBJ0:.S=.o) -ZSTD_CLI_SRC := $(filter-out benchzstd.c,$(sort $(wildcard *.c))) +ZSTD_CLI_SRC := $(filter-out benchzstd.c zstdcli_trace.c,$(sort $(wildcard *.c))) ZSTD_CLI_OBJ := $(ZSTD_CLI_SRC:.c=.o) ZSTD_ALL_SRC = $(ZSTDLIB_LOCAL_SRC) $(ZSTD_CLI_SRC) diff --git a/rust/cli/src/lib.rs b/rust/cli/src/lib.rs index 8c6fd19cc..f9b15c055 100644 --- a/rust/cli/src/lib.rs +++ b/rust/cli/src/lib.rs @@ -20,3 +20,5 @@ mod util; #[cfg(feature = "cli")] #[path = "../../src/zstd_cli.rs"] mod zstd_cli; +#[path = "../../src/zstdcli_trace.rs"] +mod zstdcli_trace; diff --git a/rust/src/zstdcli_trace.rs b/rust/src/zstdcli_trace.rs new file mode 100644 index 000000000..3558d1768 --- /dev/null +++ b/rust/src/zstdcli_trace.rs @@ -0,0 +1,229 @@ +#![allow(non_camel_case_types)] +#![allow(non_snake_case)] +#![allow(clippy::missing_safety_doc)] + +//! Command-line trace callbacks. +//! +//! `programs/zstdcli_trace.c` is a small strong-symbol override for the +//! library's weak tracing callbacks. Keep the wire-visible structures and +//! CSV format here so the CLI archive does not need a C implementation merely +//! to enable tracing. The decompression-only CLI still gets the no-op +//! callbacks below; that variant does not link the compression parameter API. + +use std::ffi::{c_char, c_int, c_uint, c_void, CStr}; +use std::fs::{File, OpenOptions}; +use std::io::Write; +use std::path::Path; +use std::sync::{Mutex, OnceLock}; +use std::time::Instant; + +#[cfg(feature = "compression")] +const ZSTD_C_COMPRESSION_LEVEL: c_int = 100; +#[cfg(feature = "compression")] +const ZSTD_C_NB_WORKERS: c_int = 400; + +#[repr(C)] +pub struct ZSTD_Trace { + pub version: c_uint, + pub streaming: c_int, + pub dictionaryID: c_uint, + pub dictionaryIsCold: c_int, + pub dictionarySize: usize, + pub uncompressedSize: usize, + pub compressedSize: usize, + pub params: *const c_void, + pub cctx: *const c_void, + pub dctx: *const c_void, +} + +#[derive(Default)] +struct TraceState { + file: Option, + enabled_at: Option, +} + +fn trace_state() -> &'static Mutex { + static STATE: OnceLock> = OnceLock::new(); + STATE.get_or_init(|| Mutex::new(TraceState::default())) +} + +#[cfg(feature = "compression")] +unsafe extern "C" { + fn ZSTD_CCtxParams_getParameter( + params: *const c_void, + param: c_int, + value: *mut c_int, + ) -> usize; +} + +fn elapsed_nanos(start: Option) -> u64 { + start + .map(|instant| instant.elapsed().as_nanos().min(u64::MAX as u128) as u64) + .unwrap_or(0) +} + +#[no_mangle] +pub unsafe extern "C" fn TRACE_enable(filename: *const c_char) { + let Some(filename) = (!filename.is_null()).then(|| unsafe { CStr::from_ptr(filename) }) else { + return; + }; + let Ok(path_string) = filename.to_str() else { + return; + }; + let path = Path::new(path_string); + let was_regular_file = std::fs::metadata(path) + .map(|metadata| metadata.is_file()) + .unwrap_or(false); + let file = OpenOptions::new().create(true).append(true).open(path).ok(); + + let mut state = trace_state() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(mut file) = file { + if !was_regular_file { + let _ = writeln!( + file, + "Algorithm, Version, Method, Mode, Level, Workers, Dictionary Size, Uncompressed Size, Compressed Size, Duration Nanos, Compression Ratio, Speed MB/s" + ); + } + state.file = Some(file); + } else { + state.file = None; + } + state.enabled_at = Some(Instant::now()); +} + +#[no_mangle] +pub extern "C" fn TRACE_finish() { + let mut state = trace_state() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + state.file = None; + state.enabled_at = None; +} + +fn trace_begin() -> u64 { + let state = trace_state() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if state.file.is_none() { + return 0; + } + elapsed_nanos(state.enabled_at).max(1) +} + +fn trace_log(method: &str, begin: u64, trace: &ZSTD_Trace, state: &mut TraceState) { + let Some(file) = state.file.as_mut() else { + return; + }; + let duration = elapsed_nanos(state.enabled_at).saturating_sub(begin); + let duration = duration.max(1); + let (level, workers) = { + #[cfg(feature = "compression")] + { + let mut level = 0; + let mut workers = 0; + if !trace.params.is_null() { + unsafe { + ZSTD_CCtxParams_getParameter( + trace.params, + ZSTD_C_COMPRESSION_LEVEL, + &mut level, + ); + ZSTD_CCtxParams_getParameter(trace.params, ZSTD_C_NB_WORKERS, &mut workers); + } + } + (level, workers) + } + #[cfg(not(feature = "compression"))] + { + (0, 0) + } + }; + let ratio = trace.uncompressedSize as f64 / trace.compressedSize.max(1) as f64; + let speed = trace.uncompressedSize as f64 * 1000.0 / duration as f64; + let _ = writeln!( + file, + "zstd, {}, {}, {}, {}, {}, {}, {}, {}, {}, {:.2}, {:.2}", + trace.version, + method, + if trace.streaming != 0 { + "streaming" + } else { + "single-pass" + }, + level, + workers, + trace.dictionarySize, + trace.uncompressedSize, + trace.compressedSize, + duration, + ratio, + speed, + ); +} + +#[no_mangle] +pub extern "C" fn ZSTD_trace_compress_begin(_cctx: *const c_void) -> u64 { + trace_begin() +} + +#[no_mangle] +pub unsafe extern "C" fn ZSTD_trace_compress_end(ctx: u64, trace: *const ZSTD_Trace) { + if ctx == 0 || trace.is_null() { + return; + } + let mut state = trace_state() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + trace_log("compress", ctx, unsafe { &*trace }, &mut state); +} + +#[no_mangle] +pub extern "C" fn ZSTD_trace_decompress_begin(_dctx: *const c_void) -> u64 { + trace_begin() +} + +#[no_mangle] +pub unsafe extern "C" fn ZSTD_trace_decompress_end(ctx: u64, trace: *const ZSTD_Trace) { + if ctx == 0 || trace.is_null() { + return; + } + let mut state = trace_state() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + trace_log("decompress", ctx, unsafe { &*trace }, &mut state); +} + +#[cfg(test)] +mod tests { + use super::*; + use std::ffi::CString; + use std::fs; + + #[test] + fn disabled_trace_callbacks_are_noops() { + TRACE_finish(); + assert_eq!(ZSTD_trace_compress_begin(std::ptr::null()), 0); + assert_eq!(ZSTD_trace_decompress_begin(std::ptr::null()), 0); + } + + #[test] + fn enabling_trace_writes_the_header_once() { + let path = std::env::temp_dir().join(format!( + "zstd-rust-trace-{}-{}.csv", + std::process::id(), + std::thread::current().name().unwrap_or("test") + )); + let path_string = path.to_string_lossy().into_owned(); + let path_c = CString::new(path_string).unwrap(); + let _ = fs::remove_file(&path); + unsafe { TRACE_enable(path_c.as_ptr()) }; + assert_ne!(ZSTD_trace_compress_begin(std::ptr::null()), 0); + TRACE_finish(); + let contents = fs::read_to_string(&path).unwrap(); + assert_eq!(contents.lines().count(), 1); + assert!(contents.starts_with("Algorithm, Version,")); + let _ = fs::remove_file(path); + } +}