Poisoning, Unwind Safety, Exception Safety, and More in HappyLock

I'm proud to announce that HappyLock 0.4 has been released. I thought that I should focus on poisoning because I thought, at the time, that it would be simple. I was very wrong. Poisoning exposed a whack-a-mole of challenges to solve, which I will illustrate here.

A Refresher on HappyLock

I recommend reading the full article on HappyLock. I've been told it's a very good read. But I'll briefly explain the idea here. Those of you who have already read the full post can skip ahead to the new stuff.

HappyLock uses the Rust borrow checker to make deadlocks impossible at compile-time. This is done by having a ThreadKey which is not Send, Copy, or Clone. But it is now Sync. Only one ThreadKey can exist on each thread at a time. In order to lock anything, you must pass either a ThreadKey or a &mut ThreadKey into the lock function. This means that in order to lock anything new, you must first release any resources you have already acquired, preventing partial allocation.

fn main() {
	// This panics if the ThreadKey has already been acquired for this thread
	let mut key = ThreadKey::get().unwrap();

	let mutex1 = Mutex::new(42);
	let mutex2 = Mutex::new(76);

	thread::spawn(|| {
		// You're allowed to have more than one ThreadKey, but not on the same thread
		let key = ThreadKey::get().unwrap();

		let guard = mutex.lock(&mut key); // Note the key. Nothing else can use it now
		*guard = 24;
		drop(guard); // Now we can use the ThreadKey again

		let guard = mutex2.lock(&mut key);
		*guard = 67;
	}).join().unwrap();

	let guard = mutex.lock(&mut key);
	assert_eq!(guard, 24);
	// the guard is implicitly dropped here
	let guard = mutex.lock(&mut key);
	assert_eq!(guard, 67);
}

If you want to lock multiple mutexes at the same time, you can use a LockCollection, which will ensure that all locks are acquired at the same time, so there's no possibility of deadlock. By default, this works by sorting all of the locks by their memory address, preventing circular wait, but there are other ways to do this with different performance qualities.

fn main() {
	let key = ThreadKey::get().unwrap();
	let collection = LockCollection::new((Mutex::new(42), Mutex::new("foo")));
	let guard = collection.lock(key);
	assert_eq!(guard.0, 42);
	assert_eq!(guard.1, "foo");
}

Lots of things can be used in a LockCollection, which have the Lockable trait. There's also the RawLock trait, which is implemented for Mutex and RwLock.

What is Unwind Safety

Unwind safety has nothing to do with the conventional notion of safety in Rust. And although object-safety has recently been rebranded to dyn-compatible, unwind safety cannot be. The notion of unwind safety is actually the entire reasoning why poisoning exists. Imagine the following incredibly contrived scenario.

struct MultipleOfFive(i8);

impl MultipleOfFive {
	fn increment(&mut self) {
		catch_unwind(|| {
			for _ in 0..5 {
				self.0 += 1;
			}
		});
	}
}

The variant that is supposed to be held here is that a MultipleOfFive is always divisible by five. But what happens if debug assertions are enabled, and our value is 125?

125 + 1 = 126
126 + 1 = 127
127 + 1 ??? integer overflow! panic!

But, we caught that panic, so now the value will always be 127, even though 127 is not divisible by five. This doesn't result in undefined behavior. It's just very bizarre. This is the purpose of the UnwindSafe trait. It asserts that unwinding won't result in unexpected behavior for the type. In fact, the code from before doesn't compile.

struct MultipleOfFive(i8);

impl MultipleOfFive {
	fn increment(&mut self) {
		catch_unwind(|| {
			for _ in 0..5 {
				self.0 += 1; // compile error: mutable references are not unwind-safe
			}
		});
	}
}

But again, this cannot be used to prevent undefined behavior. To make this point clear, there is a wrappper type in the standard library that actually lets us work around unwind-safety. We can just use AssertUnwindSafe, which is completely safe to construct.

struct MultipleOfFive(i8);

impl MultipleOfFive {
	fn increment(&mut self) {
		let this = AssertUnwindSafe(self);
		catch_unwind(move || {
			for _ in 0..5 {
				(*this).0 += 1;
			}
		});
	}
}

