use crate::{Error, Result};
use core::alloc::Layout;
use core::cell::UnsafeCell;
use core::mem::{self, MaybeUninit};
use core::ops::{Deref, DerefMut};
use core::ptr::{self, NonNull};
use core::sync::atomic::{AtomicI32, Ordering};
use hipool::{Boxed, PoolAlloc};
struct Handle<T: Send + 'static> {
tid: libc::pthread_t,
err: i32,
layout: Layout,
output: MaybeUninit<T>,
offset: usize,
}
pub struct JoinHandle<T: Send + 'static> {
handle: NonNull<Handle<T>>,
}
unsafe impl<T: Send> Send for JoinHandle<T> {}
impl<T: Send + 'static> JoinHandle<T> {
fn new<F>(f: F) -> Self
where
F: FnOnce() -> T + Send + 'static,
{
let layout = Layout::new::<Handle<T>>();
let (layout, offset) = layout.extend(Layout::new::<F>()).unwrap();
let handle = Boxed::new_buf_then(layout, |ptr| {
let addr = ptr.cast::<u8>().as_ptr();
let closure = unsafe { addr.add(offset).cast::<F>() };
unsafe { closure.write(f) };
let handle = addr.cast::<Handle<T>>();
unsafe {
handle.write(Handle::<T> {
tid: 0,
err: 0,
output: MaybeUninit::uninit(),
layout,
offset,
})
};
Ok(())
}).unwrap();
let handle = unsafe { handle.cast::<Handle<T>>().unwrap() };
Self {
handle: handle.leak().0.into(),
}
}
pub fn join(self) -> Result<T> {
let handle = unsafe { self.handle.as_ref() };
let _guard = unsafe { Boxed::from_with(self.handle, handle.layout, &PoolAlloc) };
if handle.err > 0 {
return Err(Error::new(handle.err));
}
let ret = unsafe { libc::pthread_join(handle.tid, ptr::null_mut::<*mut libc::c_void>()) };
if ret == 0 {
Ok(unsafe { handle.output.assume_init_read() })
} else {
Err(Error::new(ret))
}
}
pub fn failed(&self) -> bool {
unsafe { self.handle.as_ref().err > 0 }
}
fn set_failed(&mut self, errno: i32) {
unsafe {
self.handle.as_mut().err = errno;
}
}
fn tid_mut(&mut self) -> *mut libc::pthread_t {
unsafe { &mut self.handle.as_mut().tid }
}
}
type LibcRunnable = extern "C" fn(*mut libc::c_void) -> *mut libc::c_void;
pub fn spawn<F, T>(f: F) -> JoinHandle<T>
where
F: FnOnce() -> T + Send + 'static,
T: Send + 'static,
{
let mut handle = JoinHandle::new(f);
extern "C" fn runnable<T, F>(data: *mut libc::c_void) -> *mut libc::c_void
where
F: FnOnce() -> T + Send + 'static,
T: Send + 'static,
{
let handle = data as *mut Handle<T>;
let closure = unsafe { data.add((*handle).offset) as *mut F };
let fun = unsafe { closure.read() };
unsafe { (*handle).output.write(fun()) };
data
}
let errno = unsafe {
libc::pthread_create(
handle.tid_mut(),
ptr::null::<libc::pthread_attr_t>(),
runnable::<T, F> as LibcRunnable,
handle.handle.as_ptr() as *mut libc::c_void,
)
};
handle.set_failed(errno);
handle
}
pub fn yield_now() {
unsafe { libc::sched_yield() };
}
pub struct Semaphore {
sem: MaybeUninit<libc::sem_t>,
}
unsafe impl Send for Semaphore {}
unsafe impl Sync for Semaphore {}
impl Drop for Semaphore {
fn drop(&mut self) {
unsafe { libc::sem_destroy(self.sem.as_mut_ptr()) };
}
}
impl Semaphore {
pub fn new() -> Result<Self> {
let mut this = Self {
sem: MaybeUninit::uninit(),
};
let ret = unsafe { libc::sem_init(this.sem.as_mut_ptr(), 0, 0) };
if ret == 0 {
Ok(this)
} else {
Err(Error::last())
}
}
pub fn post(&self) {
unsafe { libc::sem_post(self.sem.as_ptr().cast_mut()) };
}
pub fn wait(&self) {
unsafe { libc::sem_wait(self.sem.as_ptr().cast_mut()) };
}
}
pub struct Mutex<T: ?Sized> {
lock: libc::pthread_mutex_t,
val: UnsafeCell<T>,
}
unsafe impl<T: Send + ?Sized> Send for Mutex<T> {}
unsafe impl<T: ?Sized> Sync for Mutex<T> {}
pub struct MutexGuard<'a, T: ?Sized + 'a> {
mutex: &'a Mutex<T>,
}
impl<T: ?Sized> Drop for Mutex<T> {
fn drop(&mut self) {
unsafe {
libc::pthread_mutex_destroy(&mut self.lock);
}
}
}
impl<T> Mutex<T> {
pub const fn new(val: T) -> Self {
Self {
lock: libc::PTHREAD_MUTEX_INITIALIZER,
val: UnsafeCell::new(val),
}
}
}
impl<T: ?Sized> Mutex<T> {
pub fn lock(&self) -> MutexGuard<'_, T> {
let ret = unsafe { libc::pthread_mutex_lock(&self.lock as *const _ as *mut _) };
assert_eq!(ret, 0);
MutexGuard::new(self)
}
pub fn try_lock(&self) -> Option<MutexGuard<'_, T>> {
let ret = unsafe { libc::pthread_mutex_trylock(&self.lock as *const _ as *mut _) };
if ret == 0 {
Some(MutexGuard::new(self))
} else {
None
}
}
}
impl<'a, T: ?Sized> MutexGuard<'a, T> {
fn new(mutex: &'a Mutex<T>) -> Self {
Self { mutex }
}
}
impl<T: ?Sized> Drop for MutexGuard<'_, T> {
fn drop(&mut self) {
unsafe {
libc::pthread_mutex_unlock(&self.mutex.lock as *const _ as *mut _);
}
}
}
impl<T: ?Sized> Deref for MutexGuard<'_, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
unsafe { &*self.mutex.val.get() }
}
}
impl<T: ?Sized> DerefMut for MutexGuard<'_, T> {
fn deref_mut(&mut self) -> &mut Self::Target {
unsafe { &mut *self.mutex.val.get() }
}
}
pub struct Cond {
cond: libc::pthread_cond_t,
mutex: libc::pthread_mutex_t,
}
unsafe impl Send for Cond {}
unsafe impl Sync for Cond {}
impl Drop for Cond {
fn drop(&mut self) {
unsafe { libc::pthread_cond_destroy(&mut self.cond) };
unsafe { libc::pthread_mutex_destroy(&mut self.mutex) };
}
}
impl Cond {
pub const fn new() -> Self {
Self {
cond: libc::PTHREAD_COND_INITIALIZER,
mutex: libc::PTHREAD_MUTEX_INITIALIZER,
}
}
pub fn signal<F>(&self, mut f: F)
where
F: FnMut(),
{
unsafe { libc::pthread_mutex_lock(&self.mutex as *const _ as *mut _) };
f();
unsafe { libc::pthread_cond_signal(&self.cond as *const _ as *mut _) };
unsafe { libc::pthread_mutex_unlock(&self.mutex as *const _ as *mut _) };
}
pub fn broadcast<F>(&self, mut f: F)
where
F: FnMut(),
{
unsafe { libc::pthread_mutex_lock(&self.mutex as *const _ as *mut _) };
f();
unsafe { libc::pthread_cond_broadcast(&self.cond as *const _ as *mut _) };
unsafe { libc::pthread_mutex_unlock(&self.mutex as *const _ as *mut _) };
}
pub fn wait<R, C, P>(&self, c: C, mut p: P) -> R
where
C: Fn() -> bool,
P: FnMut() -> R,
{
unsafe { libc::pthread_mutex_lock(&self.mutex as *const _ as *mut _) };
while !c() {
unsafe {
libc::pthread_cond_wait(
&self.cond as *const _ as *mut _,
&self.mutex as *const _ as *mut _,
)
};
}
let ret = p();
unsafe { libc::pthread_mutex_unlock(&self.mutex as *const _ as *mut _) };
ret
}
}
pub struct RWLock<T: ?Sized> {
lock: libc::pthread_rwlock_t,
val: UnsafeCell<T>,
}
unsafe impl<T: Send + ?Sized> Send for RWLock<T> {}
unsafe impl<T: ?Sized> Sync for RWLock<T> {}
impl<T> RWLock<T> {
pub const fn new(val: T) -> Self {
Self {
lock: libc::PTHREAD_RWLOCK_INITIALIZER,
val: UnsafeCell::new(val),
}
}
}
impl<T: ?Sized> RWLock<T> {
pub fn rlock(&self) -> RLockGuard<'_, T> {
unsafe { libc::pthread_rwlock_rdlock(self as *const _ as *mut _) };
RLockGuard::new(self)
}
pub fn wlock(&self) -> WLockGuard<'_, T> {
unsafe { libc::pthread_rwlock_wrlock(self as *const _ as *mut _) };
WLockGuard::new(self)
}
pub fn try_rlock(&self) -> Option<RLockGuard<'_, T>> {
let ret = unsafe { libc::pthread_rwlock_tryrdlock(self as *const _ as *mut _) };
if ret == 0 {
Some(RLockGuard::new(self))
} else {
None
}
}
pub fn try_wlock(&self) -> Option<WLockGuard<'_, T>> {
let ret = unsafe { libc::pthread_rwlock_trywrlock(self as *const _ as *mut _) };
if ret == 0 {
Some(WLockGuard::new(self))
} else {
None
}
}
}
pub struct RLockGuard<'a, T: ?Sized> {
lock: &'a RWLock<T>,
}
impl<'a, T: ?Sized> RLockGuard<'a, T> {
fn new(lock: &'a RWLock<T>) -> Self {
Self { lock }
}
}
impl<T: ?Sized> Drop for RLockGuard<'_, T> {
fn drop(&mut self) {
unsafe { libc::pthread_rwlock_unlock(&self.lock.lock as *const _ as *mut _) };
}
}
impl<T: ?Sized> Deref for RLockGuard<'_, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
unsafe { &*self.lock.val.get() }
}
}
pub struct WLockGuard<'a, T: ?Sized> {
lock: &'a RWLock<T>,
}
impl<'a, T: ?Sized> WLockGuard<'a, T> {
fn new(lock: &'a RWLock<T>) -> Self {
Self { lock }
}
}
impl<T: ?Sized> Drop for WLockGuard<'_, T> {
fn drop(&mut self) {
unsafe { libc::pthread_rwlock_unlock(&self.lock.lock as *const _ as *mut _) };
}
}
impl<T: ?Sized> Deref for WLockGuard<'_, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
unsafe { &*self.lock.val.get() }
}
}
impl<T: ?Sized> DerefMut for WLockGuard<'_, T> {
fn deref_mut(&mut self) -> &mut Self::Target {
unsafe { &mut *self.lock.val.get() }
}
}
#[derive(Debug)]
pub struct TssData {
key: libc::pthread_key_t,
}
unsafe impl Send for TssData {}
unsafe impl Sync for TssData {}
impl Drop for TssData {
fn drop(&mut self) {
unsafe { libc::pthread_key_delete(self.key) };
}
}
impl TssData {
pub fn new() -> Result<Self> {
let mut key: libc::pthread_key_t = 0;
let ret = unsafe { libc::pthread_key_create(&mut key, None) };
if ret == 0 {
Ok(Self { key })
} else {
Err(Error::new(ret))
}
}
pub fn set(&self, val: usize) {
unsafe {
libc::pthread_setspecific(self.key, val as *const libc::c_void);
}
}
pub fn get(&self) -> usize {
unsafe { libc::pthread_getspecific(self.key) as usize }
}
}
pub struct Once {
flag: AtomicI32,
lock: Mutex<()>,
}
impl Once {
pub const fn new() -> Self {
Self {
flag: AtomicI32::new(0),
lock: Mutex::new(()),
}
}
pub fn call_once<F>(&self, mut f: F)
where
F: FnMut(),
{
if self.flag.load(Ordering::Acquire) > 0 {
return;
}
let _ = self.lock.lock();
if self.flag.load(Ordering::Relaxed) == 0 {
f();
self.flag.store(1, Ordering::Release);
}
}
}
pub fn get_cpu_count() -> usize {
let mut set = MaybeUninit::<libc::cpu_set_t>::zeroed();
let size = mem::size_of_val(&set);
let ret = unsafe { libc::sched_getaffinity(libc::getpid(), size, set.as_mut_ptr()) };
if ret == 0 {
(unsafe { libc::CPU_COUNT(set.assume_init_ref()) }) as usize
} else {
1
}
}
pub fn set_cpu(cpu: usize) {
let mut set = MaybeUninit::<libc::cpu_set_t>::zeroed();
let size = mem::size_of_val(&set);
let ret = unsafe { libc::sched_getaffinity(libc::getpid(), size, set.as_mut_ptr()) };
if ret != 0 {
return;
}
let cnt = unsafe { libc::CPU_COUNT(set.assume_init_ref()) } as usize;
if cpu >= cnt {
return;
}
let mut idx = 0;
for n in 0.. {
if unsafe { libc::CPU_ISSET(n, set.assume_init_ref()) } {
if idx == cpu {
unsafe { libc::CPU_ZERO(set.assume_init_mut()) };
unsafe { libc::CPU_SET(n, set.assume_init_mut()) };
unsafe { libc::sched_setaffinity(0, size, set.as_ptr()) };
return;
}
idx += 1;
}
}
}