Made a key for locks

2 files changed, 52 insertions, 22 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 154ba8f..0c873be 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,20 +1,22 @@
 use std::any::type_name;
 use std::fmt::{self, Debug};
+use std::marker::PhantomData;
 
 mod lock;
 
-pub use lock::Lock;
+pub use lock::{Key, Lock};
+use once_cell::sync::Lazy;
+use thread_local::ThreadLocal;
 
-thread_local! {
-	pub static KEY: Lock = Lock::new();
-}
+static KEY: Lazy<ThreadLocal<Lock>> = Lazy::new(ThreadLocal::new);
 
 /// The key for the current thread.
 ///
 /// Only one of these exist per thread. To get the current thread's key, call
 /// [`ThreadKey::lock`].
 pub struct ThreadKey {
-	_priv: *const (), // this isn't Send or Sync
+	phantom: PhantomData<*const ()>, // implement !Send and !Sync
+	_key: Key<'static>,
 }
 
 impl Debug for ThreadKey {
@@ -23,27 +25,17 @@ impl Debug for ThreadKey {
 	}
 }
 
-impl Drop for ThreadKey {
-	fn drop(&mut self) {
-		KEY.with(|lock| lock.unlock());
-	}
-}
-
 impl ThreadKey {
-	/// Create a new `ThreadKey`. There should only be one `ThreadKey` per thread.
-	fn new() -> Self {
-		Self {
-			_priv: std::ptr::null(),
-		}
-	}
-
 	/// Get the current thread's `ThreadKey, if it's not already taken.
 	///
 	/// The first time this is called, it will successfully return a
 	/// `ThreadKey`. However, future calls to this function will return
 	/// [`None`], unless the key is dropped or unlocked first.
 	pub fn lock() -> Option<Self> {
-		KEY.with(|lock| lock.try_lock().then_some(Self::new()))
+		KEY.get_or_default().try_lock().map(|key| Self {
+			phantom: PhantomData,
+			_key: key,
+		})
 	}
 
 	/// Unlocks the `ThreadKey`.
diff --git a/src/lock.rs b/src/lock.rs
index 45dd1ad..fb74f25 100644
--- a/src/lock.rs
+++ b/src/lock.rs
@@ -6,18 +6,56 @@ pub struct Lock {
 	is_locked: AtomicBool,
 }
 
+#[derive(Debug)]
+pub struct Key<'a> {
+	lock: &'a Lock,
+}
+
+impl<'a> Key<'a> {
+	fn new(lock: &'a Lock) -> Self {
+		Self { lock }
+	}
+}
+
+impl<'a> Drop for Key<'a> {
+	fn drop(&mut self) {
+		// safety: this key will soon be destroyed
+		unsafe { self.lock.force_unlock() }
+	}
+}
+
 impl Lock {
+	/// Create a new unlocked `Lock`.
 	pub const fn new() -> Self {
 		Self {
 			is_locked: AtomicBool::new(false),
 		}
 	}
 
-	pub fn try_lock(&self) -> bool {
-		!self.is_locked.fetch_or(true, Ordering::Acquire)
+	/// Attempt to lock the `Lock`.
+	///
+	/// If the lock is already locked, then this'll return false. If it is
+	/// unlocked, it'll return true. If the lock is currently unlocked, then it
+	/// will be locked after this function is called.
+	///
+	/// This is not a fair lock. It is not recommended to call this function
+	/// repeatedly in a loop.
+	pub fn try_lock(&self) -> Option<Key> {
+		(!self.is_locked.fetch_or(true, Ordering::Acquire)).then_some(Key::new(self))
 	}
 
-	pub fn unlock(&self) {
+	/// Unlock the lock, without a key.
+	///
+	/// # Safety
+	///
+	/// This should only be called if the key to the lock has been "lost". That
+	/// means the program no longer has a reference to the key, but it has not
+	/// been dropped.
+	unsafe fn force_unlock(&self) {
 		self.is_locked.store(false, Ordering::Release)
 	}
+
+	pub fn unlock(key: Key) {
+		drop(key)
+	}
 }