mirror of https://github.com/xacrimon/dashmap
Use a proper lock implementation (#214)
* use a proper but reentrant safe lock impl * impl downgrade
This commit is contained in:
parent
5fbb68f188
commit
3236196d01
|
@ -18,6 +18,7 @@ raw-api = []
|
||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
lock_api = "0.4.7"
|
lock_api = "0.4.7"
|
||||||
|
parking_lot_core = "0.9.3"
|
||||||
hashbrown = { version = "0.12.0", default-features = false }
|
hashbrown = { version = "0.12.0", default-features = false }
|
||||||
serde = { version = "1.0.136", optional = true, features = ["derive"] }
|
serde = { version = "1.0.136", optional = true, features = ["derive"] }
|
||||||
cfg-if = "1.0.0"
|
cfg-if = "1.0.0"
|
||||||
|
|
318
src/lock.rs
318
src/lock.rs
|
@ -1,76 +1,300 @@
|
||||||
use lock_api::GuardSend;
|
use core::sync::atomic::{AtomicUsize, Ordering};
|
||||||
use std::hint;
|
use parking_lot_core::{ParkToken, SpinWait, UnparkToken};
|
||||||
use std::mem;
|
|
||||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
||||||
|
|
||||||
const USIZE_BITS: usize = mem::size_of::<usize>() * 8;
|
|
||||||
const EXCLUSIVE_BIT: usize = 1 << (USIZE_BITS - 1);
|
|
||||||
|
|
||||||
pub type RwLock<T> = lock_api::RwLock<RawRwLock, T>;
|
pub type RwLock<T> = lock_api::RwLock<RawRwLock, T>;
|
||||||
pub type RwLockReadGuard<'a, T> = lock_api::RwLockReadGuard<'a, RawRwLock, T>;
|
pub type RwLockReadGuard<'a, T> = lock_api::RwLockReadGuard<'a, RawRwLock, T>;
|
||||||
pub type RwLockWriteGuard<'a, T> = lock_api::RwLockWriteGuard<'a, RawRwLock, T>;
|
pub type RwLockWriteGuard<'a, T> = lock_api::RwLockWriteGuard<'a, RawRwLock, T>;
|
||||||
|
|
||||||
|
const READERS_PARKED: usize = 0b0001;
|
||||||
|
const WRITERS_PARKED: usize = 0b0010;
|
||||||
|
const ONE_READER: usize = 0b0100;
|
||||||
|
const ONE_WRITER: usize = !(READERS_PARKED | WRITERS_PARKED);
|
||||||
|
|
||||||
pub struct RawRwLock {
|
pub struct RawRwLock {
|
||||||
data: AtomicUsize,
|
state: AtomicUsize,
|
||||||
}
|
}
|
||||||
|
|
||||||
unsafe impl lock_api::RawRwLock for RawRwLock {
|
unsafe impl lock_api::RawRwLock for RawRwLock {
|
||||||
type GuardMarker = GuardSend;
|
|
||||||
|
|
||||||
#[allow(clippy::declare_interior_mutable_const)]
|
#[allow(clippy::declare_interior_mutable_const)]
|
||||||
const INIT: Self = RawRwLock {
|
const INIT: Self = Self {
|
||||||
data: AtomicUsize::new(0),
|
state: AtomicUsize::new(0),
|
||||||
};
|
};
|
||||||
|
|
||||||
fn lock_shared(&self) {
|
type GuardMarker = lock_api::GuardNoSend;
|
||||||
while !self.try_lock_shared() {
|
|
||||||
hint::spin_loop();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn try_lock_shared(&self) -> bool {
|
|
||||||
let x = self.data.load(Ordering::Acquire);
|
|
||||||
if x & EXCLUSIVE_BIT != 0 {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
|
|
||||||
let y = x + 1;
|
|
||||||
self.data
|
|
||||||
.compare_exchange(x, y, Ordering::Release, Ordering::Relaxed)
|
|
||||||
.is_ok()
|
|
||||||
}
|
|
||||||
|
|
||||||
unsafe fn unlock_shared(&self) {
|
|
||||||
self.data.fetch_sub(1, Ordering::Release);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn lock_exclusive(&self) {
|
|
||||||
while !self.try_lock_exclusive() {
|
|
||||||
hint::spin_loop();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
|
#[inline]
|
||||||
fn try_lock_exclusive(&self) -> bool {
|
fn try_lock_exclusive(&self) -> bool {
|
||||||
self.data
|
self.state
|
||||||
.compare_exchange(0, EXCLUSIVE_BIT, Ordering::Release, Ordering::Relaxed)
|
.compare_exchange(0, ONE_WRITER, Ordering::Acquire, Ordering::Relaxed)
|
||||||
.is_ok()
|
.is_ok()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
fn lock_exclusive(&self) {
|
||||||
|
if self
|
||||||
|
.state
|
||||||
|
.compare_exchange_weak(0, ONE_WRITER, Ordering::Acquire, Ordering::Relaxed)
|
||||||
|
.is_err()
|
||||||
|
{
|
||||||
|
self.lock_exclusive_slow();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline]
|
||||||
unsafe fn unlock_exclusive(&self) {
|
unsafe fn unlock_exclusive(&self) {
|
||||||
self.data.store(0, Ordering::Release)
|
if self
|
||||||
|
.state
|
||||||
|
.compare_exchange(ONE_WRITER, 0, Ordering::Release, Ordering::Relaxed)
|
||||||
|
.is_err()
|
||||||
|
{
|
||||||
|
self.unlock_exclusive_slow();
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn is_locked(&self) -> bool {
|
#[inline]
|
||||||
self.data.load(Ordering::Acquire) != 0
|
fn try_lock_shared(&self) -> bool {
|
||||||
|
self.try_lock_shared_fast() || self.try_lock_shared_slow()
|
||||||
}
|
}
|
||||||
|
|
||||||
fn is_locked_exclusive(&self) -> bool {
|
#[inline]
|
||||||
self.data.load(Ordering::Acquire) & EXCLUSIVE_BIT != 0
|
fn lock_shared(&self) {
|
||||||
|
if !self.try_lock_shared_fast() {
|
||||||
|
self.lock_shared_slow();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
unsafe fn unlock_shared(&self) {
|
||||||
|
let state = self.state.fetch_sub(ONE_READER, Ordering::Release);
|
||||||
|
|
||||||
|
if state == (ONE_READER | WRITERS_PARKED) {
|
||||||
|
self.unlock_shared_slow();
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
unsafe impl lock_api::RawRwLockDowngrade for RawRwLock {
|
unsafe impl lock_api::RawRwLockDowngrade for RawRwLock {
|
||||||
|
#[inline]
|
||||||
unsafe fn downgrade(&self) {
|
unsafe fn downgrade(&self) {
|
||||||
self.data.store(1, Ordering::SeqCst);
|
let state = self
|
||||||
|
.state
|
||||||
|
.fetch_and(ONE_READER | WRITERS_PARKED, Ordering::Release);
|
||||||
|
if state & READERS_PARKED != 0 {
|
||||||
|
parking_lot_core::unpark_all((self as *const _ as usize) + 1, UnparkToken(0));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RawRwLock {
|
||||||
|
#[cold]
|
||||||
|
fn lock_exclusive_slow(&self) {
|
||||||
|
let mut acquire_with = 0;
|
||||||
|
loop {
|
||||||
|
let mut spin = SpinWait::new();
|
||||||
|
let mut state = self.state.load(Ordering::Relaxed);
|
||||||
|
|
||||||
|
loop {
|
||||||
|
while state & ONE_WRITER == 0 {
|
||||||
|
match self.state.compare_exchange_weak(
|
||||||
|
state,
|
||||||
|
state | ONE_WRITER | acquire_with,
|
||||||
|
Ordering::Acquire,
|
||||||
|
Ordering::Relaxed,
|
||||||
|
) {
|
||||||
|
Ok(_) => return,
|
||||||
|
Err(e) => state = e,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if state & WRITERS_PARKED == 0 {
|
||||||
|
if spin.spin() {
|
||||||
|
state = self.state.load(Ordering::Relaxed);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Err(e) = self.state.compare_exchange_weak(
|
||||||
|
state,
|
||||||
|
state | WRITERS_PARKED,
|
||||||
|
Ordering::Relaxed,
|
||||||
|
Ordering::Relaxed,
|
||||||
|
) {
|
||||||
|
state = e;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let _ = unsafe {
|
||||||
|
parking_lot_core::park(
|
||||||
|
self as *const _ as usize,
|
||||||
|
|| {
|
||||||
|
let state = self.state.load(Ordering::Relaxed);
|
||||||
|
(state & ONE_WRITER != 0) && (state & WRITERS_PARKED != 0)
|
||||||
|
},
|
||||||
|
|| {},
|
||||||
|
|_, _| {},
|
||||||
|
ParkToken(0),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
};
|
||||||
|
|
||||||
|
acquire_with = WRITERS_PARKED;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cold]
|
||||||
|
fn unlock_exclusive_slow(&self) {
|
||||||
|
let state = self.state.load(Ordering::Relaxed);
|
||||||
|
assert_eq!(state & ONE_WRITER, ONE_WRITER);
|
||||||
|
|
||||||
|
let mut parked = state & (READERS_PARKED | WRITERS_PARKED);
|
||||||
|
assert_ne!(parked, 0);
|
||||||
|
|
||||||
|
if parked != (READERS_PARKED | WRITERS_PARKED) {
|
||||||
|
if let Err(new_state) =
|
||||||
|
self.state
|
||||||
|
.compare_exchange(state, 0, Ordering::Release, Ordering::Relaxed)
|
||||||
|
{
|
||||||
|
assert_eq!(new_state, ONE_WRITER | READERS_PARKED | WRITERS_PARKED);
|
||||||
|
parked = READERS_PARKED | WRITERS_PARKED;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if parked == (READERS_PARKED | WRITERS_PARKED) {
|
||||||
|
self.state.store(WRITERS_PARKED, Ordering::Release);
|
||||||
|
parked = READERS_PARKED;
|
||||||
|
}
|
||||||
|
|
||||||
|
if parked == READERS_PARKED {
|
||||||
|
return unsafe {
|
||||||
|
parking_lot_core::unpark_all((self as *const _ as usize) + 1, UnparkToken(0));
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
assert_eq!(parked, WRITERS_PARKED);
|
||||||
|
unsafe {
|
||||||
|
parking_lot_core::unpark_one(self as *const _ as usize, |_| UnparkToken(0));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline(always)]
|
||||||
|
fn try_lock_shared_fast(&self) -> bool {
|
||||||
|
let state = self.state.load(Ordering::Relaxed);
|
||||||
|
|
||||||
|
if let Some(new_state) = state.checked_add(ONE_READER) {
|
||||||
|
if new_state & ONE_WRITER != ONE_WRITER {
|
||||||
|
return self
|
||||||
|
.state
|
||||||
|
.compare_exchange_weak(state, new_state, Ordering::Acquire, Ordering::Relaxed)
|
||||||
|
.is_ok();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
false
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cold]
|
||||||
|
fn try_lock_shared_slow(&self) -> bool {
|
||||||
|
let mut state = self.state.load(Ordering::Relaxed);
|
||||||
|
|
||||||
|
while let Some(new_state) = state.checked_add(ONE_READER) {
|
||||||
|
if new_state & ONE_WRITER == ONE_WRITER {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
|
match self.state.compare_exchange_weak(
|
||||||
|
state,
|
||||||
|
new_state,
|
||||||
|
Ordering::Acquire,
|
||||||
|
Ordering::Relaxed,
|
||||||
|
) {
|
||||||
|
Ok(_) => return true,
|
||||||
|
Err(e) => state = e,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
false
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cold]
|
||||||
|
fn lock_shared_slow(&self) {
|
||||||
|
loop {
|
||||||
|
let mut spin = SpinWait::new();
|
||||||
|
let mut state = self.state.load(Ordering::Relaxed);
|
||||||
|
|
||||||
|
loop {
|
||||||
|
let mut backoff = SpinWait::new();
|
||||||
|
while let Some(new_state) = state.checked_add(ONE_READER) {
|
||||||
|
assert_ne!(
|
||||||
|
new_state & ONE_WRITER,
|
||||||
|
ONE_WRITER,
|
||||||
|
"reader count overflowed",
|
||||||
|
);
|
||||||
|
|
||||||
|
if self
|
||||||
|
.state
|
||||||
|
.compare_exchange_weak(
|
||||||
|
state,
|
||||||
|
new_state,
|
||||||
|
Ordering::Acquire,
|
||||||
|
Ordering::Relaxed,
|
||||||
|
)
|
||||||
|
.is_ok()
|
||||||
|
{
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
backoff.spin_no_yield();
|
||||||
|
state = self.state.load(Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
|
||||||
|
if state & READERS_PARKED == 0 {
|
||||||
|
if spin.spin() {
|
||||||
|
state = self.state.load(Ordering::Relaxed);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Err(e) = self.state.compare_exchange_weak(
|
||||||
|
state,
|
||||||
|
state | READERS_PARKED,
|
||||||
|
Ordering::Relaxed,
|
||||||
|
Ordering::Relaxed,
|
||||||
|
) {
|
||||||
|
state = e;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let _ = unsafe {
|
||||||
|
parking_lot_core::park(
|
||||||
|
(self as *const _ as usize) + 1,
|
||||||
|
|| {
|
||||||
|
let state = self.state.load(Ordering::Relaxed);
|
||||||
|
(state & ONE_WRITER == ONE_WRITER) && (state & READERS_PARKED != 0)
|
||||||
|
},
|
||||||
|
|| {},
|
||||||
|
|_, _| {},
|
||||||
|
ParkToken(0),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
};
|
||||||
|
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cold]
|
||||||
|
fn unlock_shared_slow(&self) {
|
||||||
|
if self
|
||||||
|
.state
|
||||||
|
.compare_exchange(WRITERS_PARKED, 0, Ordering::Relaxed, Ordering::Relaxed)
|
||||||
|
.is_ok()
|
||||||
|
{
|
||||||
|
unsafe {
|
||||||
|
parking_lot_core::unpark_one(self as *const _ as usize, |_| UnparkToken(0));
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
@ -0,0 +1,81 @@
|
||||||
|
use super::u128::AtomicU128;
|
||||||
|
use std::borrow::Borrow;
|
||||||
|
use std::cell::UnsafeCell;
|
||||||
|
use std::hash::{BuildHasher, Hash, Hasher};
|
||||||
|
use std::mem::MaybeUninit;
|
||||||
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
|
use std::sync::{Arc, Mutex};
|
||||||
|
|
||||||
|
const TOMBSTONE_BIT: u64 = 1 << 63;
|
||||||
|
const ALLOCATED_BIT: u64 = 1 << 62;
|
||||||
|
const POINTER_MASK: u64 = 0x3FFFFFFFFFFFFFFF;
|
||||||
|
|
||||||
|
fn hash<S, K>(hasher: &S, key: &K) -> u64
|
||||||
|
where
|
||||||
|
S: BuildHasher,
|
||||||
|
K: Hash,
|
||||||
|
{
|
||||||
|
let mut hasher = hasher.build_hasher();
|
||||||
|
key.hash(&mut hasher);
|
||||||
|
hasher.finish()
|
||||||
|
}
|
||||||
|
|
||||||
|
struct Slot<K, V> {
|
||||||
|
data: AtomicU64,
|
||||||
|
pair: UnsafeCell<MaybeUninit<(K, V)>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct Table<K, V, S> {
|
||||||
|
hash: Arc<S>,
|
||||||
|
slots: Box<[Slot<K, V>]>,
|
||||||
|
mask: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<K, V, S> Table<K, V, S>
|
||||||
|
where
|
||||||
|
K: Eq + Hash,
|
||||||
|
S: BuildHasher,
|
||||||
|
{
|
||||||
|
pub fn new(hasher: Arc<S>, capacity: usize) -> Self {
|
||||||
|
debug_assert!(capacity.is_power_of_two());
|
||||||
|
let slots = (0..capacity)
|
||||||
|
.map(|_| Slot {
|
||||||
|
data: AtomicU64::new(0),
|
||||||
|
pair: UnsafeCell::new(MaybeUninit::uninit()),
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
|
||||||
|
Table {
|
||||||
|
hash: hasher,
|
||||||
|
slots: slots.into_boxed_slice(),
|
||||||
|
mask: capacity - 1,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn get<Q>(&self, key: &Q) -> Option<*mut (K, V)>
|
||||||
|
where
|
||||||
|
K: Borrow<Q>,
|
||||||
|
Q: Eq + Hash,
|
||||||
|
{
|
||||||
|
let hash = hash(&*self.hash, key);
|
||||||
|
let mut idx = hash as usize & self.mask;
|
||||||
|
let mut i = 0;
|
||||||
|
|
||||||
|
loop {
|
||||||
|
let slot = &self.slots[idx];
|
||||||
|
let data = slot.data.load(Ordering::Relaxed);
|
||||||
|
let ptr = (data & POINTER_MASK) as *mut (K, V);
|
||||||
|
if !ptr.is_null() {
|
||||||
|
let stored = unsafe { (*ptr).0.borrow() };
|
||||||
|
if stored == key {
|
||||||
|
return Some(ptr);
|
||||||
|
}
|
||||||
|
} else if data & TOMBSTONE_BIT != TOMBSTONE_BIT || i > self.mask {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
|
||||||
|
idx = (idx + 1) & self.mask;
|
||||||
|
i += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
Loading…
Reference in New Issue