use crate::{Error, Result};
use core::alloc::Layout;
use core::cell::{Cell, UnsafeCell};
use core::fmt;
use core::mem::{self, ManuallyDrop, MaybeUninit};
use core::ops::{Deref, DerefMut};
use core::ptr::{self, NonNull};
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 {
done: Cell<bool>,
lock: Mutex<()>,
}
impl Once {
pub const fn new() -> Self {
Self {
done: Cell::new(false),
lock: Mutex::new(()),
}
}
pub fn call_once<F>(&self, f: F)
where
F: FnOnce(),
{
self.init(f);
}
fn init<F>(&self, f: F) -> bool
where
F: FnOnce(),
{
if self.done.get() {
return false;
}
let _guard = self.lock.lock();
if !self.done.get() {
f();
self.done.set(true);
return true;
}
false
}
fn is_called(&self) -> bool {
self.done.get()
}
fn set_called(&mut self, called: bool) {
self.done.set(called);
}
}
pub struct OnceLock<T> {
value: UnsafeCell<MaybeUninit<T>>,
once: Once,
}
unsafe impl<T: Sync> Sync for OnceLock<T> {}
unsafe impl<T: Send> Send for OnceLock<T> {}
impl<T> OnceLock<T> {
pub const fn new() -> Self {
Self {
value: UnsafeCell::new(MaybeUninit::uninit()),
once: Once::new(),
}
}
pub fn get(&self) -> Option<&T> {
if self.once.is_called() {
return Some(unsafe { (*self.value.get()).assume_init_ref() });
}
None
}
pub fn get_mut(&mut self) -> Option<&mut T> {
if self.once.is_called() {
return Some(unsafe { (*self.value.get()).assume_init_mut() });
}
None
}
pub fn get_or_init<F>(&self, f: F) -> &T
where
F: FnOnce() -> T,
{
self.once.call_once(|| unsafe {
(*self.value.get()).write(f());
});
unsafe { (*self.value.get()).assume_init_ref() }
}
pub fn get_mut_or_init<F>(&mut self, f: F) -> &mut T
where
F: FnOnce() -> T,
{
self.once.call_once(|| unsafe {
(*self.value.get()).write(f());
});
unsafe { self.value.get_mut().assume_init_mut() }
}
pub fn take(&mut self) -> Option<T> {
if self.once.is_called() {
self.once.set_called(false);
return Some(unsafe { self.value.get_mut().assume_init_read() });
}
None
}
pub fn set(&self, value: T) -> core::result::Result<(), T> {
let mut value = ManuallyDrop::new(value);
if self.once.init(|| unsafe {
(*self.value.get()).write(ManuallyDrop::take(&mut value));
}) {
return Ok(());
}
Err(unsafe { ManuallyDrop::take(&mut value) })
}
}
impl<T> Default for OnceLock<T> {
fn default() -> Self {
Self::new()
}
}
impl<T> Drop for OnceLock<T> {
fn drop(&mut self) {
if self.once.is_called() {
unsafe { self.value.get_mut().assume_init_drop() };
}
}
}
impl<T: Clone> Clone for OnceLock<T> {
fn clone(&self) -> Self {
let cloned = Self::new();
if let Some(value) = self.get() {
let _ = cloned.set(value.clone());
}
cloned
}
}
impl<T: fmt::Debug> fmt::Debug for OnceLock<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Debug::fmt(&self.get(), f)
}
}
impl<T: fmt::Display> fmt::Display for OnceLock<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if let Some(value) = self.get() {
fmt::Display::fmt(value, f)
} else {
fmt::Display::fmt("<uninit>", f)
}
}
}
union LazyData<T, F> {
data: ManuallyDrop<T>,
f: ManuallyDrop<F>,
}
pub struct LazyLock<T, F = fn() -> T> {
data: UnsafeCell<LazyData<T, F>>,
once: Once,
}
unsafe impl<T: Sync, F> Sync for LazyLock<T, F> {}
unsafe impl<T: Send, F: Send> Send for LazyLock<T, F> {}
impl<T, F: FnOnce() -> T> LazyLock<T, F> {
pub const fn new(f: F) -> Self {
Self {
once: Once::new(),
data: UnsafeCell::new(LazyData {
f: ManuallyDrop::new(f),
}),
}
}
pub fn get_mut(&mut self) -> &mut T {
self.init();
unsafe { &mut self.data.get_mut().data }
}
pub fn get(&self) -> &T {
self.init();
unsafe { &(*self.data.get()).data }
}
fn init(&self) {
self.once.call_once(|| {
let data = unsafe { &mut *self.data.get() };
let f = unsafe { ManuallyDrop::take(&mut data.f) };
data.data = ManuallyDrop::new(f());
});
}
}
impl<T, F: FnOnce() -> T> Deref for LazyLock<T, F> {
type Target = T;
fn deref(&self) -> &Self::Target {
self.get()
}
}
impl<T, F: FnOnce() -> T> DerefMut for LazyLock<T, F> {
fn deref_mut(&mut self) -> &mut Self::Target {
self.get_mut()
}
}
impl<T, F> Drop for LazyLock<T, F> {
fn drop(&mut self) {
if self.once.is_called() {
unsafe { ManuallyDrop::drop(&mut self.data.get_mut().data) }
} else {
unsafe { ManuallyDrop::drop(&mut self.data.get_mut().f) }
}
}
}
impl<T: Default> Default for LazyLock<T> {
fn default() -> Self {
Self::new(T::default)
}
}
impl<T: fmt::Debug, F: FnOnce() -> T> fmt::Debug for LazyLock<T, F> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if self.once.is_called() {
fmt::Debug::fmt(self.get(), f)
} else {
fmt::Debug::fmt("<uninit>", f)
}
}
}
impl<T: fmt::Display, F: FnOnce() -> T> fmt::Display for LazyLock<T, F> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if self.once.is_called() {
fmt::Display::fmt(self.get(), f)
} else {
fmt::Display::fmt("<uninit>", f)
}
}
}
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;
}
}
}