feat(rust): port worker pool and pthread wrappers

Replace the common worker pool and debug pthread wrappers with ABI-compatible
Rust implementations. The pool preserves bounded-queue, resize, drain-on-free,
custom-allocator, and no-thread behavior while C supplies only the build-time
multithreading configuration bit.

This moves runtime support without changing C tests or public headers. Codec
and CLI implementations remain C while their dependencies are migrated.

Test Plan:
- cargo fmt --check
- cargo test --all-targets
- cargo clippy --all-targets -- -D warnings
- cargo build --release
- make -C tests poolTests
- ./tests/poolTests
- compile pool and threading shims with -Werror in MT and no-thread modes

Refs: rust/README.md
This commit is contained in:
2026-07-10 20:05:35 +02:00
parent a0f2b2a14c
commit 26b5e202ee
6 changed files with 1037 additions and 549 deletions
+8 -364
View File
@@ -1,371 +1,15 @@
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under both the BSD-style license (found in the
* LICENSE file in the root directory of this source tree) and the GPLv2 (found
* in the COPYING file in the root directory of this source tree).
* You may select, at your option, one of the above-listed licenses.
*/
/* ====== Dependencies ======= */
#include "../common/allocations.h" /* ZSTD_customCalloc, ZSTD_customFree */
#include "zstd_deps.h" /* size_t */
#include "debug.h" /* assert */
#include "pool.h"
/* ====== Compiler specifics ====== */
#if defined(_MSC_VER)
# pragma warning(disable : 4204) /* disable: C4204: non-constant aggregate initializer */
#endif
/* The implementation lives in rust/src/pool.rs. Rust cannot observe the C
* preprocessor configuration, so expose the one bit of configuration that
* changes the pool API's behavior. */
int ZSTD_rust_pool_is_multithreaded(void);
int ZSTD_rust_pool_is_multithreaded(void)
{
#ifdef ZSTD_MULTITHREAD
#include "threading.h" /* pthread adaptation */
/* A job is a function and an opaque argument */
typedef struct POOL_job_s {
POOL_function function;
void *opaque;
} POOL_job;
struct POOL_ctx_s {
ZSTD_customMem customMem;
/* Keep track of the threads */
ZSTD_pthread_t* threads;
size_t threadCapacity;
size_t threadLimit;
/* The queue is a circular buffer */
POOL_job *queue;
size_t queueHead;
size_t queueTail;
size_t queueSize;
/* The number of threads working on jobs */
size_t numThreadsBusy;
/* Indicates if the queue is empty */
int queueEmpty;
/* The mutex protects the queue */
ZSTD_pthread_mutex_t queueMutex;
/* Condition variable for pushers to wait on when the queue is full */
ZSTD_pthread_cond_t queuePushCond;
/* Condition variables for poppers to wait on when the queue is empty */
ZSTD_pthread_cond_t queuePopCond;
/* Indicates if the queue is shutting down */
int shutdown;
};
/* POOL_thread() :
* Work thread for the thread pool.
* Waits for jobs and executes them.
* @returns : NULL on failure else non-null.
*/
static void* POOL_thread(void* opaque) {
POOL_ctx* const ctx = (POOL_ctx*)opaque;
if (!ctx) { return NULL; }
for (;;) {
/* Lock the mutex and wait for a non-empty queue or until shutdown */
ZSTD_pthread_mutex_lock(&ctx->queueMutex);
while ( ctx->queueEmpty
|| (ctx->numThreadsBusy >= ctx->threadLimit) ) {
if (ctx->shutdown) {
/* even if !queueEmpty, (possible if numThreadsBusy >= threadLimit),
* a few threads will be shutdown while !queueEmpty,
* but enough threads will remain active to finish the queue */
ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
return opaque;
}
ZSTD_pthread_cond_wait(&ctx->queuePopCond, &ctx->queueMutex);
}
/* Pop a job off the queue */
{ POOL_job const job = ctx->queue[ctx->queueHead];
ctx->queueHead = (ctx->queueHead + 1) % ctx->queueSize;
ctx->numThreadsBusy++;
ctx->queueEmpty = (ctx->queueHead == ctx->queueTail);
/* Unlock the mutex, signal a pusher, and run the job */
ZSTD_pthread_cond_signal(&ctx->queuePushCond);
ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
job.function(job.opaque);
/* If the intended queue size was 0, signal after finishing job */
ZSTD_pthread_mutex_lock(&ctx->queueMutex);
ctx->numThreadsBusy--;
ZSTD_pthread_cond_signal(&ctx->queuePushCond);
ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
}
} /* for (;;) */
assert(0); /* Unreachable */
}
/* ZSTD_createThreadPool() : public access point */
POOL_ctx* ZSTD_createThreadPool(size_t numThreads) {
return POOL_create (numThreads, 0);
}
POOL_ctx* POOL_create(size_t numThreads, size_t queueSize) {
return POOL_create_advanced(numThreads, queueSize, ZSTD_defaultCMem);
}
POOL_ctx* POOL_create_advanced(size_t numThreads, size_t queueSize,
ZSTD_customMem customMem)
{
POOL_ctx* ctx;
/* Check parameters */
if (!numThreads) { return NULL; }
/* Allocate the context and zero initialize */
ctx = (POOL_ctx*)ZSTD_customCalloc(sizeof(POOL_ctx), customMem);
if (!ctx) { return NULL; }
/* Initialize the job queue.
* It needs one extra space since one space is wasted to differentiate
* empty and full queues.
*/
ctx->queueSize = queueSize + 1;
ctx->queue = (POOL_job*)ZSTD_customCalloc(ctx->queueSize * sizeof(POOL_job), customMem);
ctx->queueHead = 0;
ctx->queueTail = 0;
ctx->numThreadsBusy = 0;
ctx->queueEmpty = 1;
{
int error = 0;
error |= ZSTD_pthread_mutex_init(&ctx->queueMutex, NULL);
error |= ZSTD_pthread_cond_init(&ctx->queuePushCond, NULL);
error |= ZSTD_pthread_cond_init(&ctx->queuePopCond, NULL);
if (error) { POOL_free(ctx); return NULL; }
}
ctx->shutdown = 0;
/* Allocate space for the thread handles */
ctx->threads = (ZSTD_pthread_t*)ZSTD_customCalloc(numThreads * sizeof(ZSTD_pthread_t), customMem);
ctx->threadCapacity = 0;
ctx->customMem = customMem;
/* Check for errors */
if (!ctx->threads || !ctx->queue) { POOL_free(ctx); return NULL; }
/* Initialize the threads */
{ size_t i;
for (i = 0; i < numThreads; ++i) {
if (ZSTD_pthread_create(&ctx->threads[i], NULL, &POOL_thread, ctx)) {
ctx->threadCapacity = i;
POOL_free(ctx);
return NULL;
} }
ctx->threadCapacity = numThreads;
ctx->threadLimit = numThreads;
}
return ctx;
}
/*! POOL_join() :
Shutdown the queue, wake any sleeping threads, and join all of the threads.
*/
static void POOL_join(POOL_ctx* ctx) {
/* Shut down the queue */
ZSTD_pthread_mutex_lock(&ctx->queueMutex);
ctx->shutdown = 1;
ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
/* Wake up sleeping threads */
ZSTD_pthread_cond_broadcast(&ctx->queuePushCond);
ZSTD_pthread_cond_broadcast(&ctx->queuePopCond);
/* Join all of the threads */
{ size_t i;
for (i = 0; i < ctx->threadCapacity; ++i) {
ZSTD_pthread_join(ctx->threads[i]); /* note : could fail */
} }
}
void POOL_free(POOL_ctx *ctx) {
if (!ctx) { return; }
POOL_join(ctx);
ZSTD_pthread_mutex_destroy(&ctx->queueMutex);
ZSTD_pthread_cond_destroy(&ctx->queuePushCond);
ZSTD_pthread_cond_destroy(&ctx->queuePopCond);
ZSTD_customFree(ctx->queue, ctx->customMem);
ZSTD_customFree(ctx->threads, ctx->customMem);
ZSTD_customFree(ctx, ctx->customMem);
}
/*! POOL_joinJobs() :
* Waits for all queued jobs to finish executing.
*/
void POOL_joinJobs(POOL_ctx* ctx) {
ZSTD_pthread_mutex_lock(&ctx->queueMutex);
while(!ctx->queueEmpty || ctx->numThreadsBusy > 0) {
ZSTD_pthread_cond_wait(&ctx->queuePushCond, &ctx->queueMutex);
}
ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
}
void ZSTD_freeThreadPool (ZSTD_threadPool* pool) {
POOL_free (pool);
}
size_t POOL_sizeof(const POOL_ctx* ctx) {
if (ctx==NULL) return 0; /* supports sizeof NULL */
return sizeof(*ctx)
+ ctx->queueSize * sizeof(POOL_job)
+ ctx->threadCapacity * sizeof(ZSTD_pthread_t);
}
/* @return : 0 on success, 1 on error */
static int POOL_resize_internal(POOL_ctx* ctx, size_t numThreads)
{
if (numThreads <= ctx->threadCapacity) {
if (!numThreads) return 1;
ctx->threadLimit = numThreads;
return 0;
}
/* numThreads > threadCapacity */
{ ZSTD_pthread_t* const threadPool = (ZSTD_pthread_t*)ZSTD_customCalloc(numThreads * sizeof(ZSTD_pthread_t), ctx->customMem);
if (!threadPool) return 1;
/* replace existing thread pool */
ZSTD_memcpy(threadPool, ctx->threads, ctx->threadCapacity * sizeof(ZSTD_pthread_t));
ZSTD_customFree(ctx->threads, ctx->customMem);
ctx->threads = threadPool;
/* Initialize additional threads */
{ size_t threadId;
for (threadId = ctx->threadCapacity; threadId < numThreads; ++threadId) {
if (ZSTD_pthread_create(&threadPool[threadId], NULL, &POOL_thread, ctx)) {
ctx->threadCapacity = threadId;
return 1;
} }
} }
/* successfully expanded */
ctx->threadCapacity = numThreads;
ctx->threadLimit = numThreads;
return 0;
}
/* @return : 0 on success, 1 on error */
int POOL_resize(POOL_ctx* ctx, size_t numThreads)
{
int result;
if (ctx==NULL) return 1;
ZSTD_pthread_mutex_lock(&ctx->queueMutex);
result = POOL_resize_internal(ctx, numThreads);
ZSTD_pthread_cond_broadcast(&ctx->queuePopCond);
ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
return result;
}
/**
* Returns 1 if the queue is full and 0 otherwise.
*
* When queueSize is 1 (pool was created with an intended queueSize of 0),
* then a queue is empty if there is a thread free _and_ no job is waiting.
*/
static int isQueueFull(POOL_ctx const* ctx) {
if (ctx->queueSize > 1) {
return ctx->queueHead == ((ctx->queueTail + 1) % ctx->queueSize);
} else {
return (ctx->numThreadsBusy == ctx->threadLimit) ||
!ctx->queueEmpty;
}
}
static void
POOL_add_internal(POOL_ctx* ctx, POOL_function function, void *opaque)
{
POOL_job job;
job.function = function;
job.opaque = opaque;
assert(ctx != NULL);
if (ctx->shutdown) return;
ctx->queueEmpty = 0;
ctx->queue[ctx->queueTail] = job;
ctx->queueTail = (ctx->queueTail + 1) % ctx->queueSize;
ZSTD_pthread_cond_signal(&ctx->queuePopCond);
}
void POOL_add(POOL_ctx* ctx, POOL_function function, void* opaque)
{
assert(ctx != NULL);
ZSTD_pthread_mutex_lock(&ctx->queueMutex);
/* Wait until there is space in the queue for the new job */
while (isQueueFull(ctx) && (!ctx->shutdown)) {
ZSTD_pthread_cond_wait(&ctx->queuePushCond, &ctx->queueMutex);
}
POOL_add_internal(ctx, function, opaque);
ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
}
int POOL_tryAdd(POOL_ctx* ctx, POOL_function function, void* opaque)
{
assert(ctx != NULL);
ZSTD_pthread_mutex_lock(&ctx->queueMutex);
if (isQueueFull(ctx)) {
ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
return 0;
}
POOL_add_internal(ctx, function, opaque);
ZSTD_pthread_mutex_unlock(&ctx->queueMutex);
return 1;
}
#else /* ZSTD_MULTITHREAD not defined */
/* ========================== */
/* No multi-threading support */
/* ========================== */
/* We don't need any data, but if it is empty, malloc() might return NULL. */
struct POOL_ctx_s {
int dummy;
};
static POOL_ctx g_poolCtx;
POOL_ctx* POOL_create(size_t numThreads, size_t queueSize) {
return POOL_create_advanced(numThreads, queueSize, ZSTD_defaultCMem);
}
POOL_ctx*
POOL_create_advanced(size_t numThreads, size_t queueSize, ZSTD_customMem customMem)
{
(void)numThreads;
(void)queueSize;
(void)customMem;
return &g_poolCtx;
}
void POOL_free(POOL_ctx* ctx) {
assert(!ctx || ctx == &g_poolCtx);
(void)ctx;
}
void POOL_joinJobs(POOL_ctx* ctx){
assert(!ctx || ctx == &g_poolCtx);
(void)ctx;
}
int POOL_resize(POOL_ctx* ctx, size_t numThreads) {
(void)ctx; (void)numThreads;
#else
return 0;
#endif
}
void POOL_add(POOL_ctx* ctx, POOL_function function, void* opaque) {
(void)ctx;
function(opaque);
}
int POOL_tryAdd(POOL_ctx* ctx, POOL_function function, void* opaque) {
(void)ctx;
function(opaque);
return 1;
}
size_t POOL_sizeof(const POOL_ctx* ctx) {
if (ctx==NULL) return 0; /* supports sizeof NULL */
assert(ctx == &g_poolCtx);
return sizeof(*ctx);
}
#endif /* ZSTD_MULTITHREAD */
+3 -179
View File
@@ -1,182 +1,6 @@
/**
* Copyright (c) 2016 Tino Reichardt
* All rights reserved.
*
* You can contact the author at:
* - zstdmt source repository: https://github.com/mcmilk/zstdmt
*
* This source code is licensed under both the BSD-style license (found in the
* LICENSE file in the root directory of this source tree) and the GPLv2 (found
* in the COPYING file in the root directory of this source tree).
* You may select, at your option, one of the above-listed licenses.
*/
/**
* This file will hold wrapper for systems, which do not support pthreads
*/
#include "threading.h"
/* create fake symbol to avoid empty translation unit warning */
/* The platform-specific functions declared by threading.h are implemented in
* rust/src/threading.rs. Keep a symbol for configurations where threading.h
* only defines macros and this translation unit would otherwise be empty. */
int g_ZSTD_threading_useless_symbol;
#if defined(ZSTD_MULTITHREAD) && defined(_WIN32)
/**
* Windows minimalist Pthread Wrapper
*/
/* === Dependencies === */
#include <process.h>
#include <errno.h>
/* === Implementation === */
typedef struct {
void* (*start_routine)(void*);
void* arg;
int initialized;
ZSTD_pthread_cond_t initialized_cond;
ZSTD_pthread_mutex_t initialized_mutex;
} ZSTD_thread_params_t;
static unsigned __stdcall worker(void *arg)
{
void* (*start_routine)(void*);
void* thread_arg;
/* Initialized thread_arg and start_routine and signal main thread that we don't need it
* to wait any longer.
*/
{
ZSTD_thread_params_t* thread_param = (ZSTD_thread_params_t*)arg;
thread_arg = thread_param->arg;
start_routine = thread_param->start_routine;
/* Signal main thread that we are running and do not depend on its memory anymore */
ZSTD_pthread_mutex_lock(&thread_param->initialized_mutex);
thread_param->initialized = 1;
ZSTD_pthread_cond_signal(&thread_param->initialized_cond);
ZSTD_pthread_mutex_unlock(&thread_param->initialized_mutex);
}
start_routine(thread_arg);
return 0;
}
int ZSTD_pthread_create(ZSTD_pthread_t* thread, const void* unused,
void* (*start_routine) (void*), void* arg)
{
ZSTD_thread_params_t thread_param;
(void)unused;
if (thread==NULL) return -1;
*thread = NULL;
thread_param.start_routine = start_routine;
thread_param.arg = arg;
thread_param.initialized = 0;
/* Setup thread initialization synchronization */
if(ZSTD_pthread_cond_init(&thread_param.initialized_cond, NULL)) {
/* Should never happen on Windows */
return -1;
}
if(ZSTD_pthread_mutex_init(&thread_param.initialized_mutex, NULL)) {
/* Should never happen on Windows */
ZSTD_pthread_cond_destroy(&thread_param.initialized_cond);
return -1;
}
/* Spawn thread */
*thread = (HANDLE)_beginthreadex(NULL, 0, worker, &thread_param, 0, NULL);
if (*thread==NULL) {
ZSTD_pthread_mutex_destroy(&thread_param.initialized_mutex);
ZSTD_pthread_cond_destroy(&thread_param.initialized_cond);
return errno;
}
/* Wait for thread to be initialized */
ZSTD_pthread_mutex_lock(&thread_param.initialized_mutex);
while(!thread_param.initialized) {
ZSTD_pthread_cond_wait(&thread_param.initialized_cond, &thread_param.initialized_mutex);
}
ZSTD_pthread_mutex_unlock(&thread_param.initialized_mutex);
ZSTD_pthread_mutex_destroy(&thread_param.initialized_mutex);
ZSTD_pthread_cond_destroy(&thread_param.initialized_cond);
return 0;
}
int ZSTD_pthread_join(ZSTD_pthread_t thread)
{
DWORD result;
if (!thread) return 0;
result = WaitForSingleObject(thread, INFINITE);
CloseHandle(thread);
switch (result) {
case WAIT_OBJECT_0:
return 0;
case WAIT_ABANDONED:
return EINVAL;
default:
return GetLastError();
}
}
#endif /* ZSTD_MULTITHREAD */
#if defined(ZSTD_MULTITHREAD) && DEBUGLEVEL >= 1 && !defined(_WIN32)
#define ZSTD_DEPS_NEED_MALLOC
#include "zstd_deps.h"
int ZSTD_pthread_mutex_init(ZSTD_pthread_mutex_t* mutex, pthread_mutexattr_t const* attr)
{
assert(mutex != NULL);
*mutex = (pthread_mutex_t*)ZSTD_malloc(sizeof(pthread_mutex_t));
if (!*mutex)
return 1;
return pthread_mutex_init(*mutex, attr);
}
int ZSTD_pthread_mutex_destroy(ZSTD_pthread_mutex_t* mutex)
{
assert(mutex != NULL);
if (!*mutex)
return 0;
{
int const ret = pthread_mutex_destroy(*mutex);
ZSTD_free(*mutex);
return ret;
}
}
int ZSTD_pthread_cond_init(ZSTD_pthread_cond_t* cond, pthread_condattr_t const* attr)
{
assert(cond != NULL);
*cond = (pthread_cond_t*)ZSTD_malloc(sizeof(pthread_cond_t));
if (!*cond)
return 1;
return pthread_cond_init(*cond, attr);
}
int ZSTD_pthread_cond_destroy(ZSTD_pthread_cond_t* cond)
{
assert(cond != NULL);
if (!*cond)
return 0;
{
int const ret = pthread_cond_destroy(*cond);
ZSTD_free(*cond);
return ret;
}
}
#endif
+9 -6
View File
@@ -20,11 +20,14 @@ zstd ABI:
- Entropy coding
- `entropy_common` reads FSE normalized counts and Huffman statistics.
- `fse_decompress` builds FSE decoding tables and decodes FSE streams.
- Runtime support
- `threading` provides platform pthread wrappers required by zstd headers.
- `pool` implements the bounded worker pool used by multithreaded compression.
The remaining compression, general decompression, dictionary, legacy, runtime,
and CLI translation units are still C. They must move before the rewrite is
complete. Keeping that boundary explicit prevents a passing hybrid build from
being mistaken for the final all-Rust result.
The remaining compression, general decompression, dictionary, legacy, and CLI
translation units are still C. They must move before the rewrite is complete.
Keeping that boundary explicit prevents a passing hybrid build from being
mistaken for the final all-Rust result.
## Compatibility boundary
@@ -53,8 +56,8 @@ Then run original compatibility tests from the repository root, starting with
the narrow target for the component being migrated. For example:
```sh
make -C tests fuzzer
./tests/fuzzer -i1 --no-big-tests
make -C tests poolTests
./tests/poolTests
```
Broader `tests/Makefile` targets remain the authoritative integration gates as
+2
View File
@@ -9,5 +9,7 @@ pub mod entropy_common;
pub mod errors;
pub mod fse_decompress;
pub mod mem;
pub mod pool;
pub mod threading;
pub mod xxhash;
pub mod zstd_common;
+733
View File
@@ -0,0 +1,733 @@
use std::cell::UnsafeCell;
use std::mem::{self, MaybeUninit};
use std::os::raw::{c_int, c_void};
use std::ptr;
use std::sync::{Condvar, Mutex};
use std::thread::{self, JoinHandle};
pub type PoolFunction = unsafe extern "C" fn(*mut c_void);
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 {
custom_alloc: Option<ZstdAllocFunction>,
custom_free: Option<ZstdFreeFunction>,
opaque: *mut c_void,
}
const DEFAULT_CUSTOM_MEM: ZstdCustomMem = ZstdCustomMem {
custom_alloc: None,
custom_free: None,
opaque: ptr::null_mut(),
};
#[derive(Clone, Copy)]
struct Job {
function: PoolFunction,
opaque: *mut c_void,
}
// A job's opaque value is owned by the caller and is required by pool.h to
// remain valid until the callback completes.
unsafe impl Send for Job {}
struct PoolState {
queue_head: usize,
queue_tail: usize,
queue_empty: bool,
num_threads_busy: usize,
thread_limit: usize,
shutdown: bool,
}
struct ThreadStorage {
slots: *mut MaybeUninit<JoinHandle<()>>,
initialized: usize,
}
pub struct PoolCtx {
custom_mem: ZstdCustomMem,
queue: *mut MaybeUninit<Job>,
queue_slots: usize,
state: Mutex<PoolState>,
queue_push_cond: Condvar,
queue_pop_cond: Condvar,
// Only POOL_resize()/POOL_free()/POOL_sizeof() access this field, and each
// does so while holding state. Workers never access thread handles.
threads: UnsafeCell<ThreadStorage>,
}
// All mutable shared state is protected by `state`. The raw queue and handle
// arrays stay allocated until every worker has been joined.
unsafe impl Send for PoolCtx {}
unsafe impl Sync for PoolCtx {}
#[derive(Clone, Copy)]
struct SendPoolPtr(*mut PoolCtx);
unsafe impl Send for SendPoolPtr {}
impl SendPoolPtr {
fn get(self) -> *mut PoolCtx {
self.0
}
}
// This has the same size and stable address as the dummy C context used by the
// original non-multithreaded implementation. It is only compared, never
// dereferenced as a PoolCtx.
static SINGLE_THREADED_POOL: c_int = 0;
fn single_threaded_pool() -> *mut PoolCtx {
ptr::from_ref(&SINGLE_THREADED_POOL).cast_mut().cast()
}
fn is_single_threaded_pool(ctx: *const PoolCtx) -> bool {
ptr::eq(ctx, single_threaded_pool().cast_const())
}
#[cfg(not(test))]
unsafe extern "C" {
fn ZSTD_rust_pool_is_multithreaded() -> c_int;
}
#[cfg(not(test))]
fn multithreading_enabled() -> bool {
// The C shim is compiled with the same preprocessor flags as pool users.
unsafe { ZSTD_rust_pool_is_multithreaded() != 0 }
}
#[cfg(test)]
fn multithreading_enabled() -> bool {
true
}
unsafe fn custom_calloc(size: usize, custom_mem: ZstdCustomMem) -> *mut c_void {
if let Some(alloc) = custom_mem.custom_alloc {
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_free(allocation: *mut c_void, custom_mem: ZstdCustomMem) {
if allocation.is_null() {
return;
}
if let Some(free) = custom_mem.custom_free {
unsafe { free(custom_mem.opaque, allocation) };
} else {
unsafe { libc::free(allocation) };
}
}
fn checked_array_size<T>(len: usize) -> Option<usize> {
mem::size_of::<T>().checked_mul(len)
}
unsafe fn create_threaded_pool(
num_threads: usize,
queue_size: usize,
custom_mem: ZstdCustomMem,
) -> *mut PoolCtx {
if num_threads == 0 {
return ptr::null_mut();
}
let Some(queue_slots) = queue_size.checked_add(1) else {
return ptr::null_mut();
};
let Some(queue_bytes) = checked_array_size::<MaybeUninit<Job>>(queue_slots) else {
return ptr::null_mut();
};
let Some(thread_bytes) = checked_array_size::<MaybeUninit<JoinHandle<()>>>(num_threads) else {
return ptr::null_mut();
};
let ctx_ptr = unsafe { custom_calloc(mem::size_of::<PoolCtx>(), custom_mem) }.cast::<PoolCtx>();
if ctx_ptr.is_null() {
return ptr::null_mut();
}
let queue = unsafe { custom_calloc(queue_bytes, custom_mem) }.cast::<MaybeUninit<Job>>();
if queue.is_null() {
unsafe { custom_free(ctx_ptr.cast(), custom_mem) };
return ptr::null_mut();
}
let thread_slots =
unsafe { custom_calloc(thread_bytes, custom_mem) }.cast::<MaybeUninit<JoinHandle<()>>>();
if thread_slots.is_null() {
unsafe {
custom_free(queue.cast(), custom_mem);
custom_free(ctx_ptr.cast(), custom_mem);
}
return ptr::null_mut();
}
unsafe {
ctx_ptr.write(PoolCtx {
custom_mem,
queue,
queue_slots,
state: Mutex::new(PoolState {
queue_head: 0,
queue_tail: 0,
queue_empty: true,
num_threads_busy: 0,
thread_limit: num_threads,
shutdown: false,
}),
queue_push_cond: Condvar::new(),
queue_pop_cond: Condvar::new(),
threads: UnsafeCell::new(ThreadStorage {
slots: thread_slots,
initialized: 0,
}),
});
}
for thread_id in 0..num_threads {
match spawn_worker(ctx_ptr) {
Ok(handle) => unsafe {
thread_slots.add(thread_id).write(MaybeUninit::new(handle));
(*(*ctx_ptr).threads.get()).initialized = thread_id + 1;
},
Err(_) => {
unsafe {
shutdown_and_join(ctx_ptr);
destroy_context(ctx_ptr);
}
return ptr::null_mut();
}
}
}
ctx_ptr
}
fn spawn_worker(ctx: *mut PoolCtx) -> std::io::Result<JoinHandle<()>> {
let send_ptr = SendPoolPtr(ctx);
thread::Builder::new().spawn(move || worker(send_ptr.get()))
}
fn worker(ctx_ptr: *mut PoolCtx) {
// The creator keeps the context alive until every handle is joined.
let ctx = unsafe { &*ctx_ptr };
loop {
let job = {
let mut state = ctx.state.lock().unwrap_or_else(|error| error.into_inner());
while state.queue_empty || state.num_threads_busy >= state.thread_limit {
if state.shutdown {
return;
}
state = ctx
.queue_pop_cond
.wait(state)
.unwrap_or_else(|error| error.into_inner());
}
let job = unsafe { ctx.queue.add(state.queue_head).read().assume_init() };
state.queue_head = (state.queue_head + 1) % ctx.queue_slots;
state.num_threads_busy += 1;
state.queue_empty = state.queue_head == state.queue_tail;
ctx.queue_push_cond.notify_one();
job
};
unsafe { (job.function)(job.opaque) };
let mut state = ctx.state.lock().unwrap_or_else(|error| error.into_inner());
debug_assert!(state.num_threads_busy > 0);
state.num_threads_busy -= 1;
ctx.queue_push_cond.notify_one();
}
}
unsafe fn shutdown_and_join(ctx_ptr: *mut PoolCtx) {
let ctx = unsafe { &*ctx_ptr };
{
let mut state = ctx.state.lock().unwrap_or_else(|error| error.into_inner());
state.shutdown = true;
}
ctx.queue_push_cond.notify_all();
ctx.queue_pop_cond.notify_all();
let threads = unsafe { &mut *ctx.threads.get() };
for thread_id in 0..threads.initialized {
let handle = unsafe { threads.slots.add(thread_id).read().assume_init() };
let _ = handle.join();
}
threads.initialized = 0;
}
unsafe fn destroy_context(ctx_ptr: *mut PoolCtx) {
let custom_mem = unsafe { (*ctx_ptr).custom_mem };
let queue = unsafe { (*ctx_ptr).queue };
let thread_slots = unsafe { (*(*ctx_ptr).threads.get()).slots };
unsafe { ptr::drop_in_place(ctx_ptr) };
unsafe {
custom_free(queue.cast(), custom_mem);
custom_free(thread_slots.cast(), custom_mem);
custom_free(ctx_ptr.cast(), custom_mem);
}
}
fn queue_is_full(ctx: &PoolCtx, state: &PoolState) -> bool {
if ctx.queue_slots > 1 {
state.queue_head == (state.queue_tail + 1) % ctx.queue_slots
} else {
state.num_threads_busy == state.thread_limit || !state.queue_empty
}
}
unsafe fn add_internal(
ctx: &PoolCtx,
state: &mut PoolState,
function: PoolFunction,
opaque: *mut c_void,
) {
if state.shutdown {
return;
}
unsafe {
ctx.queue
.add(state.queue_tail)
.write(MaybeUninit::new(Job { function, opaque }));
}
state.queue_empty = false;
state.queue_tail = (state.queue_tail + 1) % ctx.queue_slots;
ctx.queue_pop_cond.notify_one();
}
#[no_mangle]
pub extern "C" fn POOL_create(num_threads: usize, queue_size: usize) -> *mut PoolCtx {
unsafe { POOL_create_advanced(num_threads, queue_size, DEFAULT_CUSTOM_MEM) }
}
/// Public ZSTD API alias for `POOL_create` (see `ZSTD_threadPool` in zstd.h).
#[no_mangle]
pub extern "C" fn ZSTD_createThreadPool(num_threads: usize) -> *mut PoolCtx {
POOL_create(num_threads, 0)
}
/// Public ZSTD API alias for `POOL_free`.
///
/// # Safety
/// `pool` must be null or a pointer returned by a pool creation function.
#[no_mangle]
pub unsafe extern "C" fn ZSTD_freeThreadPool(pool: *mut PoolCtx) {
unsafe { POOL_free(pool) }
}
#[no_mangle]
pub unsafe extern "C" fn POOL_create_advanced(
num_threads: usize,
queue_size: usize,
custom_mem: ZstdCustomMem,
) -> *mut PoolCtx {
if !multithreading_enabled() {
return single_threaded_pool();
}
unsafe { create_threaded_pool(num_threads, queue_size, custom_mem) }
}
#[no_mangle]
pub unsafe extern "C" fn POOL_free(ctx: *mut PoolCtx) {
if ctx.is_null() || is_single_threaded_pool(ctx) {
return;
}
unsafe {
shutdown_and_join(ctx);
destroy_context(ctx);
}
}
#[no_mangle]
pub unsafe extern "C" fn POOL_joinJobs(ctx: *mut PoolCtx) {
if ctx.is_null() || is_single_threaded_pool(ctx) {
return;
}
let ctx = unsafe { &*ctx };
let mut state = ctx.state.lock().unwrap_or_else(|error| error.into_inner());
while !state.queue_empty || state.num_threads_busy > 0 {
state = ctx
.queue_push_cond
.wait(state)
.unwrap_or_else(|error| error.into_inner());
}
}
#[no_mangle]
pub unsafe extern "C" fn POOL_resize(ctx: *mut PoolCtx, num_threads: usize) -> c_int {
if is_single_threaded_pool(ctx) {
return 0;
}
if ctx.is_null() || num_threads == 0 {
return 1;
}
let ctx_ref = unsafe { &*ctx };
let mut state = ctx_ref
.state
.lock()
.unwrap_or_else(|error| error.into_inner());
let threads = unsafe { &mut *ctx_ref.threads.get() };
if num_threads <= threads.initialized {
state.thread_limit = num_threads;
drop(state);
ctx_ref.queue_pop_cond.notify_all();
return 0;
}
let Some(thread_bytes) = checked_array_size::<MaybeUninit<JoinHandle<()>>>(num_threads) else {
drop(state);
ctx_ref.queue_pop_cond.notify_all();
return 1;
};
let replacement = unsafe { custom_calloc(thread_bytes, ctx_ref.custom_mem) }
.cast::<MaybeUninit<JoinHandle<()>>>();
if replacement.is_null() {
drop(state);
ctx_ref.queue_pop_cond.notify_all();
return 1;
}
let old_initialized = threads.initialized;
for thread_id in 0..old_initialized {
let handle = unsafe { threads.slots.add(thread_id).read().assume_init() };
unsafe {
replacement.add(thread_id).write(MaybeUninit::new(handle));
}
}
unsafe { custom_free(threads.slots.cast(), ctx_ref.custom_mem) };
threads.slots = replacement;
for thread_id in old_initialized..num_threads {
match spawn_worker(ctx) {
Ok(handle) => {
unsafe {
replacement.add(thread_id).write(MaybeUninit::new(handle));
}
threads.initialized = thread_id + 1;
}
Err(_) => {
drop(state);
ctx_ref.queue_pop_cond.notify_all();
return 1;
}
}
}
state.thread_limit = num_threads;
drop(state);
ctx_ref.queue_pop_cond.notify_all();
0
}
#[no_mangle]
pub unsafe extern "C" fn POOL_sizeof(ctx: *const PoolCtx) -> usize {
if ctx.is_null() {
return 0;
}
if is_single_threaded_pool(ctx) {
return mem::size_of::<c_int>();
}
let ctx = unsafe { &*ctx };
let _state = ctx.state.lock().unwrap_or_else(|error| error.into_inner());
let thread_capacity = unsafe { (*ctx.threads.get()).initialized };
mem::size_of::<PoolCtx>()
.saturating_add(
ctx.queue_slots
.saturating_mul(mem::size_of::<MaybeUninit<Job>>()),
)
.saturating_add(
thread_capacity.saturating_mul(mem::size_of::<MaybeUninit<JoinHandle<()>>>()),
)
}
#[no_mangle]
pub unsafe extern "C" fn POOL_add(ctx: *mut PoolCtx, function: PoolFunction, opaque: *mut c_void) {
if is_single_threaded_pool(ctx) {
unsafe { function(opaque) };
return;
}
if ctx.is_null() {
return;
}
let ctx = unsafe { &*ctx };
let mut state = ctx.state.lock().unwrap_or_else(|error| error.into_inner());
while queue_is_full(ctx, &state) && !state.shutdown {
state = ctx
.queue_push_cond
.wait(state)
.unwrap_or_else(|error| error.into_inner());
}
unsafe { add_internal(ctx, &mut state, function, opaque) };
}
#[no_mangle]
pub unsafe extern "C" fn POOL_tryAdd(
ctx: *mut PoolCtx,
function: PoolFunction,
opaque: *mut c_void,
) -> c_int {
if is_single_threaded_pool(ctx) {
unsafe { function(opaque) };
return 1;
}
if ctx.is_null() {
return 0;
}
let ctx = unsafe { &*ctx };
let mut state = ctx.state.lock().unwrap_or_else(|error| error.into_inner());
if queue_is_full(ctx, &state) {
return 0;
}
unsafe { add_internal(ctx, &mut state, function, opaque) };
1
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::time::{Duration, Instant};
struct Gate {
started: AtomicUsize,
released: AtomicBool,
}
unsafe extern "C" fn gated_job(opaque: *mut c_void) {
let gate = unsafe { &*(opaque.cast::<Gate>()) };
gate.started.fetch_add(1, Ordering::Release);
while !gate.released.load(Ordering::Acquire) {
thread::yield_now();
}
}
unsafe extern "C" fn increment(opaque: *mut c_void) {
let count = unsafe { &*(opaque.cast::<AtomicUsize>()) };
count.fetch_add(1, Ordering::Relaxed);
}
fn wait_until(mut predicate: impl FnMut() -> bool) {
let deadline = Instant::now() + Duration::from_secs(5);
while !predicate() {
assert!(Instant::now() < deadline, "timed out waiting for pool job");
thread::yield_now();
}
}
#[test]
fn zero_sized_queue_only_accepts_jobs_for_free_workers() {
let ctx = POOL_create(1, 0);
assert!(!ctx.is_null());
let gate = Gate {
started: AtomicUsize::new(0),
released: AtomicBool::new(false),
};
let count = AtomicUsize::new(0);
unsafe { POOL_add(ctx, gated_job, ptr::from_ref(&gate).cast_mut().cast()) };
wait_until(|| gate.started.load(Ordering::Acquire) == 1);
assert_eq!(
unsafe { POOL_tryAdd(ctx, increment, ptr::from_ref(&count).cast_mut().cast()) },
0
);
gate.released.store(true, Ordering::Release);
unsafe { POOL_joinJobs(ctx) };
assert_eq!(
unsafe { POOL_tryAdd(ctx, increment, ptr::from_ref(&count).cast_mut().cast()) },
1
);
unsafe {
POOL_joinJobs(ctx);
POOL_free(ctx);
}
assert_eq!(count.load(Ordering::Relaxed), 1);
}
#[test]
fn bounded_queue_rejects_a_job_when_its_only_slot_is_occupied() {
let ctx = POOL_create(1, 1);
assert!(!ctx.is_null());
let gate = Gate {
started: AtomicUsize::new(0),
released: AtomicBool::new(false),
};
let count = AtomicUsize::new(0);
unsafe { POOL_add(ctx, gated_job, ptr::from_ref(&gate).cast_mut().cast()) };
wait_until(|| gate.started.load(Ordering::Acquire) == 1);
unsafe { POOL_add(ctx, increment, ptr::from_ref(&count).cast_mut().cast()) };
assert_eq!(
unsafe { POOL_tryAdd(ctx, increment, ptr::from_ref(&count).cast_mut().cast()) },
0
);
gate.released.store(true, Ordering::Release);
unsafe { POOL_free(ctx) };
assert_eq!(count.load(Ordering::Relaxed), 1);
}
#[test]
fn add_blocks_until_a_bounded_queue_slot_opens() {
let ctx = POOL_create(1, 1);
assert!(!ctx.is_null());
let gate = Gate {
started: AtomicUsize::new(0),
released: AtomicBool::new(false),
};
let count = std::sync::Arc::new(AtomicUsize::new(0));
let producer_returned = std::sync::Arc::new(AtomicBool::new(false));
unsafe { POOL_add(ctx, gated_job, ptr::from_ref(&gate).cast_mut().cast()) };
wait_until(|| gate.started.load(Ordering::Acquire) == 1);
unsafe {
POOL_add(
ctx,
increment,
std::sync::Arc::as_ptr(&count).cast_mut().cast(),
)
};
let send_ctx = SendPoolPtr(ctx);
let producer_count = std::sync::Arc::clone(&count);
let producer_returned_clone = std::sync::Arc::clone(&producer_returned);
let producer = thread::spawn(move || {
unsafe {
POOL_add(
send_ctx.get(),
increment,
std::sync::Arc::as_ptr(&producer_count).cast_mut().cast(),
)
};
producer_returned_clone.store(true, Ordering::Release);
});
thread::sleep(Duration::from_millis(10));
assert!(!producer_returned.load(Ordering::Acquire));
gate.released.store(true, Ordering::Release);
producer.join().unwrap();
unsafe { POOL_free(ctx) };
assert!(producer_returned.load(Ordering::Acquire));
assert_eq!(count.load(Ordering::Relaxed), 2);
}
#[test]
fn resize_expansion_creates_workers_and_size_tracks_capacity() {
let ctx = POOL_create(1, 3);
assert!(!ctx.is_null());
let original_size = unsafe { POOL_sizeof(ctx) };
assert_eq!(unsafe { POOL_resize(ctx, 3) }, 0);
assert!(unsafe { POOL_sizeof(ctx) } > original_size);
let gate = Gate {
started: AtomicUsize::new(0),
released: AtomicBool::new(false),
};
for _ in 0..3 {
unsafe { POOL_add(ctx, gated_job, ptr::from_ref(&gate).cast_mut().cast()) };
}
wait_until(|| gate.started.load(Ordering::Acquire) == 3);
gate.released.store(true, Ordering::Release);
unsafe { POOL_free(ctx) };
}
#[test]
fn free_drains_queued_jobs_after_reducing_the_limit() {
let ctx = POOL_create(3, 16);
assert!(!ctx.is_null());
let count = AtomicUsize::new(0);
for _ in 0..16 {
unsafe { POOL_add(ctx, increment, ptr::from_ref(&count).cast_mut().cast()) };
}
assert_eq!(unsafe { POOL_resize(ctx, 1) }, 0);
unsafe { POOL_free(ctx) };
assert_eq!(count.load(Ordering::Relaxed), 16);
}
struct AllocStats {
allocations: AtomicUsize,
frees: AtomicUsize,
}
unsafe extern "C" fn counting_alloc(opaque: *mut c_void, size: usize) -> *mut c_void {
let stats = unsafe { &*(opaque.cast::<AllocStats>()) };
stats.allocations.fetch_add(1, Ordering::Relaxed);
unsafe { libc::calloc(1, size) }
}
unsafe extern "C" fn counting_free(opaque: *mut c_void, allocation: *mut c_void) {
let stats = unsafe { &*(opaque.cast::<AllocStats>()) };
stats.frees.fetch_add(1, Ordering::Relaxed);
unsafe { libc::free(allocation) };
}
#[test]
fn advanced_pool_uses_the_custom_allocator_for_owned_storage() {
let stats = AllocStats {
allocations: AtomicUsize::new(0),
frees: AtomicUsize::new(0),
};
let custom_mem = ZstdCustomMem {
custom_alloc: Some(counting_alloc),
custom_free: Some(counting_free),
opaque: ptr::from_ref(&stats).cast_mut().cast(),
};
let ctx = unsafe { POOL_create_advanced(2, 2, custom_mem) };
assert!(!ctx.is_null());
unsafe { POOL_free(ctx) };
assert_eq!(stats.allocations.load(Ordering::Relaxed), 3);
assert_eq!(stats.frees.load(Ordering::Relaxed), 3);
}
#[test]
fn custom_memory_layout_matches_the_c_abi() {
assert_eq!(mem::size_of::<ZstdCustomMem>(), 3 * mem::size_of::<usize>());
assert_eq!(mem::align_of::<ZstdCustomMem>(), mem::align_of::<usize>());
}
#[test]
fn non_threaded_pool_executes_callbacks_synchronously() {
let ctx = single_threaded_pool();
let count = AtomicUsize::new(0);
unsafe {
POOL_add(ctx, increment, ptr::from_ref(&count).cast_mut().cast());
assert_eq!(count.load(Ordering::Relaxed), 1);
assert_eq!(
POOL_tryAdd(ctx, increment, ptr::from_ref(&count).cast_mut().cast()),
1
);
assert_eq!(POOL_resize(ctx, 0), 0);
assert_eq!(POOL_sizeof(ctx), mem::size_of::<c_int>());
POOL_free(ctx);
}
assert_eq!(count.load(Ordering::Relaxed), 2);
}
#[test]
fn rejects_zero_threads_in_threaded_mode() {
assert!(unsafe { create_threaded_pool(0, 1, DEFAULT_CUSTOM_MEM) }.is_null());
assert_eq!(unsafe { POOL_resize(ptr::null_mut(), 1) }, 1);
assert_eq!(unsafe { POOL_sizeof(ptr::null()) }, 0);
}
}
+282
View File
@@ -0,0 +1,282 @@
use std::os::raw::{c_int, c_void};
use std::ptr;
#[cfg(unix)]
use libc::{
pthread_cond_destroy, pthread_cond_init, pthread_cond_t, pthread_condattr_t, pthread_create,
pthread_join, pthread_mutex_destroy, pthread_mutex_init, pthread_mutex_t, pthread_mutexattr_t,
pthread_t,
};
/// Allocating mutex wrapper used by threading.h when `DEBUGLEVEL >= 1`.
#[cfg(unix)]
#[no_mangle]
pub unsafe extern "C" fn ZSTD_pthread_mutex_init(
mutex: *mut *mut pthread_mutex_t,
attr: *const pthread_mutexattr_t,
) -> c_int {
if mutex.is_null() {
return libc::EINVAL;
}
let allocation =
unsafe { libc::malloc(std::mem::size_of::<pthread_mutex_t>()) }.cast::<pthread_mutex_t>();
unsafe { mutex.write(allocation) };
if allocation.is_null() {
return 1;
}
unsafe { pthread_mutex_init(allocation, attr) }
}
/// Destroying mutex wrapper used by threading.h when `DEBUGLEVEL >= 1`.
#[cfg(unix)]
#[no_mangle]
pub unsafe extern "C" fn ZSTD_pthread_mutex_destroy(mutex: *mut *mut pthread_mutex_t) -> c_int {
if mutex.is_null() {
return libc::EINVAL;
}
let allocation = unsafe { mutex.read() };
if allocation.is_null() {
return 0;
}
let result = unsafe { pthread_mutex_destroy(allocation) };
unsafe { libc::free(allocation.cast()) };
result
}
/// Allocating condition-variable wrapper used by threading.h when
/// `DEBUGLEVEL >= 1`.
#[cfg(unix)]
#[no_mangle]
pub unsafe extern "C" fn ZSTD_pthread_cond_init(
cond: *mut *mut pthread_cond_t,
attr: *const pthread_condattr_t,
) -> c_int {
if cond.is_null() {
return libc::EINVAL;
}
let allocation =
unsafe { libc::malloc(std::mem::size_of::<pthread_cond_t>()) }.cast::<pthread_cond_t>();
unsafe { cond.write(allocation) };
if allocation.is_null() {
return 1;
}
unsafe { pthread_cond_init(allocation, attr) }
}
/// Destroying condition-variable wrapper used by threading.h when
/// `DEBUGLEVEL >= 1`.
#[cfg(unix)]
#[no_mangle]
pub unsafe extern "C" fn ZSTD_pthread_cond_destroy(cond: *mut *mut pthread_cond_t) -> c_int {
if cond.is_null() {
return libc::EINVAL;
}
let allocation = unsafe { cond.read() };
if allocation.is_null() {
return 0;
}
let result = unsafe { pthread_cond_destroy(allocation) };
unsafe { libc::free(allocation.cast()) };
result
}
/// POSIX-compatible thread creation entry point. On POSIX this symbol is
/// normally bypassed by threading.h's macro, but keeping it ABI-correct makes
/// the Rust object usable by configurations which reference the wrapper.
#[cfg(unix)]
#[no_mangle]
pub unsafe extern "C" fn ZSTD_pthread_create(
thread: *mut pthread_t,
attr: *const libc::pthread_attr_t,
start_routine: extern "C" fn(*mut c_void) -> *mut c_void,
arg: *mut c_void,
) -> c_int {
if thread.is_null() {
return -1;
}
unsafe { pthread_create(thread, attr, start_routine, arg) }
}
/// POSIX-compatible join which intentionally discards the callback result,
/// matching `ZSTD_pthread_join` in threading.h.
#[cfg(unix)]
#[no_mangle]
pub unsafe extern "C" fn ZSTD_pthread_join(thread: pthread_t) -> c_int {
unsafe { pthread_join(thread, ptr::null_mut()) }
}
#[cfg(windows)]
type Handle = *mut c_void;
#[cfg(windows)]
struct WindowsThreadParams {
start_routine: extern "C" fn(*mut c_void) -> *mut c_void,
arg: *mut c_void,
}
#[cfg(windows)]
unsafe impl Send for WindowsThreadParams {}
#[cfg(windows)]
unsafe extern "system" fn windows_worker(arg: *mut c_void) -> u32 {
let params = unsafe { Box::from_raw(arg.cast::<WindowsThreadParams>()) };
(params.start_routine)(params.arg);
0
}
#[cfg(windows)]
#[link(name = "msvcrt")]
unsafe extern "C" {
fn _beginthreadex(
security: *mut c_void,
stack_size: u32,
start_address: Option<unsafe extern "system" fn(*mut c_void) -> u32>,
arglist: *mut c_void,
initflag: u32,
thread_address: *mut u32,
) -> usize;
fn _errno() -> *mut c_int;
}
#[cfg(windows)]
#[link(name = "kernel32")]
unsafe extern "system" {
fn WaitForSingleObject(handle: Handle, milliseconds: u32) -> u32;
fn CloseHandle(handle: Handle) -> c_int;
fn GetLastError() -> u32;
}
#[cfg(windows)]
#[no_mangle]
pub unsafe extern "C" fn ZSTD_pthread_create(
thread: *mut Handle,
_unused: *const c_void,
start_routine: extern "C" fn(*mut c_void) -> *mut c_void,
arg: *mut c_void,
) -> c_int {
if thread.is_null() {
return -1;
}
unsafe { thread.write(ptr::null_mut()) };
let params = Box::new(WindowsThreadParams { start_routine, arg });
let params = Box::into_raw(params);
let handle = unsafe {
_beginthreadex(
ptr::null_mut(),
0,
Some(windows_worker),
params.cast(),
0,
ptr::null_mut(),
)
};
if handle == 0 {
unsafe { drop(Box::from_raw(params)) };
let errno = unsafe { _errno() };
return if errno.is_null() {
-1
} else {
unsafe { *errno }
};
}
unsafe { thread.write(handle as Handle) };
0
}
#[cfg(windows)]
#[no_mangle]
pub unsafe extern "C" fn ZSTD_pthread_join(thread: Handle) -> c_int {
const INFINITE: u32 = u32::MAX;
const WAIT_OBJECT_0: u32 = 0;
const WAIT_ABANDONED: u32 = 0x80;
if thread.is_null() {
return 0;
}
let result = unsafe { WaitForSingleObject(thread, INFINITE) };
unsafe { CloseHandle(thread) };
match result {
WAIT_OBJECT_0 => 0,
WAIT_ABANDONED => libc::EINVAL,
_ => unsafe { GetLastError() as c_int },
}
}
#[cfg(all(test, unix))]
mod tests {
use super::*;
use std::mem::MaybeUninit;
use std::sync::atomic::{AtomicUsize, Ordering};
static THREAD_CALLS: AtomicUsize = AtomicUsize::new(0);
extern "C" fn thread_main(arg: *mut c_void) -> *mut c_void {
THREAD_CALLS.fetch_add(1, Ordering::Relaxed);
arg
}
#[test]
fn debug_mutex_wrapper_allocates_initializes_and_frees() {
let mut mutex: *mut pthread_mutex_t = ptr::null_mut();
unsafe {
assert_eq!(ZSTD_pthread_mutex_init(&mut mutex, ptr::null()), 0);
assert!(!mutex.is_null());
assert_eq!(libc::pthread_mutex_lock(mutex), 0);
assert_eq!(libc::pthread_mutex_unlock(mutex), 0);
assert_eq!(ZSTD_pthread_mutex_destroy(&mut mutex), 0);
}
}
#[test]
fn debug_condition_wrapper_allocates_initializes_and_frees() {
let mut cond: *mut pthread_cond_t = ptr::null_mut();
unsafe {
assert_eq!(ZSTD_pthread_cond_init(&mut cond, ptr::null()), 0);
assert!(!cond.is_null());
assert_eq!(ZSTD_pthread_cond_destroy(&mut cond), 0);
}
}
#[test]
fn pthread_wrapper_runs_and_joins_callback() {
THREAD_CALLS.store(0, Ordering::Relaxed);
let mut thread = MaybeUninit::<pthread_t>::uninit();
unsafe {
assert_eq!(
ZSTD_pthread_create(
thread.as_mut_ptr(),
ptr::null(),
thread_main,
ptr::null_mut()
),
0
);
assert_eq!(ZSTD_pthread_join(thread.assume_init()), 0);
}
assert_eq!(THREAD_CALLS.load(Ordering::Relaxed), 1);
}
#[test]
fn null_output_pointers_are_rejected() {
unsafe {
assert_eq!(
ZSTD_pthread_mutex_init(ptr::null_mut(), ptr::null()),
libc::EINVAL
);
assert_eq!(
ZSTD_pthread_cond_init(ptr::null_mut(), ptr::null()),
libc::EINVAL
);
assert_eq!(
ZSTD_pthread_create(ptr::null_mut(), ptr::null(), thread_main, ptr::null_mut()),
-1
);
}
}
}