#![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); } }