217 lines · 5.4 KB · Rust
RawDownload
// SPDX-License-Identifier: Apache-2.0 OR MIT #![allow(clippy::undocumented_unsafe_blocks)]#![cfg_attr(feature = "alloc", feature(allocator_api))]#![allow(clippy::missing_safety_doc)] use core::{    cell::{Cell, UnsafeCell},    marker::PhantomPinned,    ops::{Deref, DerefMut},    pin::Pin,    sync::atomic::{AtomicBool, Ordering},};#[cfg(feature = "std")]use std::{    sync::Arc,    thread::{self, sleep, Builder, Thread},    time::Duration,}; use pin_init::*;#[allow(unused_attributes)]#[path = "./linked_list.rs"]pub mod linked_list;use linked_list::*; pub struct SpinLock {    inner: AtomicBool,} impl SpinLock {    #[inline]    pub fn acquire(&self) -> SpinLockGuard<'_> {        while self            .inner            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)            .is_err()        {            #[cfg(feature = "std")]            while self.inner.load(Ordering::Relaxed) {                thread::yield_now();            }        }        SpinLockGuard(self)    }     #[inline]    #[allow(clippy::new_without_default)]    pub const fn new() -> Self {        Self {            inner: AtomicBool::new(false),        }    }} pub struct SpinLockGuard<'a>(&'a SpinLock); impl Drop for SpinLockGuard<'_> {    #[inline]    fn drop(&mut self) {        self.0.inner.store(false, Ordering::Release);    }} #[pin_data]pub struct CMutex<T> {    #[pin]    wait_list: ListHead,    spin_lock: SpinLock,    locked: Cell<bool>,    #[pin]    data: UnsafeCell<T>,} impl<T> CMutex<T> {    #[inline]    pub fn new(val: impl PinInit<T>) -> impl PinInit<Self> {        pin_init!(CMutex {            wait_list <- ListHead::new(),            spin_lock: SpinLock::new(),            locked: Cell::new(false),            data <- UnsafeCell::pin_init(val),        })    }     #[inline]    pub fn lock(&self) -> Pin<CMutexGuard<'_, T>> {        let mut sguard = self.spin_lock.acquire();        if self.locked.get() {            stack_pin_init!(let _wait_entry = WaitEntry::insert_new(&self.wait_list));            // println!("wait list length: {}", self.wait_list.size());            while self.locked.get() {                drop(sguard);                #[cfg(feature = "std")]                thread::park();                sguard = self.spin_lock.acquire();            }        }        self.locked.set(true);        unsafe {            Pin::new_unchecked(CMutexGuard {                mtx: self,                _pin: PhantomPinned,            })        }    }     #[allow(dead_code)]    pub fn get_data_mut(self: Pin<&mut Self>) -> &mut T {        // SAFETY: we have an exclusive reference and thus nobody has access to data.        unsafe { &mut *self.data.get() }    }} unsafe impl<T: Send> Send for CMutex<T> {}unsafe impl<T: Send> Sync for CMutex<T> {} pub struct CMutexGuard<'a, T> {    mtx: &'a CMutex<T>,    _pin: PhantomPinned,} impl<T> Drop for CMutexGuard<'_, T> {    #[inline]    fn drop(&mut self) {        let sguard = self.mtx.spin_lock.acquire();        self.mtx.locked.set(false);        if let Some(list_field) = self.mtx.wait_list.next() {            let _wait_entry = list_field.as_ptr().cast::<WaitEntry>();            #[cfg(feature = "std")]            unsafe {                (*_wait_entry).thread.unpark()            };        }        drop(sguard);    }} impl<T> Deref for CMutexGuard<'_, T> {    type Target = T;     #[inline]    fn deref(&self) -> &Self::Target {        unsafe { &*self.mtx.data.get() }    }} impl<T> DerefMut for CMutexGuard<'_, T> {    #[inline]    fn deref_mut(&mut self) -> &mut Self::Target {        unsafe { &mut *self.mtx.data.get() }    }} #[pin_data]#[repr(C)]struct WaitEntry {    #[pin]    wait_list: ListHead,    #[cfg(feature = "std")]    thread: Thread,} impl WaitEntry {    #[inline]    fn insert_new(list: &ListHead) -> impl PinInit<Self> + '_ {        #[cfg(feature = "std")]        {            pin_init!(Self {                thread: thread::current(),                wait_list <- ListHead::insert_prev(list),            })        }        #[cfg(not(feature = "std"))]        {            pin_init!(Self {                wait_list <- ListHead::insert_prev(list),            })        }    }} #[cfg_attr(test, test)]#[allow(dead_code)]fn main() {    #[cfg(feature = "std")]    {        let mtx: Pin<Arc<CMutex<usize>>> = Arc::pin_init(CMutex::new(0)).unwrap();        let mut handles = vec![];        let thread_count = 20;        let workload = if cfg!(miri) { 100 } else { 1_000 };        for i in 0..thread_count {            let mtx = mtx.clone();            handles.push(                Builder::new()                    .name(format!("worker #{i}"))                    .spawn(move || {                        for _ in 0..workload {                            *mtx.lock() += 1;                        }                        println!("{i} halfway");                        sleep(Duration::from_millis((i as u64) * 10));                        for _ in 0..workload {                            *mtx.lock() += 1;                        }                        println!("{i} finished");                    })                    .expect("should not fail"),            );        }        for h in handles {            h.join().expect("thread panicked");        }        println!("{:?}", *mtx.lock());        assert_eq!(*mtx.lock(), workload * thread_count * 2);    }}