diff --git a/lib/compress/zstdmt_compress.c b/lib/compress/zstdmt_compress.c index 3a8833f32..6d95d8f52 100644 --- a/lib/compress/zstdmt_compress.c +++ b/lib/compress/zstdmt_compress.c @@ -695,6 +695,25 @@ size_t ZSTDMT_rust_initCStream( ZSTDMT_initResetStreamFn resetStream, ZSTDMT_initDictionaryFn updateDictionary, ZSTDMT_initSerialResetFn serialReset); +typedef struct { + void* callbackContext; + unsigned nbWorkers; +} ZSTDMT_RustResizeProjection; +typedef char ZSTDMT_rust_resize_projection_layout[ + (offsetof(ZSTDMT_RustResizeProjection, callbackContext) == 0 + && offsetof(ZSTDMT_RustResizeProjection, nbWorkers) == sizeof(void*) + && sizeof(ZSTDMT_RustResizeProjection) + == (sizeof(void*) == 8 ? 16 : 8)) ? 1 : -1]; +typedef size_t (*ZSTDMT_resizeStepFn)(void* opaque, unsigned nbWorkers); +typedef void* (*ZSTDMT_resizeResourceFn)(void* opaque, unsigned nbWorkers); +size_t ZSTDMT_rust_resize( + const ZSTDMT_RustResizeProjection* projection, + ZSTDMT_resizeStepFn resizeFactory, + ZSTDMT_resizeStepFn expandJobs, + ZSTDMT_resizeResourceFn expandBufferPool, + ZSTDMT_resizeResourceFn expandCCtxPool, + ZSTDMT_resizeResourceFn expandSeqPool, + ZSTDMT_resizeStepFn setNbWorkers); typedef struct { size_t consumed; size_t cSize; @@ -2156,20 +2175,59 @@ size_t ZSTDMT_sizeof_CCtx(ZSTDMT_CCtx* mtctx) } +static size_t ZSTDMT_rust_resizeFactory(void* opaque, unsigned nbWorkers) +{ + ZSTDMT_CCtx* const mtctx = (ZSTDMT_CCtx*)opaque; + return POOL_resize(mtctx->factory, nbWorkers) ? ERROR(memory_allocation) : 0; +} + +static size_t ZSTDMT_rust_resizeJobs(void* opaque, unsigned nbWorkers) +{ + return ZSTDMT_expandJobsTable((ZSTDMT_CCtx*)opaque, nbWorkers); +} + +static void* ZSTDMT_rust_resizeBufferPool(void* opaque, unsigned nbWorkers) +{ + ZSTDMT_CCtx* const mtctx = (ZSTDMT_CCtx*)opaque; + mtctx->bufPool = ZSTDMT_expandBufferPool( + mtctx->bufPool, BUF_POOL_MAX_NB_BUFFERS(nbWorkers)); + return mtctx->bufPool; +} + +static void* ZSTDMT_rust_resizeCCtxPool(void* opaque, unsigned nbWorkers) +{ + ZSTDMT_CCtx* const mtctx = (ZSTDMT_CCtx*)opaque; + mtctx->cctxPool = ZSTDMT_expandCCtxPool(mtctx->cctxPool, nbWorkers); + return mtctx->cctxPool; +} + +static void* ZSTDMT_rust_resizeSeqPool(void* opaque, unsigned nbWorkers) +{ + ZSTDMT_CCtx* const mtctx = (ZSTDMT_CCtx*)opaque; + mtctx->seqPool = ZSTDMT_expandSeqPool(mtctx->seqPool, nbWorkers); + return mtctx->seqPool; +} + +static size_t ZSTDMT_rust_resizeSetNbWorkers(void* opaque, unsigned nbWorkers) +{ + ZSTDMT_CCtx* const mtctx = (ZSTDMT_CCtx*)opaque; + (void)ZSTDMT_CCtxParam_setNbWorkers(&mtctx->params, nbWorkers); + return 0; +} + /* ZSTDMT_resize() : * @return : error code if fails, 0 on success */ static size_t ZSTDMT_resize(ZSTDMT_CCtx* mtctx, unsigned nbWorkers) { - if (POOL_resize(mtctx->factory, nbWorkers)) return ERROR(memory_allocation); - FORWARD_IF_ERROR( ZSTDMT_expandJobsTable(mtctx, nbWorkers) , ""); - mtctx->bufPool = ZSTDMT_expandBufferPool(mtctx->bufPool, BUF_POOL_MAX_NB_BUFFERS(nbWorkers)); - if (mtctx->bufPool == NULL) return ERROR(memory_allocation); - mtctx->cctxPool = ZSTDMT_expandCCtxPool(mtctx->cctxPool, nbWorkers); - if (mtctx->cctxPool == NULL) return ERROR(memory_allocation); - mtctx->seqPool = ZSTDMT_expandSeqPool(mtctx->seqPool, nbWorkers); - if (mtctx->seqPool == NULL) return ERROR(memory_allocation); - ZSTDMT_CCtxParam_setNbWorkers(&mtctx->params, nbWorkers); - return 0; + ZSTDMT_RustResizeProjection const projection = { mtctx, nbWorkers }; + return ZSTDMT_rust_resize( + &projection, + ZSTDMT_rust_resizeFactory, + ZSTDMT_rust_resizeJobs, + ZSTDMT_rust_resizeBufferPool, + ZSTDMT_rust_resizeCCtxPool, + ZSTDMT_rust_resizeSeqPool, + ZSTDMT_rust_resizeSetNbWorkers); } diff --git a/rust/src/zstdmt_compress.rs b/rust/src/zstdmt_compress.rs index 173ed5591..302abdeea 100644 --- a/rust/src/zstdmt_compress.rs +++ b/rust/src/zstdmt_compress.rs @@ -611,6 +611,86 @@ pub unsafe extern "C" fn ZSTDMT_rust_createCCtx( context } +pub type ZSTDMT_resizeStepFn = unsafe extern "C" fn(*mut c_void, c_uint) -> usize; +pub type ZSTDMT_resizeResourceFn = unsafe extern "C" fn(*mut c_void, c_uint) -> *mut c_void; + +/// Scalar inputs for MT pool resizing. Rust owns the resize order and +/// short-circuit policy; C retains every private pool and context mutation +/// behind these callbacks. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct ZSTDMT_RustResizeProjection { + callback_context: *mut c_void, + nb_workers: c_uint, +} + +const _: () = { + assert!(offset_of!(ZSTDMT_RustResizeProjection, callback_context) == 0); + assert!(offset_of!(ZSTDMT_RustResizeProjection, nb_workers) == size_of::()); + assert!( + size_of::() == if size_of::() == 8 { 16 } else { 8 } + ); +}; + +/// Resize MT resources in the original order: worker factory, jobs table, +/// buffer pool, CCtx pool, sequence pool, then the parameter publication. +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_resize( + projection: *const ZSTDMT_RustResizeProjection, + resize_factory: Option, + expand_jobs: Option, + expand_buffer_pool: Option, + expand_cctx_pool: Option, + expand_seq_pool: Option, + set_nb_workers: Option, +) -> usize { + let Some(projection) = (unsafe { projection.as_ref() }).copied() else { + return ERROR(ZstdErrorCode::Generic); + }; + let ( + Some(resize_factory), + Some(expand_jobs), + Some(expand_buffer_pool), + Some(expand_cctx_pool), + Some(expand_seq_pool), + Some(set_nb_workers), + ) = ( + resize_factory, + expand_jobs, + expand_buffer_pool, + expand_cctx_pool, + expand_seq_pool, + set_nb_workers, + ) + else { + return ERROR(ZstdErrorCode::Generic); + }; + if projection.callback_context.is_null() { + return ERROR(ZstdErrorCode::Generic); + } + + let context = projection.callback_context; + let workers = projection.nb_workers; + let error = unsafe { resize_factory(context, workers) }; + if ERR_isError(error) { + return error; + } + let error = unsafe { expand_jobs(context, workers) }; + if ERR_isError(error) { + return error; + } + if unsafe { expand_buffer_pool(context, workers) }.is_null() { + return ERROR(ZstdErrorCode::MemoryAllocation); + } + if unsafe { expand_cctx_pool(context, workers) }.is_null() { + return ERROR(ZstdErrorCode::MemoryAllocation); + } + if unsafe { expand_seq_pool(context, workers) }.is_null() { + return ERROR(ZstdErrorCode::MemoryAllocation); + } + unsafe { set_nb_workers(context, workers) } +} + type ZSTDMT_waitForLdmLockFn = unsafe extern "C" fn(*mut c_void); type ZSTDMT_waitForLdmOverlapFn = unsafe extern "C" fn(*mut c_void, *mut c_void, usize) -> c_int; type ZSTDMT_waitForLdmWaitFn = unsafe extern "C" fn(*mut c_void); @@ -4178,6 +4258,121 @@ mod tests { } } + #[derive(Default)] + struct ResizeTestContext { + events: Vec<&'static str>, + resource_calls: usize, + fail: Option<&'static str>, + } + + fn record_resize_event(context: *mut c_void, event: &'static str) -> bool { + let state = unsafe { &mut *context.cast::() }; + state.events.push(event); + state.fail == Some(event) + } + + unsafe extern "C" fn resize_test_step(context: *mut c_void, _workers: c_uint) -> usize { + let state = unsafe { &*context.cast::() }; + let event = match state.events.len() { + 0 => "factory", + _ => "jobs", + }; + if record_resize_event(context, event) { + ERROR(ZstdErrorCode::MemoryAllocation) + } else { + 0 + } + } + + unsafe extern "C" fn resize_test_resource( + context: *mut c_void, + _workers: c_uint, + ) -> *mut c_void { + let state = unsafe { &mut *context.cast::() }; + let event = match state.resource_calls { + 0 => "buffer-pool", + 1 => "cctx-pool", + _ => "seq-pool", + }; + state.resource_calls += 1; + if record_resize_event(context, event) { + ptr::null_mut() + } else { + ptr::dangling_mut::() + } + } + + unsafe extern "C" fn resize_test_set_workers(context: *mut c_void, _workers: c_uint) -> usize { + assert!(!record_resize_event(context, "set-workers")); + 0 + } + + fn resize_test_projection( + context: &mut ResizeTestContext, + workers: c_uint, + ) -> ZSTDMT_RustResizeProjection { + ZSTDMT_RustResizeProjection { + callback_context: context as *mut ResizeTestContext as *mut c_void, + nb_workers: workers, + } + } + + #[test] + fn resize_preserves_resource_order_and_publishes_workers_last() { + let mut context = ResizeTestContext::default(); + let projection = resize_test_projection(&mut context, 6); + let result = unsafe { + ZSTDMT_rust_resize( + &projection, + Some(resize_test_step), + Some(resize_test_step), + Some(resize_test_resource), + Some(resize_test_resource), + Some(resize_test_resource), + Some(resize_test_set_workers), + ) + }; + + assert_eq!(result, 0); + assert_eq!( + context.events, + [ + "factory", + "jobs", + "buffer-pool", + "cctx-pool", + "seq-pool", + "set-workers", + ] + ); + } + + #[test] + fn resize_stops_after_resource_failure_without_publishing_workers() { + let mut context = ResizeTestContext { + fail: Some("cctx-pool"), + ..ResizeTestContext::default() + }; + let projection = resize_test_projection(&mut context, 3); + let result = unsafe { + ZSTDMT_rust_resize( + &projection, + Some(resize_test_step), + Some(resize_test_step), + Some(resize_test_resource), + Some(resize_test_resource), + Some(resize_test_resource), + Some(resize_test_set_workers), + ) + }; + + assert_eq!(result, ERROR(ZstdErrorCode::MemoryAllocation)); + assert_eq!( + context.events, + ["factory", "jobs", "buffer-pool", "cctx-pool"] + ); + } + fn record_free_cctx_event(context: *mut c_void, event: &'static str) { unsafe { (*context.cast::()).events.push(event);