From 45cb1f509e30a429ea6d374580037400a002bc86 Mon Sep 17 00:00:00 2001 From: ddidderr Date: Sun, 12 Jul 2026 10:13:41 +0200 Subject: [PATCH] feat(compress): port multithreaded resource pools Move the reusable buffer and compression-context pools behind a Rust implementation with a narrow C ABI adapter. The scheduler, job table, serial LDM state, and stream orchestration remain in C until their private layouts are ported. Keep context types opaque across Rust modules so the pool bridge does not depend on private C layout declarations. This also keeps the existing custom allocator and pool replacement contracts intact. Test Plan: - RUSTC_WRAPPER= CARGO_BUILD_RUSTC_WRAPPER= cargo clippy --manifest-path rust/Cargo.toml --all-targets -- -D warnings - RUSTC_WRAPPER= CARGO_BUILD_RUSTC_WRAPPER= cargo test --manifest-path rust/Cargo.toml zstdmt_compress - make -B -C lib lib-mt - git diff --cached --check --- lib/compress/zstdmt_compress.c | 243 ++++--------- rust/src/dict_builder_cover.rs | 11 +- rust/src/lib.rs | 2 + rust/src/zstdmt_compress.rs | 604 +++++++++++++++++++++++++++++++++ 4 files changed, 681 insertions(+), 179 deletions(-) create mode 100644 rust/src/zstdmt_compress.rs diff --git a/lib/compress/zstdmt_compress.c b/lib/compress/zstdmt_compress.c index 0f1fe6d74..8dd2f9c02 100644 --- a/lib/compress/zstdmt_compress.c +++ b/lib/compress/zstdmt_compress.c @@ -94,28 +94,36 @@ typedef struct buffer_s { static const Buffer g_nullBuffer = { NULL, 0 }; +/* The Rust module owns the synchronized pool state. Keep a C wrapper so + * custom allocation of the containing ZSTDMT context remains unchanged. */ +typedef struct ZSTDMT_RustBufferPool_s ZSTDMT_RustBufferPool; +typedef struct { + void* start; + size_t capacity; +} ZSTDMT_RustBuffer; + +ZSTDMT_RustBufferPool* ZSTDMT_rust_buffer_pool_create(unsigned maxNbBuffers, + ZSTD_customMem cMem); +void ZSTDMT_rust_buffer_pool_free(ZSTDMT_RustBufferPool* pool); +size_t ZSTDMT_rust_buffer_pool_sizeof(const ZSTDMT_RustBufferPool* pool); +void ZSTDMT_rust_buffer_pool_set_size(ZSTDMT_RustBufferPool* pool, size_t bSize); +ZSTDMT_RustBufferPool* ZSTDMT_rust_buffer_pool_expand(ZSTDMT_RustBufferPool* pool, + unsigned maxNbBuffers); +ZSTDMT_RustBuffer ZSTDMT_rust_buffer_pool_get(ZSTDMT_RustBufferPool* pool); +void ZSTDMT_rust_buffer_pool_release(ZSTDMT_RustBufferPool* pool, + ZSTDMT_RustBuffer buffer); +ZSTDMT_RustBuffer ZSTDMT_rust_buffer_pool_resize(ZSTDMT_RustBufferPool* pool, + ZSTDMT_RustBuffer buffer); + typedef struct ZSTDMT_bufferPool_s { - ZSTD_pthread_mutex_t poolMutex; - size_t bufferSize; - unsigned totalBuffers; - unsigned nbBuffers; + ZSTDMT_RustBufferPool* rustPool; ZSTD_customMem cMem; - Buffer* buffers; } ZSTDMT_bufferPool; static void ZSTDMT_freeBufferPool(ZSTDMT_bufferPool* bufPool) { - DEBUGLOG(3, "ZSTDMT_freeBufferPool (address:%08X)", (U32)(size_t)bufPool); if (!bufPool) return; /* compatibility with free on NULL */ - if (bufPool->buffers) { - unsigned u; - for (u=0; utotalBuffers; u++) { - DEBUGLOG(4, "free buffer %2u (address:%08X)", u, (U32)(size_t)bufPool->buffers[u].start); - ZSTD_customFree(bufPool->buffers[u].start, bufPool->cMem); - } - ZSTD_customFree(bufPool->buffers, bufPool->cMem); - } - ZSTD_pthread_mutex_destroy(&bufPool->poolMutex); + ZSTDMT_rust_buffer_pool_free(bufPool->rustPool); ZSTD_customFree(bufPool, bufPool->cMem); } @@ -124,18 +132,11 @@ static ZSTDMT_bufferPool* ZSTDMT_createBufferPool(unsigned maxNbBuffers, ZSTD_cu ZSTDMT_bufferPool* const bufPool = (ZSTDMT_bufferPool*)ZSTD_customCalloc(sizeof(ZSTDMT_bufferPool), cMem); if (bufPool==NULL) return NULL; - if (ZSTD_pthread_mutex_init(&bufPool->poolMutex, NULL)) { + bufPool->rustPool = ZSTDMT_rust_buffer_pool_create(maxNbBuffers, cMem); + if (bufPool->rustPool == NULL) { ZSTD_customFree(bufPool, cMem); return NULL; } - bufPool->buffers = (Buffer*)ZSTD_customCalloc(maxNbBuffers * sizeof(Buffer), cMem); - if (bufPool->buffers==NULL) { - ZSTDMT_freeBufferPool(bufPool); - return NULL; - } - bufPool->bufferSize = 64 KB; - bufPool->totalBuffers = maxNbBuffers; - bufPool->nbBuffers = 0; bufPool->cMem = cMem; return bufPool; } @@ -143,16 +144,8 @@ static ZSTDMT_bufferPool* ZSTDMT_createBufferPool(unsigned maxNbBuffers, ZSTD_cu /* only works at initialization, not during compression */ static size_t ZSTDMT_sizeof_bufferPool(ZSTDMT_bufferPool* bufPool) { - size_t const poolSize = sizeof(*bufPool); - size_t const arraySize = bufPool->totalBuffers * sizeof(Buffer); - unsigned u; - size_t totalBufferSize = 0; - ZSTD_pthread_mutex_lock(&bufPool->poolMutex); - for (u=0; utotalBuffers; u++) - totalBufferSize += bufPool->buffers[u].capacity; - ZSTD_pthread_mutex_unlock(&bufPool->poolMutex); - - return poolSize + arraySize + totalBufferSize; + if (bufPool == NULL) return 0; + return sizeof(*bufPool) + ZSTDMT_rust_buffer_pool_sizeof(bufPool->rustPool); } /* ZSTDMT_setBufferSize() : @@ -161,28 +154,21 @@ static size_t ZSTDMT_sizeof_bufferPool(ZSTDMT_bufferPool* bufPool) * as they become freely interchangeable, reducing malloc/free usages and memory fragmentation */ static void ZSTDMT_setBufferSize(ZSTDMT_bufferPool* const bufPool, size_t const bSize) { - ZSTD_pthread_mutex_lock(&bufPool->poolMutex); DEBUGLOG(4, "ZSTDMT_setBufferSize: bSize = %u", (U32)bSize); - bufPool->bufferSize = bSize; - ZSTD_pthread_mutex_unlock(&bufPool->poolMutex); + ZSTDMT_rust_buffer_pool_set_size(bufPool->rustPool, bSize); } static ZSTDMT_bufferPool* ZSTDMT_expandBufferPool(ZSTDMT_bufferPool* srcBufPool, unsigned maxNbBuffers) { if (srcBufPool==NULL) return NULL; - if (srcBufPool->totalBuffers >= maxNbBuffers) /* good enough */ - return srcBufPool; - /* need a larger buffer pool */ - { ZSTD_customMem const cMem = srcBufPool->cMem; - size_t const bSize = srcBufPool->bufferSize; /* forward parameters */ - ZSTDMT_bufferPool* newBufPool; - ZSTDMT_freeBufferPool(srcBufPool); - newBufPool = ZSTDMT_createBufferPool(maxNbBuffers, cMem); - if (newBufPool==NULL) return newBufPool; - ZSTDMT_setBufferSize(newBufPool, bSize); - return newBufPool; + srcBufPool->rustPool = ZSTDMT_rust_buffer_pool_expand(srcBufPool->rustPool, maxNbBuffers); + if (srcBufPool->rustPool == NULL) { + ZSTD_customMem const cMem = srcBufPool->cMem; + ZSTD_customFree(srcBufPool, cMem); + return NULL; } + return srcBufPool; } /** ZSTDMT_getBuffer() : @@ -191,38 +177,10 @@ static ZSTDMT_bufferPool* ZSTDMT_expandBufferPool(ZSTDMT_bufferPool* srcBufPool, * note: allocation may fail, in this case, start==NULL and size==0 */ static Buffer ZSTDMT_getBuffer(ZSTDMT_bufferPool* bufPool) { - size_t const bSize = bufPool->bufferSize; - DEBUGLOG(5, "ZSTDMT_getBuffer: bSize = %u", (U32)bufPool->bufferSize); - ZSTD_pthread_mutex_lock(&bufPool->poolMutex); - if (bufPool->nbBuffers) { /* try to use an existing buffer */ - Buffer const buf = bufPool->buffers[--(bufPool->nbBuffers)]; - size_t const availBufferSize = buf.capacity; - bufPool->buffers[bufPool->nbBuffers] = g_nullBuffer; - if ((availBufferSize >= bSize) & ((availBufferSize>>3) <= bSize)) { - /* large enough, but not too much */ - DEBUGLOG(5, "ZSTDMT_getBuffer: provide buffer %u of size %u", - bufPool->nbBuffers, (U32)buf.capacity); - ZSTD_pthread_mutex_unlock(&bufPool->poolMutex); - return buf; - } - /* size conditions not respected : scratch this buffer, create new one */ - DEBUGLOG(5, "ZSTDMT_getBuffer: existing buffer does not meet size conditions => freeing"); - ZSTD_customFree(buf.start, bufPool->cMem); - } - ZSTD_pthread_mutex_unlock(&bufPool->poolMutex); - /* create new buffer */ - DEBUGLOG(5, "ZSTDMT_getBuffer: create a new buffer"); - { Buffer buffer; - void* const start = ZSTD_customMalloc(bSize, bufPool->cMem); - buffer.start = start; /* note : start can be NULL if malloc fails ! */ - buffer.capacity = (start==NULL) ? 0 : bSize; - if (start==NULL) { - DEBUGLOG(5, "ZSTDMT_getBuffer: buffer allocation failure !!"); - } else { - DEBUGLOG(5, "ZSTDMT_getBuffer: created buffer of size %u", (U32)bSize); - } - return buffer; - } + ZSTDMT_RustBuffer const rustBuffer = + ZSTDMT_rust_buffer_pool_get(bufPool->rustPool); + Buffer buffer = { rustBuffer.start, rustBuffer.capacity }; + return buffer; } #if ZSTD_RESIZE_SEQPOOL @@ -233,41 +191,19 @@ static Buffer ZSTDMT_getBuffer(ZSTDMT_bufferPool* bufPool) */ static Buffer ZSTDMT_resizeBuffer(ZSTDMT_bufferPool* bufPool, Buffer buffer) { - size_t const bSize = bufPool->bufferSize; - if (buffer.capacity < bSize) { - void* const start = ZSTD_customMalloc(bSize, bufPool->cMem); - Buffer newBuffer; - newBuffer.start = start; - newBuffer.capacity = start == NULL ? 0 : bSize; - if (start != NULL) { - assert(newBuffer.capacity >= buffer.capacity); - ZSTD_memcpy(newBuffer.start, buffer.start, buffer.capacity); - DEBUGLOG(5, "ZSTDMT_resizeBuffer: created buffer of size %u", (U32)bSize); - return newBuffer; - } - DEBUGLOG(5, "ZSTDMT_resizeBuffer: buffer allocation failure !!"); - } - return buffer; + ZSTDMT_RustBuffer const rustBuffer = { buffer.start, buffer.capacity }; + ZSTDMT_RustBuffer const resized = + ZSTDMT_rust_buffer_pool_resize(bufPool->rustPool, rustBuffer); + Buffer const result = { resized.start, resized.capacity }; + return result; } #endif /* store buffer for later re-use, up to pool capacity */ static void ZSTDMT_releaseBuffer(ZSTDMT_bufferPool* bufPool, Buffer buf) { - DEBUGLOG(5, "ZSTDMT_releaseBuffer"); - if (buf.start == NULL) return; /* compatible with release on NULL */ - ZSTD_pthread_mutex_lock(&bufPool->poolMutex); - if (bufPool->nbBuffers < bufPool->totalBuffers) { - bufPool->buffers[bufPool->nbBuffers++] = buf; /* stored for later use */ - DEBUGLOG(5, "ZSTDMT_releaseBuffer: stored buffer of size %u in slot %u", - (U32)buf.capacity, (U32)(bufPool->nbBuffers-1)); - ZSTD_pthread_mutex_unlock(&bufPool->poolMutex); - return; - } - ZSTD_pthread_mutex_unlock(&bufPool->poolMutex); - /* Reached bufferPool capacity (note: should not happen) */ - DEBUGLOG(5, "ZSTDMT_releaseBuffer: pool capacity reached => freeing "); - ZSTD_customFree(buf.start, bufPool->cMem); + ZSTDMT_RustBuffer const rustBuffer = { buf.start, buf.capacity }; + ZSTDMT_rust_buffer_pool_release(bufPool->rustPool, rustBuffer); } /* We need 2 output buffers per worker since each dstBuff must be flushed after it is released. @@ -308,9 +244,6 @@ static Buffer seqToBuffer(RawSeqStore_t seq) static RawSeqStore_t ZSTDMT_getSeq(ZSTDMT_seqPool* seqPool) { - if (seqPool->bufferSize == 0) { - return kNullRawSeqStore; - } return bufferToSeq(ZSTDMT_getBuffer(seqPool)); } @@ -353,25 +286,28 @@ static ZSTDMT_seqPool* ZSTDMT_expandSeqPool(ZSTDMT_seqPool* pool, U32 nbWorkers) /* ===== CCtx Pool ===== */ /* a single CCtx Pool can be invoked from multiple threads in parallel */ +typedef struct ZSTDMT_RustCCtxPool_s ZSTDMT_RustCCtxPool; + +ZSTDMT_RustCCtxPool* ZSTDMT_rust_cctx_pool_create(unsigned nbWorkers, + ZSTD_customMem cMem); +void ZSTDMT_rust_cctx_pool_free(ZSTDMT_RustCCtxPool* pool); +size_t ZSTDMT_rust_cctx_pool_sizeof(const ZSTDMT_RustCCtxPool* pool); +ZSTDMT_RustCCtxPool* ZSTDMT_rust_cctx_pool_expand(ZSTDMT_RustCCtxPool* pool, + unsigned nbWorkers); +ZSTD_CCtx* ZSTDMT_rust_cctx_pool_get(ZSTDMT_RustCCtxPool* pool); +void ZSTDMT_rust_cctx_pool_release(ZSTDMT_RustCCtxPool* pool, ZSTD_CCtx* cctx); + typedef struct { - ZSTD_pthread_mutex_t poolMutex; - int totalCCtx; - int availCCtx; + ZSTDMT_RustCCtxPool* rustPool; + int totalCCtx; /* kept for the existing MT parameter diagnostics */ ZSTD_customMem cMem; - ZSTD_CCtx** cctxs; } ZSTDMT_CCtxPool; /* note : all CCtx borrowed from the pool must be reverted back to the pool _before_ freeing the pool */ static void ZSTDMT_freeCCtxPool(ZSTDMT_CCtxPool* pool) { if (!pool) return; - ZSTD_pthread_mutex_destroy(&pool->poolMutex); - if (pool->cctxs) { - int cid; - for (cid=0; cidtotalCCtx; cid++) - ZSTD_freeCCtx(pool->cctxs[cid]); /* free compatible with NULL */ - ZSTD_customFree(pool->cctxs, pool->cMem); - } + ZSTDMT_rust_cctx_pool_free(pool->rustPool); ZSTD_customFree(pool, pool->cMem); } @@ -384,20 +320,13 @@ static ZSTDMT_CCtxPool* ZSTDMT_createCCtxPool(int nbWorkers, (ZSTDMT_CCtxPool*) ZSTD_customCalloc(sizeof(ZSTDMT_CCtxPool), cMem); assert(nbWorkers > 0); if (!cctxPool) return NULL; - if (ZSTD_pthread_mutex_init(&cctxPool->poolMutex, NULL)) { + cctxPool->rustPool = ZSTDMT_rust_cctx_pool_create((unsigned)nbWorkers, cMem); + if (cctxPool->rustPool == NULL) { ZSTD_customFree(cctxPool, cMem); return NULL; } cctxPool->totalCCtx = nbWorkers; - cctxPool->cctxs = (ZSTD_CCtx**)ZSTD_customCalloc(nbWorkers * sizeof(ZSTD_CCtx*), cMem); - if (!cctxPool->cctxs) { - ZSTDMT_freeCCtxPool(cctxPool); - return NULL; - } cctxPool->cMem = cMem; - cctxPool->cctxs[0] = ZSTD_createCCtx_advanced(cMem); - if (!cctxPool->cctxs[0]) { ZSTDMT_freeCCtxPool(cctxPool); return NULL; } - cctxPool->availCCtx = 1; /* at least one cctx for single-thread mode */ DEBUGLOG(3, "cctxPool created, with %u workers", nbWorkers); return cctxPool; } @@ -406,59 +335,33 @@ static ZSTDMT_CCtxPool* ZSTDMT_expandCCtxPool(ZSTDMT_CCtxPool* srcPool, int nbWorkers) { if (srcPool==NULL) return NULL; - if (nbWorkers <= srcPool->totalCCtx) return srcPool; /* good enough */ - /* need a larger cctx pool */ - { ZSTD_customMem const cMem = srcPool->cMem; - ZSTDMT_freeCCtxPool(srcPool); - return ZSTDMT_createCCtxPool(nbWorkers, cMem); + srcPool->rustPool = ZSTDMT_rust_cctx_pool_expand(srcPool->rustPool, + (unsigned)nbWorkers); + if (srcPool->rustPool == NULL) { + ZSTD_customMem const cMem = srcPool->cMem; + ZSTD_customFree(srcPool, cMem); + return NULL; } + srcPool->totalCCtx = nbWorkers; + return srcPool; } /* only works during initialization phase, not during compression */ static size_t ZSTDMT_sizeof_CCtxPool(ZSTDMT_CCtxPool* cctxPool) { - ZSTD_pthread_mutex_lock(&cctxPool->poolMutex); - { unsigned const nbWorkers = cctxPool->totalCCtx; - size_t const poolSize = sizeof(*cctxPool); - size_t const arraySize = cctxPool->totalCCtx * sizeof(ZSTD_CCtx*); - size_t totalCCtxSize = 0; - unsigned u; - for (u=0; ucctxs[u]); - } - ZSTD_pthread_mutex_unlock(&cctxPool->poolMutex); - assert(nbWorkers > 0); - return poolSize + arraySize + totalCCtxSize; - } + if (cctxPool == NULL) return 0; + return sizeof(*cctxPool) + ZSTDMT_rust_cctx_pool_sizeof(cctxPool->rustPool); } static ZSTD_CCtx* ZSTDMT_getCCtx(ZSTDMT_CCtxPool* cctxPool) { DEBUGLOG(5, "ZSTDMT_getCCtx"); - ZSTD_pthread_mutex_lock(&cctxPool->poolMutex); - if (cctxPool->availCCtx) { - cctxPool->availCCtx--; - { ZSTD_CCtx* const cctx = cctxPool->cctxs[cctxPool->availCCtx]; - ZSTD_pthread_mutex_unlock(&cctxPool->poolMutex); - return cctx; - } } - ZSTD_pthread_mutex_unlock(&cctxPool->poolMutex); - DEBUGLOG(5, "create one more CCtx"); - return ZSTD_createCCtx_advanced(cctxPool->cMem); /* note : can be NULL, when creation fails ! */ + return ZSTDMT_rust_cctx_pool_get(cctxPool->rustPool); } static void ZSTDMT_releaseCCtx(ZSTDMT_CCtxPool* pool, ZSTD_CCtx* cctx) { - if (cctx==NULL) return; /* compatibility with release on NULL */ - ZSTD_pthread_mutex_lock(&pool->poolMutex); - if (pool->availCCtx < pool->totalCCtx) - pool->cctxs[pool->availCCtx++] = cctx; - else { - /* pool overflow : should not happen, since totalCCtx==nbWorkers */ - DEBUGLOG(4, "CCtx pool overflow : free cctx"); - ZSTD_freeCCtx(cctx); - } - ZSTD_pthread_mutex_unlock(&pool->poolMutex); + ZSTDMT_rust_cctx_pool_release(pool->rustPool, cctx); } /* ==== Serial State ==== */ diff --git a/rust/src/dict_builder_cover.rs b/rust/src/dict_builder_cover.rs index d52ca62be..3cb0a15fc 100644 --- a/rust/src/dict_builder_cover.rs +++ b/rust/src/dict_builder_cover.rs @@ -596,15 +596,8 @@ fn best_snapshot(slot: &BestSlot) -> BestSnapshot { } } -#[repr(C)] -struct ZSTD_CCtx { - _private: [u8; 0], -} - -#[repr(C)] -struct ZSTD_CDict { - _private: [u8; 0], -} +type ZSTD_CCtx = c_void; +type ZSTD_CDict = c_void; unsafe extern "C" { fn ZDICT_finalizeDictionary( diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 9661f848a..92fed8e0b 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -63,3 +63,5 @@ pub mod zstd_opt; pub mod zstd_opt_tree; #[cfg(feature = "compression")] pub mod zstd_presplit; +#[cfg(feature = "compression")] +pub mod zstdmt_compress; diff --git a/rust/src/zstdmt_compress.rs b/rust/src/zstdmt_compress.rs new file mode 100644 index 000000000..b1c3f08ec --- /dev/null +++ b/rust/src/zstdmt_compress.rs @@ -0,0 +1,604 @@ +#![allow(non_camel_case_types)] +#![allow(non_snake_case)] +#![allow(clippy::missing_safety_doc)] +#![allow(clippy::too_many_arguments)] + +//! Rust-owned resource pools used by the multithreaded compressor. +//! +//! The job table, serial LDM state, and streaming state still use private C +//! layouts. `zstdmt_compress.c` therefore keeps those parts and projects only +//! the allocation pools into this module. The pool entry points below are a +//! narrow C ABI: buffers and `ZSTD_CCtx *` values remain opaque to Rust, while +//! allocation, reuse, expansion, and synchronization are Rust-owned. + +use std::mem::{self, MaybeUninit}; +use std::os::raw::{c_uint, c_void}; +use std::ptr; +use std::sync::Mutex; + +type ZstdAllocFunction = unsafe extern "C" fn(*mut c_void, usize) -> *mut c_void; +type ZstdFreeFunction = unsafe extern "C" fn(*mut c_void, *mut c_void); + +/// ABI-compatible representation of `ZSTD_customMem` from `zstd.h`. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct ZstdCustomMem { + pub customAlloc: Option, + pub customFree: Option, + pub opaque: *mut c_void, +} + +/// ABI-compatible buffer returned to the C adapter. +#[repr(C)] +#[derive(Clone, Copy, Default)] +pub struct ZstdMtBuffer { + pub start: *mut c_void, + pub capacity: usize, +} + +#[derive(Default)] +struct BufferPoolState { + buffer_size: usize, + nb_buffers: usize, +} + +/// The object is allocated with the caller's `ZSTD_customMem`; its mutex +/// protects all mutable pool state and the raw reusable-buffer array. +pub struct RustBufferPool { + custom_mem: ZstdCustomMem, + total_buffers: usize, + buffers: *mut MaybeUninit, + state: Mutex, +} + +// The raw array is only accessed while `state` is held. The pool itself is +// passed between C worker threads as an opaque pointer. +unsafe impl Send for RustBufferPool {} +unsafe impl Sync for RustBufferPool {} + +/// The `ZSTD_CCtx *` values are created and destroyed by the existing C API; +/// Rust owns only the synchronized reusable-pointer pool. +pub struct RustCCtxPool { + custom_mem: ZstdCustomMem, + total_cctx: usize, + cctxs: *mut MaybeUninit<*mut c_void>, + state: Mutex, +} + +unsafe impl Send for RustCCtxPool {} +unsafe impl Sync for RustCCtxPool {} + +unsafe extern "C" { + fn ZSTD_createCCtx_advanced(custom_mem: ZstdCustomMem) -> *mut c_void; + fn ZSTD_freeCCtx(cctx: *mut c_void) -> usize; + fn ZSTD_sizeof_CCtx(cctx: *const c_void) -> usize; +} + +unsafe fn custom_calloc(size: usize, custom_mem: ZstdCustomMem) -> *mut c_void { + if let Some(alloc) = custom_mem.customAlloc { + let allocation = unsafe { alloc(custom_mem.opaque, size) }; + if !allocation.is_null() { + unsafe { ptr::write_bytes(allocation, 0, size) }; + } + allocation + } else { + unsafe { libc::calloc(1, size) } + } +} + +unsafe fn custom_malloc(size: usize, custom_mem: ZstdCustomMem) -> *mut c_void { + if let Some(alloc) = custom_mem.customAlloc { + unsafe { alloc(custom_mem.opaque, size) } + } else { + unsafe { libc::malloc(size) } + } +} + +unsafe fn custom_free(allocation: *mut c_void, custom_mem: ZstdCustomMem) { + if allocation.is_null() { + return; + } + if let Some(free) = custom_mem.customFree { + unsafe { free(custom_mem.opaque, allocation) }; + } else { + unsafe { libc::free(allocation) }; + } +} + +fn checked_array_size(len: usize) -> Option { + mem::size_of::().checked_mul(len) +} + +unsafe fn create_buffer_pool( + max_nb_buffers: usize, + custom_mem: ZstdCustomMem, +) -> *mut RustBufferPool { + if max_nb_buffers == 0 { + return ptr::null_mut(); + } + let Some(buffer_bytes) = checked_array_size::>(max_nb_buffers) else { + return ptr::null_mut(); + }; + + let pool = unsafe { custom_calloc(mem::size_of::(), custom_mem) } + .cast::(); + if pool.is_null() { + return ptr::null_mut(); + } + let buffers = + unsafe { custom_calloc(buffer_bytes, custom_mem) }.cast::>(); + if buffers.is_null() { + unsafe { custom_free(pool.cast(), custom_mem) }; + return ptr::null_mut(); + } + + unsafe { + pool.write(RustBufferPool { + custom_mem, + total_buffers: max_nb_buffers, + buffers, + state: Mutex::new(BufferPoolState { + buffer_size: 64 << 10, + nb_buffers: 0, + }), + }); + } + pool +} + +unsafe fn destroy_buffer_pool(pool: *mut RustBufferPool) { + if pool.is_null() { + return; + } + let custom_mem = unsafe { (*pool).custom_mem }; + let buffers = unsafe { (*pool).buffers }; + let total_buffers = unsafe { (*pool).total_buffers }; + for index in 0..total_buffers { + let buffer = unsafe { buffers.add(index).read().assume_init() }; + unsafe { custom_free(buffer.start, custom_mem) }; + } + unsafe { ptr::drop_in_place(pool) }; + unsafe { + custom_free(buffers.cast(), custom_mem); + custom_free(pool.cast(), custom_mem); + } +} + +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_buffer_pool_create( + max_nb_buffers: c_uint, + custom_mem: ZstdCustomMem, +) -> *mut RustBufferPool { + unsafe { create_buffer_pool(max_nb_buffers as usize, custom_mem) } +} + +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_buffer_pool_free(pool: *mut RustBufferPool) { + unsafe { destroy_buffer_pool(pool) } +} + +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_buffer_pool_sizeof(pool: *const RustBufferPool) -> usize { + if pool.is_null() { + return 0; + } + let pool_ref = unsafe { &*pool }; + let _state = pool_ref + .state + .lock() + .unwrap_or_else(|error| error.into_inner()); + let mut total_buffer_size = 0usize; + for index in 0..pool_ref.total_buffers { + let buffer = unsafe { pool_ref.buffers.add(index).read().assume_init() }; + total_buffer_size = total_buffer_size.saturating_add(buffer.capacity); + } + mem::size_of::() + .saturating_add( + pool_ref + .total_buffers + .saturating_mul(mem::size_of::>()), + ) + .saturating_add(total_buffer_size) +} + +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_buffer_pool_set_size( + pool: *mut RustBufferPool, + buffer_size: usize, +) { + if pool.is_null() { + return; + } + let pool_ref = unsafe { &*pool }; + let mut state = pool_ref + .state + .lock() + .unwrap_or_else(|error| error.into_inner()); + state.buffer_size = buffer_size; +} + +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_buffer_pool_expand( + pool: *mut RustBufferPool, + max_nb_buffers: c_uint, +) -> *mut RustBufferPool { + if pool.is_null() { + return ptr::null_mut(); + } + let max_nb_buffers = max_nb_buffers as usize; + let pool_ref = unsafe { &*pool }; + let (total_buffers, buffer_size, custom_mem) = { + let state = pool_ref + .state + .lock() + .unwrap_or_else(|error| error.into_inner()); + ( + pool_ref.total_buffers, + state.buffer_size, + pool_ref.custom_mem, + ) + }; + if total_buffers >= max_nb_buffers { + return pool; + } + + // This matches the original resize contract: the old pool is consumed + // before creating the larger replacement. + unsafe { destroy_buffer_pool(pool) }; + let replacement = unsafe { create_buffer_pool(max_nb_buffers, custom_mem) }; + if !replacement.is_null() { + unsafe { ZSTDMT_rust_buffer_pool_set_size(replacement, buffer_size) }; + } + replacement +} + +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_buffer_pool_get(pool: *mut RustBufferPool) -> ZstdMtBuffer { + if pool.is_null() { + return ZstdMtBuffer::default(); + } + let pool_ref = unsafe { &*pool }; + let (buffer_size, reusable) = { + let mut state = pool_ref + .state + .lock() + .unwrap_or_else(|error| error.into_inner()); + if state.buffer_size == 0 { + return ZstdMtBuffer::default(); + } + if state.nb_buffers == 0 { + (state.buffer_size, ZstdMtBuffer::default()) + } else { + state.nb_buffers -= 1; + let index = state.nb_buffers; + let buffer = unsafe { pool_ref.buffers.add(index).read().assume_init() }; + unsafe { + pool_ref + .buffers + .add(index) + .write(MaybeUninit::new(ZstdMtBuffer::default())); + } + (state.buffer_size, buffer) + } + }; + + if !reusable.start.is_null() + && reusable.capacity >= buffer_size + && (reusable.capacity >> 3) <= buffer_size + { + return reusable; + } + if !reusable.start.is_null() { + unsafe { custom_free(reusable.start, pool_ref.custom_mem) }; + } + + let start = unsafe { custom_malloc(buffer_size, pool_ref.custom_mem) }; + ZstdMtBuffer { + start, + capacity: if start.is_null() { 0 } else { buffer_size }, + } +} + +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_buffer_pool_release( + pool: *mut RustBufferPool, + buffer: ZstdMtBuffer, +) { + if pool.is_null() || buffer.start.is_null() { + return; + } + let pool_ref = unsafe { &*pool }; + let mut state = pool_ref + .state + .lock() + .unwrap_or_else(|error| error.into_inner()); + if state.nb_buffers < pool_ref.total_buffers { + let index = state.nb_buffers; + unsafe { + pool_ref.buffers.add(index).write(MaybeUninit::new(buffer)); + } + state.nb_buffers += 1; + return; + } + drop(state); + unsafe { custom_free(buffer.start, pool_ref.custom_mem) }; +} + +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_buffer_pool_resize( + pool: *mut RustBufferPool, + buffer: ZstdMtBuffer, +) -> ZstdMtBuffer { + if pool.is_null() || buffer.start.is_null() { + return buffer; + } + let pool_ref = unsafe { &*pool }; + let buffer_size = { + let state = pool_ref + .state + .lock() + .unwrap_or_else(|error| error.into_inner()); + state.buffer_size + }; + if buffer.capacity >= buffer_size { + return buffer; + } + let start = unsafe { custom_malloc(buffer_size, pool_ref.custom_mem) }; + if start.is_null() { + return buffer; + } + unsafe { + ptr::copy_nonoverlapping( + buffer.start.cast::(), + start.cast::(), + buffer.capacity, + ); + custom_free(buffer.start, pool_ref.custom_mem); + } + ZstdMtBuffer { + start, + capacity: buffer_size, + } +} + +unsafe fn create_cctx_pool(nb_workers: usize, custom_mem: ZstdCustomMem) -> *mut RustCCtxPool { + if nb_workers == 0 { + return ptr::null_mut(); + } + let Some(cctx_bytes) = checked_array_size::>(nb_workers) else { + return ptr::null_mut(); + }; + let pool = + unsafe { custom_calloc(mem::size_of::(), custom_mem) }.cast::(); + if pool.is_null() { + return ptr::null_mut(); + } + let cctxs = unsafe { custom_calloc(cctx_bytes, custom_mem) }.cast::>(); + if cctxs.is_null() { + unsafe { custom_free(pool.cast(), custom_mem) }; + return ptr::null_mut(); + } + let first = unsafe { ZSTD_createCCtx_advanced(custom_mem) }; + if first.is_null() { + unsafe { + custom_free(cctxs.cast(), custom_mem); + custom_free(pool.cast(), custom_mem); + } + return ptr::null_mut(); + } + unsafe { cctxs.write(MaybeUninit::new(first)) }; + unsafe { + pool.write(RustCCtxPool { + custom_mem, + total_cctx: nb_workers, + cctxs, + state: Mutex::new(1), + }); + } + pool +} + +unsafe fn destroy_cctx_pool(pool: *mut RustCCtxPool) { + if pool.is_null() { + return; + } + let custom_mem = unsafe { (*pool).custom_mem }; + let cctxs = unsafe { (*pool).cctxs }; + let total_cctx = unsafe { (*pool).total_cctx }; + for index in 0..total_cctx { + let cctx = unsafe { cctxs.add(index).read().assume_init() }; + unsafe { ZSTD_freeCCtx(cctx) }; + } + unsafe { ptr::drop_in_place(pool) }; + unsafe { + custom_free(cctxs.cast(), custom_mem); + custom_free(pool.cast(), custom_mem); + } +} + +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_cctx_pool_create( + nb_workers: c_uint, + custom_mem: ZstdCustomMem, +) -> *mut RustCCtxPool { + unsafe { create_cctx_pool(nb_workers as usize, custom_mem) } +} + +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_cctx_pool_free(pool: *mut RustCCtxPool) { + unsafe { destroy_cctx_pool(pool) } +} + +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_cctx_pool_sizeof(pool: *const RustCCtxPool) -> usize { + if pool.is_null() { + return 0; + } + let pool_ref = unsafe { &*pool }; + let _state = pool_ref + .state + .lock() + .unwrap_or_else(|error| error.into_inner()); + let mut total_cctx_size = 0usize; + for index in 0..pool_ref.total_cctx { + let cctx = unsafe { pool_ref.cctxs.add(index).read().assume_init() }; + total_cctx_size = + total_cctx_size.saturating_add(unsafe { ZSTD_sizeof_CCtx(cctx.cast_const()) }); + } + mem::size_of::() + .saturating_add( + pool_ref + .total_cctx + .saturating_mul(mem::size_of::>()), + ) + .saturating_add(total_cctx_size) +} + +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_cctx_pool_expand( + pool: *mut RustCCtxPool, + nb_workers: c_uint, +) -> *mut RustCCtxPool { + if pool.is_null() { + return ptr::null_mut(); + } + let nb_workers = nb_workers as usize; + let pool_ref = unsafe { &*pool }; + if pool_ref.total_cctx >= nb_workers { + return pool; + } + let custom_mem = pool_ref.custom_mem; + unsafe { destroy_cctx_pool(pool) }; + unsafe { create_cctx_pool(nb_workers, custom_mem) } +} + +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_cctx_pool_get(pool: *mut RustCCtxPool) -> *mut c_void { + if pool.is_null() { + return ptr::null_mut(); + } + let pool_ref = unsafe { &*pool }; + let cctx = { + let mut available = pool_ref + .state + .lock() + .unwrap_or_else(|error| error.into_inner()); + if *available == 0 { + ptr::null_mut() + } else { + *available -= 1; + let index = *available; + unsafe { pool_ref.cctxs.add(index).read().assume_init() } + } + }; + if !cctx.is_null() { + cctx + } else { + unsafe { ZSTD_createCCtx_advanced(pool_ref.custom_mem) } + } +} + +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_cctx_pool_release(pool: *mut RustCCtxPool, cctx: *mut c_void) { + if pool.is_null() || cctx.is_null() { + return; + } + let pool_ref = unsafe { &*pool }; + let mut available = pool_ref + .state + .lock() + .unwrap_or_else(|error| error.into_inner()); + if *available < pool_ref.total_cctx { + let index = *available; + unsafe { + pool_ref.cctxs.add(index).write(MaybeUninit::new(cctx)); + } + *available += 1; + return; + } + drop(available); + unsafe { ZSTD_freeCCtx(cctx) }; +} + +#[cfg(test)] +mod tests { + use super::*; + + const DEFAULT_MEM: ZstdCustomMem = ZstdCustomMem { + customAlloc: None, + customFree: None, + opaque: ptr::null_mut(), + }; + + #[test] + fn buffer_pool_reuses_and_resizes_buffers() { + let pool = unsafe { ZSTDMT_rust_buffer_pool_create(2, DEFAULT_MEM) }; + assert!(!pool.is_null()); + + let first = unsafe { ZSTDMT_rust_buffer_pool_get(pool) }; + assert!(!first.start.is_null()); + assert_eq!(first.capacity, 64 << 10); + unsafe { ZSTDMT_rust_buffer_pool_release(pool, first) }; + + let reused = unsafe { ZSTDMT_rust_buffer_pool_get(pool) }; + assert_eq!(reused.start, first.start); + assert_eq!(reused.capacity, first.capacity); + unsafe { ZSTDMT_rust_buffer_pool_release(pool, reused) }; + + unsafe { ZSTDMT_rust_buffer_pool_set_size(pool, 128 << 10) }; + let resized = unsafe { ZSTDMT_rust_buffer_pool_get(pool) }; + assert_eq!(resized.capacity, 128 << 10); + unsafe { + ZSTDMT_rust_buffer_pool_release(pool, resized); + ZSTDMT_rust_buffer_pool_free(pool); + } + } + + #[test] + fn zero_size_pool_does_not_allocate_sequence_storage() { + let pool = unsafe { ZSTDMT_rust_buffer_pool_create(1, DEFAULT_MEM) }; + assert!(!pool.is_null()); + unsafe { ZSTDMT_rust_buffer_pool_set_size(pool, 0) }; + let buffer = unsafe { ZSTDMT_rust_buffer_pool_get(pool) }; + assert!(buffer.start.is_null()); + assert_eq!(buffer.capacity, 0); + unsafe { ZSTDMT_rust_buffer_pool_free(pool) }; + } + + #[test] + fn expansion_preserves_requested_buffer_size() { + let pool = unsafe { ZSTDMT_rust_buffer_pool_create(1, DEFAULT_MEM) }; + assert!(!pool.is_null()); + unsafe { ZSTDMT_rust_buffer_pool_set_size(pool, 96 << 10) }; + let expanded = unsafe { ZSTDMT_rust_buffer_pool_expand(pool, 3) }; + assert!(!expanded.is_null()); + let buffer = unsafe { ZSTDMT_rust_buffer_pool_get(expanded) }; + assert_eq!(buffer.capacity, 96 << 10); + unsafe { + ZSTDMT_rust_buffer_pool_release(expanded, buffer); + ZSTDMT_rust_buffer_pool_free(expanded); + } + } + + #[test] + fn resize_preserves_existing_buffer_contents() { + let pool = unsafe { ZSTDMT_rust_buffer_pool_create(1, DEFAULT_MEM) }; + assert!(!pool.is_null()); + let buffer = unsafe { ZSTDMT_rust_buffer_pool_get(pool) }; + assert!(!buffer.start.is_null()); + let sample = b"multithreaded-buffer-pool"; + unsafe { + ptr::copy_nonoverlapping(sample.as_ptr(), buffer.start.cast(), sample.len()); + ZSTDMT_rust_buffer_pool_set_size(pool, buffer.capacity + 1); + } + let resized = unsafe { ZSTDMT_rust_buffer_pool_resize(pool, buffer) }; + assert_eq!(resized.capacity, buffer.capacity + 1); + let contents = + unsafe { std::slice::from_raw_parts(resized.start.cast::(), sample.len()) }; + assert_eq!(contents, sample); + unsafe { + ZSTDMT_rust_buffer_pool_release(pool, resized); + ZSTDMT_rust_buffer_pool_free(pool); + } + } +}