So that begs the question: if mutable references are not unwind safe, why is Mutex unwind safe? That's because of poisoning. Any time when a panic is caught while a MutexGuard is in scope, the mutex gets poisoned. Any future calls to lock on that mutex will return a PoisonError. Then, it becomes the responsibility of the caller to decide how to handle that error. Most people don't care and just `unwrap` the result. Most people seem to think that providing poisoning by default in the standard library was a mistake. Several locking crates, like parking lot and spin, don't include poisoning at all. There's even some work being done to provide non-poisoning variations of the locking primitive in the standard library. For HappyLock, I want to provide as many of the standard library's features as possible, so I worked out an alternative approach where you can get poisoning, but not by default.

How to Poison a Mutex

I decided to implement poisoning in the simplest way I could think of. HappyLock is already filled with wrappers, so why not use one here too? Put simply, the implementation looks like this:

struct Poisonable<L: RawLock + Lockable> { /* ... */ }

impl Poisonable<L> {
	// this is way over-simplified but you probably get the point
	pub fn lock(&self) -> PoisonGuard<L::Guard>;
}

Unlike with LockCollection, Poisonable requires that L be a RawLock. That's because Poisonable can't decide on its own how to deal with multiple locks without deadlocking. That's the purpose of LockCollection. But, RawLock already has a lock method, so it's easy to just call that to lock whatever we need.

At the same time, it would be nice if we could put a LockCollection inside of a Poisonable. So, LockCollection now implements RawLock, so you can create a Poisonable<LockCollection<Vec<Mutex<String>>>> and not unwrap every single mutex.

There's another method that the standard library's Mutex has that is much harder to implement: get_mut. This is already implemented on HappyLock's Mutex. It also exists on the lock collections, but instead of getting mutable references to the data being locked, it actually returns the collection. For example, LockCollection<Vec<Mutex<i32>>>::get_mut returns a &mut Vec<Mutex<i32>>. Doing the same thing on a Poisonable would require calling get_mut twice where the standard library only requires it to be called once.

So, I've made a couple more traits: LockableGetMut and LockableIntoInner

trait LockableGetMut {
	type Inner<'a>;

	fn get_mut(&mut self) -> Self::Inner<'_>;
}

trait LockableIntoInner {
	type Inner;

	fn into_inner(self) -> Self::Inner;
}

You might wonder, why not just have the same Inner type on both traits, and then we can have get_mut return &mut Self::Inner. But consider the case of a tuple. In order to have a method like get_mut, we have to create a new tuple with a different type that has no mutexes. But we can't return a mutable reference to a tuple we just created! But, we can implement the function quite simply with this system.

impl<A: LockableGetMut, B: LockableGetMut> LockableGetMut for (A, B) {
	type Inner<'a> = (A::Inner<'a>, B::Inner<'a>);

	fn get_mut(&mut self) -> Self::Inner<'_> {
		(self.0.get_mut(), self.1.get_mut())
	}
}

For consistency, lock collections now also use this functionality, and the previous into_inner method is now called into_child. This is a breaking change, but there are other breaking changes that we'll also need to make here, as we'll soon see.

The Exception Safety of HappyLock

Exception safety is a less optional form of unwind safety. The Rustonomicon actually refers to unwind-safety as "maximal exception safety". But even if we don't always guarantee that safe code will have the correct behavior, we need to at least make sure that our types are exception-safe to the point where unwinding does not cause potential unwind safety. Here's an example of how this could be violated.

impl<T: Clone> Vec<T> {
	fn push(&mut self, x: &T) {
		self.reserve(1);
		unsafe {
			self.set_len(self.len() + 1);
			self.ptr().add(1).write(x.clone());
		}
	}
}

The problem is that clone can panic, because you can implement it however you want. If it panics, then the length will be incorrect forever. This is something that you'll always need to be careful about in unsafe code. Now, look at this function in HappyLock.

