feat(mt): move produced-output flush into Rust

Move the multithreaded produced-output state machine into Rust while keeping
C-owned job descriptors, synchronization, checksum state, and buffer-pool
cleanup behind projection callbacks. Rust now decides publication progress,
job completion, pending-work results, and frame-completion reporting without
crossing the private ZSTDMT context layout. Add focused tests for partial
output, worker errors, checksum projection, and frame-end state.

Test Plan:
- cargo test --manifest-path rust/Cargo.toml --all-targets -- --test-threads=1
- cargo test --manifest-path rust/Cargo.toml --no-default-features --features compression,decompression,dict-builder,legacy-v01,legacy-v02,legacy-v03,legacy-v04,legacy-v05,legacy-v06,legacy-v07 --all-targets -- --test-threads=1
- make -B -C lib -j2 lib
- make -B -C tests -j2 test-zstd test-rust-lib-smoke
- cargo clippy --manifest-path rust/Cargo.toml --tests -- -D warnings
This commit is contained in:
2026-07-18 23:07:10 +02:00
parent 726a32fa79
commit 575360c6b1
2 changed files with 604 additions and 84 deletions
+457
View File
@@ -65,6 +65,63 @@ pub struct ZSTDMT_flushPublicationResult {
pub dstFlushed: usize,
}
/// Scalar snapshot of the oldest active C-owned job used by the flush state
/// machine. C keeps the descriptor, mutex, buffer pool, and checksum state;
/// Rust decides when this snapshot can be published and completed.
#[repr(C)]
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct ZSTDMT_flushJobProjection {
pub consumed: usize,
pub cSize: usize,
pub srcSize: usize,
pub dstStart: *const c_void,
pub dstCapacity: usize,
pub dstFlushed: usize,
pub frameChecksumNeeded: c_uint,
}
/// Scalar view of the C-owned MT state needed after the current job is
/// published. The C adapter applies descriptor and context mutations made by
/// Rust's decisions through callbacks.
#[repr(C)]
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct ZSTDMT_flushContextProjection {
pub doneJobID: c_uint,
pub nextJobID: c_uint,
pub jobIDMask: c_uint,
pub jobReady: c_uint,
pub frameEnded: c_uint,
pub inBuffFilled: usize,
}
#[repr(C)]
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct ZSTDMT_flushProducedResult {
pub result: usize,
pub outputPos: usize,
pub updateAllJobsCompleted: c_uint,
pub allJobsCompleted: c_uint,
}
pub type ZSTDMT_flushProjectJobFn = unsafe extern "C" fn(
opaque: *mut c_void,
job_id: c_uint,
block_to_flush: c_uint,
projection: *mut ZSTDMT_flushJobProjection,
);
pub type ZSTDMT_flushChecksumFn = unsafe extern "C" fn(
opaque: *mut c_void,
job_id: c_uint,
projection: *mut ZSTDMT_flushJobProjection,
);
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);
const ZSTD_E_END: c_uint = 2;
/// Scalar view of the C-owned job state needed to finish a streaming frame
/// with an empty last block. The job descriptor, buffer pool, and mutex stay
/// private to C; Rust receives only this projection and a buffer callback.
@@ -238,6 +295,120 @@ unsafe fn publish_job_output_with(
}
}
#[inline]
fn flush_produced_result(result: usize, output_pos: usize) -> ZSTDMT_flushProducedResult {
ZSTDMT_flushProducedResult {
result,
outputPos: output_pos,
updateAllJobsCompleted: 0,
allJobsCompleted: 0,
}
}
#[inline]
unsafe fn flush_produced_with<P, C, U, F, E>(
context: ZSTDMT_flushContextProjection,
output_dst: *mut c_void,
output_size: usize,
output_pos: usize,
block_to_flush: c_uint,
end: c_uint,
mut project_job: P,
mut add_frame_checksum: C,
mut update_job: U,
mut complete_job: F,
mut on_error: E,
) -> 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(),
{
let mut output_pos = output_pos;
let mut done_job_id = context.doneJobID;
if done_job_id < context.nextJobID {
let job_id = done_job_id & context.jobIDMask;
let mut projection = project_job(job_id, block_to_flush);
let src_consumed = projection.consumed;
let src_size = projection.srcSize;
if ERR_isError(projection.cSize) {
on_error();
return flush_produced_result(projection.cSize, output_pos);
}
/* The checksum is added only once, after the worker has consumed the
* whole source range. C retains the serial XXH state and updates the
* projected cSize through this callback. */
if src_consumed == src_size && projection.frameChecksumNeeded != 0 {
add_frame_checksum(job_id, &mut projection);
}
let c_size = projection.cSize;
let mut dst_flushed = projection.dstFlushed;
if c_size > 0 {
let publication = unsafe {
publish_job_output_with(
output_dst,
output_size,
output_pos,
projection.dstStart,
projection.dstCapacity,
c_size,
dst_flushed,
)
};
if publication.status != FLUSH_PUBLICATION_OK {
return flush_produced_result(ERROR(ZstdErrorCode::Generic), output_pos);
}
output_pos = publication.outputPos;
dst_flushed = publication.dstFlushed;
update_job(job_id, dst_flushed);
if src_consumed == src_size && dst_flushed == c_size {
complete_job(job_id, src_size, c_size);
done_job_id = done_job_id.wrapping_add(1);
}
}
/* Match the C API's distinction between known pending output and a
* live worker whose output size is not known yet. */
if c_size > dst_flushed {
return flush_produced_result(c_size - dst_flushed, output_pos);
}
if src_size > src_consumed {
return flush_produced_result(1, output_pos);
}
}
if done_job_id < context.nextJobID {
return flush_produced_result(1, output_pos);
}
if context.jobReady != 0 {
return flush_produced_result(1, output_pos);
}
if context.inBuffFilled > 0 {
return flush_produced_result(1, output_pos);
}
/* All jobs are entirely flushed. For ZSTD_e_end the caller asks whether
* the frame itself has ended, rather than whether buffers are empty. */
ZSTDMT_flushProducedResult {
result: if end == ZSTD_E_END {
(context.frameEnded == 0) as usize
} else {
0
},
outputPos: output_pos,
updateAllJobsCompleted: 1,
allJobsCompleted: context.frameEnded,
}
}
#[cfg(not(test))]
unsafe extern "C" {
fn ZSTD_compressContinue_public(
@@ -391,6 +562,71 @@ pub unsafe extern "C" fn ZSTDMT_rust_publishJobOutput(
}
}
#[cfg(not(test))]
#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_flushProduced(
context: *const ZSTDMT_flushContextProjection,
output_dst: *mut c_void,
output_size: usize,
output_pos: usize,
block_to_flush: c_uint,
end: c_uint,
opaque: *mut c_void,
projectJob: Option<ZSTDMT_flushProjectJobFn>,
addFrameChecksum: Option<ZSTDMT_flushChecksumFn>,
updateJob: Option<ZSTDMT_flushUpdateJobFn>,
completeJob: Option<ZSTDMT_flushCompleteJobFn>,
onError: Option<ZSTDMT_flushErrorFn>,
) -> ZSTDMT_flushProducedResult {
let Some(context) = (unsafe { context.as_ref() }).copied() else {
return flush_produced_result(ERROR(ZstdErrorCode::Generic), output_pos);
};
let (
Some(project_job),
Some(add_frame_checksum),
Some(update_job),
Some(complete_job),
Some(on_error),
) = (
projectJob,
addFrameChecksum,
updateJob,
completeJob,
onError,
)
else {
return flush_produced_result(ERROR(ZstdErrorCode::Generic), output_pos);
};
unsafe {
flush_produced_with(
context,
output_dst,
output_size,
output_pos,
block_to_flush,
end,
|job_id, block_to_flush| {
let mut projection = ZSTDMT_flushJobProjection::default();
project_job(opaque, job_id, block_to_flush, &mut projection);
projection
},
|job_id, projection| {
add_frame_checksum(opaque, job_id, projection);
},
|job_id, dst_flushed| {
update_job(opaque, job_id, dst_flushed);
},
|job_id, src_size, c_size| {
complete_job(opaque, job_id, src_size, c_size);
},
|| {
on_error(opaque);
},
)
}
}
#[inline]
fn get_input_data_in_use_with<F>(
first_job_id: c_uint,
@@ -1919,6 +2155,227 @@ mod tests {
assert_eq!(output, original);
}
#[test]
fn flush_state_machine_preserves_offsets_and_completes_job() {
let job = [0x60u8, 0x61, 0x62, 0x63, 0x64, 0x65];
let mut output = [0xa5u8; 8];
let projection = ZSTDMT_flushJobProjection {
consumed: 4,
cSize: job.len(),
srcSize: 4,
dstStart: job.as_ptr().cast(),
dstCapacity: job.len(),
dstFlushed: 2,
frameChecksumNeeded: 0,
};
let context = ZSTDMT_flushContextProjection {
doneJobID: 0,
nextJobID: 1,
jobIDMask: 0,
..ZSTDMT_flushContextProjection::default()
};
let mut updated = None;
let mut completed = None;
let result = unsafe {
flush_produced_with(
context,
output.as_mut_ptr().cast(),
output.len(),
1,
0,
1,
|job_id, block_to_flush| {
assert_eq!(job_id, 0);
assert_eq!(block_to_flush, 0);
projection
},
|_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"),
)
};
assert_eq!(result.result, 0);
assert_eq!(result.outputPos, 5);
assert_eq!(result.updateAllJobsCompleted, 1);
assert_eq!(result.allJobsCompleted, 0);
assert_eq!(updated, Some(6));
assert_eq!(completed, Some((4, 6)));
assert_eq!(&output[1..5], &job[2..]);
assert_eq!(&output[..1], &[0xa5]);
}
#[test]
fn flush_state_machine_keeps_pending_output_without_space() {
let job = [0x70u8, 0x71, 0x72, 0x73, 0x74, 0x75];
let mut output = [0xa5u8; 4];
let original = output;
let projection = ZSTDMT_flushJobProjection {
consumed: 2,
cSize: job.len(),
srcSize: 5,
dstStart: job.as_ptr().cast(),
dstCapacity: job.len(),
dstFlushed: 2,
frameChecksumNeeded: 0,
};
let context = ZSTDMT_flushContextProjection {
doneJobID: 3,
nextJobID: 4,
jobIDMask: 0,
..ZSTDMT_flushContextProjection::default()
};
let mut updated = None;
let mut completed = false;
let result = unsafe {
flush_produced_with(
context,
output.as_mut_ptr().cast(),
output.len(),
output.len(),
1,
1,
|_job_id, _block_to_flush| projection,
|_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"),
)
};
assert_eq!(result.result, 4);
assert_eq!(result.outputPos, output.len());
assert_eq!(result.updateAllJobsCompleted, 0);
assert_eq!(updated, Some(2));
assert!(!completed);
assert_eq!(output, original);
}
#[test]
fn flush_state_machine_normalizes_worker_error_and_checksum_projection() {
let error_projection = ZSTDMT_flushJobProjection {
cSize: ERROR(ZstdErrorCode::Generic),
..ZSTDMT_flushJobProjection::default()
};
let context = ZSTDMT_flushContextProjection {
doneJobID: 0,
nextJobID: 1,
jobIDMask: 0,
..ZSTDMT_flushContextProjection::default()
};
let mut error_cleanup = false;
let result = unsafe {
flush_produced_with(
context,
ptr::null_mut(),
0,
0,
1,
1,
|_job_id, _block_to_flush| error_projection,
|_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,
)
};
assert_eq!(result.result, ERROR(ZstdErrorCode::Generic));
assert_eq!(result.outputPos, 0);
assert!(error_cleanup);
let job = [0x80u8, 0x81, 0x82, 0x83, 0x84, 0x85, 0x86];
let mut output = [0xa5u8; 8];
let projection = ZSTDMT_flushJobProjection {
consumed: 4,
cSize: 3,
srcSize: 4,
dstStart: job.as_ptr().cast(),
dstCapacity: job.len(),
dstFlushed: 0,
frameChecksumNeeded: 1,
};
let mut checksum_calls = 0;
let mut completed_size = None;
let result = unsafe {
flush_produced_with(
ZSTDMT_flushContextProjection {
doneJobID: 0,
nextJobID: 1,
jobIDMask: 0,
..ZSTDMT_flushContextProjection::default()
},
output.as_mut_ptr().cast(),
output.len(),
0,
0,
1,
|_job_id, _block_to_flush| projection,
|_job_id, projected| {
checksum_calls += 1;
projected.cSize = 7;
projected.frameChecksumNeeded = 0;
},
|_job_id, _dst_flushed| {},
|_job_id, _src_size, c_size| completed_size = Some(c_size),
|| panic!("unexpected error callback"),
)
};
assert_eq!(checksum_calls, 1);
assert_eq!(completed_size, Some(7));
assert_eq!(result.result, 0);
assert_eq!(&output[..7], &job);
}
#[test]
fn flush_state_machine_reports_end_of_frame_state() {
let not_ended = unsafe {
flush_produced_with(
ZSTDMT_flushContextProjection {
frameEnded: 0,
..ZSTDMT_flushContextProjection::default()
},
ptr::null_mut(),
0,
0,
0,
ZSTD_E_END,
|_job_id, _block_to_flush| panic!("no job expected"),
|_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"),
)
};
assert_eq!(not_ended.result, 1);
assert_eq!(not_ended.updateAllJobsCompleted, 1);
assert_eq!(not_ended.allJobsCompleted, 0);
let ended = unsafe {
flush_produced_with(
ZSTDMT_flushContextProjection {
frameEnded: 1,
..ZSTDMT_flushContextProjection::default()
},
ptr::null_mut(),
0,
0,
0,
ZSTD_E_END,
|_job_id, _block_to_flush| panic!("no job expected"),
|_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"),
)
};
assert_eq!(ended.result, 0);
assert_eq!(ended.updateAllJobsCompleted, 1);
assert_eq!(ended.allJobsCompleted, 1);
}
#[test]
fn empty_last_block_serializes_after_buffer_callback() {
let mut storage = vec![0xa5u8; 3];