diff --git a/programs/fileio.c b/programs/fileio.c index fa6812a45..232c96d8e 100644 --- a/programs/fileio.c +++ b/programs/fileio.c @@ -315,6 +315,9 @@ void FIO_rust_setOutBuffer(ZSTD_outBuffer* output, void* buf, size_t s, size_t p const char* FIO_rust_determineCompressedName(const char* srcFileName, const char* outDirName, const char* suffix); const char* FIO_rust_determineDstName(const char* srcFileName, const char* outDirName, const char* const* suffixList, const char* suffixListStr); +int FIO_rust_adjustMemLimitForPatchFromMode(FIO_prefs_t* prefs, + unsigned long long dictSize, + unsigned long long maxSrcFileSize); #ifdef ZSTD_LZ4COMPRESS int FIO_rust_LZ4_GetBlockSize_FromBlockId(int id); #endif @@ -707,14 +710,19 @@ static void FIO_adjustMemLimitForPatchFromMode(FIO_prefs_t* const prefs, unsigned long long const dictSize, unsigned long long const maxSrcFileSize) { - unsigned long long maxSize = MAX(prefs->memLimit, MAX(dictSize, maxSrcFileSize)); + enum { + FIO_PATCH_MEM_LIMIT_SUCCESS = 0, + FIO_PATCH_MEM_LIMIT_UNKNOWN_SIZE = 1, + FIO_PATCH_MEM_LIMIT_TOO_LARGE = 2 + }; unsigned const maxWindowSize = (1U << ZSTD_WINDOWLOG_MAX); - if (maxSize == UTIL_FILESIZE_UNKNOWN) + + int const status = FIO_rust_adjustMemLimitForPatchFromMode(prefs, dictSize, maxSrcFileSize); + if (status == FIO_PATCH_MEM_LIMIT_UNKNOWN_SIZE) EXM_THROW(42, "Using --patch-from with stdin requires --stream-size"); - assert(maxSize != UTIL_FILESIZE_UNKNOWN); - if (maxSize > maxWindowSize) + if (status == FIO_PATCH_MEM_LIMIT_TOO_LARGE) EXM_THROW(42, "Can't handle files larger than %u GB\n", maxWindowSize/(1 GB)); - FIO_setMemLimit(prefs, (unsigned)maxSize); + assert(status == FIO_PATCH_MEM_LIMIT_SUCCESS); } /* FIO_multiFilesConcatWarning() : diff --git a/rust/src/fileio_prefs.rs b/rust/src/fileio_prefs.rs index 0e5b1b4a4..d32e5d6bb 100644 --- a/rust/src/fileio_prefs.rs +++ b/rust/src/fileio_prefs.rs @@ -19,6 +19,11 @@ use std::ptr; const FIO_ZSTD_COMPRESSION: c_int = 0; const FIO_OVERLAP_LOG_NOTSET: c_int = 9999; const FIO_LDM_PARAM_NOTSET: c_int = 9999; +const FIO_PATCH_MEM_LIMIT_SUCCESS: c_int = 0; +const FIO_PATCH_MEM_LIMIT_UNKNOWN_SIZE: c_int = 1; +const FIO_PATCH_MEM_LIMIT_TOO_LARGE: c_int = 2; +const UTIL_FILESIZE_UNKNOWN: u64 = u64::MAX; +const ZSTD_WINDOWLOG_MAX: u32 = if size_of::() == 4 { 30 } else { 31 }; static STDOUT_MARK: &[u8] = b"/*stdout*\\\0"; static mut COMPRESSED_NAME_CAPACITY: usize = 0; @@ -295,6 +300,44 @@ pub unsafe extern "C" fn FIO_setMemLimit(prefs: *mut FIO_prefs_t, mem_limit: c_u unsafe { ptr::addr_of_mut!((*prefs).memLimit).write(mem_limit) }; } +fn adjusted_patch_mem_limit( + current_mem_limit: c_uint, + dict_size: u64, + max_src_file_size: u64, +) -> Result { + let max_size = u64::from(current_mem_limit) + .max(dict_size) + .max(max_src_file_size); + + if max_size == UTIL_FILESIZE_UNKNOWN { + return Err(FIO_PATCH_MEM_LIMIT_UNKNOWN_SIZE); + } + if max_size > (1u64 << ZSTD_WINDOWLOG_MAX) { + return Err(FIO_PATCH_MEM_LIMIT_TOO_LARGE); + } + Ok(max_size as c_uint) +} + +/// Adjust the patch-from memory limit without changing it on a rejected size. +/// +/// The status values are mirrored by the C diagnostic wrapper immediately +/// below its ABI declarations. +#[no_mangle] +pub unsafe extern "C" fn FIO_rust_adjustMemLimitForPatchFromMode( + prefs: *mut FIO_prefs_t, + dict_size: u64, + max_src_file_size: u64, +) -> c_int { + let current_mem_limit = unsafe { (*prefs).memLimit }; + match adjusted_patch_mem_limit(current_mem_limit, dict_size, max_src_file_size) { + Ok(mem_limit) => { + unsafe { ptr::addr_of_mut!((*prefs).memLimit).write(mem_limit) }; + FIO_PATCH_MEM_LIMIT_SUCCESS + } + Err(status) => status, + } +} + #[cfg(feature = "compression")] #[no_mangle] pub unsafe extern "C" fn FIO_setNbWorkers(prefs: *mut FIO_prefs_t, nb_workers: c_int) { @@ -936,12 +979,73 @@ mod tests { use std::fs; use std::mem::{align_of, offset_of, size_of}; - const UTIL_FILESIZE_UNKNOWN: u64 = u64::MAX; - fn temporary_file_path(name: &str) -> std::path::PathBuf { std::env::temp_dir().join(format!("zstd-fileio-prefs-{}-{name}", std::process::id())) } + fn prefs_with_mem_limit(mem_limit: c_uint) -> FIO_prefs_t { + let mut prefs = unsafe { std::mem::zeroed::() }; + prefs.memLimit = mem_limit; + prefs + } + + #[test] + fn patch_mem_limit_preserves_existing_limit() { + let mut prefs = prefs_with_mem_limit(4096); + let status = unsafe { FIO_rust_adjustMemLimitForPatchFromMode(&mut prefs, 1024, 2048) }; + + assert_eq!(status, FIO_PATCH_MEM_LIMIT_SUCCESS); + assert_eq!(prefs.memLimit, 4096); + } + + #[test] + fn patch_mem_limit_uses_dictionary_and_source_maxima() { + let mut prefs = prefs_with_mem_limit(128); + let status = unsafe { FIO_rust_adjustMemLimitForPatchFromMode(&mut prefs, 4096, 1024) }; + assert_eq!(status, FIO_PATCH_MEM_LIMIT_SUCCESS); + assert_eq!(prefs.memLimit, 4096); + + let mut prefs = prefs_with_mem_limit(128); + let status = unsafe { FIO_rust_adjustMemLimitForPatchFromMode(&mut prefs, 1024, 8192) }; + assert_eq!(status, FIO_PATCH_MEM_LIMIT_SUCCESS); + assert_eq!(prefs.memLimit, 8192); + } + + #[test] + fn patch_mem_limit_rejects_unknown_sizes_without_update() { + for (dict_size, max_src_file_size) in + [(UTIL_FILESIZE_UNKNOWN, 0), (0, UTIL_FILESIZE_UNKNOWN)] + { + let mut prefs = prefs_with_mem_limit(4096); + let status = unsafe { + FIO_rust_adjustMemLimitForPatchFromMode(&mut prefs, dict_size, max_src_file_size) + }; + + assert_eq!(status, FIO_PATCH_MEM_LIMIT_UNKNOWN_SIZE); + assert_eq!(prefs.memLimit, 4096); + } + } + + #[test] + fn patch_mem_limit_rejects_sizes_over_the_window_without_update() { + let mut prefs = prefs_with_mem_limit(4096); + let status = unsafe { + FIO_rust_adjustMemLimitForPatchFromMode(&mut prefs, (1u64 << ZSTD_WINDOWLOG_MAX) + 1, 0) + }; + + assert_eq!(status, FIO_PATCH_MEM_LIMIT_TOO_LARGE); + assert_eq!(prefs.memLimit, 4096); + } + + #[test] + fn patch_mem_limit_updates_on_success() { + let mut prefs = prefs_with_mem_limit(128); + let status = unsafe { FIO_rust_adjustMemLimitForPatchFromMode(&mut prefs, 1024, 8192) }; + + assert_eq!(status, FIO_PATCH_MEM_LIMIT_SUCCESS); + assert_eq!(prefs.memLimit, 8192); + } + #[test] fn largest_file_size_scan_handles_empty_input() { assert_eq!(unsafe { FIO_rust_getLargestFileSize(ptr::null(), 0) }, 0);