impl<L: OwnedLockable> OwnedLockCollection<L> {
	pub fn lock<'g, 'key, Key: Keyable + 'key>(
		&'g self,
		key: Key,
	) -> LockGuard<'key, L::Guard<'g>, Key> {
		let locks = get_locks(&self.data);
		for lock in locks {
			// safety: we have the thread key, and these locks happen in a
			//         predetermined order
			unsafe { lock.lock() };
		}

		let guard = unsafe { self.data.guard() };
		LockGuard {
			guard,
			key,
			_phantom: PhantomData,
		}
	}
}

As we can see, it is considered undefined behavior in HappyLock to cause a deadlock. Now, what would happen if we had a collection of three locks, and the second one panicked. The first mutex would never unlock, and then we could call it again, causing a deadlock, like this.

fn main() {
	let collection = OwnedLockCollection::new((
		Mutex::new(1),
		Mutex::new(2),
		Mutex::new(3)
	));

	catch_unwind(|| {
		let key = ThreadKey::get().unwrap();
		collection.lock();
		// We'll say for the sake of argument that our second mutex panics here
		// The first mutex was never unlocked
	});

	// The ThreadKey was still dropped, so we can re-acquire it.
	let key = ThreadKey::get().unwrap();

	// But wait, the first mutex is still locked by this thread.
	collection.lock(); // DEADLOCK!!!
}

Note that this only becomes a problem if, somehow, parking_lot panics when calling the lock function. This would be considered a bug, but since you can use whatever RawMutex you want with HappyLock, I can't guarantee that this will never happen. This could be solved by just requiring a specific underlying API to be used, but I'm not ready to impose that restriction (yet). Somebody at RustConf suggested using a fancy guard insertion strategy so that the mutexes will automatically unlock on panic. That might work, but I have a simpler solution.

We can use the scopeguard library to automatically run code when unwinding. What we'll need to do is keep track of which locks have been locked so far, and then, if we start unwinding, unlock them.

Now, if you use the scopeguard crate, this will actually result in a stack overflow, depending on how rigourously you check for this. It won't check to make sure that the panic was caused by your function, so it will keep calling itself, and you'll get a stack overflow. So I made my own utility to check that the unwinding was caused by a specific piece of code

fn handle_unwind<R, F: FnOnce -> R, G: FnOnce()>(try_fn: F, catch: G) -> R {
	let try_fn = AssertUnwindSafe(try_fn);
	catch_unwind(try_fn).unwrap_or_else(|e| {
		catch();
		resume_unwind(e)
	})
}

Then we can use it to unlock on failure, like so

unsafe fn ordered_lock(locks: &[&dyn RawLock]) {
	let locked = Cell::new(0);

	handle_unwind(
		|| {
			for lock in locks {
				lock.raw_lock();
				locked.set(locked.get() + 1);
			}
		},
		|| attempt_to_recover_locks_from_panic(&locked[0..locked.get()]),
	)
}

unsafe fn attempt_to_recover_locks_from_panic(locked: &[&dyn RawLock]) {
	handle_unwind(
		|| locked.for_each(|lock| lock.raw_unlock()),
		|| locked.for_each(|l| l.poison()),
	)
}

Ok, so I lied when I said there's no poisoning by default. It's not documented anywhere, but yeah. If the unlock function panics, when we're kinda screwed, so all RawLocks now have to be poisonable. This is a different kind of poisoning, which causes all future calls to the mutex to panic, similar to the poisoning of Once. It's not exposed in the signatures for any of the functions. I don't want you to ever have to think about this kind of poisoning. I emphasize, this should never happen. If it does, then the RawMutex probably contains a bug that already needs to be fixed.

The lock that panicked is poisoned, regardless of whether it was locked or not. That's because I have zero clue whether or not it actually locked. I would hope that if the lock function panics, it didn't lock, but I don't know that for sure. I also don't have a good way to check, without excluding support for certain mutexes in the future. Someone might suggest using the result of try_lock to check if it's locked or not, and then unlock it in either case. That person might be surprised to learn that, in C, it is undefined behavior to call try_lock on a mutex that is already acquired by the calling thread. So there's not much I can do other than panic here.

Miscellaneous Changes

