From 61cc00d23871b9d2e4c4f23af6917dc5ca7efb58 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?John=20K=C3=A5re=20Alsaker?= Date: Sat, 2 Sep 2023 23:19:34 +0200 Subject: [PATCH] Refactor `Lock` implementation --- compiler/rustc_data_structures/src/sharded.rs | 11 +- compiler/rustc_data_structures/src/sync.rs | 3 +- .../rustc_data_structures/src/sync/lock.rs | 494 ++++++++---------- 3 files changed, 239 insertions(+), 269 deletions(-) diff --git a/compiler/rustc_data_structures/src/sharded.rs b/compiler/rustc_data_structures/src/sharded.rs index 3706f62383f..c56fce02af6 100644 --- a/compiler/rustc_data_structures/src/sharded.rs +++ b/compiler/rustc_data_structures/src/sharded.rs @@ -1,7 +1,7 @@ use crate::fx::{FxHashMap, FxHasher}; #[cfg(parallel_compiler)] use crate::sync::{is_dyn_thread_safe, CacheAligned}; -use crate::sync::{Lock, LockGuard}; +use crate::sync::{Assume, Lock, LockGuard}; #[cfg(parallel_compiler)] use itertools::Either; use std::borrow::Borrow; @@ -75,6 +75,7 @@ impl Sharded { /// The shard is selected by hashing `val` with `FxHasher`. #[inline] + #[track_caller] pub fn lock_shard_by_value(&self, _val: &K) -> LockGuard<'_, T> { match self { Self::Single(single) => { @@ -83,7 +84,7 @@ impl Sharded { // SAFETY: We know `is_dyn_thread_safe` was false when creating the lock thus // `might_be_dyn_thread_safe` was also false. - unsafe { single.lock_assume_no_sync() } + unsafe { single.lock_assume(Assume::NoSync) } } #[cfg(parallel_compiler)] Self::Shards(..) => self.lock_shard_by_hash(make_hash(_val)), @@ -91,11 +92,13 @@ impl Sharded { } #[inline] + #[track_caller] pub fn lock_shard_by_hash(&self, hash: u64) -> LockGuard<'_, T> { self.lock_shard_by_index(get_shard_hash(hash)) } #[inline] + #[track_caller] pub fn lock_shard_by_index(&self, _i: usize) -> LockGuard<'_, T> { match self { Self::Single(single) => { @@ -104,7 +107,7 @@ impl Sharded { // SAFETY: We know `is_dyn_thread_safe` was false when creating the lock thus // `might_be_dyn_thread_safe` was also false. - unsafe { single.lock_assume_no_sync() } + unsafe { single.lock_assume(Assume::NoSync) } } #[cfg(parallel_compiler)] Self::Shards(shards) => { @@ -115,7 +118,7 @@ impl Sharded { // always inbounds. // SAFETY (lock_assume_sync): We know `is_dyn_thread_safe` was true when creating // the lock thus `might_be_dyn_thread_safe` was also true. - unsafe { shards.get_unchecked(_i & (SHARDS - 1)).0.lock_assume_sync() } + unsafe { shards.get_unchecked(_i & (SHARDS - 1)).0.lock_assume(Assume::Sync) } } } } diff --git a/compiler/rustc_data_structures/src/sync.rs b/compiler/rustc_data_structures/src/sync.rs index 8eda9043e35..a30ffcc120b 100644 --- a/compiler/rustc_data_structures/src/sync.rs +++ b/compiler/rustc_data_structures/src/sync.rs @@ -49,7 +49,7 @@ use std::ops::{Deref, DerefMut}; use std::panic::{catch_unwind, resume_unwind, AssertUnwindSafe}; mod lock; -pub use lock::{Lock, LockGuard}; +pub use lock::{Assume, Lock, LockGuard}; mod worker_local; pub use worker_local::{Registry, WorkerLocal}; @@ -86,7 +86,6 @@ mod mode { // Whether thread safety might be enabled. #[inline] - #[cfg(parallel_compiler)] pub fn might_be_dyn_thread_safe() -> bool { DYN_THREAD_SAFE_MODE.load(Ordering::Relaxed) != DYN_NOT_THREAD_SAFE } diff --git a/compiler/rustc_data_structures/src/sync/lock.rs b/compiler/rustc_data_structures/src/sync/lock.rs index 0a8c0306fe5..0188eb114fb 100644 --- a/compiler/rustc_data_structures/src/sync/lock.rs +++ b/compiler/rustc_data_structures/src/sync/lock.rs @@ -3,283 +3,257 @@ //! //! When `cfg(parallel_compiler)` is not set, the lock is instead a wrapper around `RefCell`. -#[cfg(not(parallel_compiler))] -use std::cell::RefCell; -#[cfg(parallel_compiler)] -use { - crate::cold_path, - crate::sync::DynSend, - crate::sync::DynSync, - parking_lot::lock_api::RawMutex, - std::cell::Cell, - std::cell::UnsafeCell, - std::fmt, - std::intrinsics::{likely, unlikely}, - std::marker::PhantomData, - std::mem::ManuallyDrop, - std::ops::{Deref, DerefMut}, -}; +#![allow(dead_code)] + +use std::fmt; #[cfg(not(parallel_compiler))] -pub use std::cell::RefMut as LockGuard; +pub use disabled::*; +#[cfg(parallel_compiler)] +pub use enabled::*; -#[cfg(not(parallel_compiler))] -#[derive(Debug)] -pub struct Lock(RefCell); - -#[cfg(not(parallel_compiler))] -impl Lock { - #[inline(always)] - pub fn new(inner: T) -> Self { - Lock(RefCell::new(inner)) - } - - #[inline(always)] - pub fn into_inner(self) -> T { - self.0.into_inner() - } - - #[inline(always)] - pub fn get_mut(&mut self) -> &mut T { - self.0.get_mut() - } - - #[inline(always)] - pub fn try_lock(&self) -> Option> { - self.0.try_borrow_mut().ok() - } - - #[inline(always)] - #[track_caller] - // This is unsafe to match the API for the `parallel_compiler` case. - pub unsafe fn lock_assume_no_sync(&self) -> LockGuard<'_, T> { - self.0.borrow_mut() - } - - #[inline(always)] - #[track_caller] - // This is unsafe to match the API for the `parallel_compiler` case. - pub unsafe fn lock_assume_sync(&self) -> LockGuard<'_, T> { - self.0.borrow_mut() - } - - #[inline(always)] - #[track_caller] - pub fn lock(&self) -> LockGuard<'_, T> { - self.0.borrow_mut() - } +#[derive(Clone, Copy, PartialEq)] +pub enum Assume { + NoSync, + Sync, } -/// A guard holding mutable access to a `Lock` which is in a locked state. -#[cfg(parallel_compiler)] -#[must_use = "if unused the Lock will immediately unlock"] -pub struct LockGuard<'a, T> { - lock: &'a Lock, - marker: PhantomData<&'a mut T>, -} +mod enabled { + use super::Assume; + use crate::sync::mode; + #[cfg(parallel_compiler)] + use crate::sync::{DynSend, DynSync}; + use parking_lot::lock_api::RawMutex as _; + use parking_lot::RawMutex; + use std::cell::Cell; + use std::cell::UnsafeCell; + use std::hint::unreachable_unchecked; + use std::intrinsics::unlikely; + use std::marker::PhantomData; + use std::ops::{Deref, DerefMut}; -#[cfg(parallel_compiler)] -impl<'a, T: 'a> Deref for LockGuard<'a, T> { - type Target = T; - #[inline] - fn deref(&self) -> &T { - // SAFETY: We have shared access to the mutable access owned by this type, - // so we can give out a shared reference. - unsafe { &*self.lock.data.get() } + /// A guard holding mutable access to a `Lock` which is in a locked state. + #[must_use = "if unused the Lock will immediately unlock"] + pub struct LockGuard<'a, T> { + lock: &'a Lock, + marker: PhantomData<&'a mut T>, + + /// The syncronization mode of the lock. This is explicitly passed to let LLVM relate it + /// to the original lock operation. + assume: Assume, } -} -#[cfg(parallel_compiler)] -impl<'a, T: 'a> DerefMut for LockGuard<'a, T> { - #[inline] - fn deref_mut(&mut self) -> &mut T { - // SAFETY: We have mutable access to the data so we can give out a mutable reference. - unsafe { &mut *self.lock.data.get() } - } -} - -#[cfg(parallel_compiler)] -impl<'a, T: 'a> Drop for LockGuard<'a, T> { - #[inline] - fn drop(&mut self) { - // SAFETY: We know that the lock is in a locked - // state because it is a invariant of this type. - unsafe { self.lock.raw.unlock() }; - } -} - -#[cfg(parallel_compiler)] -union LockRawUnion { - /// Indicates if the cell is locked. Only used if `LockRaw.sync` is false. - cell: ManuallyDrop>, - - /// A lock implementation that's only used if `LockRaw.sync` is true. - lock: ManuallyDrop, -} - -/// A raw lock which only uses synchronization if `might_be_dyn_thread_safe` is true. -/// It contains no associated data and is used in the implementation of `Lock` which does have such data. -/// -/// A manual implementation of a tagged union is used with the `sync` field and the `LockRawUnion` instead -/// of using enums as it results in better code generation. -#[cfg(parallel_compiler)] -struct LockRaw { - /// Indicates if synchronization is used via `opt.lock` if true, - /// or if a non-thread safe cell is used via `opt.cell`. This is set on initialization and never changed. - sync: bool, - opt: LockRawUnion, -} - -#[cfg(parallel_compiler)] -impl LockRaw { - fn new() -> Self { - if unlikely(super::mode::might_be_dyn_thread_safe()) { - // Create the lock with synchronization enabled using the `RawMutex` type. - LockRaw { - sync: true, - opt: LockRawUnion { lock: ManuallyDrop::new(parking_lot::RawMutex::INIT) }, - } - } else { - // Create the lock with synchronization disabled. - LockRaw { sync: false, opt: LockRawUnion { cell: ManuallyDrop::new(Cell::new(false)) } } + impl<'a, T: 'a> Deref for LockGuard<'a, T> { + type Target = T; + #[inline] + fn deref(&self) -> &T { + // SAFETY: We have shared access to the mutable access owned by this type, + // so we can give out a shared reference. + unsafe { &*self.lock.data.get() } } } - #[inline(always)] - fn try_lock(&self) -> bool { - // SAFETY: This is safe since the union fields are used in accordance with `self.sync`. - unsafe { - if likely(!self.sync) { - if self.opt.cell.get() { - false + impl<'a, T: 'a> DerefMut for LockGuard<'a, T> { + #[inline] + fn deref_mut(&mut self) -> &mut T { + // SAFETY: We have mutable access to the data so we can give out a mutable reference. + unsafe { &mut *self.lock.data.get() } + } + } + + impl<'a, T: 'a> Drop for LockGuard<'a, T> { + #[inline] + fn drop(&mut self) { + // SAFETY (dispatch): We get `self.assume` from the lock operation so it is consistent + // with the lock state. + // SAFETY (unlock): We know that the lock is locked as this type is a proof of that. + unsafe { + self.lock.dispatch( + self.assume, + |cell| { + debug_assert_eq!(cell.get(), true); + cell.set(false); + Some(()) + }, + |lock| lock.unlock(), + ); + }; + } + } + + enum LockMode { + NoSync(Cell), + Sync(RawMutex), + } + + impl LockMode { + #[inline(always)] + fn to_assume(&self) -> Assume { + match self { + LockMode::NoSync(..) => Assume::NoSync, + LockMode::Sync(..) => Assume::Sync, + } + } + } + + /// The value representing a locked state for the `Cell`. + const LOCKED: bool = true; + + /// A lock which only uses synchronization if `might_be_dyn_thread_safe` is true. + /// It implements `DynSend` and `DynSync` instead of the typical `Send` and `Sync`. + pub struct Lock { + mode: LockMode, + data: UnsafeCell, + } + + impl Lock { + #[inline(always)] + pub fn new(inner: T) -> Self { + Lock { + mode: if unlikely(mode::might_be_dyn_thread_safe()) { + // Create the lock with synchronization enabled using the `RawMutex` type. + LockMode::Sync(RawMutex::INIT) } else { - self.opt.cell.set(true); - true + // Create the lock with synchronization disabled. + LockMode::NoSync(Cell::new(!LOCKED)) + }, + data: UnsafeCell::new(inner), + } + } + + #[inline(always)] + pub fn into_inner(self) -> T { + self.data.into_inner() + } + + #[inline(always)] + pub fn get_mut(&mut self) -> &mut T { + self.data.get_mut() + } + + /// This dispatches on the `LockMode` and gives access to its variants depending on + /// `assume`. If `no_sync` returns `None` this will panic. + /// + /// Safety + /// This method must only be called if `might_be_dyn_thread_safe` on lock creation matches + /// matches the `assume` argument. + #[inline(always)] + #[track_caller] + unsafe fn dispatch( + &self, + assume: Assume, + no_sync: impl FnOnce(&Cell) -> Option, + sync: impl FnOnce(&RawMutex) -> R, + ) -> R { + #[inline(never)] + #[track_caller] + #[cold] + fn lock_held() -> ! { + panic!("lock was already held") + } + + match assume { + Assume::NoSync => { + let LockMode::NoSync(cell) = &self.mode else { + unsafe { unreachable_unchecked() } + }; + if let Some(v) = no_sync(cell) { + v + } else { + // Call this here instead of in `no_sync` so `track_caller` gets properly + // passed along. + lock_held() + } + } + Assume::Sync => { + let LockMode::Sync(lock) = &self.mode else { + unsafe { unreachable_unchecked() } + }; + sync(lock) } - } else { - self.opt.lock.try_lock() } } + + #[inline(always)] + pub fn try_lock(&self) -> Option> { + let assume = self.mode.to_assume(); + unsafe { + self.dispatch( + assume, + |cell| Some((cell.get() != LOCKED).then(|| cell.set(LOCKED)).is_some()), + RawMutex::try_lock, + ) + .then(|| LockGuard { lock: self, marker: PhantomData, assume }) + } + } + + #[inline(always)] + #[track_caller] + pub unsafe fn lock_assume(&self, assume: Assume) -> LockGuard<'_, T> { + unsafe { + self.dispatch( + assume, + |cell| (cell.replace(LOCKED) != LOCKED).then(|| ()), + RawMutex::lock, + ); + LockGuard { lock: self, marker: PhantomData, assume } + } + } + + #[inline(always)] + #[track_caller] + pub fn lock(&self) -> LockGuard<'_, T> { + unsafe { self.lock_assume(self.mode.to_assume()) } + } } - #[inline(always)] - fn lock(&self) { - // SAFETY: This is safe since `self.sync` is used in accordance with the preconditions of - // `lock_assume_no_sync` and `lock_assume_sync`. - unsafe { - if likely(!self.sync) { - self.lock_assume_no_sync() - } else { - self.lock_assume_sync(); - } - } - } - - /// This acquires the lock assuming no syncronization is required. - /// - /// Safety - /// This method must only be called if `might_be_dyn_thread_safe` was false on lock creation. - #[inline(always)] - unsafe fn lock_assume_no_sync(&self) { - // SAFETY: This is safe since `self.opt.cell` is the union field used due to the - // precondition on this function. - unsafe { - if unlikely(self.opt.cell.replace(true)) { - cold_path(|| panic!("lock was already held")) - } - } - } - - /// This acquires the lock assuming syncronization is required. - /// - /// Safety - /// This method must only be called if `might_be_dyn_thread_safe` was true on lock creation. - #[inline(always)] - unsafe fn lock_assume_sync(&self) { - // SAFETY: This is safe since `self.opt.lock` is the union field used due to the - // precondition on this function. - unsafe { - self.opt.lock.lock(); - } - } - - /// This unlocks the lock. - /// - /// Safety - /// This method may only be called if the lock is currently held. - #[inline(always)] - unsafe fn unlock(&self) { - // SAFETY: The union use is safe since the union fields are used in accordance with - // `self.sync` and the `unlock` method precondition is upheld by the caller. - unsafe { - if likely(!self.sync) { - debug_assert_eq!(self.opt.cell.get(), true); - self.opt.cell.set(false); - } else { - self.opt.lock.unlock(); - } - } - } + #[cfg(parallel_compiler)] + unsafe impl DynSend for Lock {} + #[cfg(parallel_compiler)] + unsafe impl DynSync for Lock {} } -/// A lock which only uses synchronization if `might_be_dyn_thread_safe` is true. -/// It implements `DynSend` and `DynSync` instead of the typical `Send` and `Sync`. -#[cfg(parallel_compiler)] -pub struct Lock { - raw: LockRaw, - data: UnsafeCell, -} +mod disabled { + use super::Assume; + use std::cell::RefCell; -#[cfg(parallel_compiler)] -impl Lock { - #[inline(always)] - pub fn new(inner: T) -> Self { - Lock { raw: LockRaw::new(), data: UnsafeCell::new(inner) } - } + pub use std::cell::RefMut as LockGuard; - #[inline(always)] - pub fn into_inner(self) -> T { - self.data.into_inner() - } + pub struct Lock(RefCell); - #[inline(always)] - pub fn get_mut(&mut self) -> &mut T { - self.data.get_mut() - } - - #[inline(always)] - pub fn try_lock(&self) -> Option> { - if self.raw.try_lock() { Some(LockGuard { lock: self, marker: PhantomData }) } else { None } - } - - /// This acquires the lock assuming no syncronization is required. - /// - /// Safety - /// This method must only be called if `might_be_dyn_thread_safe` was false on lock creation. - #[inline(always)] - pub(crate) unsafe fn lock_assume_no_sync(&self) -> LockGuard<'_, T> { - unsafe { - self.raw.lock_assume_no_sync(); + impl Lock { + #[inline(always)] + pub fn new(inner: T) -> Self { + Lock(RefCell::new(inner)) } - LockGuard { lock: self, marker: PhantomData } - } - /// This acquires the lock assuming syncronization is required. - /// - /// Safety - /// This method must only be called if `might_be_dyn_thread_safe` was true on lock creation. - #[inline(always)] - pub(crate) unsafe fn lock_assume_sync(&self) -> LockGuard<'_, T> { - unsafe { - self.raw.lock_assume_sync(); + #[inline(always)] + pub fn into_inner(self) -> T { + self.0.into_inner() } - LockGuard { lock: self, marker: PhantomData } - } - #[inline(always)] - pub fn lock(&self) -> LockGuard<'_, T> { - self.raw.lock(); - LockGuard { lock: self, marker: PhantomData } + #[inline(always)] + pub fn get_mut(&mut self) -> &mut T { + self.0.get_mut() + } + + #[inline(always)] + pub fn try_lock(&self) -> Option> { + self.0.try_borrow_mut().ok() + } + + #[inline(always)] + #[track_caller] + // This is unsafe to match the API for the `parallel_compiler` case. + pub unsafe fn lock_assume(&self, _assume: Assume) -> LockGuard<'_, T> { + self.0.borrow_mut() + } + + #[inline(always)] + #[track_caller] + pub fn lock(&self) -> LockGuard<'_, T> { + self.0.borrow_mut() + } } } @@ -303,12 +277,13 @@ impl Lock { } } -#[cfg(parallel_compiler)] -unsafe impl DynSend for Lock {} -#[cfg(parallel_compiler)] -unsafe impl DynSync for Lock {} +impl Default for Lock { + #[inline] + fn default() -> Self { + Lock::new(T::default()) + } +} -#[cfg(parallel_compiler)] impl fmt::Debug for Lock { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self.try_lock() { @@ -326,10 +301,3 @@ impl fmt::Debug for Lock { } } } - -impl Default for Lock { - #[inline] - fn default() -> Self { - Lock::new(T::default()) - } -}