diff --git a/lib/compress/zstdmt_compress.c b/lib/compress/zstdmt_compress.c index 157b43b7e..f3c69b24f 100644 --- a/lib/compress/zstdmt_compress.c +++ b/lib/compress/zstdmt_compress.c @@ -666,6 +666,9 @@ typedef int (*ZSTDMT_jobTableInitFn)(void* jobTable, unsigned nbJobs, size_t jobSize); typedef void (*ZSTDMT_jobTableDestroyFn)(void* jobTable, unsigned nbJobs, size_t jobSize); +void* ZSTDMT_rust_createJobsTable( + unsigned* nbJobsPtr, size_t jobSize, ZSTD_customMem cMem, + ZSTDMT_jobTableInitFn initSync, ZSTDMT_jobTableDestroyFn destroySync); size_t ZSTDMT_rust_expandJobsTable( void** jobTablePtr, unsigned* jobIDMaskPtr, unsigned nbWorkers, size_t jobSize, ZSTD_customMem cMem, @@ -1753,14 +1756,9 @@ void ZSTDMT_job_table_destroy_sync(void* jobTable, unsigned nbJobs, size_t jobSi * update *nbJobsPtr to next power of 2 value, as size of table */ static ZSTDMT_jobDescription* ZSTDMT_createJobsTable(U32* nbJobsPtr, ZSTD_customMem cMem) { - ZSTDMT_jobDescription* const jobTable = (ZSTDMT_jobDescription*) - ZSTDMT_rust_job_table_create(nbJobsPtr, sizeof(ZSTDMT_jobDescription), cMem); - if (jobTable == NULL) return NULL; - if (ZSTDMT_job_table_init_sync(jobTable, *nbJobsPtr, sizeof(*jobTable)) != 0) { - ZSTDMT_freeJobsTable(jobTable, *nbJobsPtr, cMem); - return NULL; - } - return jobTable; + return (ZSTDMT_jobDescription*) ZSTDMT_rust_createJobsTable( + nbJobsPtr, sizeof(ZSTDMT_jobDescription), cMem, + ZSTDMT_job_table_init_sync, ZSTDMT_job_table_destroy_sync); } static size_t ZSTDMT_expandJobsTable (ZSTDMT_CCtx* mtctx, U32 nbWorkers) { diff --git a/rust/src/zstdmt_compress.rs b/rust/src/zstdmt_compress.rs index b4c920331..26cce4a0c 100644 --- a/rust/src/zstdmt_compress.rs +++ b/rust/src/zstdmt_compress.rs @@ -3197,6 +3197,27 @@ unsafe fn free_job_table_with_sync( unsafe { ZSTDMT_rust_job_table_free(job_table, nb_jobs, job_size, custom_mem) }; } +unsafe fn create_job_table_with_sync( + nb_jobs_ptr: *mut c_uint, + job_size: usize, + custom_mem: ZstdCustomMem, + init_sync: JobTableInitSync, + destroy_sync: JobTableDestroySync, +) -> *mut c_void { + let job_table = unsafe { create_job_table(nb_jobs_ptr, job_size, custom_mem) }; + if job_table.is_null() { + return ptr::null_mut(); + } + + let nb_jobs = unsafe { *nb_jobs_ptr }; + if unsafe { init_sync(job_table, nb_jobs, job_size) } != 0 { + unsafe { free_job_table_with_sync(job_table, nb_jobs, job_size, custom_mem, destroy_sync) }; + return ptr::null_mut(); + } + + job_table +} + unsafe fn expand_job_table( job_table_ptr: *mut *mut c_void, job_id_mask_ptr: *mut c_uint, @@ -3234,16 +3255,12 @@ unsafe fn expand_job_table( } let mut nb_jobs = requested_jobs; - let new_job_table = unsafe { ZSTDMT_rust_job_table_create(&mut nb_jobs, job_size, custom_mem) }; + let new_job_table = unsafe { + create_job_table_with_sync(&mut nb_jobs, job_size, custom_mem, init_sync, destroy_sync) + }; 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 { @@ -3262,6 +3279,19 @@ pub unsafe extern "C" fn ZSTDMT_rust_job_table_create( unsafe { create_job_table(nb_jobs_ptr, job_size, custom_mem) } } +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_createJobsTable( + nb_jobs_ptr: *mut c_uint, + job_size: usize, + custom_mem: ZstdCustomMem, + init_sync: JobTableInitSync, + destroy_sync: JobTableDestroySync, +) -> *mut c_void { + unsafe { + create_job_table_with_sync(nb_jobs_ptr, job_size, custom_mem, init_sync, destroy_sync) + } +} + #[no_mangle] pub unsafe extern "C" fn ZSTDMT_rust_job_table_free( job_table: *mut c_void, @@ -6770,6 +6800,32 @@ mod tests { assert_eq!(rounded_job_count(256), Some(512)); } + #[test] + fn job_table_creation_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 job_table = unsafe { + ZSTDMT_rust_createJobsTable( + &mut requested_jobs, + 1, + DEFAULT_MEM, + probe_job_table_init, + probe_job_table_destroy, + ) + }; + + assert!(job_table.is_null()); + assert_eq!(requested_jobs, 2); + assert_eq!(JOB_TABLE_INIT_CALLS.load(Ordering::Relaxed), 1); + assert_eq!(JOB_TABLE_DESTROY_CALLS.load(Ordering::Relaxed), 1); + + JOB_TABLE_FAIL_INIT.store(false, Ordering::Relaxed); + } + #[test] fn job_table_expansion_preserves_mask_and_sync_lifecycle() { let _lock = JOB_TABLE_TEST_LOCK.lock().unwrap();