From 6bc0332f376a5685699e32e3ae2b2275de2be1d8 Mon Sep 17 00:00:00 2001
From: Matthieu Le brazidec <r3v2d0g@jesus.gg>
Date: Wed, 20 Mar 2024 15:33:20 +0100
Subject: [PATCH] Add a `rwlock()` method to owned `RwLock` guards

Adds a `rwlock()` method returning a reference to the original `Arc<RwLock>` to the `Owned*` guards.
---
 tokio/src/sync/rwlock/owned_read_guard.rs     | 30 ++++++++++++++++--
 tokio/src/sync/rwlock/owned_write_guard.rs    | 31 +++++++++++++++----
 .../sync/rwlock/owned_write_guard_mapped.rs   | 25 +++++++++++++--
 tokio/src/sync/rwlock/read_guard.rs           |  3 +-
 4 files changed, 75 insertions(+), 14 deletions(-)

diff --git a/tokio/src/sync/rwlock/owned_read_guard.rs b/tokio/src/sync/rwlock/owned_read_guard.rs
index 273e7b86f2f..67b1f8541c0 100644
--- a/tokio/src/sync/rwlock/owned_read_guard.rs
+++ b/tokio/src/sync/rwlock/owned_read_guard.rs
@@ -1,7 +1,5 @@
 use crate::sync::rwlock::RwLock;
-use std::marker::PhantomData;
-use std::sync::Arc;
-use std::{fmt, mem, ops, ptr};
+use std::{fmt, marker::PhantomData, mem, ops, ptr, sync::Arc};
 
 /// Owned RAII structure used to release the shared read access of a lock when
 /// dropped.
@@ -138,6 +136,32 @@ impl<T: ?Sized, U: ?Sized> OwnedRwLockReadGuard<T, U> {
             resource_span: this.resource_span,
         })
     }
+
+    /// Returns a reference to the original `Arc<RwLock>`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use std::sync::Arc;
+    /// use tokio::sync::{RwLock, OwnedRwLockReadGuard};
+    ///
+    /// #[derive(Debug, Clone, Copy, PartialEq, Eq)]
+    /// struct Foo(u32);
+    ///
+    /// # #[tokio::main]
+    /// # async fn main() {
+    /// let lock = Arc::new(RwLock::new(Foo(1)));
+    ///
+    /// let guard = lock.clone().read_owned().await;
+    /// assert!(Arc::ptr_eq(&lock, guard.rwlock()));
+    ///
+    /// let guard = OwnedRwLockReadGuard::map(guard, |f| &f.0);
+    /// assert!(Arc::ptr_eq(&lock, guard.rwlock()));
+    /// # }
+    /// ```
+    pub fn rwlock(&self) -> &Arc<RwLock<T>> {
+        &self.lock
+    }
 }
 
 impl<T: ?Sized, U: ?Sized> ops::Deref for OwnedRwLockReadGuard<T, U> {
diff --git a/tokio/src/sync/rwlock/owned_write_guard.rs b/tokio/src/sync/rwlock/owned_write_guard.rs
index a8ce4a1603f..e1584f50339 100644
--- a/tokio/src/sync/rwlock/owned_write_guard.rs
+++ b/tokio/src/sync/rwlock/owned_write_guard.rs
@@ -1,9 +1,8 @@
-use crate::sync::rwlock::owned_read_guard::OwnedRwLockReadGuard;
-use crate::sync::rwlock::owned_write_guard_mapped::OwnedRwLockMappedWriteGuard;
-use crate::sync::rwlock::RwLock;
-use std::marker::PhantomData;
-use std::sync::Arc;
-use std::{fmt, mem, ops, ptr};
+use crate::sync::rwlock::{
+    owned_read_guard::OwnedRwLockReadGuard, owned_write_guard_mapped::OwnedRwLockMappedWriteGuard,
+    RwLock,
+};
+use std::{fmt, marker::PhantomData, mem, ops, ptr, sync::Arc};
 
 /// Owned RAII structure used to release the exclusive write access of a lock when
 /// dropped.
@@ -390,6 +389,26 @@ impl<T: ?Sized> OwnedRwLockWriteGuard<T> {
 
         guard
     }
+
+    /// Returns a reference to the original `Arc<RwLock>`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use std::sync::Arc;
+    /// use tokio::sync::RwLock;
+    ///
+    /// # #[tokio::main]
+    /// # async fn main() {
+    /// let lock = Arc::new(RwLock::new(1));
+    ///
+    /// let guard = lock.clone().write_owned().await;
+    /// assert!(Arc::ptr_eq(&lock, guard.rwlock()));
+    /// # }
+    /// ```
+    pub fn rwlock(&self) -> &Arc<RwLock<T>> {
+        &self.lock
+    }
 }
 
 impl<T: ?Sized> ops::Deref for OwnedRwLockWriteGuard<T> {
diff --git a/tokio/src/sync/rwlock/owned_write_guard_mapped.rs b/tokio/src/sync/rwlock/owned_write_guard_mapped.rs
index 9f4952100a5..164fe8feb99 100644
--- a/tokio/src/sync/rwlock/owned_write_guard_mapped.rs
+++ b/tokio/src/sync/rwlock/owned_write_guard_mapped.rs
@@ -1,7 +1,5 @@
 use crate::sync::rwlock::RwLock;
-use std::marker::PhantomData;
-use std::sync::Arc;
-use std::{fmt, mem, ops, ptr};
+use std::{fmt, marker::PhantomData, mem, ops, ptr, sync::Arc};
 
 /// Owned RAII structure used to release the exclusive write access of a lock when
 /// dropped.
@@ -155,6 +153,27 @@ impl<T: ?Sized, U: ?Sized> OwnedRwLockMappedWriteGuard<T, U> {
             resource_span: this.resource_span,
         })
     }
+
+    /// Returns a reference to the original `Arc<RwLock>`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use std::sync::Arc;
+    /// use tokio::sync::{RwLock, OwnedRwLockMappedWriteGuard};
+    ///
+    /// # #[tokio::main]
+    /// # async fn main() {
+    /// let lock = Arc::new(RwLock::new(1));
+    ///
+    /// let guard = lock.clone().write_owned().await;
+    /// let guard = OwnedRwLockMappedWriteGuard::map(guard, |x| x);
+    /// assert!(Arc::ptr_eq(&lock, guard.rwlock()));
+    /// # }
+    /// ```
+    pub fn rwlock(&self) -> &Arc<RwLock<T>> {
+        &self.lock
+    }
 }
 
 impl<T: ?Sized, U: ?Sized> ops::Deref for OwnedRwLockMappedWriteGuard<T, U> {
diff --git a/tokio/src/sync/rwlock/read_guard.rs b/tokio/src/sync/rwlock/read_guard.rs
index a04b59588d5..ec892ce2972 100644
--- a/tokio/src/sync/rwlock/read_guard.rs
+++ b/tokio/src/sync/rwlock/read_guard.rs
@@ -1,6 +1,5 @@
 use crate::sync::batch_semaphore::Semaphore;
-use std::marker::PhantomData;
-use std::{fmt, mem, ops};
+use std::{fmt, marker::PhantomData, mem, ops};
 
 /// RAII structure used to release the shared read access of a lock when
 /// dropped.