use std::cell::UnsafeCell;
use std::collections::HashMap;
use std::fmt;
use std::marker::PhantomData;
use std::mem;
use std::ops::{Deref, DerefMut};
use std::panic::{RefUnwindSafe, UnwindSafe};
use std::sync::{LockResult, PoisonError, TryLockError, TryLockResult};
use std::sync::{Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard};
use std::thread::{self, ThreadId};
use CachePadded;
const NUM_SHARDS: usize = 8;
struct Shard {
lock: RwLock<()>,
write_guard: UnsafeCell<Option<RwLockWriteGuard<'static, ()>>>,
}
pub struct ShardedLock<T: ?Sized> {
shards: Box<[CachePadded<Shard>]>,
value: UnsafeCell<T>,
}
unsafe impl<T: ?Sized + Send> Send for ShardedLock<T> {}
unsafe impl<T: ?Sized + Send + Sync> Sync for ShardedLock<T> {}
impl<T: ?Sized> UnwindSafe for ShardedLock<T> {}
impl<T: ?Sized> RefUnwindSafe for ShardedLock<T> {}
impl<T> ShardedLock<T> {
pub fn new(value: T) -> ShardedLock<T> {
ShardedLock {
shards: (0..NUM_SHARDS)
.map(|_| {
CachePadded::new(Shard {
lock: RwLock::new(()),
write_guard: UnsafeCell::new(None),
})
})
.collect::<Vec<_>>()
.into_boxed_slice(),
value: UnsafeCell::new(value),
}
}
pub fn into_inner(self) -> LockResult<T> {
let is_poisoned = self.is_poisoned();
let inner = self.value.into_inner();
if is_poisoned {
Err(PoisonError::new(inner))
} else {
Ok(inner)
}
}
}
impl<T: ?Sized> ShardedLock<T> {
pub fn is_poisoned(&self) -> bool {
self.shards[0].lock.is_poisoned()
}
pub fn get_mut(&mut self) -> LockResult<&mut T> {
let is_poisoned = self.is_poisoned();
let inner = unsafe { &mut *self.value.get() };
if is_poisoned {
Err(PoisonError::new(inner))
} else {
Ok(inner)
}
}
pub fn try_read(&self) -> TryLockResult<ShardedLockReadGuard<T>> {
let current_index = current_index().unwrap_or(0);
let shard_index = current_index & (self.shards.len() - 1);
match self.shards[shard_index].lock.try_read() {
Ok(guard) => Ok(ShardedLockReadGuard {
lock: self,
_guard: guard,
_marker: PhantomData,
}),
Err(TryLockError::Poisoned(err)) => {
let guard = ShardedLockReadGuard {
lock: self,
_guard: err.into_inner(),
_marker: PhantomData,
};
Err(TryLockError::Poisoned(PoisonError::new(guard)))
}
Err(TryLockError::WouldBlock) => Err(TryLockError::WouldBlock),
}
}
pub fn read(&self) -> LockResult<ShardedLockReadGuard<T>> {
let current_index = current_index().unwrap_or(0);
let shard_index = current_index & (self.shards.len() - 1);
match self.shards[shard_index].lock.read() {
Ok(guard) => Ok(ShardedLockReadGuard {
lock: self,
_guard: guard,
_marker: PhantomData,
}),
Err(err) => Err(PoisonError::new(ShardedLockReadGuard {
lock: self,
_guard: err.into_inner(),
_marker: PhantomData,
})),
}
}
pub fn try_write(&self) -> TryLockResult<ShardedLockWriteGuard<T>> {
let mut poisoned = false;
let mut blocked = None;
for (i, shard) in self.shards.iter().enumerate() {
let guard = match shard.lock.try_write() {
Ok(guard) => guard,
Err(TryLockError::Poisoned(err)) => {
poisoned = true;
err.into_inner()
}
Err(TryLockError::WouldBlock) => {
blocked = Some(i);
break;
}
};
unsafe {
let guard: RwLockWriteGuard<'static, ()> = mem::transmute(guard);
let dest: *mut _ = shard.write_guard.get();
*dest = Some(guard);
}
}
if let Some(i) = blocked {
for shard in self.shards[0..i].iter().rev() {
unsafe {
let dest: *mut _ = shard.write_guard.get();
let guard = mem::replace(&mut *dest, None);
drop(guard);
}
}
Err(TryLockError::WouldBlock)
} else if poisoned {
let guard = ShardedLockWriteGuard {
lock: self,
_marker: PhantomData,
};
Err(TryLockError::Poisoned(PoisonError::new(guard)))
} else {
Ok(ShardedLockWriteGuard {
lock: self,
_marker: PhantomData,
})
}
}
pub fn write(&self) -> LockResult<ShardedLockWriteGuard<T>> {
let mut poisoned = false;
for shard in self.shards.iter() {
let guard = match shard.lock.write() {
Ok(guard) => guard,
Err(err) => {
poisoned = true;
err.into_inner()
}
};
unsafe {
let guard: RwLockWriteGuard<'_, ()> = guard;
let guard: RwLockWriteGuard<'static, ()> = mem::transmute(guard);
let dest: *mut _ = shard.write_guard.get();
*dest = Some(guard);
}
}
if poisoned {
Err(PoisonError::new(ShardedLockWriteGuard {
lock: self,
_marker: PhantomData,
}))
} else {
Ok(ShardedLockWriteGuard {
lock: self,
_marker: PhantomData,
})
}
}
}
impl<T: ?Sized + fmt::Debug> fmt::Debug for ShardedLock<T> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self.try_read() {
Ok(guard) => f
.debug_struct("ShardedLock")
.field("data", &&*guard)
.finish(),
Err(TryLockError::Poisoned(err)) => f
.debug_struct("ShardedLock")
.field("data", &&**err.get_ref())
.finish(),
Err(TryLockError::WouldBlock) => {
struct LockedPlaceholder;
impl fmt::Debug for LockedPlaceholder {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.write_str("<locked>")
}
}
f.debug_struct("ShardedLock")
.field("data", &LockedPlaceholder)
.finish()
}
}
}
}
impl<T: Default> Default for ShardedLock<T> {
fn default() -> ShardedLock<T> {
ShardedLock::new(Default::default())
}
}
impl<T> From<T> for ShardedLock<T> {
fn from(t: T) -> Self {
ShardedLock::new(t)
}
}
pub struct ShardedLockReadGuard<'a, T: ?Sized + 'a> {
lock: &'a ShardedLock<T>,
_guard: RwLockReadGuard<'a, ()>,
_marker: PhantomData<RwLockReadGuard<'a, T>>,
}
unsafe impl<'a, T: ?Sized + Sync> Sync for ShardedLockReadGuard<'a, T> {}
impl<'a, T: ?Sized> Deref for ShardedLockReadGuard<'a, T> {
type Target = T;
fn deref(&self) -> &T {
unsafe { &*self.lock.value.get() }
}
}
impl<'a, T: fmt::Debug> fmt::Debug for ShardedLockReadGuard<'a, T> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("ShardedLockReadGuard")
.field("lock", &self.lock)
.finish()
}
}
impl<'a, T: ?Sized + fmt::Display> fmt::Display for ShardedLockReadGuard<'a, T> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
(**self).fmt(f)
}
}
pub struct ShardedLockWriteGuard<'a, T: ?Sized + 'a> {
lock: &'a ShardedLock<T>,
_marker: PhantomData<RwLockWriteGuard<'a, T>>,
}
unsafe impl<'a, T: ?Sized + Sync> Sync for ShardedLockWriteGuard<'a, T> {}
impl<'a, T: ?Sized> Drop for ShardedLockWriteGuard<'a, T> {
fn drop(&mut self) {
for shard in self.lock.shards.iter().rev() {
unsafe {
let dest: *mut _ = shard.write_guard.get();
let guard = mem::replace(&mut *dest, None);
drop(guard);
}
}
}
}
impl<'a, T: fmt::Debug> fmt::Debug for ShardedLockWriteGuard<'a, T> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("ShardedLockWriteGuard")
.field("lock", &self.lock)
.finish()
}
}
impl<'a, T: ?Sized + fmt::Display> fmt::Display for ShardedLockWriteGuard<'a, T> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
(**self).fmt(f)
}
}
impl<'a, T: ?Sized> Deref for ShardedLockWriteGuard<'a, T> {
type Target = T;
fn deref(&self) -> &T {
unsafe { &*self.lock.value.get() }
}
}
impl<'a, T: ?Sized> DerefMut for ShardedLockWriteGuard<'a, T> {
fn deref_mut(&mut self) -> &mut T {
unsafe { &mut *self.lock.value.get() }
}
}
#[inline]
fn current_index() -> Option<usize> {
REGISTRATION.try_with(|reg| reg.index).ok()
}
struct ThreadIndices {
mapping: HashMap<ThreadId, usize>,
free_list: Vec<usize>,
next_index: usize,
}
lazy_static! {
static ref THREAD_INDICES: Mutex<ThreadIndices> = Mutex::new(ThreadIndices {
mapping: HashMap::new(),
free_list: Vec::new(),
next_index: 0,
});
}
struct Registration {
index: usize,
thread_id: ThreadId,
}
impl Drop for Registration {
fn drop(&mut self) {
let mut indices = THREAD_INDICES.lock().unwrap();
indices.mapping.remove(&self.thread_id);
indices.free_list.push(self.index);
}
}
thread_local! {
static REGISTRATION: Registration = {
let thread_id = thread::current().id();
let mut indices = THREAD_INDICES.lock().unwrap();
let index = match indices.free_list.pop() {
Some(i) => i,
None => {
let i = indices.next_index;
indices.next_index += 1;
i
}
};
indices.mapping.insert(thread_id, index);
Registration {
index,
thread_id,
}
};
}