diff --git a/lib/compress/zstdmt_compress.c b/lib/compress/zstdmt_compress.c index 7021698d4..efbff5159 100644 --- a/lib/compress/zstdmt_compress.c +++ b/lib/compress/zstdmt_compress.c @@ -357,7 +357,8 @@ typedef void (*ZSTDMT_flushUpdateJobFn)( void* opaque, unsigned jobID, size_t dstFlushed); typedef void (*ZSTDMT_flushCompleteJobFn)( void* opaque, unsigned jobID, size_t srcSize, size_t cSize); -typedef void (*ZSTDMT_flushErrorFn)(void* opaque); +typedef void (*ZSTDMT_flushWaitForAllJobsFn)(void* opaque); +typedef void (*ZSTDMT_flushReleaseAllJobResourcesFn)(void* opaque); ZSTDMT_RustFlushProducedResult ZSTDMT_rust_flushProduced( const ZSTDMT_RustFlushContextProjection* context, void* outputDst, size_t outputSize, size_t outputPos, @@ -366,7 +367,8 @@ ZSTDMT_RustFlushProducedResult ZSTDMT_rust_flushProduced( ZSTDMT_flushChecksumFn addFrameChecksum, ZSTDMT_flushUpdateJobFn updateJob, ZSTDMT_flushCompleteJobFn completeJob, - ZSTDMT_flushErrorFn onError); + ZSTDMT_flushWaitForAllJobsFn waitForAllJobs, + ZSTDMT_flushReleaseAllJobResourcesFn releaseAllJobResources); size_t ZSTDMT_rust_writeFrameChecksum( const ZSTDMT_RustFrameChecksumProjection* projection); @@ -2937,10 +2939,15 @@ static void ZSTDMT_completeFlushJob(void* opaque, unsigned jobID, mtctx->doneJobID++; } -static void ZSTDMT_flushError(void* opaque) +static void ZSTDMT_flushWaitForAllJobs(void* opaque) { ZSTDMT_CCtx* const mtctx = (ZSTDMT_CCtx*)opaque; ZSTDMT_waitForAllJobsCompleted(mtctx); +} + +static void ZSTDMT_flushReleaseAllJobResources(void* opaque) +{ + ZSTDMT_CCtx* const mtctx = (ZSTDMT_CCtx*)opaque; ZSTDMT_releaseAllJobResources(mtctx); } @@ -2968,7 +2975,8 @@ static size_t ZSTDMT_flushProduced(ZSTDMT_CCtx* mtctx, ZSTD_outBuffer* output, u ZSTDMT_addFrameChecksum, ZSTDMT_updateFlushJob, ZSTDMT_completeFlushJob, - ZSTDMT_flushError); + ZSTDMT_flushWaitForAllJobs, + ZSTDMT_flushReleaseAllJobResources); assert(output->size >= output->pos); output->pos = result.outputPos; diff --git a/rust/src/zstdmt_compress.rs b/rust/src/zstdmt_compress.rs index 6999a0ea2..6354ceeef 100644 --- a/rust/src/zstdmt_compress.rs +++ b/rust/src/zstdmt_compress.rs @@ -1017,7 +1017,8 @@ pub type ZSTDMT_flushUpdateJobFn = unsafe extern "C" fn(opaque: *mut c_void, job_id: c_uint, dst_flushed: usize); pub type ZSTDMT_flushCompleteJobFn = unsafe extern "C" fn(opaque: *mut c_void, job_id: c_uint, src_size: usize, c_size: usize); -pub type ZSTDMT_flushErrorFn = unsafe extern "C" fn(opaque: *mut c_void); +pub type ZSTDMT_flushWaitForAllJobsFn = unsafe extern "C" fn(opaque: *mut c_void); +pub type ZSTDMT_flushReleaseAllJobResourcesFn = unsafe extern "C" fn(opaque: *mut c_void); /// Scalar MT state used by the outer streaming scheduler. The C context, /// input-range ownership, and synchronization objects remain private to C; @@ -2337,7 +2338,7 @@ fn flush_produced_result(result: usize, output_pos: usize) -> ZSTDMT_flushProduc } #[inline] -unsafe fn flush_produced_with( +unsafe fn flush_produced_with( context: ZSTDMT_flushContextProjection, output_dst: *mut c_void, output_size: usize, @@ -2348,14 +2349,16 @@ unsafe fn flush_produced_with( mut add_frame_checksum: C, mut update_job: U, mut complete_job: F, - mut on_error: E, + mut wait_for_all_jobs: W, + mut release_all_job_resources: R, ) -> ZSTDMT_flushProducedResult where P: FnMut(c_uint, c_uint) -> ZSTDMT_flushJobProjection, C: FnMut(c_uint, &mut ZSTDMT_flushJobProjection), U: FnMut(c_uint, usize), F: FnMut(c_uint, usize, usize), - E: FnMut(), + W: FnMut(), + R: FnMut(), { let mut output_pos = output_pos; let mut done_job_id = context.doneJobID; @@ -2367,7 +2370,8 @@ where let src_size = projection.srcSize; if ERR_isError(projection.cSize) { - on_error(); + wait_for_all_jobs(); + release_all_job_resources(); return flush_produced_result(projection.cSize, output_pos); } @@ -2379,7 +2383,8 @@ where } if ERR_isError(projection.cSize) { - on_error(); + wait_for_all_jobs(); + release_all_job_resources(); return flush_produced_result(projection.cSize, output_pos); } @@ -2612,7 +2617,8 @@ pub unsafe extern "C" fn ZSTDMT_rust_flushProduced( addFrameChecksum: Option, updateJob: Option, completeJob: Option, - onError: Option, + waitForAllJobs: Option, + releaseAllJobResources: Option, ) -> ZSTDMT_flushProducedResult { let Some(context) = (unsafe { context.as_ref() }).copied() else { return flush_produced_result(ERROR(ZstdErrorCode::Generic), output_pos); @@ -2622,13 +2628,15 @@ pub unsafe extern "C" fn ZSTDMT_rust_flushProduced( Some(add_frame_checksum), Some(update_job), Some(complete_job), - Some(on_error), + Some(wait_for_all_jobs), + Some(release_all_job_resources), ) = ( projectJob, addFrameChecksum, updateJob, completeJob, - onError, + waitForAllJobs, + releaseAllJobResources, ) else { return flush_produced_result(ERROR(ZstdErrorCode::Generic), output_pos); @@ -2657,7 +2665,10 @@ pub unsafe extern "C" fn ZSTDMT_rust_flushProduced( complete_job(opaque, job_id, src_size, c_size); }, || { - on_error(opaque); + wait_for_all_jobs(opaque); + }, + || { + release_all_job_resources(opaque); }, ) } @@ -6735,7 +6746,8 @@ mod tests { |_job_id, _projection| panic!("unexpected checksum callback"), |_job_id, dst_flushed| updated = Some(dst_flushed), |_job_id, src_size, c_size| completed = Some((src_size, c_size)), - || panic!("unexpected error callback"), + || panic!("unexpected wait callback"), + || panic!("unexpected release callback"), ) }; @@ -6784,7 +6796,8 @@ mod tests { |_job_id, _projection| panic!("unexpected checksum callback"), |_job_id, dst_flushed| updated = Some(dst_flushed), |_job_id, _src_size, _c_size| completed = true, - || panic!("unexpected error callback"), + || panic!("unexpected wait callback"), + || panic!("unexpected release callback"), ) }; @@ -6808,7 +6821,7 @@ mod tests { jobIDMask: 0, ..ZSTDMT_flushContextProjection::default() }; - let mut error_cleanup = false; + let mut error_cleanup = Vec::new(); let result = unsafe { flush_produced_with( context, @@ -6821,12 +6834,13 @@ mod tests { |_job_id, _projection| panic!("unexpected checksum callback"), |_job_id, _dst_flushed| panic!("unexpected update callback"), |_job_id, _src_size, _c_size| panic!("unexpected completion callback"), - || error_cleanup = true, + || error_cleanup.push("wait"), + || error_cleanup.push("release"), ) }; assert_eq!(result.result, ERROR(ZstdErrorCode::Generic)); assert_eq!(result.outputPos, 0); - assert!(error_cleanup); + assert_eq!(error_cleanup, ["wait", "release"]); let job = [0x80u8, 0x81, 0x82, 0x83, 0x84, 0x85, 0x86]; let mut output = [0xa5u8; 8]; @@ -6862,7 +6876,8 @@ mod tests { }, |_job_id, _dst_flushed| {}, |_job_id, _src_size, c_size| completed_size = Some(c_size), - || panic!("unexpected error callback"), + || panic!("unexpected wait callback"), + || panic!("unexpected release callback"), ) }; assert_eq!(checksum_calls, 1); @@ -6884,7 +6899,7 @@ mod tests { frameChecksumNeeded: 1, ..ZSTDMT_flushJobProjection::default() }; - let mut error_cleanup = false; + let mut error_cleanup = Vec::new(); let mut completed = false; let result = unsafe { @@ -6906,13 +6921,14 @@ mod tests { }, |_job_id, _dst_flushed| panic!("unexpected update callback"), |_job_id, _src_size, _c_size| completed = true, - || error_cleanup = true, + || error_cleanup.push("wait"), + || error_cleanup.push("release"), ) }; assert_eq!(result.result, ERROR(ZstdErrorCode::DstSizeTooSmall)); assert_eq!(result.outputPos, 0); - assert!(error_cleanup); + assert_eq!(error_cleanup, ["wait", "release"]); assert!(!completed); assert_eq!(output, [0xa5; 8]); } @@ -6990,7 +7006,8 @@ mod tests { |_job_id, _projection| panic!("no checksum expected"), |_job_id, _dst_flushed| panic!("no update expected"), |_job_id, _src_size, _c_size| panic!("no completion expected"), - || panic!("no error expected"), + || panic!("no wait expected"), + || panic!("no release expected"), ) }; assert_eq!(not_ended.result, 1); @@ -7012,7 +7029,8 @@ mod tests { |_job_id, _projection| panic!("no checksum expected"), |_job_id, _dst_flushed| panic!("no update expected"), |_job_id, _src_size, _c_size| panic!("no completion expected"), - || panic!("no error expected"), + || panic!("no wait expected"), + || panic!("no release expected"), ) }; assert_eq!(ended.result, 0);