refactor(mt): move context construction policy to Rust

Move MT context validation, worker clamping, resource-construction order, and
partial-failure cleanup policy into Rust. Keep the private context layout,
allocator-owned resources, pools, jobs table, and synchronization primitives in
C behind callbacks.

Test Plan:
- ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo +nightly fmt --manifest-path rust/Cargo.toml -- --check
- ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo clippy --manifest-path rust/Cargo.toml --all-targets -- -D warnings
- ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo test --manifest-path rust/Cargo.toml --all-targets (782 passed)
- ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo clippy --manifest-path rust/cli/Cargo.toml --all-targets -- -D warnings
- ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo test --manifest-path rust/cli/Cargo.toml --all-targets (179 passed)
- ulimit -v 41943040; make -j1
- ulimit -v 41943040; make -j1 -C tests test (all tests completed successfully)
This commit is contained in:
2026-07-20 06:44:03 +02:00
parent a00ba4c2b7
commit 2b2f9df5de
2 changed files with 520 additions and 39 deletions
+365
View File
@@ -468,6 +468,149 @@ const _: () = {
assert!(size_of::<ZSTDMT_RustFreeCCtxState>() == size_of::<[usize; 13]>());
};
pub type ZSTDMT_createCCtxAllocateFn = unsafe extern "C" fn(*mut c_void, usize) -> *mut c_void;
pub type ZSTDMT_createCCtxSetWorkersFn = unsafe extern "C" fn(*mut c_void, c_uint) -> usize;
pub type ZSTDMT_createCCtxSetInitialStateFn = unsafe extern "C" fn(*mut c_void);
pub type ZSTDMT_createCCtxFactoryFn = unsafe extern "C" fn(*mut c_void, c_uint) -> *mut c_void;
pub type ZSTDMT_createCCtxJobsFn = unsafe extern "C" fn(*mut c_void, *mut c_uint) -> *mut c_void;
pub type ZSTDMT_createCCtxResourceFn = unsafe extern "C" fn(*mut c_void, c_uint) -> *mut c_void;
pub type ZSTDMT_createCCtxSerialInitFn = unsafe extern "C" fn(*mut c_void) -> c_int;
pub type ZSTDMT_createCCtxFreeFn = unsafe extern "C" fn(*mut c_void);
/// Scalar inputs and callbacks for MT context construction. Rust owns the
/// validation, worker clamp, and resource-construction order; C retains the
/// private context layout, allocator calls, pools, and synchronization.
#[repr(C)]
#[derive(Clone, Copy)]
pub struct ZSTDMT_RustCreateCCtxProjection {
callback_context: *mut c_void,
requested_nb_workers: c_uint,
max_nb_workers: c_uint,
context_size: usize,
custom_mem: ZstdCustomMem,
}
const _: () = {
assert!(offset_of!(ZSTDMT_RustCreateCCtxProjection, callback_context) == 0);
assert!(
offset_of!(ZSTDMT_RustCreateCCtxProjection, requested_nb_workers) == size_of::<usize>()
);
assert!(
offset_of!(ZSTDMT_RustCreateCCtxProjection, max_nb_workers)
== size_of::<usize>() + size_of::<c_uint>()
);
assert!(
offset_of!(ZSTDMT_RustCreateCCtxProjection, context_size)
== if size_of::<usize>() == 8 { 16 } else { 12 }
);
assert!(
offset_of!(ZSTDMT_RustCreateCCtxProjection, custom_mem)
== if size_of::<usize>() == 8 { 24 } else { 16 }
);
assert!(
size_of::<ZSTDMT_RustCreateCCtxProjection>()
== if size_of::<usize>() == 8 { 48 } else { 28 }
);
};
fn custom_mem_is_valid(custom_mem: ZstdCustomMem) -> bool {
custom_mem.customAlloc.is_some() == custom_mem.customFree.is_some()
}
/// Construct an MT context through the original C resource callbacks.
/// Callbacks are all invoked through the allocation and cleanup sequence even
/// after an intermediate resource returns NULL, matching the C constructor's
/// cleanup contract for partially initialized contexts.
#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_createCCtx(
projection: *const ZSTDMT_RustCreateCCtxProjection,
allocate: Option<ZSTDMT_createCCtxAllocateFn>,
set_workers: Option<ZSTDMT_createCCtxSetWorkersFn>,
set_initial_state: Option<ZSTDMT_createCCtxSetInitialStateFn>,
create_factory: Option<ZSTDMT_createCCtxFactoryFn>,
create_jobs: Option<ZSTDMT_createCCtxJobsFn>,
create_buffer_pool: Option<ZSTDMT_createCCtxResourceFn>,
create_cctx_pool: Option<ZSTDMT_createCCtxResourceFn>,
create_seq_pool: Option<ZSTDMT_createCCtxResourceFn>,
init_serial: Option<ZSTDMT_createCCtxSerialInitFn>,
free_context: Option<ZSTDMT_createCCtxFreeFn>,
) -> *mut c_void {
let Some(projection) = (unsafe { projection.as_ref() }).copied() else {
return ptr::null_mut();
};
let (
Some(allocate),
Some(set_workers),
Some(set_initial_state),
Some(create_factory),
Some(create_jobs),
Some(create_buffer_pool),
Some(create_cctx_pool),
Some(create_seq_pool),
Some(init_serial),
Some(free_context),
) = (
allocate,
set_workers,
set_initial_state,
create_factory,
create_jobs,
create_buffer_pool,
create_cctx_pool,
create_seq_pool,
init_serial,
free_context,
)
else {
return ptr::null_mut();
};
if projection.callback_context.is_null()
|| projection.requested_nb_workers == 0
|| projection.max_nb_workers == 0
|| projection.context_size == 0
|| !custom_mem_is_valid(projection.custom_mem)
{
return ptr::null_mut();
}
let nb_workers = projection
.requested_nb_workers
.min(projection.max_nb_workers);
let mut nb_jobs = projection.requested_nb_workers.wrapping_add(2);
let context = unsafe { allocate(projection.callback_context, projection.context_size) };
if context.is_null() {
return ptr::null_mut();
}
// The C constructor ignores this setter's result for a valid worker
// count; preserve that exact policy before publishing the custom memory.
unsafe {
let _ = set_workers(projection.callback_context, nb_workers);
set_initial_state(projection.callback_context);
}
let factory = unsafe { create_factory(projection.callback_context, nb_workers) };
let jobs = unsafe { create_jobs(projection.callback_context, &mut nb_jobs) };
let buffer_pool = unsafe { create_buffer_pool(projection.callback_context, nb_workers) };
let cctx_pool = unsafe { create_cctx_pool(projection.callback_context, nb_workers) };
let seq_pool = unsafe { create_seq_pool(projection.callback_context, nb_workers) };
let serial_error = unsafe { init_serial(projection.callback_context) };
if factory.is_null()
|| jobs.is_null()
|| buffer_pool.is_null()
|| cctx_pool.is_null()
|| seq_pool.is_null()
|| serial_error != 0
{
unsafe { free_context(projection.callback_context) };
return ptr::null_mut();
}
context
}
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);
@@ -3932,6 +4075,109 @@ mod tests {
events: Vec<&'static str>,
}
#[derive(Default)]
struct CreateCCtxTestContext {
events: Vec<&'static str>,
workers: c_uint,
jobs: c_uint,
factory_ok: bool,
resource_calls: usize,
}
fn record_create_cctx_event(context: *mut c_void, event: &'static str) {
unsafe {
(*context.cast::<CreateCCtxTestContext>())
.events
.push(event);
}
}
unsafe extern "C" fn create_cctx_test_allocate(
context: *mut c_void,
_size: usize,
) -> *mut c_void {
record_create_cctx_event(context, "allocate");
ptr::dangling_mut::<c_void>()
}
unsafe extern "C" fn create_cctx_test_set_workers(
context: *mut c_void,
workers: c_uint,
) -> usize {
record_create_cctx_event(context, "workers");
unsafe { (*context.cast::<CreateCCtxTestContext>()).workers = workers };
0
}
unsafe extern "C" fn create_cctx_test_initial_state(context: *mut c_void) {
record_create_cctx_event(context, "initial-state");
}
unsafe extern "C" fn create_cctx_test_factory(
context: *mut c_void,
_workers: c_uint,
) -> *mut c_void {
record_create_cctx_event(context, "factory");
let context = unsafe { &*context.cast::<CreateCCtxTestContext>() };
if context.factory_ok {
ptr::dangling_mut::<c_void>()
} else {
ptr::null_mut()
}
}
unsafe extern "C" fn create_cctx_test_jobs(
context: *mut c_void,
jobs: *mut c_uint,
) -> *mut c_void {
record_create_cctx_event(context, "jobs");
unsafe {
(*context.cast::<CreateCCtxTestContext>()).jobs = 8;
*jobs = 8;
}
ptr::dangling_mut::<c_void>()
}
unsafe extern "C" fn create_cctx_test_resource(
context: *mut c_void,
_workers: c_uint,
) -> *mut c_void {
let event = unsafe {
let state = &mut *context.cast::<CreateCCtxTestContext>();
let event = match state.resource_calls {
0 => "buffer-pool",
1 => "cctx-pool",
_ => "seq-pool",
};
state.resource_calls += 1;
event
};
record_create_cctx_event(context, event);
ptr::dangling_mut::<c_void>()
}
unsafe extern "C" fn create_cctx_test_serial_init(context: *mut c_void) -> c_int {
record_create_cctx_event(context, "serial");
0
}
unsafe extern "C" fn create_cctx_test_free(context: *mut c_void) {
record_create_cctx_event(context, "free");
}
fn create_cctx_test_projection(
context: &mut CreateCCtxTestContext,
custom_mem: ZstdCustomMem,
) -> ZSTDMT_RustCreateCCtxProjection {
ZSTDMT_RustCreateCCtxProjection {
callback_context: context as *mut CreateCCtxTestContext as *mut c_void,
requested_nb_workers: 99,
max_nb_workers: 4,
context_size: 128,
custom_mem,
}
}
fn record_free_cctx_event(context: *mut c_void, event: &'static str) {
unsafe {
(*context.cast::<FreeCCtxTestContext>()).events.push(event);
@@ -4005,6 +4251,125 @@ mod tests {
assert_eq!(unsafe { ZSTDMT_rust_freeCCtx(ptr::null()) }, 0);
}
#[test]
fn create_cctx_preserves_resource_order_and_clamps_workers() {
let mut context = CreateCCtxTestContext {
factory_ok: true,
..CreateCCtxTestContext::default()
};
let projection = create_cctx_test_projection(&mut context, DEFAULT_MEM);
let result = unsafe {
ZSTDMT_rust_createCCtx(
&projection,
Some(create_cctx_test_allocate),
Some(create_cctx_test_set_workers),
Some(create_cctx_test_initial_state),
Some(create_cctx_test_factory),
Some(create_cctx_test_jobs),
Some(create_cctx_test_resource),
Some(create_cctx_test_resource),
Some(create_cctx_test_resource),
Some(create_cctx_test_serial_init),
Some(create_cctx_test_free),
)
};
assert_eq!(result, ptr::dangling_mut::<c_void>());
assert_eq!(context.workers, 4);
assert_eq!(context.jobs, 8);
assert_eq!(
context.events,
vec![
"allocate",
"workers",
"initial-state",
"factory",
"jobs",
"buffer-pool",
"cctx-pool",
"seq-pool",
"serial",
]
);
}
#[test]
fn create_cctx_runs_remaining_callbacks_before_failure_cleanup() {
let mut context = CreateCCtxTestContext::default();
let projection = create_cctx_test_projection(&mut context, DEFAULT_MEM);
let result = unsafe {
ZSTDMT_rust_createCCtx(
&projection,
Some(create_cctx_test_allocate),
Some(create_cctx_test_set_workers),
Some(create_cctx_test_initial_state),
Some(create_cctx_test_factory),
Some(create_cctx_test_jobs),
Some(create_cctx_test_resource),
Some(create_cctx_test_resource),
Some(create_cctx_test_resource),
Some(create_cctx_test_serial_init),
Some(create_cctx_test_free),
)
};
assert!(result.is_null());
assert_eq!(
context.events,
vec![
"allocate",
"workers",
"initial-state",
"factory",
"jobs",
"buffer-pool",
"cctx-pool",
"seq-pool",
"serial",
"free",
]
);
}
unsafe extern "C" fn create_cctx_test_alloc(_opaque: *mut c_void, _size: usize) -> *mut c_void {
ptr::null_mut()
}
#[test]
fn create_cctx_rejects_mismatched_custom_allocator() {
let mut context = CreateCCtxTestContext {
factory_ok: true,
..CreateCCtxTestContext::default()
};
let invalid_mem = ZstdCustomMem {
customAlloc: Some(create_cctx_test_alloc),
customFree: None,
opaque: ptr::null_mut(),
};
let projection = create_cctx_test_projection(&mut context, invalid_mem);
let result = unsafe {
ZSTDMT_rust_createCCtx(
&projection,
Some(create_cctx_test_allocate),
Some(create_cctx_test_set_workers),
Some(create_cctx_test_initial_state),
Some(create_cctx_test_factory),
Some(create_cctx_test_jobs),
Some(create_cctx_test_resource),
Some(create_cctx_test_resource),
Some(create_cctx_test_resource),
Some(create_cctx_test_serial_init),
Some(create_cctx_test_free),
)
};
assert!(result.is_null());
assert!(context.events.is_empty());
}
#[test]
fn free_cctx_preserves_callback_order_and_ownership_conditions() {
let mut context = FreeCCtxTestContext::default();