feat(cli): move trace callbacks into Rust

Implement the CLI trace callbacks and CSV logger in Rust, preserving the ZSTD_Trace ABI, timing fields, parameter reporting, mutex-protected writes, and no-op disabled behavior. Add the module to every Rust CLI archive and exclude the duplicate C trace object from Rust-backed program targets. Test Plan: - cargo test --manifest-path rust/cli/Cargo.toml (92 passed) - cargo clippy --manifest-path rust/cli/Cargo.toml --all-targets - make -B -C programs -j2 zstd - Rust trace smoke: programs/zstd -q --trace=/tmp/zstd-rust-trace-smoke.csv -c /tmp/levels-input - make -B -C tests -j2 test-cli-tests (41 passed)
2026-07-18 02:50:25 +02:00
parent 49b3405d63
commit 9f85a7ebd7
3 changed files with 233 additions and 1 deletions
@@ -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)
@@ -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;
@@ -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<File>,
+    enabled_at: Option<Instant>,
+}
+
+fn trace_state() -> &'static Mutex<TraceState> {
+    static STATE: OnceLock<Mutex<TraceState>> = 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<Instant>) -> 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);
+    }
+}