diff --git a/lib/compress/zstdmt_compress.c b/lib/compress/zstdmt_compress.c index cc5b8d028..29b6e4b01 100644 --- a/lib/compress/zstdmt_compress.c +++ b/lib/compress/zstdmt_compress.c @@ -163,6 +163,15 @@ ZSTD_frameProgression ZSTDMT_rust_frameProgression( ZSTD_frameProgression ZSTDMT_rust_frameProgressionAddJob( ZSTD_frameProgression progression, size_t srcSize, size_t consumed, size_t produced, size_t flushed); +typedef int (*ZSTDMT_jobTableInitFn)(void* jobTable, unsigned nbJobs, + size_t jobSize); +typedef void (*ZSTDMT_jobTableDestroyFn)(void* jobTable, unsigned nbJobs, + size_t jobSize); +size_t ZSTDMT_rust_expandJobsTable( + void** jobTablePtr, unsigned* jobIDMaskPtr, unsigned nbWorkers, + size_t jobSize, ZSTD_customMem cMem, + ZSTDMT_jobTableInitFn initSync, + ZSTDMT_jobTableDestroyFn destroySync); typedef struct ZSTDMT_bufferPool_s { ZSTDMT_RustBufferPool* rustPool; @@ -911,16 +920,14 @@ static ZSTDMT_jobDescription* ZSTDMT_createJobsTable(U32* nbJobsPtr, ZSTD_custom } static size_t ZSTDMT_expandJobsTable (ZSTDMT_CCtx* mtctx, U32 nbWorkers) { - U32 nbJobs = nbWorkers + 2; - if (nbJobs > mtctx->jobIDMask+1) { /* need more job capacity */ - ZSTDMT_freeJobsTable(mtctx->jobs, mtctx->jobIDMask+1, mtctx->cMem); - mtctx->jobIDMask = 0; - mtctx->jobs = ZSTDMT_createJobsTable(&nbJobs, mtctx->cMem); - if (mtctx->jobs==NULL) return ERROR(memory_allocation); - assert((nbJobs != 0) && ((nbJobs & (nbJobs - 1)) == 0)); /* ensure nbJobs is a power of 2 */ - mtctx->jobIDMask = nbJobs - 1; - } - return 0; + void* jobs = mtctx->jobs; + U32 jobIDMask = mtctx->jobIDMask; + size_t const error = ZSTDMT_rust_expandJobsTable( + &jobs, &jobIDMask, nbWorkers, sizeof(ZSTDMT_jobDescription), + mtctx->cMem, ZSTDMT_job_table_init_sync, ZSTDMT_job_table_destroy_sync); + mtctx->jobs = (ZSTDMT_jobDescription*)jobs; + mtctx->jobIDMask = jobIDMask; + return error; } diff --git a/rust/src/zstdmt_compress.rs b/rust/src/zstdmt_compress.rs index 57d595f35..4659f4963 100644 --- a/rust/src/zstdmt_compress.rs +++ b/rust/src/zstdmt_compress.rs @@ -10,14 +10,16 @@ //! those operations and projects only allocation/lifecycle pieces and pure //! sizing policy into this module. The entry points below are narrow C ABIs: //! buffers, `ZSTD_CCtx *` values, and job descriptors remain opaque to Rust, -//! while allocation, reuse, expansion, synchronization, and sizing policy are -//! Rust-owned. +//! while allocation, reuse, table-capacity orchestration, and sizing policy +//! are Rust-owned. C retains ownership of private descriptor fields and +//! platform synchronization. use std::mem::{self, MaybeUninit}; use std::os::raw::{c_int, c_uint, c_void}; use std::ptr; use std::sync::Mutex; +use crate::errors::{ZstdErrorCode, ERROR}; use crate::zstd_compress::ZSTD_frameProgression; const ZSTDMT_JOBLOG_MAX: c_uint = if mem::size_of::() == 4 { 29 } else { 30 }; @@ -653,6 +655,79 @@ unsafe fn free_job_table_storage(job_table: *mut c_void, custom_mem: ZstdCustomM unsafe { custom_free(job_table, custom_mem) }; } +type JobTableInitSync = unsafe extern "C" fn(*mut c_void, c_uint, usize) -> c_int; +type JobTableDestroySync = unsafe extern "C" fn(*mut c_void, c_uint, usize); + +unsafe fn free_job_table_with_sync( + job_table: *mut c_void, + nb_jobs: c_uint, + job_size: usize, + custom_mem: ZstdCustomMem, + destroy_sync: JobTableDestroySync, +) { + if job_table.is_null() { + return; + } + unsafe { destroy_sync(job_table, nb_jobs, job_size) }; + unsafe { ZSTDMT_rust_job_table_free(job_table, nb_jobs, job_size, custom_mem) }; +} + +unsafe fn expand_job_table( + job_table_ptr: *mut *mut c_void, + job_id_mask_ptr: *mut c_uint, + nb_workers: c_uint, + job_size: usize, + custom_mem: ZstdCustomMem, + init_sync: JobTableInitSync, + destroy_sync: JobTableDestroySync, +) -> usize { + if job_table_ptr.is_null() || job_id_mask_ptr.is_null() { + return ERROR(ZstdErrorCode::Generic); + } + + let requested_jobs = nb_workers.wrapping_add(2); + let current_capacity = unsafe { (*job_id_mask_ptr).wrapping_add(1) }; + if requested_jobs <= current_capacity { + return 0; + } + + let old_job_table = unsafe { *job_table_ptr }; + unsafe { + free_job_table_with_sync( + old_job_table, + current_capacity, + job_size, + custom_mem, + destroy_sync, + ); + // Match ZSTDMT_expandJobsTable(): after the old table is released, a + // failed replacement leaves the C context with no table and a zero + // mask. The C wrapper copies these scalar results back to its private + // context only after this adapter returns. + *job_table_ptr = ptr::null_mut(); + *job_id_mask_ptr = 0; + } + + let mut nb_jobs = requested_jobs; + let new_job_table = unsafe { ZSTDMT_rust_job_table_create(&mut nb_jobs, job_size, custom_mem) }; + if new_job_table.is_null() { + return ERROR(ZstdErrorCode::MemoryAllocation); + } + if unsafe { init_sync(new_job_table, nb_jobs, job_size) } != 0 { + unsafe { + free_job_table_with_sync(new_job_table, nb_jobs, job_size, custom_mem, destroy_sync) + }; + return ERROR(ZstdErrorCode::MemoryAllocation); + } + + debug_assert!(nb_jobs.is_power_of_two()); + unsafe { + *job_table_ptr = new_job_table; + *job_id_mask_ptr = nb_jobs.wrapping_sub(1); + } + 0 +} + #[no_mangle] pub unsafe extern "C" fn ZSTDMT_rust_job_table_create( nb_jobs_ptr: *mut c_uint, @@ -676,6 +751,33 @@ pub unsafe extern "C" fn ZSTDMT_rust_job_table_free( unsafe { free_job_table_storage(job_table, custom_mem) } } +/// Expand the C-owned MT job table when the requested worker count outgrows +/// its masked capacity. The table storage and custom allocator calls stay in +/// Rust; C supplies callbacks for the private descriptor synchronization +/// lifecycle. +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_expandJobsTable( + job_table_ptr: *mut *mut c_void, + job_id_mask_ptr: *mut c_uint, + nb_workers: c_uint, + job_size: usize, + custom_mem: ZstdCustomMem, + init_sync: JobTableInitSync, + destroy_sync: JobTableDestroySync, +) -> usize { + unsafe { + expand_job_table( + job_table_ptr, + job_id_mask_ptr, + nb_workers, + job_size, + custom_mem, + init_sync, + destroy_sync, + ) + } +} + unsafe fn create_buffer_pool( max_nb_buffers: usize, custom_mem: ZstdCustomMem, @@ -1110,6 +1212,7 @@ pub unsafe extern "C" fn ZSTDMT_rust_cctx_pool_release(pool: *mut RustCCtxPool, #[cfg(test)] mod tests { use super::*; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; const DEFAULT_MEM: ZstdCustomMem = ZstdCustomMem { customAlloc: None, @@ -1117,6 +1220,32 @@ mod tests { opaque: ptr::null_mut(), }; + static JOB_TABLE_INIT_CALLS: AtomicUsize = AtomicUsize::new(0); + static JOB_TABLE_DESTROY_CALLS: AtomicUsize = AtomicUsize::new(0); + static JOB_TABLE_FAIL_INIT: AtomicBool = AtomicBool::new(false); + static JOB_TABLE_TEST_LOCK: Mutex<()> = Mutex::new(()); + + unsafe extern "C" fn probe_job_table_init( + _job_table: *mut c_void, + _nb_jobs: c_uint, + _job_size: usize, + ) -> c_int { + JOB_TABLE_INIT_CALLS.fetch_add(1, Ordering::Relaxed); + if JOB_TABLE_FAIL_INIT.load(Ordering::Relaxed) { + -1 + } else { + 0 + } + } + + unsafe extern "C" fn probe_job_table_destroy( + _job_table: *mut c_void, + _nb_jobs: c_uint, + _job_size: usize, + ) { + JOB_TABLE_DESTROY_CALLS.fetch_add(1, Ordering::Relaxed); + } + fn rsync_prime_power() -> u64 { (0..RSYNC_LENGTH - 1).fold(1, |power, _| power.wrapping_mul(PRIME8_BYTES)) } @@ -1703,6 +1832,99 @@ mod tests { assert_eq!(rounded_job_count(256), Some(512)); } + #[test] + fn job_table_expansion_preserves_mask_and_sync_lifecycle() { + let _lock = JOB_TABLE_TEST_LOCK.lock().unwrap(); + JOB_TABLE_INIT_CALLS.store(0, Ordering::Relaxed); + JOB_TABLE_DESTROY_CALLS.store(0, Ordering::Relaxed); + JOB_TABLE_FAIL_INIT.store(false, Ordering::Relaxed); + + let mut requested_jobs = 1; + let initial_table = + unsafe { ZSTDMT_rust_job_table_create(&mut requested_jobs, 1, DEFAULT_MEM) }; + assert!(!initial_table.is_null()); + assert_eq!(requested_jobs, 2); + + let mut job_table = initial_table; + let mut job_id_mask = 1; + assert_eq!( + unsafe { + ZSTDMT_rust_expandJobsTable( + &mut job_table, + &mut job_id_mask, + 2, + 1, + DEFAULT_MEM, + probe_job_table_init, + probe_job_table_destroy, + ) + }, + 0 + ); + assert!(!job_table.is_null()); + assert_eq!(job_id_mask, 7); + assert_eq!(JOB_TABLE_INIT_CALLS.load(Ordering::Relaxed), 1); + assert_eq!(JOB_TABLE_DESTROY_CALLS.load(Ordering::Relaxed), 1); + + let expanded_table = job_table; + assert_eq!( + unsafe { + ZSTDMT_rust_expandJobsTable( + &mut job_table, + &mut job_id_mask, + 2, + 1, + DEFAULT_MEM, + probe_job_table_init, + probe_job_table_destroy, + ) + }, + 0 + ); + assert_eq!(job_table, expanded_table); + assert_eq!(job_id_mask, 7); + assert_eq!(JOB_TABLE_INIT_CALLS.load(Ordering::Relaxed), 1); + assert_eq!(JOB_TABLE_DESTROY_CALLS.load(Ordering::Relaxed), 1); + + unsafe { ZSTDMT_rust_job_table_free(job_table, 8, 1, DEFAULT_MEM) }; + } + + #[test] + fn job_table_expansion_cleans_up_after_sync_init_failure() { + let _lock = JOB_TABLE_TEST_LOCK.lock().unwrap(); + JOB_TABLE_INIT_CALLS.store(0, Ordering::Relaxed); + JOB_TABLE_DESTROY_CALLS.store(0, Ordering::Relaxed); + JOB_TABLE_FAIL_INIT.store(true, Ordering::Relaxed); + + let mut requested_jobs = 1; + let initial_table = + unsafe { ZSTDMT_rust_job_table_create(&mut requested_jobs, 1, DEFAULT_MEM) }; + assert!(!initial_table.is_null()); + + let mut job_table = initial_table; + let mut job_id_mask = 1; + assert_eq!( + unsafe { + ZSTDMT_rust_expandJobsTable( + &mut job_table, + &mut job_id_mask, + 2, + 1, + DEFAULT_MEM, + probe_job_table_init, + probe_job_table_destroy, + ) + }, + ERROR(ZstdErrorCode::MemoryAllocation) + ); + assert!(job_table.is_null()); + assert_eq!(job_id_mask, 0); + assert_eq!(JOB_TABLE_INIT_CALLS.load(Ordering::Relaxed), 1); + assert_eq!(JOB_TABLE_DESTROY_CALLS.load(Ordering::Relaxed), 2); + + JOB_TABLE_FAIL_INIT.store(false, Ordering::Relaxed); + } + #[test] fn target_job_log_preserves_ldm_and_non_ldm_policy() { assert_eq!(