zstd-rs/rust/src/zstdcli_trace.rs

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