From 0e0d87ddf50814ffbe0f52725682059506f4c840 Mon Sep 17 00:00:00 2001 From: ddidderr Date: Sat, 18 Jul 2026 22:40:56 +0200 Subject: [PATCH] feat(mt): move empty-job emission into Rust The multithreaded compressor still handled the terminal empty-job case in C, including the decision to acquire a buffer, emit the final empty block, and clear the consumed source range. That left a small but stateful branch outside the Rust job helpers and made allocation-failure behavior difficult to test. Project the terminal job and result through a stable C ABI, then let Rust own that branch's assertions, buffer callback, block-header serialization, and failure-preserving result. C remains responsible for its private job layout and applies the returned buffer, source, and compressed-size fields exactly as before. Test Plan: - `cargo test --manifest-path rust/Cargo.toml --all-targets -- --test-threads=1` -- 473 passed. - Native multithreaded cases in `test-fuzzer`, `test-zstream`, `test-cli-tests`, and `test-zstd` -- passed. - `make -B -C lib -j2 lib ZSTD_FORCE_DECOMPRESS_SEQUENCES_SHORT=1 ZSTD_FORCE_DECOMPRESS_SEQUENCES_LONG=0` -- passed. - `make -B -C lib -j2 lib ZSTD_FORCE_DECOMPRESS_SEQUENCES_SHORT=0 ZSTD_FORCE_DECOMPRESS_SEQUENCES_LONG=1` -- passed. --- lib/compress/zstdmt_compress.c | 60 ++++++++++++--- rust/src/zstdmt_compress.rs | 135 +++++++++++++++++++++++++++++++++ 2 files changed, 186 insertions(+), 9 deletions(-) diff --git a/lib/compress/zstdmt_compress.c b/lib/compress/zstdmt_compress.c index 8cda31afa..539620a72 100644 --- a/lib/compress/zstdmt_compress.c +++ b/lib/compress/zstdmt_compress.c @@ -140,6 +140,24 @@ ZSTDMT_flushPublicationResult ZSTDMT_rust_publishJobOutput( const void* jobDst, size_t jobCapacity, size_t cSize, size_t dstFlushed); +typedef struct { + unsigned lastJob; + size_t srcSize; + unsigned firstJob; + void* dstStart; + size_t dstCapacity; + size_t consumed; +} ZSTDMT_RustEmptyBlockJobProjection; +typedef struct { + ZSTDMT_RustBuffer buffer; + size_t cSize; + unsigned clearSource; +} ZSTDMT_RustEmptyBlockResult; +typedef ZSTDMT_RustBuffer (*ZSTDMT_bufferGetFn)(void* opaque); +ZSTDMT_RustEmptyBlockResult ZSTDMT_rust_writeLastEmptyBlock( + const ZSTDMT_RustEmptyBlockJobProjection* projection, + void* opaque, ZSTDMT_bufferGetFn getBuffer); + typedef struct { rawSeq* seq; size_t pos; @@ -288,6 +306,13 @@ static Buffer ZSTDMT_getBuffer(ZSTDMT_bufferPool* bufPool) return buffer; } +static ZSTDMT_RustBuffer ZSTDMT_getBufferForRust(void* opaque) +{ + Buffer const buffer = ZSTDMT_getBuffer((ZSTDMT_bufferPool*)opaque); + ZSTDMT_RustBuffer const result = { buffer.start, buffer.capacity }; + return result; +} + #if ZSTD_RESIZE_SEQPOOL /** ZSTDMT_resizeBuffer() : * assumption : bufPool must be valid @@ -1375,20 +1400,37 @@ size_t ZSTDMT_initCStream_internal( */ static void ZSTDMT_writeLastEmptyBlock(ZSTDMT_jobDescription* job) { + ZSTDMT_RustEmptyBlockJobProjection projection; + ZSTDMT_RustEmptyBlockResult result; + assert(job->lastJob == 1); assert(job->src.size == 0); /* last job is empty -> will be simplified into a last empty block */ assert(job->firstJob == 0); /* cannot be first job, as it also needs to create frame header */ assert(job->dstBuff.start == NULL); /* invoked from streaming variant only (otherwise, dstBuff might be user's output) */ - job->dstBuff = ZSTDMT_getBuffer(job->bufPool); - if (job->dstBuff.start == NULL) { - job->cSize = ERROR(memory_allocation); - return; - } - assert(job->dstBuff.capacity >= ZSTD_blockHeaderSize); /* no buffer should ever be that small */ - job->src = kNullRange; - job->cSize = ZSTD_writeLastEmptyBlock(job->dstBuff.start, job->dstBuff.capacity); - assert(!ZSTD_isError(job->cSize)); assert(job->consumed == 0); + + projection = (ZSTDMT_RustEmptyBlockJobProjection){ + job->lastJob, + job->src.size, + job->firstJob, + job->dstBuff.start, + job->dstBuff.capacity, + job->consumed + }; + result = ZSTDMT_rust_writeLastEmptyBlock( + &projection, job->bufPool, ZSTDMT_getBufferForRust); + + job->dstBuff = (Buffer){ result.buffer.start, result.buffer.capacity }; + if (job->dstBuff.start == NULL) { + assert(!result.clearSource); + job->cSize = result.cSize; + return; + } + assert(result.clearSource); + assert(job->dstBuff.capacity >= ZSTD_blockHeaderSize); /* no buffer should ever be that small */ + if (result.clearSource) job->src = kNullRange; + job->cSize = result.cSize; + assert(!ZSTD_isError(job->cSize)); } static size_t ZSTDMT_createCompressionJob(ZSTDMT_CCtx* mtctx, size_t srcSize, ZSTD_EndDirective endOp) diff --git a/rust/src/zstdmt_compress.rs b/rust/src/zstdmt_compress.rs index d89da97f6..9bb68334e 100644 --- a/rust/src/zstdmt_compress.rs +++ b/rust/src/zstdmt_compress.rs @@ -65,6 +65,30 @@ pub struct ZSTDMT_flushPublicationResult { pub dstFlushed: usize, } +/// 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. +#[repr(C)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct ZSTDMT_emptyBlockJobProjection { + pub lastJob: c_uint, + pub srcSize: usize, + pub firstJob: c_uint, + pub dstStart: *mut c_void, + pub dstCapacity: usize, + pub consumed: usize, +} + +#[repr(C)] +#[derive(Clone, Copy, Default)] +pub struct ZSTDMT_emptyBlockResult { + pub buffer: ZstdMtBuffer, + pub cSize: usize, + pub clearSource: c_uint, +} + +pub type ZSTDMT_bufferGetFn = unsafe extern "C" fn(*mut c_void) -> ZstdMtBuffer; + /// Scalar snapshot of one C-owned job descriptor. /// /// The C adapter fills this projection while holding the descriptor mutex. @@ -98,6 +122,67 @@ pub type ZSTDMT_jobProjectionFn = unsafe extern "C" fn( const FLUSH_PUBLICATION_OK: c_int = 0; const FLUSH_PUBLICATION_INVALID_BOUNDS: c_int = 1; +#[inline] +fn empty_block_allocation_error() -> ZSTDMT_emptyBlockResult { + ZSTDMT_emptyBlockResult { + buffer: ZstdMtBuffer::default(), + cSize: ERROR(ZstdErrorCode::MemoryAllocation), + clearSource: 0, + } +} + +#[inline] +unsafe fn write_last_empty_block_with( + projection: ZSTDMT_emptyBlockJobProjection, + mut get_buffer: F, +) -> ZSTDMT_emptyBlockResult +where + F: FnMut() -> ZstdMtBuffer, +{ + debug_assert_eq!(projection.lastJob, 1); + debug_assert_eq!(projection.srcSize, 0); + debug_assert_eq!(projection.firstJob, 0); + debug_assert!(projection.dstStart.is_null()); + debug_assert_eq!(projection.dstCapacity, 0); + debug_assert_eq!(projection.consumed, 0); + + let buffer = get_buffer(); + if buffer.start.is_null() { + return empty_block_allocation_error(); + } + + let c_size = unsafe { + crate::zstd_compress_frame::ZSTD_writeLastEmptyBlock(buffer.start, buffer.capacity) + }; + ZSTDMT_emptyBlockResult { + buffer, + cSize: c_size, + /* The C implementation clears src immediately after acquiring the + * buffer, before serializing the block. */ + clearSource: 1, + } +} + +/// Acquire and serialize the terminal empty block for a C-owned MT job. +/// +/// C keeps the private descriptor and applies the returned state after this +/// callback-driven operation. In particular, an allocation failure leaves +/// the source range untouched just as the original C helper did. +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_writeLastEmptyBlock( + projection: *const ZSTDMT_emptyBlockJobProjection, + opaque: *mut c_void, + getBuffer: Option, +) -> ZSTDMT_emptyBlockResult { + let Some(projection) = (unsafe { projection.as_ref() }).copied() else { + return empty_block_allocation_error(); + }; + let Some(get_buffer) = getBuffer else { + return empty_block_allocation_error(); + }; + unsafe { write_last_empty_block_with(projection, || get_buffer(opaque)) } +} + #[inline] fn invalid_flush_publication( output_pos: usize, @@ -1834,6 +1919,56 @@ mod tests { assert_eq!(output, original); } + #[test] + fn empty_last_block_serializes_after_buffer_callback() { + let mut storage = vec![0xa5u8; 3]; + let expected_start = storage.as_mut_ptr().cast::(); + let mut callback_calls = 0; + let projection = ZSTDMT_emptyBlockJobProjection { + lastJob: 1, + srcSize: 0, + firstJob: 0, + dstStart: ptr::null_mut(), + dstCapacity: 0, + consumed: 0, + }; + + let result = unsafe { + write_last_empty_block_with(projection, || { + callback_calls += 1; + ZstdMtBuffer { + start: expected_start, + capacity: storage.len(), + } + }) + }; + + assert_eq!(callback_calls, 1); + assert_eq!(result.buffer.start, expected_start); + assert_eq!(result.buffer.capacity, storage.len()); + assert_eq!(result.cSize, 3); + assert_eq!(result.clearSource, 1); + assert_eq!(storage, [1, 0, 0]); + } + + #[test] + fn empty_last_block_keeps_source_on_buffer_allocation_failure() { + let projection = ZSTDMT_emptyBlockJobProjection { + lastJob: 1, + srcSize: 0, + firstJob: 0, + dstStart: ptr::null_mut(), + dstCapacity: 0, + consumed: 0, + }; + let result = unsafe { write_last_empty_block_with(projection, ZstdMtBuffer::default) }; + + assert!(result.buffer.start.is_null()); + assert_eq!(result.buffer.capacity, 0); + assert_eq!(result.cSize, ERROR(ZstdErrorCode::MemoryAllocation)); + assert_eq!(result.clearSource, 0); + } + #[test] fn chunk_loop_handles_empty_jobs_without_compression() { let mut compressor = MockChunkCompressor {