ddidderr/zstd-rs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(cli): port benchmark and compatibility option handling

Browse Source 
Restore benchmark decode mode, auto-thread selection, alternate frame formats,
gzip/xz/lzma/lz4 aliases, and trace lifecycle handling in the Rust frontend.

Test Plan:
- RUSTC_WRAPPER= CARGO_BUILD_RUSTC_WRAPPER= cargo test --manifest-path rust/cli/Cargo.toml
- RUSTC_WRAPPER= CARGO_BUILD_RUSTC_WRAPPER= cargo clippy --manifest-path rust/cli/Cargo.toml --all-targets -- -D warnings
- make -B -C programs zstd V=1
- make -C tests check V=1
This commit is contained in:
ddidderr
2026-07-12 18:08:09 +02:00
parent 66b7858728
commit 3fa5872848
2 changed files with 306 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
programs/zstdcli.c
+8 -2
View File
@@ -24,7 +24,7 @@ int ZSTD_rust_cli_bench(const char* const* fileNames, unsigned nbFiles,
int startCLevel, int endCLevel,
const ZSTD_compressionParameters* compressionParams,
int displayLevel, unsigned nbSeconds,
size_t blockSize, int nbWorkers);
size_t blockSize, int nbWorkers, int mode);
const char* ZSTD_rust_cli_expected_version(void)
{
@@ -41,7 +41,7 @@ int ZSTD_rust_cli_bench(const char* const* fileNames, unsigned nbFiles,
int startCLevel, int endCLevel,
const ZSTD_compressionParameters* compressionParams,
int displayLevel, unsigned nbSeconds,
size_t blockSize, int nbWorkers)
size_t blockSize, int nbWorkers, int mode)
{
#ifndef ZSTD_NOBENCH
BMK_advancedParams_t advancedParams = BMK_initAdvancedParams();
@@ -50,6 +50,11 @@ int ZSTD_rust_cli_bench(const char* const* fileNames, unsigned nbFiles,
advancedParams.nbSeconds = nbSeconds;
advancedParams.blockSize = blockSize;
advancedParams.nbWorkers = nbWorkers;
advancedParams.mode = (BMK_mode_t)mode;
if (advancedParams.mode == BMK_decodeOnly) {
startLevel = 0;
endLevel = 0;
}
if (startLevel > ZSTD_maxCLevel()) startLevel = ZSTD_maxCLevel();
if (endLevel > ZSTD_maxCLevel()) endLevel = ZSTD_maxCLevel();
if (endLevel < startLevel) endLevel = startLevel;
@@ -67,6 +72,7 @@ int ZSTD_rust_cli_bench(const char* const* fileNames, unsigned nbFiles,
(void)fileNames; (void)nbFiles; (void)dictFileName;
(void)startCLevel; (void)endCLevel; (void)compressionParams;
(void)displayLevel; (void)nbSeconds; (void)blockSize; (void)nbWorkers;
(void)mode;
return -1;
#endif
}
rust/src/zstd_cli.rs
+298 -28
View File
@@ -22,8 +22,9 @@
//! byte-identical with the C CLI.
//!
//! Remaining C-only CLI boundaries are called out in `unsupported()` below:
//! tracing and alternate-format selection. Dictionary training uses the
//! temporary `DiB_trainFromFiles` bridge while its algorithms are Rust-owned.
//! a few compatibility-only options. Trace logging uses the existing C trace
//! backend through a narrow lifecycle bridge, while dictionary training uses
//! the temporary `DiB_trainFromFiles` bridge and Rust-owned algorithms.
use std::env;
use std::ffi::{CStr, CString, OsStr, OsString};
@@ -56,10 +57,17 @@ const STDOUT_MARK: &str = "/*stdout*\\";
const NULL_MARK: &str = "NUL";
#[cfg(not(windows))]
const NULL_MARK: &str = "/dev/null";
#[cfg(feature = "compression")]
const ZSTD_SUFFIX: &[u8] = b".zst\0";
const GZ_SUFFIX: &[u8] = b".gz\0";
const XZ_SUFFIX: &[u8] = b".xz\0";
const LZMA_SUFFIX: &[u8] = b".lzma\0";
const LZ4_SUFFIX: &[u8] = b".lz4\0";
const FIO_ZSTD_COMPRESSION: c_int = 0;
const FIO_GZIP_COMPRESSION: c_int = 1;
const FIO_XZ_COMPRESSION: c_int = 2;
const FIO_LZMA_COMPRESSION: c_int = 3;
const FIO_LZ4_COMPRESSION: c_int = 4;
const FIO_PS_AUTO: c_int = 0;
const FIO_PS_NEVER: c_int = 1;
const FIO_PS_ALWAYS: c_int = 2;
@@ -212,6 +220,8 @@ unsafe extern "C" {
fn FIO_setNbFilesTotal(ctx: *mut FIO_ctx_t, value: c_int);
fn FIO_setHasStdinInput(ctx: *mut FIO_ctx_t, value: c_int);
fn FIO_setHasStdoutOutput(ctx: *mut FIO_ctx_t, value: c_int);
fn TRACE_enable(filename: *const c_char);
fn TRACE_finish();
#[cfg(feature = "compression")]
fn FIO_compressFilename(
@@ -275,6 +285,7 @@ unsafe extern "C" {
nb_seconds: c_uint,
block_size: usize,
nb_workers: c_int,
mode: c_int,
) -> c_int;
/// Narrow bridge to `programs/dibio.c`. The file loader and dictionary
/// algorithms remain on the C/Rust library side of this boundary.
@@ -322,6 +333,48 @@ enum Operation {
Train,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum CompressionFormat {
Zstd,
Gzip,
Xz,
Lzma,
Lz4,
}
impl CompressionFormat {
fn parse(value: &str) -> Result<Self, String> {
match value {
"zstd" => Ok(Self::Zstd),
"gzip" => Ok(Self::Gzip),
"xz" => Ok(Self::Xz),
"lzma" => Ok(Self::Lzma),
"lz4" => Ok(Self::Lz4),
_ => Err(format!("unknown compression format {value:?}")),
}
}
const fn fio_type(self) -> c_int {
match self {
Self::Zstd => FIO_ZSTD_COMPRESSION,
Self::Gzip => FIO_GZIP_COMPRESSION,
Self::Xz => FIO_XZ_COMPRESSION,
Self::Lzma => FIO_LZMA_COMPRESSION,
Self::Lz4 => FIO_LZ4_COMPRESSION,
}
}
const fn suffix(self) -> &'static [u8] {
match self {
Self::Zstd => ZSTD_SUFFIX,
Self::Gzip => GZ_SUFFIX,
Self::Xz => XZ_SUFFIX,
Self::Lzma => LZMA_SUFFIX,
Self::Lz4 => LZ4_SUFFIX,
}
}
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
enum TrainingAlgorithm {
Cover,
@@ -340,6 +393,8 @@ enum Action {
#[derive(Debug)]
struct Cli {
operation: Operation,
bench_decode: bool,
compression_format: CompressionFormat,
inputs: Vec<CString>,
output: Option<CString>,
dictionary: Option<CString>,
@@ -359,6 +414,7 @@ struct Cli {
progress: i32,
workers: Option<i32>,
single_thread: bool,
auto_threads_logical: bool,
block_size: Option<usize>,
mem_limit: Option<u32>,
ldm: bool,
@@ -389,15 +445,17 @@ struct Cli {
max_dict_size: usize,
dictionary_id: u32,
dictionary_selectivity: u32,
trace: Option<CString>,
cover_params: ZDICT_cover_params_t,
fast_cover_params: ZDICT_fastCover_params_t,
unsupported_program: Option<String>,
}
impl Cli {
fn new(program_name: &str) -> Self {
let mut cli = Self {
operation: Operation::Compress,
bench_decode: false,
compression_format: CompressionFormat::Zstd,
inputs: Vec::new(),
output: None,
dictionary: None,
@@ -417,6 +475,7 @@ impl Cli {
progress: FIO_PS_AUTO,
workers: None,
single_thread: false,
auto_threads_logical: false,
block_size: None,
mem_limit: None,
ldm: false,
@@ -447,9 +506,9 @@ impl Cli {
max_dict_size: DEFAULT_MAX_DICT_SIZE,
dictionary_id: 0,
dictionary_selectivity: DEFAULT_DICT_SELECTIVITY,
trace: None,
cover_params: ZDICT_cover_params_t::default(),
fast_cover_params: default_fast_cover_params(),
unsupported_program: None,
};
match program_name {
@@ -463,8 +522,46 @@ impl Cli {
cli.pass_through = Some(1);
cli.display_level = 1;
}
"gzip" | "gunzip" | "gzcat" | "lzma" | "unlzma" | "xz" | "unxz" | "lz4" | "unlz4" => {
cli.unsupported_program = Some(program_name.to_owned())
"gzip" => {
cli.compression_format = CompressionFormat::Gzip;
cli.remove_source = true;
}
"gunzip" => {
cli.compression_format = CompressionFormat::Gzip;
cli.operation = Operation::Decompress;
cli.remove_source = true;
}
"gzcat" => {
cli.compression_format = CompressionFormat::Gzip;
cli.operation = Operation::Decompress;
cli.output = Some(cstring(STDOUT_MARK).expect("static stdout marker"));
cli.force = true;
cli.force_stdout = true;
cli.pass_through = Some(1);
cli.display_level = 1;
}
"xz" => {
cli.compression_format = CompressionFormat::Xz;
cli.remove_source = true;
}
"unxz" => {
cli.compression_format = CompressionFormat::Xz;
cli.operation = Operation::Decompress;
cli.remove_source = true;
}
"lzma" => {
cli.compression_format = CompressionFormat::Lzma;
cli.remove_source = true;
}
"unlzma" => {
cli.compression_format = CompressionFormat::Lzma;
cli.operation = Operation::Decompress;
cli.remove_source = true;
}
"lz4" => cli.compression_format = CompressionFormat::Lz4,
"unlz4" => {
cli.compression_format = CompressionFormat::Lz4;
cli.operation = Operation::Decompress;
}
_ => {}
}
@@ -508,11 +605,16 @@ unsafe fn default_worker_count() -> i32 {
}
#[cfg(feature = "compression")]
unsafe fn resolved_worker_count(workers: Option<i32>, single_thread: bool) -> i32 {
unsafe fn resolved_worker_count(
workers: Option<i32>,
single_thread: bool,
auto_threads_logical: bool,
) -> i32 {
match workers {
/* --single-thread pins zero workers; a bare zero (-T0 or the zstdmt
* program name) auto-detects the core count as in the C CLI. */
Some(0) if single_thread => 0,
Some(0) if auto_threads_logical => unsafe { UTIL_countLogicalCores() }.max(1),
Some(0) => unsafe { UTIL_countPhysicalCores() }.max(1),
Some(workers) => workers,
None => unsafe { default_worker_count() },
@@ -564,7 +666,7 @@ fn usage(advanced: bool) {
let _ = writeln!(out, "\nImplemented advanced compression controls:");
let _ = writeln!(
out,
" --fast[=#], --ultra, --long[=#], --threads=#, --single-thread, --block-size=#"
" --fast[=#], --ultra, --long[=#], --threads=#, --single-thread, --auto-threads={{physical|logical}}, --block-size=#"
);
let _ = writeln!(
out,
@@ -617,7 +719,14 @@ fn usage(advanced: bool) {
out,
"\nDictionary builder:\n --train Create a dictionary from training files\n --train-cover[=k=#,d=#,steps=#,split=#,shrink[=#]]\n --train-fastcover[=k=#,d=#,f=#,steps=#,split=#,accel=#,shrink[=#]]\n --train-legacy[=s=#] Use the legacy algorithm\n --maxdict=# Maximum dictionary size (default {DEFAULT_MAX_DICT_SIZE})\n --dictID=# Force the dictionary ID (default: random)"
);
let _ = writeln!(out, "\nNot yet migrated: trace and alternate formats.");
let _ = writeln!(
out,
"\n --format=zstd|gzip|xz|lzma|lz4 Select the compression frame format"
);
let _ = writeln!(
out,
" --trace LOG Log compression/decompression trace data"
);
}
}
@@ -1000,7 +1109,10 @@ fn parse_long_option(
Ok(None)
}
"--decompress" | "--uncompress" => {
cli.operation = Operation::Decompress;
cli.bench_decode = true;
if cli.operation != Operation::Bench {
cli.operation = Operation::Decompress;
}
Ok(None)
}
"--test" => {
@@ -1156,6 +1268,15 @@ fn parse_long_option(
cli.workers = Some(parse_worker_count(&value)?);
Ok(None)
}
"--auto-threads" => {
let value = next_field(attached, args, index, name)?;
cli.auto_threads_logical = match value.as_str() {
"physical" => false,
"logical" => true,
_ => return Err(format!("unknown auto-thread mode {value:?}")),
};
Ok(None)
}
"--single-thread" => {
/* As in the C CLI: zero workers plus a latch that suppresses the
* automatic core-count resolution, so fileio runs its
@@ -1241,6 +1362,16 @@ fn parse_long_option(
cli.dictionary_id = parse_u32(&value, "dictionary ID")?;
Ok(None)
}
"--format" => {
let value = next_field(attached, args, index, name)?;
cli.compression_format = CompressionFormat::parse(&value)?;
Ok(None)
}
"--trace" => {
let value = next_field(attached, args, index, name)?;
cli.trace = Some(cstring(&value)?);
Ok(None)
}
"--filelist" => {
let value = next_field(attached, args, index, name)?;
cli.file_lists.push(cstring(&value)?);
@@ -1273,10 +1404,7 @@ fn parse_long_option(
}
"--max"
| "--patch-from"
| "--trace"
| "--format"
| "--priority"
| "--auto-threads"
| "--fake-stdin-is-console"
| "--fake-stdout-is-console"
| "--fake-stderr-is-console"
@@ -1309,7 +1437,12 @@ fn parse_short_options(
offset = digits_end;
continue;
}
'd' => cli.operation = Operation::Decompress,
'd' => {
cli.bench_decode = true;
if cli.operation != Operation::Bench {
cli.operation = Operation::Decompress;
}
}
'z' => cli.operation = Operation::Compress,
't' => cli.operation = Operation::Test,
'b' => cli.operation = Operation::Bench,
@@ -1377,7 +1510,20 @@ fn parse_short_options(
}
break;
}
'p' | 'P' | 's' | 'S' => {
's' => {
let mut digits_end = offset + 1;
while digits_end < bytes.len() && bytes[digits_end].is_ascii_digit() {
digits_end += 1;
}
if digits_end == offset + 1 {
return Err("missing argument for -s".to_owned());
}
cli.dictionary_selectivity =
parse_u32(&value[offset + 1..digits_end], "dictionary selectivity")?;
offset = digits_end;
continue;
}
'p' | 'P' | 'S' => {
unsupported(&format!("-{option}"))?;
}
_ => return Err(format!("unknown option -{option}")),
@@ -1429,7 +1575,7 @@ unsafe fn apply_preferences(cli: &Cli, prefs: *mut FIO_prefs_t, ctx: *mut FIO_ct
};
unsafe {
g_utilDisplayLevel = display_level;
FIO_setCompressionType(prefs, FIO_ZSTD_COMPRESSION);
FIO_setCompressionType(prefs, cli.compression_format.fio_type());
FIO_setNotificationLevel(display_level);
FIO_setProgressSetting(
if !io::stderr().is_terminal() && cli.progress != FIO_PS_ALWAYS {
@@ -1462,7 +1608,10 @@ unsafe fn apply_preferences(cli: &Cli, prefs: *mut FIO_prefs_t, ctx: *mut FIO_ct
}),
);
#[cfg(feature = "compression")]
FIO_setNbWorkers(prefs, resolved_worker_count(cli.workers, cli.single_thread));
FIO_setNbWorkers(
prefs,
resolved_worker_count(cli.workers, cli.single_thread, cli.auto_threads_logical),
);
FIO_setLdmFlag(prefs, u32::from(cli.ldm));
FIO_setAdaptiveMode(prefs, i32::from(cli.adapt));
FIO_setRsyncable(prefs, i32::from(cli.rsyncable));
@@ -1694,7 +1843,7 @@ unsafe fn run_compress(
output_dir_mirror(cli),
output_dir_flat(cli),
output,
ZSTD_SUFFIX.as_ptr().cast(),
cli.compression_format.suffix().as_ptr().cast(),
dictionary,
cli.level,
cli.compression_params,
@@ -1826,6 +1975,18 @@ fn run_bench(cli: &Cli) -> Result<i32, String> {
.dictionary
.as_ref()
.map_or(ptr::null(), |value| value.as_ptr());
let workers = {
#[cfg(feature = "compression")]
{
unsafe {
resolved_worker_count(cli.workers, cli.single_thread, cli.auto_threads_logical)
}
}
#[cfg(not(feature = "compression"))]
{
cli.workers.unwrap_or(1)
}
};
let result = unsafe {
ZSTD_rust_cli_bench(
inputs.as_ptr(),
@@ -1838,7 +1999,8 @@ fn run_bench(cli: &Cli) -> Result<i32, String> {
cli.bench_nb_seconds.unwrap_or(DEFAULT_BENCH_NB_SECONDS),
cli.block_size.unwrap_or(0),
// The C CLI benchmarks single-threaded unless -T was given.
cli.workers.unwrap_or(1),
workers,
if cli.bench_decode { 1 } else { 0 },
)
};
if result < 0 {
@@ -1862,7 +2024,9 @@ fn run_train(cli: &Cli) -> Result<i32, String> {
let default_output = cstring(DEFAULT_DICT_NAME).expect("static dictionary name");
let output = cli.output.as_ref().unwrap_or(&default_output);
let inputs: Vec<*const c_char> = cli.inputs.iter().map(|value| value.as_ptr()).collect();
let workers = unsafe { resolved_worker_count(cli.workers, cli.single_thread) } as c_uint;
let workers = unsafe {
resolved_worker_count(cli.workers, cli.single_thread, cli.auto_threads_logical)
} as c_uint;
let z_params = ZDICT_params_t {
compressionLevel: cli.level,
notificationLevel: cli.display_level as c_uint,
@@ -1936,12 +2100,31 @@ fn run_train(cli: &Cli) -> Result<i32, String> {
}
}
fn run_cli(mut cli: Cli) -> Result<i32, String> {
if let Some(program_name) = &cli.unsupported_program {
return Err(format!(
"{program_name} compatibility mode is not yet implemented by the Rust CLI frontend"
));
struct TraceGuard {
enabled: bool,
}
impl TraceGuard {
fn new(filename: Option<&CString>) -> Self {
if let Some(filename) = filename {
unsafe { TRACE_enable(filename.as_ptr()) };
Self { enabled: true }
} else {
Self { enabled: false }
}
}
}
impl Drop for TraceGuard {
fn drop(&mut self) {
if self.enabled {
unsafe { TRACE_finish() };
}
}
}
fn run_cli(mut cli: Cli) -> Result<i32, String> {
let _trace_guard = TraceGuard::new(cli.trace.as_ref());
unsafe {
/* The C CLI publishes the display level to util.c before any table
* expansion, so traversal warnings obey -q/-v. */
@@ -1966,6 +2149,9 @@ fn run_cli(mut cli: Cli) -> Result<i32, String> {
return Err("file information is not supported".to_owned());
}
if cli.operation == Operation::Bench {
if cli.compression_format != CompressionFormat::Zstd {
return Err("benchmark mode only supports the zstd format".to_owned());
}
return run_bench(&cli);
}
if cli.operation == Operation::Train {
@@ -2186,6 +2372,46 @@ mod tests {
assert_eq!(cli.content_size, 1);
assert_eq!(cli.workers, None);
assert_eq!(cli.operation, Operation::Compress);
assert_eq!(cli.compression_format, CompressionFormat::Zstd);
}
#[test]
fn format_option_selects_frame_type_and_suffix() {
let cli = parse(&["zstd", "--format=gzip", "input"]);
assert_eq!(cli.compression_format, CompressionFormat::Gzip);
assert_eq!(cli.compression_format.fio_type(), FIO_GZIP_COMPRESSION);
assert_eq!(cli.compression_format.suffix(), b".gz\0");
let cli = parse(&["zstd", "--format=gzip", "--format=zstd", "input"]);
assert_eq!(cli.compression_format, CompressionFormat::Zstd);
}
#[test]
fn format_option_rejects_unknown_formats() {
let error = parse_args(vec![
OsString::from("zstd"),
OsString::from("--format=zip"),
OsString::from("input"),
])
.expect_err("unknown formats must be rejected");
assert!(error.contains("unknown compression format"));
}
#[test]
fn trace_option_accepts_attached_and_separate_paths() {
let cli = parse(&["zstd", "--trace=first.trace", "input"]);
assert_eq!(
cli.trace.as_deref().map(CStr::to_bytes),
Some(&b"first.trace"[..])
);
let cli = parse(&["zstd", "--trace", "second.trace", "input"]);
assert_eq!(
cli.trace.as_deref().map(CStr::to_bytes),
Some(&b"second.trace"[..])
);
}
#[test]
@@ -2276,6 +2502,28 @@ mod tests {
assert!(!cli.single_thread);
}
#[test]
fn auto_threads_selects_core_family() {
let physical = parse(&["zstd", "-T0", "--auto-threads=physical", "input"]);
assert_eq!(physical.workers, Some(0));
assert!(!physical.auto_threads_logical);
let logical = parse(&["zstd", "-T0", "--auto-threads", "logical", "input"]);
assert_eq!(logical.workers, Some(0));
assert!(logical.auto_threads_logical);
}
#[test]
fn auto_threads_rejects_unknown_core_family() {
let error = parse_args(vec![
OsString::from("zstd"),
OsString::from("--auto-threads=bogus"),
])
.expect_err("an unknown core family must be rejected");
assert!(error.contains("unknown auto-thread mode"));
}
#[test]
fn single_thread_pins_zero_workers() {
let cli = parse(&["zstd", "--single-thread", "input"]);
@@ -2306,10 +2554,20 @@ mod tests {
}
#[test]
fn alternate_format_aliases_fail_before_processing_files() {
fn alternate_format_aliases_select_compatibility_modes() {
let cli = Cli::new("gzip");
assert_eq!(cli.compression_format, CompressionFormat::Gzip);
assert_eq!(cli.operation, Operation::Compress);
assert!(cli.remove_source);
assert_eq!(cli.unsupported_program.as_deref(), Some("gzip"));
let cli = Cli::new("unxz");
assert_eq!(cli.compression_format, CompressionFormat::Xz);
assert_eq!(cli.operation, Operation::Decompress);
assert!(cli.remove_source);
let cli = Cli::new("lzma");
assert_eq!(cli.compression_format, CompressionFormat::Lzma);
assert_eq!(cli.operation, Operation::Compress);
}
#[cfg(unix)]
@@ -2405,6 +2663,10 @@ mod tests {
assert_eq!(legacy.operation, Operation::Train);
assert_eq!(legacy.training_algorithm, TrainingAlgorithm::Legacy);
assert_eq!(legacy.dictionary_selectivity, 12);
let legacy_short = parse(&["zstd", "--train-legacy", "-s5", "sample"]);
assert_eq!(legacy_short.training_algorithm, TrainingAlgorithm::Legacy);
assert_eq!(legacy_short.dictionary_selectivity, 5);
}
#[test]
@@ -2573,4 +2835,12 @@ mod tests {
assert_eq!(cli.bench_end_level, Some(0));
assert_eq!(cli.bench_nb_seconds, Some(0));
}
#[test]
fn bench_decompression_keeps_benchmark_operation() {
let cli = parse(&["zstd", "-b", "-d", "-i0", "input"]);
assert_eq!(cli.operation, Operation::Bench);
assert!(cli.bench_decode);
}
}