use std::hash::{Hash, Hasher};
use std::collections::hash_map::DefaultHasher;
use std::sync::{Mutex, MutexGuard};
pub struct StripedLock<const N: usize> {
locks: [Mutex<()>; N],
}
impl<const N: usize> StripedLock<N> {
pub fn new() -> Self {
Self {
locks: std::array::from_fn(|_| Mutex::new(())),
}
}
pub fn lock(&self, key: &str) -> MutexGuard<'_, ()> {
let index = Self::stripe_index(key);
self.locks[index].lock().unwrap_or_else(|e| e.into_inner())
}
pub fn lock_two(&self, a: &str, b: &str) -> (MutexGuard<'_, ()>, Option<MutexGuard<'_, ()>>) {
let ia = Self::stripe_index(a);
let ib = Self::stripe_index(b);
if ia == ib {
return (self.locks[ia].lock().unwrap_or_else(|e| e.into_inner()), None);
}
let (lo, hi) = if ia < ib { (ia, ib) } else { (ib, ia) };
let g_lo = self.locks[lo].lock().unwrap_or_else(|e| e.into_inner());
let g_hi = self.locks[hi].lock().unwrap_or_else(|e| e.into_inner());
(g_lo, Some(g_hi))
}
fn stripe_index(key: &str) -> usize {
let mut hasher = DefaultHasher::new();
key.hash(&mut hasher);
(hasher.finish() as usize) % N
}
}
impl<const N: usize> Default for StripedLock<N> {
fn default() -> Self {
Self::new()
}
}
unsafe impl<const N: usize> Sync for StripedLock<N> {}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_striped_lock_different_keys() {
let lock = StripedLock::<64>::new();
let _g1 = lock.lock("parent_a");
}
#[test]
fn test_striped_lock_same_key_serializes() {
let lock = StripedLock::<64>::new();
{
let _g = lock.lock("same_parent");
}
let _g2 = lock.lock("same_parent");
}
}