refactor(mt): move stream input copy into Rust

The MT streaming scheduler previously delegated each bounded input-buffer copy
through a C-only callback. That callback mutated the C context as a side
effect, which kept a trivial data movement operation outside the Rust stream
path and made the result projection incomplete.

Copy the checked input range directly in Rust with copy_nonoverlapping, after
the existing available-input, destination-capacity, and non-null-source checks.
The C input-range callback remains responsible for exposing the reusable buffer,
while the Rust result now returns its updated filled count. C publishes that
count back to mtctx after the scheduler call. The result ABI has explicit
cross-language layout assertions, and the scheduler test verifies both copied
bytes and fill accounting. Job scheduling and all other MT callbacks are
unchanged.

Test Plan:
- `cc -fsyntax-only -Werror=incompatible-pointer-types -Ilib -Ilib/common -Ilib/compress -Ilib/decompress -Ilib/dict -Ilib/legacy lib/compress/zstdmt_compress.c` -- passed
- `cargo +nightly fmt --manifest-path rust/Cargo.toml --all -- --check` -- passed
- `git diff --check` and `git diff --cached --check` -- passed
- Cargo build/test, make, fuzzers, and other heavy checks were not run per task instructions
This commit is contained in:
2026-07-20 14:54:44 +02:00
parent f3dad84d86
commit 40137b28bf
2 changed files with 81 additions and 60 deletions
+12 -15
View File
@@ -17,7 +17,7 @@
/* ====== Dependencies ====== */
#include "../common/allocations.h" /* ZSTD_customMalloc, ZSTD_customCalloc, ZSTD_customFree */
#include "../common/zstd_deps.h" /* ZSTD_memcpy, ZSTD_memset, INT_MAX, UINT_MAX */
#include "../common/zstd_deps.h" /* ZSTD_memmove, ZSTD_memset, INT_MAX, UINT_MAX */
#include "../common/mem.h" /* MEM_STATIC */
#include "../common/pool.h" /* threadpool */
#include "../common/threading.h" /* mutex */
@@ -381,10 +381,18 @@ typedef struct {
size_t result;
size_t inputPos;
size_t outputPos;
size_t inBuffFilled;
} ZSTDMT_RustCompressStreamResult;
typedef char ZSTDMT_compress_stream_result_layout[
(offsetof(ZSTDMT_RustCompressStreamResult, result) == 0
&& offsetof(ZSTDMT_RustCompressStreamResult, inputPos) == sizeof(size_t)
&& offsetof(ZSTDMT_RustCompressStreamResult, outputPos) == 2 * sizeof(size_t)
&& offsetof(ZSTDMT_RustCompressStreamResult, inBuffFilled)
== 3 * sizeof(size_t)
&& sizeof(ZSTDMT_RustCompressStreamResult) == 4 * sizeof(size_t))
? 1 : -1];
typedef int (*ZSTDMT_streamTryGetInputRangeFn)(
void* opaque, ZSTDMT_RustStreamInputRangeProjection* projection);
typedef int (*ZSTDMT_streamLoadInputFn)(void* opaque, const void* src, size_t size);
typedef size_t (*ZSTDMT_streamCreateJobFn)(void* opaque, size_t srcSize, unsigned end);
typedef ZSTDMT_RustStreamFlushResult (*ZSTDMT_streamFlushProducedFn)(
void* opaque, void* outputDst, size_t outputSize, size_t outputPos,
@@ -395,7 +403,6 @@ ZSTDMT_RustCompressStreamResult ZSTDMT_rust_compressStreamGeneric(
const ZSTDMT_RustStreamOutputProjection* output,
unsigned end, void* opaque,
ZSTDMT_streamTryGetInputRangeFn tryGetInputRange,
ZSTDMT_streamLoadInputFn loadInput,
ZSTDMT_streamCreateJobFn createJob,
ZSTDMT_streamFlushProducedFn flushProduced);
@@ -3041,17 +3048,6 @@ static int ZSTDMT_streamTryGetInputRange(
return ready;
}
static int ZSTDMT_streamLoadInput(void* opaque, const void* src, size_t size)
{
ZSTDMT_CCtx* const mtctx = (ZSTDMT_CCtx*)opaque;
assert(mtctx->inBuff.buffer.start != NULL);
assert(mtctx->inBuff.filled <= mtctx->inBuff.buffer.capacity);
assert(size <= mtctx->inBuff.buffer.capacity - mtctx->inBuff.filled);
ZSTD_memcpy((char*)mtctx->inBuff.buffer.start + mtctx->inBuff.filled, src, size);
mtctx->inBuff.filled += size;
return 1;
}
size_t ZSTDMT_nextInputSizeHint(const ZSTDMT_CCtx* mtctx)
{
return ZSTDMT_rust_nextInputSizeHint(mtctx->targetSectionSize,
@@ -3091,13 +3087,14 @@ size_t ZSTDMT_compressStream_generic(ZSTDMT_CCtx* mtctx,
ZSTDMT_RustCompressStreamResult const result =
ZSTDMT_rust_compressStreamGeneric(
&context, &inputProjection, &outputProjection, (unsigned)endOp,
mtctx, ZSTDMT_streamTryGetInputRange, ZSTDMT_streamLoadInput,
mtctx, ZSTDMT_streamTryGetInputRange,
ZSTDMT_createCompressionJob, ZSTDMT_streamFlushProduced);
DEBUGLOG(5, "ZSTDMT_compressStream_generic (endOp=%u, srcSize=%u)",
(U32)endOp, (U32)(input->size - input->pos));
assert(output->pos <= output->size);
assert(input->pos <= input->size);
mtctx->inBuff.filled = result.inBuffFilled;
input->pos = result.inputPos;
output->pos = result.outputPos;
DEBUGLOG(5, "end of ZSTDMT_compressStream_generic: remainingToFlush = %u", (U32)result.result);
+69 -45
View File
@@ -1042,14 +1042,21 @@ pub struct ZSTDMT_compressStreamResult {
pub result: usize,
pub inputPos: usize,
pub outputPos: usize,
pub inBuffFilled: usize,
}
const _: () = {
assert!(offset_of!(ZSTDMT_compressStreamResult, result) == 0);
assert!(offset_of!(ZSTDMT_compressStreamResult, inputPos) == size_of::<usize>());
assert!(offset_of!(ZSTDMT_compressStreamResult, outputPos) == 2 * size_of::<usize>());
assert!(offset_of!(ZSTDMT_compressStreamResult, inBuffFilled) == 3 * size_of::<usize>());
assert!(size_of::<ZSTDMT_compressStreamResult>() == 4 * size_of::<usize>());
};
pub type ZSTDMT_streamTryGetInputRangeFn = unsafe extern "C" fn(
opaque: *mut c_void,
projection: *mut ZSTDMT_streamInputRangeProjection,
) -> c_int;
pub type ZSTDMT_streamLoadInputFn =
unsafe extern "C" fn(opaque: *mut c_void, src: *const c_void, size: usize) -> c_int;
pub type ZSTDMT_streamCreateJobFn =
unsafe extern "C" fn(opaque: *mut c_void, src_size: usize, end: c_uint) -> usize;
pub type ZSTDMT_streamFlushProducedFn = unsafe extern "C" fn(
@@ -2547,37 +2554,37 @@ fn compress_stream_result(
result: usize,
input_pos: usize,
output_pos: usize,
in_buff_filled: usize,
) -> ZSTDMT_compressStreamResult {
ZSTDMT_compressStreamResult {
result,
inputPos: input_pos,
outputPos: output_pos,
inBuffFilled: in_buff_filled,
}
}
/// Run one outer MT streaming pass while C retains the context and worker
/// implementation details.
///
/// The callbacks are intentionally narrow: input-range acquisition and input
/// copying can update C-owned buffers, job creation can post the prepared job,
/// and flushing can wait on C-owned job synchronization. Rust owns only the
/// ordering of those operations, the end-directive adjustments, and the
/// public progress/return policy.
/// The callbacks are intentionally narrow: input-range acquisition exposes a
/// C-owned buffer, job creation can post the prepared job, and flushing can
/// wait on C-owned job synchronization. Rust owns the bounded input copy, the
/// ordering of those operations, the end-directive adjustments, and the public
/// progress/return policy.
#[allow(clippy::too_many_arguments)]
unsafe fn compress_stream_generic_with<T, L, J, F, S>(
unsafe fn compress_stream_generic_with<T, J, F, S>(
context: ZSTDMT_compressStreamContextProjection,
input: ZSTDMT_streamInputProjection,
output: ZSTDMT_streamOutputProjection,
end_op: c_uint,
mut try_get_input_range: T,
mut load_input: L,
mut create_job: J,
mut flush_produced: F,
mut find_sync_point: S,
) -> ZSTDMT_compressStreamResult
where
T: FnMut(&mut ZSTDMT_streamInputRangeProjection) -> c_int,
L: FnMut(*const c_void, usize) -> c_int,
J: FnMut(usize, c_uint) -> usize,
F: FnMut(*mut c_void, usize, usize, c_uint, c_uint) -> ZSTDMT_streamFlushResult,
S: FnMut(
@@ -2594,11 +2601,21 @@ where
|| context.inBuffFilled > context.inBuffCapacity
|| (context.inBuffStart.is_null() && context.inBuffFilled != 0)
{
return compress_stream_result(ERROR(ZstdErrorCode::Generic), input_pos, output_pos);
return compress_stream_result(
ERROR(ZstdErrorCode::Generic),
input_pos,
output_pos,
context.inBuffFilled,
);
}
if context.frameEnded != 0 && end_op == ZSTD_E_CONTINUE {
return compress_stream_result(ERROR(ZstdErrorCode::StageWrong), input_pos, output_pos);
return compress_stream_result(
ERROR(ZstdErrorCode::StageWrong),
input_pos,
output_pos,
context.inBuffFilled,
);
}
let mut end_op = end_op;
@@ -2634,6 +2651,7 @@ where
ERROR(ZstdErrorCode::Generic),
input_pos,
output_pos,
input_range.bufferFilled,
);
};
if sync_point.toLoad > available
@@ -2644,6 +2662,7 @@ where
ERROR(ZstdErrorCode::Generic),
input_pos,
output_pos,
input_range.bufferFilled,
);
}
@@ -2653,11 +2672,14 @@ where
.cast::<u8>()
.wrapping_add(input_pos)
.cast::<c_void>();
if load_input(source, sync_point.toLoad) == 0 {
return compress_stream_result(
ERROR(ZstdErrorCode::Generic),
input_pos,
output_pos,
unsafe {
ptr::copy_nonoverlapping(
source.cast::<u8>(),
input_range
.bufferStart
.cast::<u8>()
.add(input_range.bufferFilled),
sync_point.toLoad,
);
}
input_pos += sync_point.toLoad;
@@ -2681,7 +2703,12 @@ where
debug_assert!(input_range.bufferFilled <= context.targetSectionSize);
let create_result = create_job(input_range.bufferFilled, end_op);
if ERR_isError(create_result) {
return compress_stream_result(create_result, input_pos, output_pos);
return compress_stream_result(
create_result,
input_pos,
output_pos,
input_range.bufferFilled,
);
}
}
@@ -2698,7 +2725,12 @@ where
} else {
remaining_to_flush
};
compress_stream_result(result, input_pos, flush_result.outputPos)
compress_stream_result(
result,
input_pos,
flush_result.outputPos,
input_range.bufferFilled,
)
}
/// C ABI entry point for the MT outer scheduling/flush policy.
@@ -2710,23 +2742,32 @@ pub unsafe extern "C" fn ZSTDMT_rust_compressStreamGeneric(
end_op: c_uint,
opaque: *mut c_void,
tryGetInputRange: Option<ZSTDMT_streamTryGetInputRangeFn>,
loadInput: Option<ZSTDMT_streamLoadInputFn>,
createJob: Option<ZSTDMT_streamCreateJobFn>,
flushProduced: Option<ZSTDMT_streamFlushProducedFn>,
) -> ZSTDMT_compressStreamResult {
let Some(context) = (unsafe { context.as_ref() }).copied() else {
return compress_stream_result(ERROR(ZstdErrorCode::Generic), 0, 0);
return compress_stream_result(ERROR(ZstdErrorCode::Generic), 0, 0, 0);
};
let Some(input) = (unsafe { input.as_ref() }).copied() else {
return compress_stream_result(ERROR(ZstdErrorCode::Generic), 0, 0);
return compress_stream_result(ERROR(ZstdErrorCode::Generic), 0, 0, context.inBuffFilled);
};
let Some(output) = (unsafe { output.as_ref() }).copied() else {
return compress_stream_result(ERROR(ZstdErrorCode::Generic), input.pos, 0);
return compress_stream_result(
ERROR(ZstdErrorCode::Generic),
input.pos,
0,
context.inBuffFilled,
);
};
let (Some(try_get_input_range), Some(load_input), Some(create_job), Some(flush_produced)) =
(tryGetInputRange, loadInput, createJob, flushProduced)
let (Some(try_get_input_range), Some(create_job), Some(flush_produced)) =
(tryGetInputRange, createJob, flushProduced)
else {
return compress_stream_result(ERROR(ZstdErrorCode::Generic), input.pos, output.pos);
return compress_stream_result(
ERROR(ZstdErrorCode::Generic),
input.pos,
output.pos,
context.inBuffFilled,
);
};
unsafe {
@@ -2736,7 +2777,6 @@ pub unsafe extern "C" fn ZSTDMT_rust_compressStreamGeneric(
output,
end_op,
|projection| try_get_input_range(opaque, projection),
|src, size| load_input(opaque, src, size),
|src_size, end| create_job(opaque, src_size, end),
|dst, size, pos, block_to_flush, end| {
flush_produced(opaque, dst, size, pos, block_to_flush, end)
@@ -6750,7 +6790,7 @@ mod tests {
}
#[test]
fn outer_scheduler_loads_input_posts_job_and_does_not_block_after_progress() {
fn outer_scheduler_copies_input_and_reports_filled_after_progress() {
let input_bytes = [1u8, 2, 3, 4, 5];
let mut round_buffer = [0xa5u8; 4];
let mut output_buffer = [0xb6u8; 8];
@@ -6775,7 +6815,6 @@ mod tests {
pos: 1,
};
let mut range_calls = 0;
let mut load_calls = 0;
let mut create_calls = Vec::new();
let mut flush_calls = Vec::new();
let round_buffer_ptr = round_buffer.as_mut_ptr();
@@ -6794,12 +6833,6 @@ mod tests {
projection.bufferFilled = 0;
1
},
|src, size| {
load_calls += 1;
assert_eq!(size, round_capacity);
ptr::copy_nonoverlapping(src.cast::<u8>(), round_buffer_ptr, size);
1
},
|src_size, end| {
create_calls.push((src_size, end));
0
@@ -6819,13 +6852,13 @@ mod tests {
};
assert_eq!(range_calls, 1);
assert_eq!(load_calls, 1);
assert_eq!(round_buffer, input_bytes[..round_buffer.len()]);
assert_eq!(create_calls, vec![(round_buffer.len(), ZSTD_E_CONTINUE)]);
assert_eq!(flush_calls, vec![(0, ZSTD_E_CONTINUE)]);
assert_eq!(result.result, 1);
assert_eq!(result.inputPos, round_buffer.len());
assert_eq!(result.outputPos, 3);
assert_eq!(result.inBuffFilled, round_buffer.len());
assert_eq!(output_buffer[0], 0xb6);
}
@@ -6857,10 +6890,6 @@ mod tests {
output,
ZSTD_E_END,
|_projection| panic!("input range should already be available"),
|_src, size| {
assert_eq!(size, round_buffer.len());
1
},
|src_size, end| {
create_end = Some((src_size, end));
0
@@ -6912,7 +6941,6 @@ mod tests {
output,
ZSTD_E_CONTINUE,
|_projection| panic!("a ready job must skip input acquisition"),
|_src, _size| panic!("a ready job must skip input copying"),
|src_size, end| {
create_call = Some((src_size, end));
0
@@ -6959,10 +6987,6 @@ mod tests {
output,
ZSTD_E_CONTINUE,
|_projection| panic!("input range should already be available"),
|_src, size| {
assert_eq!(size, round_buffer.len());
1
},
|_src_size, _end| ERROR(ZstdErrorCode::DstSizeTooSmall),
|_dst, _size, output_pos, _block_to_flush, _end| {
flush_called = true;