use super::RawRwLock;
#[cfg(feature = "threading")]
use core::cell::Cell;
use lock_api::{
RawRwLock as RawRwLockTrait, RawRwLockDowngrade, RawRwLockUpgrade as RawRwLockUpgradeTrait,
RawRwLockUpgradeDowngrade,
};
#[cfg(feature = "threading")]
use std::sync::OnceLock;
#[cfg(feature = "threading")]
pub type BlockingWaitHook = fn(wait: &dyn Fn());
#[cfg(feature = "threading")]
static BLOCKING_WAIT: OnceLock<BlockingWaitHook> = OnceLock::new();
#[cfg(feature = "threading")]
pub fn set_blocking_wait_hook(hook: BlockingWaitHook) {
let _ = BLOCKING_WAIT.set(hook);
}
#[cfg(feature = "threading")]
std::thread_local! {
static IN_HOOK: Cell<bool> = const { Cell::new(false) };
}
#[cfg(feature = "threading")]
struct HookGuard;
#[cfg(feature = "threading")]
impl Drop for HookGuard {
fn drop(&mut self) {
let _ = IN_HOOK.try_with(|in_hook| in_hook.set(false));
}
}
#[cfg(all(feature = "threading", debug_assertions))]
std::thread_local! {
static WORLD_STOPPED: Cell<bool> = const { Cell::new(false) };
}
#[cfg(feature = "threading")]
#[inline]
pub fn set_world_stopped(stopped: bool) {
#[cfg(debug_assertions)]
let _ = WORLD_STOPPED.try_with(|flag| flag.set(stopped));
#[cfg(not(debug_assertions))]
let _ = stopped;
}
#[cfg(all(feature = "threading", debug_assertions))]
#[track_caller]
fn assert_not_stopping_the_world() {
let stopped = WORLD_STOPPED.try_with(Cell::get).unwrap_or(false);
assert!(
!stopped,
"a stop-the-world section took a detaching lock, which a parked thread \
may be holding and only this section can release"
);
}
#[cfg(not(all(feature = "threading", debug_assertions)))]
#[inline(always)]
fn assert_not_stopping_the_world() {}
#[cfg(feature = "threading")]
#[cold]
#[inline(never)]
fn wait_detached(wait: impl Fn()) {
let Some(hook) = BLOCKING_WAIT.get() else {
wait();
return;
};
let entered = IN_HOOK
.try_with(|in_hook| !in_hook.replace(true))
.unwrap_or(false);
if !entered {
wait();
return;
}
let _guard = HookGuard;
hook(&wait);
}
#[cfg(not(feature = "threading"))]
#[inline]
fn wait_detached(wait: impl Fn()) {
wait();
}
#[repr(transparent)]
pub struct RawDetachingRwLock(RawRwLock);
unsafe impl RawRwLockTrait for RawDetachingRwLock {
#[allow(
clippy::declare_interior_mutable_const,
reason = "raw lock initializer, as in the type it wraps"
)]
const INIT: Self = Self(<RawRwLock as RawRwLockTrait>::INIT);
type GuardMarker = <RawRwLock as RawRwLockTrait>::GuardMarker;
#[inline]
fn lock_shared(&self) {
assert_not_stopping_the_world();
if !self.0.try_lock_shared() {
wait_detached(|| self.0.lock_shared());
}
}
#[inline]
fn try_lock_shared(&self) -> bool {
self.0.try_lock_shared()
}
#[inline]
unsafe fn unlock_shared(&self) {
unsafe { self.0.unlock_shared() }
}
#[inline]
fn lock_exclusive(&self) {
assert_not_stopping_the_world();
if !self.0.try_lock_exclusive() {
wait_detached(|| self.0.lock_exclusive());
}
}
#[inline]
fn try_lock_exclusive(&self) -> bool {
self.0.try_lock_exclusive()
}
#[inline]
unsafe fn unlock_exclusive(&self) {
unsafe { self.0.unlock_exclusive() }
}
#[inline]
fn is_locked(&self) -> bool {
self.0.is_locked()
}
#[inline]
fn is_locked_exclusive(&self) -> bool {
self.0.is_locked_exclusive()
}
}
unsafe impl RawRwLockDowngrade for RawDetachingRwLock {
#[inline]
unsafe fn downgrade(&self) {
unsafe { self.0.downgrade() }
}
}
unsafe impl RawRwLockUpgradeTrait for RawDetachingRwLock {
#[inline]
fn lock_upgradable(&self) {
assert_not_stopping_the_world();
if !self.0.try_lock_upgradable() {
wait_detached(|| self.0.lock_upgradable());
}
}
#[inline]
fn try_lock_upgradable(&self) -> bool {
self.0.try_lock_upgradable()
}
#[inline]
unsafe fn unlock_upgradable(&self) {
unsafe { self.0.unlock_upgradable() }
}
#[inline]
unsafe fn upgrade(&self) {
unsafe { self.0.upgrade() }
}
#[inline]
unsafe fn try_upgrade(&self) -> bool {
unsafe { self.0.try_upgrade() }
}
}
unsafe impl RawRwLockUpgradeDowngrade for RawDetachingRwLock {
#[inline]
unsafe fn downgrade_upgradable(&self) {
unsafe { self.0.downgrade_upgradable() }
}
#[inline]
unsafe fn downgrade_to_upgradable(&self) {
unsafe { self.0.downgrade_to_upgradable() }
}
}
#[cfg(test)]
mod tests {
#[cfg(all(feature = "threading", debug_assertions))]
use super::set_world_stopped;
#[cfg(all(feature = "threading", debug_assertions))]
use crate::lock::PyDetachingRwLock;
#[cfg(all(feature = "threading", debug_assertions))]
#[test]
fn taking_one_while_stopping_the_world_is_caught() {
let lock = PyDetachingRwLock::new(());
drop(lock.write());
set_world_stopped(true);
let taken = std::panic::catch_unwind(core::panic::AssertUnwindSafe(|| {
let _guard = lock.read();
}));
set_world_stopped(false);
assert!(
taken.is_err(),
"a stop-the-world section took a detaching lock and nothing complained"
);
drop(lock.write());
}
}