zstd-rs/rust/src/threading.rs

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
            );
        }
    }
}