The try_lock function now returns a Result instead of an Option. This is something I've been thinking about for a while now, and the fact that Poisonable now exists makes an even more compelling argument for using Result. If something is already locked, then you can have the key back without having to call ThreadKey::get.

I've kinda shown this already in some of the examples, but the RawLock trait methods have been renamed. For example, RawLock::lock is now RawLock::raw_lock.

Some of the reading APIs have been moved from Lockable to Sharable. The reason they were in Lockable before was because I had an idea where you could mix sharables with non-sharables in a lock collection, and then calling read would just get an exclusive lock if there was no shared lock available for an entry. Eventually I decided this was a bad idea, but I still had the read_guard method in Lockable. So now it's moved to Sharable.

I mentioned this already in the introduction to HappyLock, but ThreadKey is now Sync. This is possible for the same reason that the nightly Exclusive type is Sync. A &ThreadKey is completely useless, so no undefined behavior can be caused because of it. This also allows guards which hold a key (i.e. MutexGuard) to be Sync. They still can't be send because, among other things, dropping a ThreadKey on the wrong thread can lead to having multiple ThreadKeys on the same thread, making deadlock possible. The Rust standard library also doesn't implement Send on its guard types, so I think this is acceptable.

I'm now using cargo-diet before I publish HappyLock, so examples are no longer included when you download it.

The ThreadKey now uses a Cell<bool> instead of an AtomicBool when deciding if you can get a key or not. It's a thread-local variable, so there was no synchronization needed in order to access the value.

Speaking of which, I've removed the thread-local crate as a dependency and learned enough to just use the standard library's thread_local macro. It's more efficient anyway. Also, now that LazyCell is in the standard library, I've dropped the once_cell dependency. But this does increase the MSRV substantially.

The guard and ref types now implement more traits like Eq, Ord and Hash. This would be a breaking change in the standard library, but this is already a breaking change here.

I've been using cargo-mutants to help me write better unit tests. This helped me catch one bug before it was even released. Although, I did still find a bug, so it's not perfect. HappyLock is in this weird state where it's being used and I'm scared of obvious bugs, but it's not used enough to the point where obvious bugs are immediately found and fixed. I should probably start using tarpaulin as well.

The Future

I'm still deciding what my next course of action will be for this crate. I should write more tests for it before I do another big release. The documentation examples could definitely be improved. Some of my plans for the future have changed. I still would like to try copying the standard library lock implementations at some point. I no longer think compile-time duplicate checks are possible due to the lack of const trait methods in Rust. LockCell and ReadOnlyLockCollection are still just ideas in the back of my mind. But I have acquired a few new ideas since last year.

Once

It is definitely possible to make a deadlock-free Once, by having call_once require a ThreadKey. I tried somehow fitting it into a LockCollection, but I don't think that's practical, and I don't know why you would ever want more than a single Once at a time. Although, I can see a point to using a Mutex inside of a call_once, so I can also provide a way to pass in a RawLock + Lockable with it.

I'm still not sure if I want to include Condvar. It is useful but I don't think I would be able to prevent it from deadlocking. I'll keep it in mind though.

Async

During my time at RustConf, I realized that the biggest barrier to HappyLock adoption is probably the inability to use it with async Rust. The problem is that you can have multiple tasks running on the same thread, so ThreadKey::get() will return None very often. But a good solution to this would be to have TaskKeys instead, which would only be useful on the async mutexes. But this is something that the async runtime would need to be able to provide. But if the spawn function gave you a TaskKey, then there would be no need to even call a get function for it. The biggest problem here would be that you can't use a blocking lock, because if it's already locked in a task that runs on the same thread, then this would cause a deadlock.

Expanding Cyclic Wait

A reddit comment gave me an idea on how to pursue this. I don't think I'll go into too much detail right now, but it would involve a IntoLockIterator trait, several new types, adding a &ThreadKey to all of the ref guard types, and quite a bit of type-state analysis. I'll probably talk about this over at the Miami Rust User Group, so join us online if you're interested. I think we can turn HappyLock from a library that primarily prevents partial allocation and prevents cyclic wait as an optimization, to a library that is used to prevent cyclic wait and uses remnants of total allocation to prevent you from screwing up.