diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c index a858fb47a..444b5529d 100644 --- a/lib/compress/zstd_compress.c +++ b/lib/compress/zstd_compress.c @@ -1978,6 +1978,9 @@ ZSTD_inBuffer ZSTD_rust_inBufferForEndFlush(int inBufferMode, size_t expectedPos); size_t ZSTD_rust_endStreamRemaining(size_t remainingToFlush, int frameEnded, int checksumFlag); +size_t ZSTD_rust_endStreamRemainingPolicy(size_t remainingToFlush, + int nbWorkers, + int frameEnded, int checksumFlag); size_t ZSTD_rust_checkBufferStability( int inBufferMode, int outBufferMode, const void* expectedInSrc, size_t expectedInPos, @@ -8676,12 +8679,12 @@ size_t ZSTD_endStream(ZSTD_CStream* zcs, ZSTD_outBuffer* output) ZSTD_inBuffer input = inBuffer_forEndFlush(zcs); size_t const remainingToFlush = ZSTD_compressStream2(zcs, output, &input, ZSTD_e_end); FORWARD_IF_ERROR(remainingToFlush , "ZSTD_compressStream2(,,ZSTD_e_end) failed"); - if (zcs->appliedParams.nbWorkers > 0) return remainingToFlush; /* minimal estimation */ - /* single thread mode : attempt to calculate remaining to flush more precisely */ - { size_t const toFlush = ZSTD_rust_endStreamRemaining( - remainingToFlush, zcs->frameEnded, - zcs->appliedParams.fParams.checksumFlag); - DEBUGLOG(4, "ZSTD_endStream : remaining to flush : %u", (unsigned)toFlush); + { size_t const toFlush = ZSTD_rust_endStreamRemainingPolicy( + remainingToFlush, zcs->appliedParams.nbWorkers, + zcs->frameEnded, zcs->appliedParams.fParams.checksumFlag); + if (zcs->appliedParams.nbWorkers <= 0) { + DEBUGLOG(4, "ZSTD_endStream : remaining to flush : %u", (unsigned)toFlush); + } return toFlush; } } diff --git a/rust/src/zstd_compress.rs b/rust/src/zstd_compress.rs index ff951ab48..42f7c56af 100644 --- a/rust/src/zstd_compress.rs +++ b/rust/src/zstd_compress.rs @@ -6495,6 +6495,33 @@ pub extern "C" fn ZSTD_rust_endStreamRemaining( end_stream_remaining(remaining_to_flush, frame_ended, checksum_flag) } +#[inline] +fn end_stream_remaining_policy( + remaining_to_flush: usize, + nb_workers: c_int, + frame_ended: c_int, + checksum_flag: c_int, +) -> usize { + if nb_workers > 0 { + return remaining_to_flush; + } + + ZSTD_rust_endStreamRemaining(remaining_to_flush, frame_ended, checksum_flag) +} + +/// Select the end-stream remaining-output result without crossing C context +/// state. Multithreaded streams keep the minimal estimate; single-threaded +/// streams use the more precise scalar arithmetic above. +#[no_mangle] +pub extern "C" fn ZSTD_rust_endStreamRemainingPolicy( + remaining_to_flush: usize, + nb_workers: c_int, + frame_ended: c_int, + checksum_flag: c_int, +) -> usize { + end_stream_remaining_policy(remaining_to_flush, nb_workers, frame_ended, checksum_flag) +} + #[inline] fn check_buffer_stability( in_buffer_mode: c_int, @@ -15253,6 +15280,31 @@ mod tests { assert_eq!(ZSTD_rust_endStreamRemaining(usize::MAX, 0, 1), 6); } + #[test] + fn end_stream_remaining_policy_returns_minimal_estimate_for_workers() { + for frame_ended in [0, 1] { + for checksum_flag in [0, 1] { + assert_eq!( + end_stream_remaining_policy(17, 1, frame_ended, checksum_flag), + 17 + ); + assert_eq!( + ZSTD_rust_endStreamRemainingPolicy(17, 1, frame_ended, checksum_flag), + 17 + ); + } + } + } + + #[test] + fn end_stream_remaining_policy_estimates_single_thread_combinations() { + assert_eq!(end_stream_remaining_policy(17, 0, 1, 1), 17); + assert_eq!(end_stream_remaining_policy(17, 0, 1, 0), 17); + assert_eq!(end_stream_remaining_policy(17, 0, 0, 0), 20); + assert_eq!(end_stream_remaining_policy(17, 0, 0, 1), 24); + assert_eq!(ZSTD_rust_endStreamRemainingPolicy(17, 0, 0, 1), 24); + } + #[test] fn check_buffer_stability_accepts_matching_stable_input() { let expected_src = b"input".as_ptr().cast::();