happylock/

mutex.rs

use std::cell::UnsafeCell;
use std::marker::PhantomData;

use lock_api::RawMutex;

use crate::poisonable::PoisonFlag;
use crate::ThreadKey;

mod guard;
mod mutex;

/// A spinning mutex
#[cfg(feature = "spin")]
pub type SpinLock<T> = Mutex<T, spin::Mutex<()>>;

/// A parking lot mutex
#[cfg(feature = "parking_lot")]
pub type ParkingMutex<T> = Mutex<T, parking_lot::RawMutex>;

/// A mutual exclusion primitive useful for protecting shared data, which
/// cannot deadlock.
///
/// This mutex will block threads waiting for the lock to become available. The
/// mutex can be created via a `new` constructor. Each mutex has a type
/// parameter which represents the data that it is protecting. The data can
/// only be accessed through the [`MutexGuard`]s returned from [`lock`] and
/// [`try_lock`], which guarantees that the data is only ever accessed when
/// the mutex is locked.
///
/// Locking the mutex on a thread that already locked it is impossible, due to
/// the requirement of the [`ThreadKey`]. Therefore, this will never deadlock.
///
/// # Examples
///
/// ```
/// use std::sync::Arc;
/// use std::thread;
/// use std::sync::mpsc;
///
/// use happylock::{Mutex, ThreadKey};
///
/// // Spawn a few threads to increment a shared variable (non-atomically),
/// // and let the main thread know once all increments are done.
/// //
/// // Here we're using an Arc to share memory among threads, and the data
/// // inside the Arc is protected with a mutex.
/// const N: usize = 10;
///
/// let data = Arc::new(Mutex::new(0));
///
/// let (tx, rx) = mpsc::channel();
/// for _ in 0..N {
///     let (data, tx) = (Arc::clone(&data), tx.clone());
///     thread::spawn(move || {
///         let key = ThreadKey::get().unwrap();
///         let mut data = data.lock(key);
///         *data += 1;
///         if *data == N {
///             tx.send(()).unwrap();
///         }
///         // the lock is unlocked
///     });
/// }
///
/// rx.recv().unwrap();
/// ```
///
/// To unlock a mutex guard sooner than the end of the enclosing scope, either
/// create an inner scope, drop the guard manually, or call [`Mutex::unlock`].
///
/// ```
/// use std::sync::Arc;
/// use std::thread;
///
/// use happylock::{Mutex, ThreadKey};
///
/// const N: usize = 3;
///
/// let data_mutex = Arc::new(Mutex::new(vec![1, 2, 3, 4]));
/// let res_mutex = Arc::new(Mutex::new(0));
///
/// let mut threads = Vec::with_capacity(N);
/// (0..N).for_each(|_| {
///     let data_mutex_clone = Arc::clone(&data_mutex);
///     let res_mutex_clone = Arc::clone(&res_mutex);
///
///     threads.push(thread::spawn(move || {
///         let mut key = ThreadKey::get().unwrap();
///
///         // Here we use a block to limit the lifetime of the lock guard.
///         let result = data_mutex_clone.scoped_lock(&mut key, |data| {
///             let result = data.iter().fold(0, |acc, x| acc + x * 2);
///             data.push(result);
///             result
///             // The mutex guard gets dropped here, so the lock is released
///         });
///         // The thread key is available again
///         *res_mutex_clone.lock(key) += result;
///     }));
/// });
///
/// let key = ThreadKey::get().unwrap();
/// let mut data = data_mutex.lock(key);
/// let result = data.iter().fold(0, |acc, x| acc + x * 2);
/// data.push(result);
///
/// // We drop the `data` explicitly because it's not necessary anymore. This
/// // allows other threads to start working on the data immediately. Dropping
/// // the data also gives us access to the thread key, so we can lock
/// // another mutex.
/// let key = Mutex::unlock(data);
///
/// // Here the mutex guard is not assigned to a variable and so, even if the
/// // scope does not end after this line, the mutex is still released: there is
/// // no deadlock.
/// *res_mutex.lock(key) += result;
///
/// threads.into_iter().for_each(|thread| {
///     thread
///         .join()
///         .expect("The thread creating or execution failed !")
/// });
///
/// let key = ThreadKey::get().unwrap();
/// assert_eq!(*res_mutex.lock(key), 800);
/// ```
///
/// [`lock`]: `Mutex::lock`
/// [`try_lock`]: `Mutex::try_lock`
/// [`ThreadKey`]: `crate::ThreadKey`
pub struct Mutex<T: ?Sized, R> {
	raw: R,
	poison: PoisonFlag,
	data: UnsafeCell<T>,
}

