recursive_mutex.rs

#![feature(optin_builtin_traits,unsafe_destructor)]

pub mod recursive_mutex {
    #![allow(unstable)]

    use std::cell::UnsafeCell;
    use std::sync::Semaphore;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::thread::Thread;

    thread_local!( static THREAD_ID: Box<Thread> = Box::new(Thread::current()) );
    
    pub struct RecursiveMutex {
        counter: AtomicUsize,
        owner: AtomicUsize,
        recursion: UnsafeCell<usize>,
        semaphore: Semaphore,
    }

    unsafe impl Sync for RecursiveMutex {}
    unsafe impl Send for RecursiveMutex {}

    struct RecursiveMutexGuardInner<'a> {
        drop_safe: bool,
        mutex: &'a RecursiveMutex,
    }
    
    pub struct RecursiveMutexGuard<'a> {
        inner: RecursiveMutexGuardInner<'a>,
    }
    
    #[unsafe_destructor]
    impl<'a> Drop for RecursiveMutexGuardInner<'a> {
        fn drop(&mut self) {
            let res = self.mutex.counter.fetch_sub(1, Ordering::Release);
            if res > 1 && self.drop_safe {
                let recur = unsafe { *self.mutex.recursion.get() };
                if recur == 0 {
                    self.mutex.semaphore.release()
                }
            }
        }
    }

    // Cannot implement Send because we check the thread ID on destruction.
    impl<'a> !Send for RecursiveMutexGuardInner<'a> {}

    impl RecursiveMutex {
        pub fn new() -> Self {
            RecursiveMutex {
                counter: AtomicUsize::new(0),
                owner: AtomicUsize::new(0),
                recursion: UnsafeCell::new(0),
                semaphore: Semaphore::new(0),
            }
        }
        pub fn lock(&self) -> RecursiveMutexGuard {
            let tid = THREAD_ID.with( |tid| &**tid as *const _ as usize );
            let res = self.counter.fetch_add(1, Ordering::Acquire);
            let mut guard = RecursiveMutexGuardInner { mutex: self, drop_safe: false };
            if res > 0 && tid != self.owner.load(Ordering::Relaxed) {
                self.semaphore.acquire();
            }
            guard.drop_safe = true;
            self.owner.store(tid, Ordering::Relaxed);
            unsafe {
                *self.recursion.get() += 1;
            }
            RecursiveMutexGuard {
                inner: guard
            }
        }
    }
    
    #[unsafe_destructor]
    impl<'a> Drop for RecursiveMutexGuard<'a> {
        fn drop(&mut self) {
            // We can avoid the assertion here because Rust can statically guarantee
            // we are not running the destructor in the wrong thread.
            let recur = unsafe {
                let recur = self.inner.mutex.recursion.get();
                let recursion = *recur - 1;
                *recur = recursion;
                recursion
            };
            if recur == 0 {
                self.inner.mutex.owner.store(0, Ordering::Relaxed)
            }
        }
    }
}

fn main() {
    let mutex = ::std::sync::Arc::new(recursive_mutex::RecursiveMutex::new());
    let mutex_ = mutex.clone();
    let _ = {
        let _guard = mutex.lock(); 
        println!("Thread 1 acquired the lock");
        let thread = ::std::thread::Thread::scoped( move || {
            let mutex = mutex_;
            let _guard = mutex.lock();
            println!("Thread 2 acquired the lock");
            let _guard = mutex.lock();
            println!("Thread 2 acquired the lock recursively.");
        });
        let _guard = mutex.lock();
        println!("Thread 1 acquired the lock recursively.");
        thread
    };
}

static_recursive_mutex.rs

#![feature(thread_local,optin_builtin_traits,unsafe_destructor)]

pub mod recursive_mutex {
    #![allow(unstable)]

    use std::cell::UnsafeCell;
    use std::mem;
    use std::num::Int;
    use std::sync::{self, MutexGuard, StaticMutex};
    use std::sync::atomic::{self, AtomicUsize, Ordering};

    // This may seem useless but provided that each thread has a unique
    // thread local address, and this is created once per thread, it will
    // always be unique.
    #[thread_local] static THREAD_ID: () = ();

    #[allow(missing_copy_implementations)]
    #[derive(Show)]
    pub enum LockError {
        /// Mutex was poisoned,
        Poisoned,
        /// Mutex would block due to exceeded recursion limits.
        WouldBlockRecursive,
    }

    #[allow(missing_copy_implementations)]
    #[derive(Show)]
    pub enum TryLockError {
        /// Mutex was poisoned
        Poisoned,
        /// Mutex would block because it is taken by another thread.
        WouldBlockExclusive,
        /// Mutex would block due to exceeded recursion limits.
        WouldBlockRecursive,
    }

    pub struct RecursiveMutex {
        owner: AtomicUsize,
        recursion: UnsafeCell<u64>,
        mutex: StaticMutex,
        guard: UnsafeCell<*mut MutexGuard<'static, ()>>,
    }

