ddidderr/zstd-rs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(compress): move MT stream initialization policy into Rust

Browse Source 
Move the high-level ZSTDMT_initCStream_internal setup policy into Rust. Rust
now owns worker-count resizing decisions, job-size normalization, unfinished-
job draining order, overlap and section sizing, rsync setup, buffer sizing,
and stream reset sequencing through a scalar projection and callbacks.

Keep MT contexts, pools, job resources, dictionaries, buffers, synchronization,
and serial state private to C. C callbacks perform those private mutations while
Rust controls the transparent initialization flow and can test its normalization
and ordering independently of the private layouts.

Test Plan:
- cargo test --manifest-path rust/Cargo.toml --all-targets -- --test-threads=1
- cargo test --manifest-path rust/cli/Cargo.toml --all-targets -- --test-threads=1
- run the legacy Rust feature matrix and all six library/CLI clippy gates with
  -D warnings
- run lib and program native rebuilds plus test-cli-tests,
  test-rust-lib-smoke, and test-zstd with make -j1
- run fuzzer, zstream, and decode-corpus stress gates serially with
  ulimit -v 41943040

Commit is intentionally unsigned because GPG pinentry hangs in this
non-interactive environment.
This commit is contained in:
ddidderr
2026-07-19 09:43:28 +02:00
parent a9589d2d7d
commit 15d42bcc95
2 changed files with 704 additions and 130 deletions
Download Patch File Download Diff File Expand all files Collapse all files
lib/compress/zstdmt_compress.c
+222 -130
View File
@@ -371,6 +371,44 @@ void ZSTDMT_rust_findSynchronizationPoint(const void* inputSrc, size_t inputSize
U64 ZSTDMT_rust_rollingHashPrimePower(U32 length);
size_t ZSTDMT_rust_nextInputSizeHint(size_t targetSectionSize,
size_t inBuffFilled);
typedef struct {
unsigned requestedNbWorkers;
unsigned currentNbWorkers;
size_t jobSize;
size_t jobSizeMin;
size_t jobSizeMax;
int enableLdm;
unsigned windowLog;
unsigned chainLog;
int strategy;
int overlapLog;
int rsyncable;
size_t roundBuffCapacity;
unsigned allJobsCompleted;
} ZSTDMT_RustInitCStreamProjection;
typedef size_t (*ZSTDMT_initResizeFn)(void* opaque, unsigned nbWorkers);
typedef void (*ZSTDMT_initDrainFn)(void* opaque);
typedef void (*ZSTDMT_initApplyParametersFn)(void* opaque, size_t jobSize);
typedef size_t (*ZSTDMT_initDictionaryFn)(void* opaque);
typedef void (*ZSTDMT_initSetSizeFn)(void* opaque, size_t size);
typedef void (*ZSTDMT_initSetRsyncFn)(void* opaque, U64 hitMask, U64 primePower);
typedef void (*ZSTDMT_initSetBufferSizeFn)(void* opaque, size_t size);
typedef size_t (*ZSTDMT_initResizeRoundBufferFn)(void* opaque, size_t capacity);
typedef void (*ZSTDMT_initResetStreamFn)(void* opaque);
typedef size_t (*ZSTDMT_initSerialResetFn)(void* opaque, size_t targetSectionSize);
size_t ZSTDMT_rust_initCStream(
const ZSTDMT_RustInitCStreamProjection* projection, void* opaque,
ZSTDMT_initResizeFn resize, ZSTDMT_initDrainFn drain,
ZSTDMT_initApplyParametersFn applyParameters,
ZSTDMT_initDictionaryFn prepareDictionary,
ZSTDMT_initSetSizeFn setTargetPrefixSize,
ZSTDMT_initSetSizeFn setTargetSectionSize,
ZSTDMT_initSetRsyncFn setRsync,
ZSTDMT_initSetBufferSizeFn setBufferSize,
ZSTDMT_initResizeRoundBufferFn resizeRoundBuffer,
ZSTDMT_initResetStreamFn resetStream,
ZSTDMT_initDictionaryFn updateDictionary,
ZSTDMT_initSerialResetFn serialReset);
typedef struct {
size_t consumed;
size_t cSize;
@@ -1524,26 +1562,159 @@ size_t ZSTDMT_toFlushNow(ZSTDMT_CCtx* mtctx)
/* ===== Multi-threaded compression ===== */
/* ------------------------------------------ */
static unsigned ZSTDMT_computeTargetJobLog(const ZSTD_CCtx_params* params)
typedef struct {
ZSTDMT_CCtx* mtctx;
ZSTD_CCtx_params params;
const void* dict;
size_t dictSize;
ZSTD_dictContentType_e dictContentType;
const ZSTD_CDict* cdict;
unsigned long long pledgedSrcSize;
} ZSTDMT_initCStreamState;
static size_t ZSTDMT_initCStreamResize(void* opaque, unsigned nbWorkers)
{
return ZSTDMT_rust_computeTargetJobLog(params->cParams.windowLog,
params->cParams.chainLog,
(int)params->cParams.strategy,
(int)params->ldmParams.enableLdm);
ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque;
return ZSTDMT_resize(state->mtctx, nbWorkers);
}
static size_t ZSTDMT_computeOverlapSize(const ZSTD_CCtx_params* params)
static void ZSTDMT_initCStreamDrain(void* opaque)
{
size_t overlapSize;
assert(0 <= params->overlapLog && params->overlapLog <= 9);
overlapSize = ZSTDMT_rust_computeOverlapSize(params->cParams.windowLog,
params->cParams.chainLog,
(int)params->cParams.strategy,
params->overlapLog,
(int)params->ldmParams.enableLdm);
DEBUGLOG(4, "overlapLog : %i", params->overlapLog);
DEBUGLOG(4, "overlap size : %i", (int)(overlapSize == 0 ? 1 : overlapSize));
return overlapSize;
ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque;
ZSTDMT_CCtx* const mtctx = state->mtctx;
ZSTDMT_waitForAllJobsCompleted(mtctx);
ZSTDMT_releaseAllJobResources(mtctx);
mtctx->allJobsCompleted = 1;
}
static void ZSTDMT_initCStreamApplyParameters(void* opaque, size_t jobSize)
{
ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque;
state->params.jobSize = jobSize;
state->mtctx->params = state->params;
state->mtctx->frameContentSize = state->pledgedSrcSize;
}
static size_t ZSTDMT_initCStreamPrepareDictionary(void* opaque)
{
ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque;
ZSTDMT_CCtx* const mtctx = state->mtctx;
ZSTD_freeCDict(mtctx->cdictLocal);
if (state->dict) {
mtctx->cdictLocal = ZSTD_createCDict_advanced(
state->dict, state->dictSize, ZSTD_dlm_byCopy,
state->dictContentType, state->params.cParams, mtctx->cMem);
mtctx->cdict = mtctx->cdictLocal;
if (mtctx->cdictLocal == NULL) return ERROR(memory_allocation);
} else {
mtctx->cdictLocal = NULL;
mtctx->cdict = state->cdict;
}
return 0;
}
static void ZSTDMT_initCStreamSetTargetPrefixSize(void* opaque, size_t size)
{
ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque;
state->mtctx->targetPrefixSize = size;
DEBUGLOG(4, "overlapLog=%i => %u KB", state->params.overlapLog, (U32)(size >> 10));
}
static void ZSTDMT_initCStreamSetTargetSectionSize(void* opaque, size_t size)
{
ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque;
state->mtctx->targetSectionSize = size;
}
static void ZSTDMT_initCStreamSetRsync(void* opaque, U64 hitMask, U64 primePower)
{
ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque;
state->mtctx->rsync.hash = 0;
state->mtctx->rsync.hitMask = hitMask;
state->mtctx->rsync.primePower = primePower;
DEBUGLOG(4, "rsyncLog = %u", ZSTD_highbit32((U32)(hitMask + 1)));
}
static void ZSTDMT_initCStreamSetBufferSize(void* opaque, size_t size)
{
ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque;
DEBUGLOG(4, "Job Size : %u KB (note : set to %u)",
(U32)(state->mtctx->targetSectionSize >> 10),
(U32)state->params.jobSize);
DEBUGLOG(4, "inBuff Size : %u KB", (U32)(state->mtctx->targetSectionSize >> 10));
ZSTDMT_setBufferSize(state->mtctx->bufPool, size);
}
static size_t ZSTDMT_initCStreamResizeRoundBuffer(void* opaque, size_t capacity)
{
ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque;
ZSTDMT_CCtx* const mtctx = state->mtctx;
if (mtctx->roundBuff.capacity < capacity) {
if (mtctx->roundBuff.buffer)
ZSTD_customFree(mtctx->roundBuff.buffer, mtctx->cMem);
mtctx->roundBuff.buffer = (BYTE*)ZSTD_customMalloc(capacity, mtctx->cMem);
if (mtctx->roundBuff.buffer == NULL) {
mtctx->roundBuff.capacity = 0;
return ERROR(memory_allocation);
}
mtctx->roundBuff.capacity = capacity;
}
return 0;
}
static void ZSTDMT_initCStreamResetStream(void* opaque)
{
ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque;
ZSTDMT_CCtx* const mtctx = state->mtctx;
DEBUGLOG(4, "roundBuff capacity : %u KB", (U32)(mtctx->roundBuff.capacity >> 10));
mtctx->roundBuff.pos = 0;
mtctx->inBuff.buffer = g_nullBuffer;
mtctx->inBuff.filled = 0;
mtctx->inBuff.prefix = kNullRange;
mtctx->doneJobID = 0;
mtctx->nextJobID = 0;
mtctx->frameEnded = 0;
mtctx->allJobsCompleted = 0;
mtctx->consumed = 0;
mtctx->produced = 0;
}
static size_t ZSTDMT_initCStreamUpdateDictionary(void* opaque)
{
ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque;
ZSTDMT_CCtx* const mtctx = state->mtctx;
ZSTD_freeCDict(mtctx->cdictLocal);
mtctx->cdictLocal = NULL;
mtctx->cdict = NULL;
if (state->dict) {
if (state->dictContentType == ZSTD_dct_rawContent) {
mtctx->inBuff.prefix.start = (const BYTE*)state->dict;
mtctx->inBuff.prefix.size = state->dictSize;
} else {
/* note : a loadPrefix becomes an internal CDict */
mtctx->cdictLocal = ZSTD_createCDict_advanced(
state->dict, state->dictSize, ZSTD_dlm_byRef,
state->dictContentType, state->params.cParams, mtctx->cMem);
mtctx->cdict = mtctx->cdictLocal;
if (mtctx->cdictLocal == NULL) return ERROR(memory_allocation);
}
} else {
mtctx->cdict = state->cdict;
}
return 0;
}
static size_t ZSTDMT_initCStreamSerialReset(void* opaque, size_t targetSectionSize)
{
ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque;
if (ZSTDMT_serialState_reset(
&state->mtctx->serial, state->mtctx->seqPool, state->params,
targetSectionSize, state->dict, state->dictSize,
state->dictContentType))
return ERROR(memory_allocation);
return 0;
}
/* ====================================== */
@@ -1556,6 +1727,9 @@ size_t ZSTDMT_initCStream_internal(
const ZSTD_CDict* cdict, ZSTD_CCtx_params params,
unsigned long long pledgedSrcSize)
{
ZSTDMT_initCStreamState state;
ZSTDMT_RustInitCStreamProjection projection;
DEBUGLOG(4, "ZSTDMT_initCStream_internal (pledgedSrcSize=%u, nbWorkers=%u, cctxPool=%u)",
(U32)pledgedSrcSize, params.nbWorkers, mtctx->cctxPool->totalCCtx);
@@ -1563,121 +1737,39 @@ size_t ZSTDMT_initCStream_internal(
assert(!ZSTD_isError(ZSTD_checkCParams(params.cParams)));
assert(!((dict) && (cdict))); /* either dict or cdict, not both */
/* init */
if (params.nbWorkers != mtctx->params.nbWorkers)
FORWARD_IF_ERROR( ZSTDMT_resize(mtctx, (unsigned)params.nbWorkers) , "");
state = (ZSTDMT_initCStreamState){
mtctx, params, dict, dictSize, dictContentType, cdict, pledgedSrcSize
};
projection = (ZSTDMT_RustInitCStreamProjection){
(unsigned)params.nbWorkers,
(unsigned)mtctx->params.nbWorkers,
params.jobSize,
ZSTDMT_JOBSIZE_MIN,
(size_t)ZSTDMT_JOBSIZE_MAX,
params.ldmParams.enableLdm,
params.cParams.windowLog,
params.cParams.chainLog,
params.cParams.strategy,
params.overlapLog,
params.rsyncable,
mtctx->roundBuff.capacity,
mtctx->allJobsCompleted
};
if (params.jobSize != 0 && params.jobSize < ZSTDMT_JOBSIZE_MIN) params.jobSize = ZSTDMT_JOBSIZE_MIN;
if (params.jobSize > (size_t)ZSTDMT_JOBSIZE_MAX) params.jobSize = (size_t)ZSTDMT_JOBSIZE_MAX;
if (mtctx->allJobsCompleted == 0) { /* previous compression not correctly finished */
ZSTDMT_waitForAllJobsCompleted(mtctx);
ZSTDMT_releaseAllJobResources(mtctx);
mtctx->allJobsCompleted = 1;
}
mtctx->params = params;
mtctx->frameContentSize = pledgedSrcSize;
ZSTD_freeCDict(mtctx->cdictLocal);
if (dict) {
mtctx->cdictLocal = ZSTD_createCDict_advanced(dict, dictSize,
ZSTD_dlm_byCopy, dictContentType, /* note : a loadPrefix becomes an internal CDict */
params.cParams, mtctx->cMem);
mtctx->cdict = mtctx->cdictLocal;
if (mtctx->cdictLocal == NULL) return ERROR(memory_allocation);
} else {
mtctx->cdictLocal = NULL;
mtctx->cdict = cdict;
}
mtctx->targetPrefixSize = ZSTDMT_computeOverlapSize(&params);
DEBUGLOG(4, "overlapLog=%i => %u KB", params.overlapLog, (U32)(mtctx->targetPrefixSize>>10));
mtctx->targetSectionSize = params.jobSize;
if (mtctx->targetSectionSize == 0) {
mtctx->targetSectionSize = 1ULL << ZSTDMT_computeTargetJobLog(&params);
}
assert(mtctx->targetSectionSize <= (size_t)ZSTDMT_JOBSIZE_MAX);
if (params.rsyncable) {
/* Aim for the targetsectionSize as the average job size. */
U32 const jobSizeKB = (U32)(mtctx->targetSectionSize >> 10);
U32 const rsyncBits = (assert(jobSizeKB >= 1), ZSTD_highbit32(jobSizeKB) + 10);
/* We refuse to create jobs < RSYNC_MIN_BLOCK_SIZE bytes, so make sure our
* expected job size is at least 4x larger. */
assert(rsyncBits >= RSYNC_MIN_BLOCK_LOG + 2);
DEBUGLOG(4, "rsyncLog = %u", rsyncBits);
mtctx->rsync.hash = 0;
mtctx->rsync.hitMask = (1ULL << rsyncBits) - 1;
mtctx->rsync.primePower = ZSTDMT_rust_rollingHashPrimePower(RSYNC_LENGTH);
}
if (mtctx->targetSectionSize < mtctx->targetPrefixSize) mtctx->targetSectionSize = mtctx->targetPrefixSize; /* job size must be >= overlap size */
DEBUGLOG(4, "Job Size : %u KB (note : set to %u)", (U32)(mtctx->targetSectionSize>>10), (U32)params.jobSize);
DEBUGLOG(4, "inBuff Size : %u KB", (U32)(mtctx->targetSectionSize>>10));
ZSTDMT_setBufferSize(mtctx->bufPool, ZSTD_compressBound(mtctx->targetSectionSize));
{
/* If ldm is enabled we need windowSize space. */
size_t const windowSize = mtctx->params.ldmParams.enableLdm == ZSTD_ps_enable ? (1U << mtctx->params.cParams.windowLog) : 0;
/* Two buffers of slack, plus extra space for the overlap
* This is the minimum slack that LDM works with. One extra because
* flush might waste up to targetSectionSize-1 bytes. Another extra
* for the overlap (if > 0), then one to fill which doesn't overlap
* with the LDM window.
*/
size_t const nbSlackBuffers = 2 + (mtctx->targetPrefixSize > 0);
size_t const slackSize = mtctx->targetSectionSize * nbSlackBuffers;
/* Compute the total size, and always have enough slack */
size_t const nbWorkers = MAX(mtctx->params.nbWorkers, 1);
size_t const sectionsSize = mtctx->targetSectionSize * nbWorkers;
size_t const capacity = MAX(windowSize, sectionsSize) + slackSize;
if (mtctx->roundBuff.capacity < capacity) {
if (mtctx->roundBuff.buffer)
ZSTD_customFree(mtctx->roundBuff.buffer, mtctx->cMem);
mtctx->roundBuff.buffer = (BYTE*)ZSTD_customMalloc(capacity, mtctx->cMem);
if (mtctx->roundBuff.buffer == NULL) {
mtctx->roundBuff.capacity = 0;
return ERROR(memory_allocation);
}
mtctx->roundBuff.capacity = capacity;
}
}
DEBUGLOG(4, "roundBuff capacity : %u KB", (U32)(mtctx->roundBuff.capacity>>10));
mtctx->roundBuff.pos = 0;
mtctx->inBuff.buffer = g_nullBuffer;
mtctx->inBuff.filled = 0;
mtctx->inBuff.prefix = kNullRange;
mtctx->doneJobID = 0;
mtctx->nextJobID = 0;
mtctx->frameEnded = 0;
mtctx->allJobsCompleted = 0;
mtctx->consumed = 0;
mtctx->produced = 0;
/* update dictionary */
ZSTD_freeCDict(mtctx->cdictLocal);
mtctx->cdictLocal = NULL;
mtctx->cdict = NULL;
if (dict) {
if (dictContentType == ZSTD_dct_rawContent) {
mtctx->inBuff.prefix.start = (const BYTE*)dict;
mtctx->inBuff.prefix.size = dictSize;
} else {
/* note : a loadPrefix becomes an internal CDict */
mtctx->cdictLocal = ZSTD_createCDict_advanced(dict, dictSize,
ZSTD_dlm_byRef, dictContentType,
params.cParams, mtctx->cMem);
mtctx->cdict = mtctx->cdictLocal;
if (mtctx->cdictLocal == NULL) return ERROR(memory_allocation);
}
} else {
mtctx->cdict = cdict;
}
if (ZSTDMT_serialState_reset(&mtctx->serial, mtctx->seqPool, params, mtctx->targetSectionSize,
dict, dictSize, dictContentType))
return ERROR(memory_allocation);
return 0;
return ZSTDMT_rust_initCStream(
&projection, &state,
ZSTDMT_initCStreamResize,
ZSTDMT_initCStreamDrain,
ZSTDMT_initCStreamApplyParameters,
ZSTDMT_initCStreamPrepareDictionary,
ZSTDMT_initCStreamSetTargetPrefixSize,
ZSTDMT_initCStreamSetTargetSectionSize,
ZSTDMT_initCStreamSetRsync,
ZSTDMT_initCStreamSetBufferSize,
ZSTDMT_initCStreamResizeRoundBuffer,
ZSTDMT_initCStreamResetStream,
ZSTDMT_initCStreamUpdateDictionary,
ZSTDMT_initCStreamSerialReset);
}
rust/src/zstdmt_compress.rs
+482
View File
@@ -19,11 +19,20 @@ use std::os::raw::{c_int, c_uint, c_void};
use std::ptr;
use std::sync::Mutex;
use crate::bits::ZSTD_highbit32;
use crate::errors::{ERR_isError, ZstdErrorCode, ERROR};
use crate::zstd_compress::ZSTD_frameProgression;
const ZSTDMT_JOBLOG_MAX: c_uint = if mem::size_of::<usize>() == 4 { 29 } else { 30 };
const ZSTD_WINDOWLOG_MAX: c_uint = if mem::size_of::<usize>() == 4 { 30 } else { 31 };
#[cfg(test)]
const DEFAULT_ZSTDMT_JOBSIZE_MIN: usize = 512 << 10;
#[cfg(test)]
const DEFAULT_ZSTDMT_JOBSIZE_MAX: usize = if mem::size_of::<usize>() == 4 {
512 << 20
} else {
1024 << 20
};
const ZSTD_FAST: c_int = 1;
const ZSTD_DFAST: c_int = 2;
@@ -75,6 +84,39 @@ pub type ZSTDMT_compressionJobCompressFn =
pub type ZSTDMT_compressionJobErrorFn = unsafe extern "C" fn(*mut c_void, usize);
pub type ZSTDMT_compressionJobFinishFn = unsafe extern "C" fn(*mut c_void, usize);
/// Scalar inputs for the MT streaming initializer. The full parameter
/// object, dictionary handles, pools, buffers, and synchronization remain
/// private to C. Rust owns the order in which the C callbacks are invoked and
/// the pure normalization/sizing decisions between those callbacks.
#[repr(C)]
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct ZSTDMT_initCStreamProjection {
pub requestedNbWorkers: c_uint,
pub currentNbWorkers: c_uint,
pub jobSize: usize,
pub jobSizeMin: usize,
pub jobSizeMax: usize,
pub enableLdm: c_int,
pub windowLog: c_uint,
pub chainLog: c_uint,
pub strategy: c_int,
pub overlapLog: c_int,
pub rsyncable: c_int,
pub roundBuffCapacity: usize,
pub allJobsCompleted: c_uint,
}
pub type ZSTDMT_initResizeFn = unsafe extern "C" fn(*mut c_void, c_uint) -> usize;
pub type ZSTDMT_initDrainFn = unsafe extern "C" fn(*mut c_void);
pub type ZSTDMT_initApplyParametersFn = unsafe extern "C" fn(*mut c_void, usize);
pub type ZSTDMT_initDictionaryFn = unsafe extern "C" fn(*mut c_void) -> usize;
pub type ZSTDMT_initSetSizeFn = unsafe extern "C" fn(*mut c_void, usize);
pub type ZSTDMT_initSetRsyncFn = unsafe extern "C" fn(*mut c_void, u64, u64);
pub type ZSTDMT_initSetBufferSizeFn = unsafe extern "C" fn(*mut c_void, usize);
pub type ZSTDMT_initResizeRoundBufferFn = unsafe extern "C" fn(*mut c_void, usize) -> usize;
pub type ZSTDMT_initResetStreamFn = unsafe extern "C" fn(*mut c_void);
pub type ZSTDMT_initSerialResetFn = unsafe extern "C" fn(*mut c_void, usize) -> usize;
#[repr(C)]
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct ZSTDMT_flushPublicationResult {
@@ -708,6 +750,244 @@ pub unsafe extern "C" fn ZSTDMT_rust_compressionJob(
);
}
#[inline]
fn normalize_mt_job_size(job_size: usize, job_size_min: usize, job_size_max: usize) -> usize {
if job_size != 0 && job_size < job_size_min {
job_size_min
} else if job_size > job_size_max {
job_size_max
} else {
job_size
}
}
#[inline]
fn init_round_buffer_capacity(
projection: ZSTDMT_initCStreamProjection,
target_prefix_size: usize,
target_section_size: usize,
) -> usize {
let window_size = if projection.enableLdm == ZSTD_PS_ENABLE {
1usize << projection.windowLog as usize
} else {
0
};
let nb_slack_buffers = 2 + usize::from(target_prefix_size > 0);
let slack_size = target_section_size.wrapping_mul(nb_slack_buffers);
let nb_workers = projection.requestedNbWorkers.max(1) as usize;
let sections_size = target_section_size.wrapping_mul(nb_workers);
window_size.max(sections_size).wrapping_add(slack_size)
}
#[inline]
fn init_rsync_parameters(target_section_size: usize) -> (u64, u64) {
let job_size_kb = (target_section_size >> 10) as u32;
debug_assert!(job_size_kb >= 1);
let rsync_bits = ZSTD_highbit32(job_size_kb) + 10;
debug_assert!(rsync_bits >= (RSYNC_MIN_BLOCK_LOG + 2) as u32);
let hit_mask = 1u64.wrapping_shl(rsync_bits).wrapping_sub(1);
(hit_mask, rolling_hash_prime_power(RSYNC_LENGTH as c_uint))
}
/// Run the high-level MT streaming initialization policy while C retains all
/// private state and side effects behind callbacks. The callback order is
/// intentionally the order of the original C initializer: resize, normalize,
/// drain, attach the first dictionary, size buffers, reset stream state,
/// update the active dictionary, and finally reset serial state.
#[inline]
fn init_c_stream_with<
Resize,
Drain,
ApplyParams,
PrepareDictionary,
SetPrefixSize,
SetSectionSize,
SetRsync,
SetBufferSize,
ResizeRoundBuffer,
ResetStream,
UpdateDictionary,
SerialReset,
>(
projection: ZSTDMT_initCStreamProjection,
mut resize: Resize,
mut drain: Drain,
mut apply_parameters: ApplyParams,
mut prepare_dictionary: PrepareDictionary,
mut set_target_prefix_size: SetPrefixSize,
mut set_target_section_size: SetSectionSize,
mut set_rsync: SetRsync,
mut set_buffer_size: SetBufferSize,
mut resize_round_buffer: ResizeRoundBuffer,
mut reset_stream: ResetStream,
mut update_dictionary: UpdateDictionary,
mut serial_reset: SerialReset,
) -> usize
where
Resize: FnMut(c_uint) -> usize,
Drain: FnMut(),
ApplyParams: FnMut(usize),
PrepareDictionary: FnMut() -> usize,
SetPrefixSize: FnMut(usize),
SetSectionSize: FnMut(usize),
SetRsync: FnMut(u64, u64),
SetBufferSize: FnMut(usize),
ResizeRoundBuffer: FnMut(usize) -> usize,
ResetStream: FnMut(),
UpdateDictionary: FnMut() -> usize,
SerialReset: FnMut(usize) -> usize,
{
if projection.requestedNbWorkers != projection.currentNbWorkers {
let error = resize(projection.requestedNbWorkers);
if ERR_isError(error) {
return error;
}
}
let job_size = normalize_mt_job_size(
projection.jobSize,
projection.jobSizeMin,
projection.jobSizeMax,
);
if projection.allJobsCompleted == 0 {
drain();
}
apply_parameters(job_size);
let error = prepare_dictionary();
if ERR_isError(error) {
return error;
}
let target_prefix_size = compute_overlap_size(
projection.windowLog,
projection.chainLog,
projection.strategy,
projection.overlapLog,
projection.enableLdm,
);
set_target_prefix_size(target_prefix_size);
let initial_target_section_size = if job_size == 0 {
1usize
<< compute_target_job_log(
projection.windowLog,
projection.chainLog,
projection.strategy,
projection.enableLdm,
) as usize
} else {
job_size
};
debug_assert!(initial_target_section_size <= projection.jobSizeMax);
set_target_section_size(initial_target_section_size);
if projection.rsyncable != 0 {
let (hit_mask, prime_power) = init_rsync_parameters(initial_target_section_size);
set_rsync(hit_mask, prime_power);
}
let target_section_size = initial_target_section_size.max(target_prefix_size);
if target_section_size != initial_target_section_size {
set_target_section_size(target_section_size);
}
set_buffer_size(crate::zstd_compress_api::ZSTD_compressBound(
target_section_size,
));
let round_capacity =
init_round_buffer_capacity(projection, target_prefix_size, target_section_size);
if projection.roundBuffCapacity < round_capacity {
let error = resize_round_buffer(round_capacity);
if ERR_isError(error) {
return error;
}
}
reset_stream();
let error = update_dictionary();
if ERR_isError(error) {
return error;
}
serial_reset(target_section_size)
}
/// C ABI entry point for the MT streaming initializer. C owns every
/// allocation, dictionary handle, synchronization object, and private context
/// mutation; this wrapper only connects those operations to the Rust policy.
#[cfg(not(test))]
#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_initCStream(
projection: *const ZSTDMT_initCStreamProjection,
opaque: *mut c_void,
resize: Option<ZSTDMT_initResizeFn>,
drain: Option<ZSTDMT_initDrainFn>,
applyParameters: Option<ZSTDMT_initApplyParametersFn>,
prepareDictionary: Option<ZSTDMT_initDictionaryFn>,
setTargetPrefixSize: Option<ZSTDMT_initSetSizeFn>,
setTargetSectionSize: Option<ZSTDMT_initSetSizeFn>,
setRsync: Option<ZSTDMT_initSetRsyncFn>,
setBufferSize: Option<ZSTDMT_initSetBufferSizeFn>,
resizeRoundBuffer: Option<ZSTDMT_initResizeRoundBufferFn>,
resetStream: Option<ZSTDMT_initResetStreamFn>,
updateDictionary: Option<ZSTDMT_initDictionaryFn>,
serialReset: Option<ZSTDMT_initSerialResetFn>,
) -> usize {
let Some(projection) = (unsafe { projection.as_ref() }).copied() else {
return ERROR(ZstdErrorCode::Generic);
};
let (
Some(resize),
Some(drain),
Some(apply_parameters),
Some(prepare_dictionary),
Some(set_target_prefix_size),
Some(set_target_section_size),
Some(set_rsync),
Some(set_buffer_size),
Some(resize_round_buffer),
Some(reset_stream),
Some(update_dictionary),
Some(serial_reset),
) = (
resize,
drain,
applyParameters,
prepareDictionary,
setTargetPrefixSize,
setTargetSectionSize,
setRsync,
setBufferSize,
resizeRoundBuffer,
resetStream,
updateDictionary,
serialReset,
)
else {
return ERROR(ZstdErrorCode::Generic);
};
init_c_stream_with(
projection,
|nb_workers| unsafe { resize(opaque, nb_workers) },
|| unsafe { drain(opaque) },
|job_size| unsafe { apply_parameters(opaque, job_size) },
|| unsafe { prepare_dictionary(opaque) },
|size| unsafe { set_target_prefix_size(opaque, size) },
|size| unsafe { set_target_section_size(opaque, size) },
|hit_mask, prime_power| unsafe { set_rsync(opaque, hit_mask, prime_power) },
|size| unsafe { set_buffer_size(opaque, size) },
|capacity| unsafe { resize_round_buffer(opaque, capacity) },
|| unsafe { reset_stream(opaque) },
|| unsafe { update_dictionary(opaque) },
|target_section_size| unsafe { serial_reset(opaque, target_section_size) },
)
}
#[inline]
fn invalid_flush_publication(
output_pos: usize,
@@ -2875,6 +3155,208 @@ mod tests {
assert_eq!(state.finished, vec![0]);
}
fn init_projection() -> ZSTDMT_initCStreamProjection {
ZSTDMT_initCStreamProjection {
requestedNbWorkers: 4,
currentNbWorkers: 2,
jobSize: 1,
jobSizeMin: DEFAULT_ZSTDMT_JOBSIZE_MIN,
jobSizeMax: DEFAULT_ZSTDMT_JOBSIZE_MAX,
enableLdm: 0,
windowLog: 20,
chainLog: 16,
strategy: ZSTD_FAST,
overlapLog: 1,
rsyncable: 1,
roundBuffCapacity: 0,
allJobsCompleted: 0,
}
}
#[test]
fn init_c_stream_preserves_success_order_and_normalization() {
let projection = init_projection();
let events = Rc::new(RefCell::new(Vec::<&'static str>::new()));
let values = Rc::new(RefCell::new(Vec::<(&'static str, usize)>::new()));
let resize_events = Rc::clone(&events);
let drain_events = Rc::clone(&events);
let apply_events = Rc::clone(&events);
let apply_values = Rc::clone(&values);
let prepare_events = Rc::clone(&events);
let prefix_events = Rc::clone(&events);
let prefix_values = Rc::clone(&values);
let section_events = Rc::clone(&events);
let section_values = Rc::clone(&values);
let rsync_events = Rc::clone(&events);
let buffer_events = Rc::clone(&events);
let buffer_values = Rc::clone(&values);
let round_events = Rc::clone(&events);
let round_values = Rc::clone(&values);
let reset_events = Rc::clone(&events);
let update_events = Rc::clone(&events);
let serial_events = Rc::clone(&events);
let serial_values = Rc::clone(&values);
let result = init_c_stream_with(
projection,
move |workers| {
assert_eq!(workers, 4);
resize_events.borrow_mut().push("resize");
0
},
move || drain_events.borrow_mut().push("drain"),
move |job_size| {
apply_events.borrow_mut().push("apply");
apply_values.borrow_mut().push(("job", job_size));
},
move || {
prepare_events.borrow_mut().push("prepare-dict");
0
},
move |size| {
prefix_events.borrow_mut().push("prefix");
prefix_values.borrow_mut().push(("prefix", size));
},
move |size| {
section_events.borrow_mut().push("section");
section_values.borrow_mut().push(("section", size));
},
move |_hit_mask, _prime_power| rsync_events.borrow_mut().push("rsync"),
move |size| {
buffer_events.borrow_mut().push("buffer");
buffer_values.borrow_mut().push(("buffer", size));
},
move |capacity| {
round_events.borrow_mut().push("round");
round_values.borrow_mut().push(("round", capacity));
0
},
move || reset_events.borrow_mut().push("reset"),
move || {
update_events.borrow_mut().push("update-dict");
0
},
move |size| {
serial_events.borrow_mut().push("serial");
serial_values.borrow_mut().push(("serial", size));
0
},
);
assert_eq!(result, 0);
assert_eq!(
events.borrow().as_slice(),
&[
"resize",
"drain",
"apply",
"prepare-dict",
"prefix",
"section",
"rsync",
"buffer",
"round",
"reset",
"update-dict",
"serial",
]
);
assert_eq!(values.borrow()[0], ("job", DEFAULT_ZSTDMT_JOBSIZE_MIN));
assert_eq!(values.borrow()[1], ("prefix", 0));
assert_eq!(values.borrow()[2], ("section", DEFAULT_ZSTDMT_JOBSIZE_MIN));
assert_eq!(values.borrow()[5], ("serial", DEFAULT_ZSTDMT_JOBSIZE_MIN));
assert!(values.borrow()[3].1 > 0);
assert!(values.borrow()[4].1 > 0);
}
#[test]
fn init_c_stream_stops_before_normalization_when_resize_fails() {
let expected_error = ERROR(ZstdErrorCode::MemoryAllocation);
let result = init_c_stream_with(
init_projection(),
|_| expected_error,
|| panic!("a failed resize must stop initialization"),
|_| panic!("a failed resize must stop initialization"),
|| panic!("a failed resize must stop initialization"),
|_| panic!("a failed resize must stop initialization"),
|_| panic!("a failed resize must stop initialization"),
|_, _| panic!("a failed resize must stop initialization"),
|_| panic!("a failed resize must stop initialization"),
|_| panic!("a failed resize must stop initialization"),
|| panic!("a failed resize must stop initialization"),
|| panic!("a failed resize must stop initialization"),
|_| panic!("a failed resize must stop initialization"),
);
assert_eq!(result, expected_error);
}
#[test]
fn init_c_stream_stops_after_round_buffer_error() {
let expected_error = ERROR(ZstdErrorCode::MemoryAllocation);
let mut projection = init_projection();
projection.currentNbWorkers = projection.requestedNbWorkers;
projection.allJobsCompleted = 1;
let events = Rc::new(RefCell::new(Vec::<&'static str>::new()));
let result = init_c_stream_with(
projection,
|_| panic!("workers already match"),
|| panic!("all jobs are complete"),
{
let events = Rc::clone(&events);
move |_| events.borrow_mut().push("apply")
},
{
let events = Rc::clone(&events);
move || {
events.borrow_mut().push("prepare");
0
}
},
{
let events = Rc::clone(&events);
move |_| events.borrow_mut().push("prefix")
},
{
let events = Rc::clone(&events);
move |_| events.borrow_mut().push("section")
},
{
let events = Rc::clone(&events);
move |_, _| events.borrow_mut().push("rsync")
},
{
let events = Rc::clone(&events);
move |_| events.borrow_mut().push("buffer")
},
{
let events = Rc::clone(&events);
move |_| {
events.borrow_mut().push("round");
expected_error
}
},
{
let events = Rc::clone(&events);
move || events.borrow_mut().push("reset")
},
{
let events = Rc::clone(&events);
move || {
events.borrow_mut().push("update");
0
}
},
|_| panic!("serial reset must follow successful dictionary update"),
);
assert_eq!(result, expected_error);
assert_eq!(
events.borrow().as_slice(),
&["apply", "prepare", "prefix", "section", "rsync", "buffer", "round"]
);
}
unsafe extern "C" fn mock_compress_continue(
cctx: *mut c_void,
_dst: *mut c_void,