[core] Add force_unlock_read methods to lock primitives

Author: cwfitzgerald

wgpu-core/src/lock/mod.rs

@@ -52,4 +52,4 @@ use observing as chosen;
 #[cfg(not(any(wgpu_validate_locks, feature = "observe_locks")))]
 use vanilla as chosen;
 
-pub use chosen::{Mutex, MutexGuard, RwLock, RwLockReadGuard, RwLockWriteGuard};
+pub use chosen::{Mutex, MutexGuard, RankData, RwLock, RwLockReadGuard, RwLockWriteGuard};

wgpu-core/src/lock/observing.rs

@@ -38,6 +38,8 @@ use std::{
 use super::rank::{LockRank, LockRankSet};
 use crate::FastHashSet;
 
+pub type RankData = Option<HeldLock>;
+
 /// A `Mutex` instrumented for lock acquisition order observation.
 ///
 /// This is just a wrapper around a [`parking_lot::Mutex`], along with
@@ -160,6 +162,27 @@ impl<T> RwLock<T> {
             _state: LockStateGuard { saved },
         }
     }
+
+    /// Force an read-unlock operation on this lock.
+    ///
+    /// Safety:
+    /// - A read lock must be held which is not held by a guard.
+    pub unsafe fn force_unlock_read(&self, data: RankData) {
+        release(data);
+        unsafe { self.inner.force_unlock_read() };
+    }
+}
+
+impl<'a, T> RwLockReadGuard<'a, T> {
+    // Forget the read guard, leaving the lock in a locked state with no guard.
+    //
+    // Equivalent to std::mem::forget, but preserves the information about the lock
+    // rank.
+    pub fn forget(this: Self) -> RankData {
+        core::mem::forget(this.inner);
+
+        this._state.saved
+    }
 }
 
 impl<'a, T> RwLockWriteGuard<'a, T> {
@@ -316,7 +339,7 @@ enum ThreadState {
 
 /// Information about a currently held lock.
 #[derive(Debug, Copy, Clone)]
-struct HeldLock {
+pub struct HeldLock {
     /// The lock's rank.
     rank: LockRank,
 

wgpu-core/src/lock/ranked.rs

@@ -59,6 +59,8 @@ use core::{cell::Cell, fmt, ops, panic::Location};
 
 use super::rank::LockRank;
 
+pub use LockState as RankData;
+
 /// A `Mutex` instrumented for deadlock prevention.
 ///
 /// This is just a wrapper around a [`parking_lot::Mutex`], along with
@@ -87,7 +89,7 @@ std::thread_local! {
 
 /// Per-thread state for the deadlock checker.
 #[derive(Debug, Copy, Clone)]
-struct LockState {
+pub struct LockState {
     /// The last lock we acquired, and where.
     last_acquired: Option<(LockRank, &'static Location<'static>)>,
 
@@ -270,6 +272,27 @@ impl<T> RwLock<T> {
             saved: LockStateGuard(saved),
         }
     }
+
+    /// Force an read-unlock operation on this lock.
+    ///
+    /// Safety:
+    /// - A read lock must be held which is not held by a guard.
+    pub unsafe fn force_unlock_read(&self, data: RankData) {
+        release(data);
+        unsafe { self.inner.force_unlock_read() };
+    }
+}
+
+impl<'a, T> RwLockReadGuard<'a, T> {
+    // Forget the read guard, leaving the lock in a locked state with no guard.
+    //
+    // Equivalent to std::mem::forget, but preserves the information about the lock
+    // rank.
+    pub fn forget(this: Self) -> RankData {
+        core::mem::forget(this.inner);
+
+        this.saved.0
+    }
 }
 
 impl<'a, T> RwLockWriteGuard<'a, T> {

wgpu-core/src/lock/vanilla.rs

@@ -5,6 +5,10 @@
 
 use core::{fmt, ops};
 
+use crate::lock::rank::LockRank;
+
+pub struct RankData;
+
 /// A plain wrapper around [`parking_lot::Mutex`].
 ///
 /// This is just like [`parking_lot::Mutex`], except that our [`new`]
@@ -23,7 +27,7 @@ pub struct Mutex<T>(parking_lot::Mutex<T>);
 pub struct MutexGuard<'a, T>(parking_lot::MutexGuard<'a, T>);
 
 impl<T> Mutex<T> {
-    pub fn new(_rank: super::rank::LockRank, value: T) -> Mutex<T> {
+    pub fn new(_rank: LockRank, value: T) -> Mutex<T> {
         Mutex(parking_lot::Mutex::new(value))
     }
 
@@ -79,7 +83,7 @@ pub struct RwLockReadGuard<'a, T>(parking_lot::RwLockReadGuard<'a, T>);
 pub struct RwLockWriteGuard<'a, T>(parking_lot::RwLockWriteGuard<'a, T>);
 
 impl<T> RwLock<T> {
-    pub fn new(_rank: super::rank::LockRank, value: T) -> RwLock<T> {
+    pub fn new(_rank: LockRank, value: T) -> RwLock<T> {
         RwLock(parking_lot::RwLock::new(value))
     }
 
@@ -90,6 +94,26 @@ impl<T> RwLock<T> {
     pub fn write(&self) -> RwLockWriteGuard<T> {
         RwLockWriteGuard(self.0.write())
     }
+
+    /// Force an read-unlock operation on this lock.
+    ///
+    /// Safety:
+    /// - A read lock must be held which is not held by a guard.
+    pub unsafe fn force_unlock_read(&self, _data: RankData) {
+        unsafe { self.0.force_unlock_read() };
+    }
+}
+
+impl<'a, T> RwLockReadGuard<'a, T> {
+    // Forget the read guard, leaving the lock in a locked state with no guard.
+    //
+    // Equivalent to std::mem::forget, but preserves the information about the lock
+    // rank.
+    pub fn forget(this: Self) -> RankData {
+        core::mem::forget(this.0);
+
+        RankData
+    }
 }
 
 impl<'a, T> RwLockWriteGuard<'a, T> {

wgpu-core/src/snatch.rs

@@ -1,9 +1,9 @@
-use core::{cell::UnsafeCell, fmt};
+use core::{cell::UnsafeCell, fmt, mem::ManuallyDrop};
 
-use crate::lock::{rank, RwLock, RwLockReadGuard, RwLockWriteGuard};
+use crate::lock::{rank, RankData, RwLock, RwLockReadGuard, RwLockWriteGuard};
 
 /// A guard that provides read access to snatchable data.
-pub struct SnatchGuard<'a>(#[expect(dead_code)] RwLockReadGuard<'a, ()>);
+pub struct SnatchGuard<'a>(RwLockReadGuard<'a, ()>);
 /// A guard that allows snatching the snatchable data.
 pub struct ExclusiveSnatchGuard<'a>(#[expect(dead_code)] RwLockWriteGuard<'a, ()>);
 
@@ -158,6 +158,34 @@ impl SnatchLock {
         LockTrace::enter("write");
         ExclusiveSnatchGuard(self.lock.write())
     }
+
+    #[track_caller]
+    #[expect(unused)]
+    pub unsafe fn force_unlock_read(&self, data: RankData) {
+        // This is unsafe because it can cause deadlocks if the lock is held.
+        // It should only be used in very specific cases, like when a resource
+        // needs to be snatched in a panic handler.
+        LockTrace::exit();
+        unsafe { self.lock.force_unlock_read(data) };
+    }
+}
+
+impl SnatchGuard<'_> {
+    /// Forget the guard, leaving the lock in a locked state with no guard.
+    ///
+    /// This is equivalent to `std::mem::forget`, but preserves the information about the lock
+    /// rank.
+    pub fn forget(this: Self) -> RankData {
+        // Cancel the drop implementation of the current guard.
+        let manually_drop = ManuallyDrop::new(this);
+
+        // As we are unable to destructure out of this guard due to the drop implementation,
+        // so we manually read the inner value.
+        // SAFETY: This is safe because we never access the original guard again.
+        let inner_guard = unsafe { core::ptr::read(&manually_drop.0) };
+
+        RwLockReadGuard::forget(inner_guard)
+    }
 }
 
 impl Drop for SnatchGuard<'_> {