    pub const RECURSIVE_MUTEX_INIT: RecursiveMutex = RecursiveMutex {
        owner: atomic::ATOMIC_USIZE_INIT,
        recursion: UnsafeCell { value: 0 },
        mutex: sync::MUTEX_INIT,
        guard: UnsafeCell { value: 0 as *mut _ },
    };

    unsafe impl Sync for RecursiveMutex {}
    unsafe impl Send for RecursiveMutex {}

    #[must_use]
    pub struct RecursiveMutexGuard<'a> {
        mutex: &'a RecursiveMutex,
    }

    // Cannot implement Send because we rely on the guard being dropped in the
    // same thread (otherwise we can't use Relaxed).  We might be able to allow
    // it with Acquire / Release?
    impl<'a> !Send for RecursiveMutexGuard<'a> {}

    impl RecursiveMutex {
        pub fn lock(&'static self) -> Result<RecursiveMutexGuard, LockError> {
            let tid = &THREAD_ID as *const _ as usize;
            // Relaxed is sufficient.  If tid == self.owner, it must have been set in the
            // same thread, and nothing else could have taken the lock in another thread;
            // hence, it is synchronized.  Similarly, if tid != self.owner, either the
            // lock was never taken by this thread, or the lock was taken by this thread
            // and then dropped in the same thread (known because the guard is not Send),
            // so that is synchronized as well.  The only reason it needs to be atomic at
            // all is to ensure it doesn't see partial data, and to make sure the load and
            // store aren't reordered around the acquire incorrectly (I believe this is why
            // Unordered is not suitable here, but I may be wrong since acquire() provides
            // a memory fence).
            if tid != self.owner.load(Ordering::Relaxed) {
                match self.mutex.lock() {
                    Ok(guard) => unsafe {
                        self.owner.store(tid, Ordering::Relaxed);
                        *self.guard.get() = mem::transmute(Box::new(guard));
                    },
                    Err(_) => return Err(LockError::Poisoned),
                }
            }
            unsafe {
                let r = self.recursion.get();
                match (*r).checked_add(1) {
                    Some(n) => {
                        *r = n;
                    },
                    None => return Err(LockError::WouldBlockRecursive)
                }
            }
            Ok(RecursiveMutexGuard {
                mutex: self
            })
        }

        pub fn try_lock(&'static self) -> Result<RecursiveMutexGuard, TryLockError> {
            let tid = &THREAD_ID as *const _ as usize;
            // Relaxed is sufficient.  If tid == self.owner, it must have been set in the
            // same thread, and nothing else could have taken the lock in another thread;
            // hence, it is synchronized.  Similarly, if tid != self.owner, either the
            // lock was never taken by this thread, or the lock was taken by this thread
            // and then dropped in the same thread (known because the guard is not Send),
            // so that is synchronized as well.  The only reason it needs to be atomic at
            // all is to ensure it doesn't see partial data, and to make sure the load and
            // store aren't reordered around the acquire incorrectly (I believe this is why
            // Unordered is not suitable here, but I may be wrong since acquire() provides
            // a memory fence).
            if tid != self.owner.load(Ordering::Relaxed) {
                match self.mutex.try_lock() {
                    Ok(guard) => unsafe {
                        self.owner.store(tid, Ordering::Relaxed);
                        *self.guard.get() = mem::transmute(Box::new(guard));
                    },
                    Err(sync::TryLockError::Poisoned(_)) => return Err(TryLockError::Poisoned),
                    Err(sync::TryLockError::WouldBlock) => return Err(TryLockError::WouldBlockExclusive), 
                }
            }
            unsafe {
                let r = self.recursion.get();
                match (*r).checked_add(1) {
                    Some(n) => {
                        *r = n;
                    },
                    None => return Err(TryLockError::WouldBlockRecursive)
                }
            }
            Ok(RecursiveMutexGuard {
                mutex: self
            })
        }
    }
    
    #[unsafe_destructor]
    impl<'a> Drop for RecursiveMutexGuard<'a> {
        fn drop(&mut self) {
            // We can avoid the assertion here because Rust can statically guarantee
            // we are not running the destructor in the wrong thread.
            unsafe {
                let recur = self.mutex.recursion.get();
                *recur -= 1;
                if *recur == 0 {
                    self.mutex.owner.store(0, Ordering::Relaxed);
                    mem::transmute::<_,Box<MutexGuard<()>>>(*self.mutex.guard.get());
                }
            }
        }
    }
}

fn main() {
    static MUTEX: recursive_mutex::RecursiveMutex = recursive_mutex::RECURSIVE_MUTEX_INIT;
    let _thread;
    {
        let _guard = MUTEX.lock().unwrap();
        println!("Thread 1 acquired the lock");
        _thread = ::std::thread::Thread::scoped( || {
            let _guard = MUTEX.lock().unwrap();
            println!("Thread 2 acquired the lock");
            let _guard = MUTEX.lock().unwrap();
            println!("Thread 2 acquired the lock recursively.");
        });
        let _guard = MUTEX.lock().unwrap();
        println!("Thread 1 acquired the lock recursively.");
    };
}