diff --git a/lib/compress/zstdmt_compress.c b/lib/compress/zstdmt_compress.c index aa6b0fd1b..267bfe4c5 100644 --- a/lib/compress/zstdmt_compress.c +++ b/lib/compress/zstdmt_compress.c @@ -756,7 +756,8 @@ typedef struct { unsigned allJobsCompleted; } ZSTDMT_RustInitCStreamProjection; typedef size_t (*ZSTDMT_initResizeFn)(void* opaque, unsigned nbWorkers); -typedef void (*ZSTDMT_initDrainFn)(void* opaque); +typedef void (*ZSTDMT_initWaitForAllJobsFn)(void* opaque); +typedef void (*ZSTDMT_initReleaseAllJobResourcesFn)(void* opaque); typedef void (*ZSTDMT_initApplyParametersFn)(void* opaque, size_t jobSize); typedef size_t (*ZSTDMT_initDictionaryFn)(void* opaque); typedef void (*ZSTDMT_initSetSizeFn)(void* opaque, size_t size); @@ -938,7 +939,9 @@ void* ZSTDMT_rust_createCCtx( ZSTDMT_createCCtxFreeFn freeContext); size_t ZSTDMT_rust_initCStream( const ZSTDMT_RustInitCStreamProjection* projection, void* opaque, - ZSTDMT_initResizeFn resize, ZSTDMT_initDrainFn drain, + ZSTDMT_initResizeFn resize, + ZSTDMT_initWaitForAllJobsFn waitForAllJobs, + ZSTDMT_initReleaseAllJobResourcesFn releaseAllJobResources, ZSTDMT_initApplyParametersFn applyParameters, ZSTDMT_initDictionaryFn prepareDictionary, ZSTDMT_initSetSizeFn setTargetPrefixSize, @@ -2782,13 +2785,16 @@ static size_t ZSTDMT_initCStreamResize(void* opaque, unsigned nbWorkers) return ZSTDMT_resize(state->mtctx, nbWorkers); } -static void ZSTDMT_initCStreamDrain(void* opaque) +static void ZSTDMT_initCStreamWaitForAllJobs(void* opaque) { ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque; - ZSTDMT_CCtx* const mtctx = state->mtctx; - ZSTDMT_waitForAllJobsCompleted(mtctx); - ZSTDMT_releaseAllJobResources(mtctx); - mtctx->allJobsCompleted = 1; + ZSTDMT_waitForAllJobsCompleted(state->mtctx); +} + +static void ZSTDMT_initCStreamReleaseAllJobResources(void* opaque) +{ + ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque; + ZSTDMT_releaseAllJobResources(state->mtctx); } static void ZSTDMT_initCStreamApplyParameters(void* opaque, size_t jobSize) @@ -3031,7 +3037,8 @@ size_t ZSTDMT_initCStream_internal( return ZSTDMT_rust_initCStream( &projection, &state, ZSTDMT_initCStreamResize, - ZSTDMT_initCStreamDrain, + ZSTDMT_initCStreamWaitForAllJobs, + ZSTDMT_initCStreamReleaseAllJobResources, ZSTDMT_initCStreamApplyParameters, ZSTDMT_initCStreamPrepareDictionary, ZSTDMT_initCStreamSetTargetPrefixSize, diff --git a/rust/src/zstdmt_compress.rs b/rust/src/zstdmt_compress.rs index e584aff46..5a01a53e8 100644 --- a/rust/src/zstdmt_compress.rs +++ b/rust/src/zstdmt_compress.rs @@ -1166,7 +1166,8 @@ pub struct ZSTDMT_initCStreamProjection { } pub type ZSTDMT_initResizeFn = unsafe extern "C" fn(*mut c_void, c_uint) -> usize; -pub type ZSTDMT_initDrainFn = unsafe extern "C" fn(*mut c_void); +pub type ZSTDMT_initWaitForAllJobsFn = unsafe extern "C" fn(*mut c_void); +pub type ZSTDMT_initReleaseAllJobResourcesFn = unsafe extern "C" fn(*mut c_void); pub type ZSTDMT_initApplyParametersFn = unsafe extern "C" fn(*mut c_void, usize); pub type ZSTDMT_initDictionaryFn = unsafe extern "C" fn(*mut c_void) -> usize; pub type ZSTDMT_initSetSizeFn = unsafe extern "C" fn(*mut c_void, usize); @@ -2766,12 +2767,14 @@ fn init_rsync_parameters(target_section_size: usize) -> (u64, u64) { /// Run the high-level MT streaming initialization policy while C retains all /// private state and side effects behind callbacks. The callback order is /// intentionally the order of the original C initializer: resize, normalize, -/// drain, attach the first dictionary, size buffers, reset stream state, +/// wait for all jobs, release their resources, attach the first dictionary, +/// size buffers, reset stream state, /// update the active dictionary, and finally reset serial state. #[inline] fn init_c_stream_with< Resize, - Drain, + WaitForAllJobs, + ReleaseAllJobResources, ApplyParams, PrepareDictionary, SetPrefixSize, @@ -2785,7 +2788,8 @@ fn init_c_stream_with< >( projection: ZSTDMT_initCStreamProjection, mut resize: Resize, - mut drain: Drain, + mut wait_for_all_jobs: WaitForAllJobs, + mut release_all_job_resources: ReleaseAllJobResources, mut apply_parameters: ApplyParams, mut prepare_dictionary: PrepareDictionary, mut set_target_prefix_size: SetPrefixSize, @@ -2799,7 +2803,8 @@ fn init_c_stream_with< ) -> usize where Resize: FnMut(c_uint) -> usize, - Drain: FnMut(), + WaitForAllJobs: FnMut(), + ReleaseAllJobResources: FnMut(), ApplyParams: FnMut(usize), PrepareDictionary: FnMut() -> usize, SetPrefixSize: FnMut(usize), @@ -2825,7 +2830,8 @@ where ); if projection.allJobsCompleted == 0 { - drain(); + wait_for_all_jobs(); + release_all_job_resources(); } apply_parameters(job_size); @@ -3181,7 +3187,8 @@ pub unsafe extern "C" fn ZSTDMT_rust_initCStream( projection: *const ZSTDMT_initCStreamProjection, opaque: *mut c_void, resize: Option, - drain: Option, + waitForAllJobs: Option, + releaseAllJobResources: Option, applyParameters: Option, prepareDictionary: Option, setTargetPrefixSize: Option, @@ -3198,7 +3205,8 @@ pub unsafe extern "C" fn ZSTDMT_rust_initCStream( }; let ( Some(resize), - Some(drain), + Some(wait_for_all_jobs), + Some(release_all_job_resources), Some(apply_parameters), Some(prepare_dictionary), Some(set_target_prefix_size), @@ -3211,7 +3219,8 @@ pub unsafe extern "C" fn ZSTDMT_rust_initCStream( Some(serial_reset), ) = ( resize, - drain, + waitForAllJobs, + releaseAllJobResources, applyParameters, prepareDictionary, setTargetPrefixSize, @@ -3230,7 +3239,8 @@ pub unsafe extern "C" fn ZSTDMT_rust_initCStream( init_c_stream_with( projection, |nb_workers| unsafe { resize(opaque, nb_workers) }, - || unsafe { drain(opaque) }, + || unsafe { wait_for_all_jobs(opaque) }, + || unsafe { release_all_job_resources(opaque) }, |job_size| unsafe { apply_parameters(opaque, job_size) }, || unsafe { prepare_dictionary(opaque) }, |size| unsafe { set_target_prefix_size(opaque, size) }, @@ -7768,7 +7778,8 @@ mod tests { let events = Rc::new(RefCell::new(Vec::<&'static str>::new())); let values = Rc::new(RefCell::new(Vec::<(&'static str, usize)>::new())); let resize_events = Rc::clone(&events); - let drain_events = Rc::clone(&events); + let wait_events = Rc::clone(&events); + let release_events = Rc::clone(&events); let apply_events = Rc::clone(&events); let apply_values = Rc::clone(&values); let prepare_events = Rc::clone(&events); @@ -7793,7 +7804,8 @@ mod tests { resize_events.borrow_mut().push("resize"); 0 }, - move || drain_events.borrow_mut().push("drain"), + move || wait_events.borrow_mut().push("wait"), + move || release_events.borrow_mut().push("release"), move |job_size| { apply_events.borrow_mut().push("apply"); apply_values.borrow_mut().push(("job", job_size)); @@ -7837,7 +7849,8 @@ mod tests { events.borrow().as_slice(), &[ "resize", - "drain", + "wait", + "release", "apply", "prepare-dict", "prefix", @@ -7865,6 +7878,7 @@ mod tests { init_projection(), |_| expected_error, || panic!("a failed resize must stop initialization"), + || panic!("a failed resize must stop initialization"), |_| panic!("a failed resize must stop initialization"), || panic!("a failed resize must stop initialization"), |_| panic!("a failed resize must stop initialization"), @@ -7891,6 +7905,7 @@ mod tests { projection, |_| panic!("workers already match"), || panic!("all jobs are complete"), + || panic!("all jobs are complete"), { let events = Rc::clone(&events); move |_| events.borrow_mut().push("apply")