/// An RAII implementation of a “scoped lock” of a mutex. When this structure
/// is dropped (falls out of scope), the lock will be unlocked.
///
/// The data protected by the mutex can be accessed through this guard via its
/// [`Deref`] and [`DerefMut`] implementations.
///
/// This is created by calling the [`lock`] and [`try_lock`] methods on [`Mutex`]
///
/// This is similar to the [`MutexGuard`] type, except it does not hold a
/// [`ThreadKey`].
///
/// [`lock`]: `Mutex::lock`
/// [`try_lock`]: `Mutex::try_lock`
/// [`Deref`]: `std::ops::Deref`
/// [`DerefMut`]: `std::ops::DerefMut`
pub struct MutexRef<'a, T: ?Sized + 'a, R: RawMutex>(&'a Mutex<T, R>, PhantomData<R::GuardMarker>);

/// An RAII implementation of a “scoped lock” of a mutex. When this structure
/// is dropped (falls out of scope), the lock will be unlocked.
///
/// The data protected by the mutex can be accessed through this guard via its
/// [`Deref`] and [`DerefMut`] implementations.
///
/// This is created by calling the [`lock`] and [`try_lock`] methods on [`Mutex`]
///
/// This guard holds on to a [`ThreadKey`], which ensures that nothing else is
/// locked until this guard is dropped. The [`ThreadKey`] can be reacquired
/// using [`Mutex::unlock`].
///
/// [`Deref`]: `std::ops::Deref`
/// [`DerefMut`]: `std::ops::DerefMut`
/// [`lock`]: `Mutex::lock`
/// [`try_lock`]: `Mutex::try_lock`
//
// This is the most lifetime-intensive thing I've ever written. Can I graduate
// from borrow checker university now?
pub struct MutexGuard<'a, T: ?Sized + 'a, R: RawMutex> {
	mutex: MutexRef<'a, T, R>, // this way we don't need to re-implement Drop
	thread_key: ThreadKey,
}

#[cfg(test)]
mod tests {
	use crate::{LockCollection, ThreadKey};

	use super::*;

	#[test]
	fn unlocked_when_initialized() {
		let lock: crate::Mutex<_> = Mutex::new("Hello, world!");

		assert!(!lock.is_locked());
	}

	#[test]
	fn locked_after_read() {
		let key = ThreadKey::get().unwrap();
		let lock: crate::Mutex<_> = Mutex::new("Hello, world!");

		let guard = lock.lock(key);

		assert!(lock.is_locked());
		drop(guard)
	}

	#[test]
	fn from_works() {
		let key = ThreadKey::get().unwrap();
		let mutex: crate::Mutex<_> = Mutex::from("Hello, world!");

		let guard = mutex.lock(key);
		assert_eq!(*guard, "Hello, world!");
	}

	#[test]
	fn as_mut_works() {
		let key = ThreadKey::get().unwrap();
		let mut mutex = crate::Mutex::from(42);

		let mut_ref = mutex.as_mut();
		*mut_ref = 24;

		mutex.scoped_lock(key, |guard| assert_eq!(*guard, 24))
	}

	#[test]
	fn display_works_for_guard() {
		let key = ThreadKey::get().unwrap();
		let mutex: crate::Mutex<_> = Mutex::new("Hello, world!");
		let guard = mutex.lock(key);
		assert_eq!(guard.to_string(), "Hello, world!".to_string());
	}

	#[test]
	fn display_works_for_ref() {
		let mutex: crate::Mutex<_> = Mutex::new("Hello, world!");
		let guard = unsafe { mutex.try_lock_no_key().unwrap() };
		assert_eq!(guard.to_string(), "Hello, world!".to_string());
	}

	#[test]
	fn ref_as_mut() {
		let key = ThreadKey::get().unwrap();
		let collection = LockCollection::new(crate::Mutex::new(0));
		let mut guard = collection.lock(key);
		let guard_mut = guard.as_mut().as_mut();

		*guard_mut = 3;
		let key = LockCollection::<crate::Mutex<_>>::unlock(guard);

		let guard = collection.lock(key);

		assert_eq!(guard.as_ref().as_ref(), &3);
	}

	#[test]
	fn guard_as_mut() {
		let key = ThreadKey::get().unwrap();
		let mutex = crate::Mutex::new(0);
		let mut guard = mutex.lock(key);
		let guard_mut = guard.as_mut();

		*guard_mut = 3;
		let key = Mutex::unlock(guard);

		let guard = mutex.lock(key);

		assert_eq!(guard.as_ref(), &3);
	}

	#[test]
	fn dropping_guard_releases_mutex() {
		let key = ThreadKey::get().unwrap();
		let mutex: crate::Mutex<_> = Mutex::new("Hello, world!");

		let guard = mutex.lock(key);
		drop(guard);

		assert!(!mutex.is_locked());
	}

	#[test]
	fn dropping_ref_releases_mutex() {
		let mutex: crate::Mutex<_> = Mutex::new("Hello, world!");

		let guard = unsafe { mutex.try_lock_no_key().unwrap() };
		drop(guard);

		assert!(!mutex.is_locked());
	}
}