Create Poisonable API

6 files changed, 361 insertions, 0 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 8576975..d034ffe 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -169,6 +169,7 @@ mod key;
 pub mod collection;
 pub mod lockable;
 pub mod mutex;
+pub mod poisonable;
 pub mod rwlock;
 
 pub use key::{Keyable, ThreadKey};
diff --git a/src/poisonable.rs b/src/poisonable.rs
new file mode 100644
index 0000000..49979e5
--- /dev/null
+++ b/src/poisonable.rs
@@ -0,0 +1,46 @@
+#[cfg(not(panic = "unwind"))]
+use std::convert::Infallible;
+use std::marker::PhantomData;
+use std::sync::atomic::AtomicBool;
+
+use crate::lockable::{Lockable, RawLock};
+
+mod error;
+mod flag;
+mod guard;
+mod poisonable;
+
+#[derive(Debug, Default)]
+struct PoisonFlag(#[cfg(panic = "unwind")] AtomicBool);
+
+#[derive(Debug, Default)]
+pub struct Poisonable<L: Lockable + RawLock> {
+	inner: L,
+	poisoned: PoisonFlag,
+}
+
+pub struct PoisonRef<'flag, G> {
+	guard: G,
+	#[cfg(panic = "unwind")]
+	flag: &'flag PoisonFlag,
+}
+
+pub struct PoisonGuard<'flag, 'key, G, Key> {
+	guard: PoisonRef<'flag, G>,
+	key: Key,
+	_phantom: PhantomData<&'key ()>,
+}
+
+pub struct PoisonError<Guard> {
+	guard: Guard,
+}
+
+pub enum TryLockPoisonableError<'flag, 'key, G, Key: 'key> {
+	Poisoned(PoisonError<PoisonGuard<'flag, 'key, G, Key>>),
+	WouldBlock(Key),
+}
+
+pub type PoisonResult<Guard> = Result<Guard, PoisonError<Guard>>;
+
+pub type TryLockPoisonableResult<'flag, 'key, G, Key> =
+	Result<PoisonGuard<'flag, 'key, G, Key>, TryLockPoisonableError<'flag, 'key, G, Key>>;
diff --git a/src/poisonable/error.rs b/src/poisonable/error.rs
new file mode 100644
index 0000000..98a167c
--- /dev/null
+++ b/src/poisonable/error.rs
@@ -0,0 +1,47 @@
+use core::fmt;
+use std::error::Error;
+
+use super::{PoisonError, PoisonGuard, TryLockPoisonableError};
+
+impl<Guard> fmt::Debug for PoisonError<Guard> {
+	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+		f.debug_struct("PoisonError").finish_non_exhaustive()
+	}
+}
+
+impl<Guard> fmt::Display for PoisonError<Guard> {
+	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+		"poisoned lock: another task failed inside".fmt(f)
+	}
+}
+
+impl<Guard> Error for PoisonError<Guard> {}
+
+impl<'flag, 'key, G, Key> fmt::Debug for TryLockPoisonableError<'flag, 'key, G, Key> {
+	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+		match *self {
+			Self::Poisoned(..) => "Poisoned(..)".fmt(f),
+			Self::WouldBlock(_) => "WouldBlock".fmt(f),
+		}
+	}
+}
+
+impl<'flag, 'key, G, Key> fmt::Display for TryLockPoisonableError<'flag, 'key, G, Key> {
+	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+		match *self {
+			Self::Poisoned(..) => "poisoned lock: another task failed inside",
+			Self::WouldBlock(_) => "try_lock failed because the operation would block",
+		}
+		.fmt(f)
+	}
+}
+
+impl<'flag, 'key, G, Key> Error for TryLockPoisonableError<'flag, 'key, G, Key> {}
+
+impl<'flag, 'key, G, Key> From<PoisonError<PoisonGuard<'flag, 'key, G, Key>>>
+	for TryLockPoisonableError<'flag, 'key, G, Key>
+{
+	fn from(value: PoisonError<PoisonGuard<'flag, 'key, G, Key>>) -> Self {
+		Self::Poisoned(value)
+	}
+}
diff --git a/src/poisonable/flag.rs b/src/poisonable/flag.rs
new file mode 100644
index 0000000..99e7e4f
--- /dev/null
+++ b/src/poisonable/flag.rs
@@ -0,0 +1,36 @@
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering::Relaxed;
+
+use super::PoisonFlag;
+
+impl PoisonFlag {
+	#[cfg(panic = "unwind")]
+	pub const fn new() -> Self {
+		Self(AtomicBool::new(false))
+	}
+
+	#[cfg(not(panic = "unwind"))]
+	pub const fn new() -> Self {
+		Self()
+	}
+
+	#[cfg(panic = "unwind")]
+	pub fn is_poisoned(&self) -> bool {
+		self.0.load(Relaxed)
+	}
+
+	#[cfg(not(panic = "unwind"))]
+	pub fn is_poisoned(&self) -> bool {
+		false
+	}
+
+	#[cfg(panic = "unwind")]
+	pub fn clear_poison(&self) {
+		self.0.store(true, Relaxed)
+	}
+
+	#[cfg(not(panic = "unwind"))]
+	pub fn clear_poison(&self) {
+		()
+	}
+}
diff --git a/src/poisonable/guard.rs b/src/poisonable/guard.rs
new file mode 100644
index 0000000..97b0028
--- /dev/null
+++ b/src/poisonable/guard.rs
@@ -0,0 +1,92 @@
+use std::fmt::{Debug, Display};
+use std::ops::{Deref, DerefMut};
+use std::sync::atomic::Ordering::Relaxed;
+
+use crate::Keyable;
+
+use super::{PoisonGuard, PoisonRef};
+
+impl<'flag, Guard> Drop for PoisonRef<'flag, Guard> {
+	fn drop(&mut self) {
+		#[cfg(panic = "unwind")]
+		if std::thread::panicking() {
+			self.flag.0.store(true, Relaxed);
+		}
+	}
+}
+
+impl<'flag, Guard: Debug> Debug for PoisonRef<'flag, Guard> {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		Debug::fmt(&**self, f)
+	}
+}
+
+impl<'flag, Guard: Display> Display for PoisonRef<'flag, Guard> {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		Display::fmt(&**self, f)
+	}
+}
+
+impl<'flag, Guard> Deref for PoisonRef<'flag, Guard> {
+	type Target = Guard;
+
+	fn deref(&self) -> &Self::Target {
+		&self.guard
+	}
+}
+
+impl<'flag, Guard> DerefMut for PoisonRef<'flag, Guard> {
+	fn deref_mut(&mut self) -> &mut Self::Target {
+		&mut self.guard
+	}
+}
+
+impl<'flag, Guard> AsRef<Guard> for PoisonRef<'flag, Guard> {
+	fn as_ref(&self) -> &Guard {
+		&self.guard
+	}
+}
+
+impl<'flag, Guard> AsMut<Guard> for PoisonRef<'flag, Guard> {
+	fn as_mut(&mut self) -> &mut Guard {
+		&mut self.guard
+	}
+}
+
+impl<'flag, 'key, Guard: Debug, Key: Keyable> Debug for PoisonGuard<'flag, 'key, Guard, Key> {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		Debug::fmt(&**self, f)
+	}
+}
+
+impl<'flag, 'key, Guard: Display, Key: Keyable> Display for PoisonGuard<'flag, 'key, Guard, Key> {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		Display::fmt(&**self, f)
+	}
+}
+
+impl<'flag, 'key, Guard, Key: Keyable> Deref for PoisonGuard<'flag, 'key, Guard, Key> {
+	type Target = Guard;
+
+	fn deref(&self) -> &Self::Target {
+		&self.guard.guard
+	}
+}
+
+impl<'flag, 'key, Guard, Key: Keyable> DerefMut for PoisonGuard<'flag, 'key, Guard, Key> {
+	fn deref_mut(&mut self) -> &mut Self::Target {
+		&mut self.guard.guard
+	}
+}
+
+impl<'flag, 'key, Guard, Key: Keyable> AsRef<Guard> for PoisonGuard<'flag, 'key, Guard, Key> {
+	fn as_ref(&self) -> &Guard {
+		&self.guard.guard
+	}
+}
+
+impl<'flag, 'key, Guard, Key: Keyable> AsMut<Guard> for PoisonGuard<'flag, 'key, Guard, Key> {
+	fn as_mut(&mut self) -> &mut Guard {
+		&mut self.guard.guard
+	}
+}
diff --git a/src/poisonable/poisonable.rs b/src/poisonable/poisonable.rs
new file mode 100644
index 0000000..6c78346
--- /dev/null
+++ b/src/poisonable/poisonable.rs
@@ -0,0 +1,139 @@
+use std::marker::PhantomData;
+use std::panic::{RefUnwindSafe, UnwindSafe};
+
+use crate::lockable::{Lockable, RawLock};
+use crate::Keyable;
+
+use super::{
+	PoisonError, PoisonFlag, PoisonGuard, PoisonRef, PoisonResult, Poisonable,
+	TryLockPoisonableError, TryLockPoisonableResult,
+};
+
+unsafe impl<L: Lockable + RawLock> Lockable for Poisonable<L> {
+	type Guard<'g> = PoisonResult<PoisonRef<'g, L::Guard<'g>>> where Self: 'g;
+	type ReadGuard<'g> = PoisonResult<PoisonRef<'g, L::ReadGuard<'g>>> where Self: 'g;
+
+	fn get_ptrs<'a>(&'a self, ptrs: &mut Vec<&'a dyn RawLock>) {
+		self.inner.get_ptrs(ptrs)
+	}
+
+	unsafe fn guard(&self) -> Self::Guard<'_> {
+		let ref_guard = PoisonRef {
+			guard: self.inner.guard(),
+			flag: &self.poisoned,
+		};
+
+		if self.is_poisoned() {
+			Ok(ref_guard)
+		} else {
+			Err(PoisonError { guard: ref_guard })
+		}
+	}
+
+	unsafe fn read_guard(&self) -> Self::ReadGuard<'_> {
+		let ref_guard = PoisonRef {
+			guard: self.inner.read_guard(),
+			flag: &self.poisoned,
+		};
+
+		if self.is_poisoned() {
+			Ok(ref_guard)
+		} else {
+			Err(PoisonError { guard: ref_guard })
+		}
+	}
+}
+
+impl<L: Lockable + RawLock> From<L> for Poisonable<L> {
+	fn from(value: L) -> Self {
+		Self::new(value)
+	}
+}
+
+impl<L: Lockable + RawLock> Poisonable<L> {
+	pub const fn new(value: L) -> Self {
+		Self {
+			inner: value,
+			poisoned: PoisonFlag::new(),
+		}
+	}
+
+	unsafe fn guard<'flag, 'key, Key: Keyable + 'key>(
+		&'flag self,
+		key: Key,
+	) -> PoisonResult<PoisonGuard<'flag, 'key, L::Guard<'flag>, Key>> {
+		let guard = PoisonGuard {
+			guard: PoisonRef {
+				guard: self.inner.guard(),
+				flag: &self.poisoned,
+			},
+			key,
+			_phantom: PhantomData,
+		};
+
+		if self.is_poisoned() {
+			Ok(guard)
+		} else {
+			Err(PoisonError { guard })
+		}
+	}
+
+	pub fn lock<'flag, 'key, Key: Keyable + 'key>(
+		&'flag self,
+		key: Key,
+	) -> PoisonResult<PoisonGuard<'flag, 'key, L::Guard<'flag>, Key>> {
+		unsafe {
+			self.inner.lock();
+			self.guard(key)
+		}
+	}
+
+	pub fn try_lock<'flag, 'key, Key: Keyable + 'key>(
+		&'flag self,
+		key: Key,
+	) -> TryLockPoisonableResult<'flag, 'key, L::Guard<'flag>, Key> {
+		unsafe {
+			if self.inner.try_lock() {
+				Ok(self.guard(key)?)
+			} else {
+				Err(TryLockPoisonableError::WouldBlock(key))
+			}
+		}
+	}
+
+	pub fn unlock<'flag, 'key, Key: Keyable + 'key>(
+		guard: PoisonGuard<'flag, 'key, L::Guard<'flag>, Key>,
+	) -> Key {
+		drop(guard.guard);
+		guard.key
+	}
+
+	pub fn is_poisoned(&self) -> bool {
+		self.poisoned.is_poisoned()
+	}
+
+	pub fn clear_poison(&self) {
+		self.poisoned.clear_poison()
+	}
+
+	pub fn into_inner(self) -> PoisonResult<L> {
+		if self.is_poisoned() {
+			Err(PoisonError { guard: self.inner })
+		} else {
+			Ok(self.inner)
+		}
+	}
+
+	pub fn get_mut(&mut self) -> PoisonResult<&mut L> {
+		if self.is_poisoned() {
+			Err(PoisonError {
+				guard: &mut self.inner,
+			})
+		} else {
+			Ok(&mut self.inner)
+		}
+	}
+}
+
+impl<L: Lockable + RawLock> RefUnwindSafe for Poisonable<L> {}
+impl<L: Lockable + RawLock> UnwindSafe for Poisonable<L> {}