From 55f28cdd71ff4f98f5c5384ba557bf4e93377621 Mon Sep 17 00:00:00 2001 From: ddidderr Date: Sun, 19 Jul 2026 13:55:21 +0200 Subject: [PATCH] feat(compress): move thread-pool attachment policy into Rust Move ZSTD_CCtx_refThreadPool init-stage validation and pool-slot assignment behind a Rust-owned ABI bridge while keeping the opaque pool and CCtx storage in C. Preserve the stage error and avoid mutating the pool on rejection. Test Plan: - cargo test --manifest-path rust/Cargo.toml --lib - cargo clippy --manifest-path rust/Cargo.toml --all-targets -- -D warnings - make -B -C programs -j1 zstd - make -C tests -j1 test-zstream ZSTREAM_TESTTIME=-T1s - focused ref_thread_pool unit tests --- lib/compress/zstd_compress.c | 23 +++++++++-- rust/src/zstd_compress.rs | 80 ++++++++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+), 4 deletions(-) diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c index d8b5e5eaf..9dee14ea8 100644 --- a/lib/compress/zstd_compress.c +++ b/lib/compress/zstd_compress.c @@ -287,6 +287,20 @@ typedef char ZSTD_rust_set_parameter_state_layout[ == 4 * sizeof(void*) && sizeof(ZSTD_rust_setParameterState) == 5 * sizeof(void*)) ? 1 : -1]; +/* Thread-pool attachment policy lives in Rust. The pool slot remains + * C-owned; this projection passes its storage location and opaque value. */ +typedef struct { + void** pool; + void* requestedPool; + int streamStage; +} ZSTD_rust_refThreadPoolState; +size_t ZSTD_rust_refThreadPool(const ZSTD_rust_refThreadPoolState* state); +typedef char ZSTD_rust_ref_thread_pool_state_layout[ + (offsetof(ZSTD_rust_refThreadPoolState, pool) == 0 + && offsetof(ZSTD_rust_refThreadPoolState, requestedPool) == sizeof(void*) + && offsetof(ZSTD_rust_refThreadPoolState, streamStage) == 2 * sizeof(void*) + && sizeof(ZSTD_rust_refThreadPoolState) == 3 * sizeof(void*)) + ? 1 : -1]; typedef void (*ZSTD_rust_resetCCtxClearAllDicts_f)(void* context); typedef size_t (*ZSTD_rust_resetCCtxResetParams_f)(void* context); typedef struct { @@ -2084,10 +2098,11 @@ size_t ZSTD_CCtx_refCDict(ZSTD_CCtx* cctx, const ZSTD_CDict* cdict) size_t ZSTD_CCtx_refThreadPool(ZSTD_CCtx* cctx, ZSTD_threadPool* pool) { - RETURN_ERROR_IF(cctx->streamStage != zcss_init, stage_wrong, - "Can't ref a pool when ctx not in init stage."); - cctx->pool = pool; - return 0; + ZSTD_rust_refThreadPoolState state; + state.pool = (void**)&cctx->pool; + state.requestedPool = pool; + state.streamStage = (int)cctx->streamStage; + return ZSTD_rust_refThreadPool(&state); } size_t ZSTD_CCtx_refPrefix(ZSTD_CCtx* cctx, const void* prefix, size_t prefixSize) diff --git a/rust/src/zstd_compress.rs b/rust/src/zstd_compress.rs index a8dcf4f87..541acac10 100644 --- a/rust/src/zstd_compress.rs +++ b/rust/src/zstd_compress.rs @@ -1012,6 +1012,43 @@ pub unsafe extern "C" fn ZSTD_rust_setParameter( result } +/// Explicit projection for `ZSTD_CCtx_refThreadPool`. +/// +/// Rust owns the init-stage check and assignment policy while the pool slot +/// itself remains storage owned by the C context. The pool value is opaque to +/// Rust because its private `POOL_ctx_s` layout is not part of this bridge. +#[repr(C)] +pub struct ZSTD_rust_refThreadPoolState { + pool: *mut *mut c_void, + requested_pool: *mut c_void, + stream_stage: c_int, +} + +const _: () = { + assert!(offset_of!(ZSTD_rust_refThreadPoolState, pool) == 0); + assert!(offset_of!(ZSTD_rust_refThreadPoolState, requested_pool) == size_of::()); + assert!(offset_of!(ZSTD_rust_refThreadPoolState, stream_stage) == 2 * size_of::()); + assert!(size_of::() == 3 * size_of::()); +}; + +/// Attach an opaque thread pool through the C-owned context slot. +#[no_mangle] +pub unsafe extern "C" fn ZSTD_rust_refThreadPool( + state: *const ZSTD_rust_refThreadPoolState, +) -> usize { + if state.is_null() { + return ERROR(ZstdErrorCode::Generic); + } + let state = unsafe { &*state }; + if state.stream_stage != ZSTD_CSTREAM_STAGE_INIT { + return ERROR(ZstdErrorCode::StageWrong); + } + unsafe { + *state.pool = state.requested_pool; + } + 0 +} + type ResetCCtxClearAllDictsFn = unsafe extern "C" fn(*mut c_void); type ResetCCtxResetParamsFn = unsafe extern "C" fn(*mut c_void) -> usize; @@ -10478,6 +10515,49 @@ mod tests { assert_eq!(max_block_size_set, 0); } + fn ref_thread_pool_test_state( + pool: &mut *mut c_void, + requested_pool: *mut c_void, + stream_stage: c_int, + ) -> ZSTD_rust_refThreadPoolState { + ZSTD_rust_refThreadPoolState { + pool, + requested_pool, + stream_stage, + } + } + + #[test] + fn ref_thread_pool_assigns_the_requested_pool_at_init_stage() { + let mut pool = ptr::null_mut::(); + let requested_pool = ptr::dangling_mut::(); + let state = ref_thread_pool_test_state(&mut pool, requested_pool, ZSTD_CSTREAM_STAGE_INIT); + + let result = unsafe { ZSTD_rust_refThreadPool(&state) }; + + assert_eq!(result, 0); + assert_eq!(pool, requested_pool); + } + + #[test] + fn ref_thread_pool_rejects_non_init_stage_without_mutating_the_pool() { + let original_pool = ptr::dangling_mut::(); + let mut pool = original_pool; + let state = ref_thread_pool_test_state(&mut pool, ptr::null_mut(), ZSTD_CSTREAM_STAGE_LOAD); + + let result = unsafe { ZSTD_rust_refThreadPool(&state) }; + + assert_eq!(result, ERROR(ZstdErrorCode::StageWrong)); + assert_eq!(pool, original_pool); + } + + #[test] + fn ref_thread_pool_rejects_null_state() { + let result = unsafe { ZSTD_rust_refThreadPool(ptr::null()) }; + + assert_eq!(result, ERROR(ZstdErrorCode::Generic)); + } + #[derive(Default)] struct SetCParamsTestContext { events: Vec<&'static str>,