1 use crate::sync::batch_semaphore::Semaphore;
2 use std::marker::PhantomData;
3 use std::{fmt, mem, ops};
4 
5 /// RAII structure used to release the exclusive write access of a lock when
6 /// dropped.
7 ///
8 /// This structure is created by [mapping] an [`RwLockWriteGuard`]. It is a
9 /// separate type from `RwLockWriteGuard` to disallow downgrading a mapped
10 /// guard, since doing so can cause undefined behavior.
11 ///
12 /// [mapping]: method@crate::sync::RwLockWriteGuard::map
13 /// [`RwLockWriteGuard`]: struct@crate::sync::RwLockWriteGuard
14 #[clippy::has_significant_drop]
15 pub struct RwLockMappedWriteGuard<'a, T: ?Sized> {
16     // When changing the fields in this struct, make sure to update the
17     // `skip_drop` method.
18     #[cfg(all(tokio_unstable, feature = "tracing"))]
19     pub(super) resource_span: tracing::Span,
20     pub(super) permits_acquired: u32,
21     pub(super) s: &'a Semaphore,
22     pub(super) data: *mut T,
23     pub(super) marker: PhantomData<&'a mut T>,
24 }
25 
26 #[allow(dead_code)] // Unused fields are still used in Drop.
27 struct Inner<'a, T: ?Sized> {
28     #[cfg(all(tokio_unstable, feature = "tracing"))]
29     resource_span: tracing::Span,
30     permits_acquired: u32,
31     s: &'a Semaphore,
32     data: *mut T,
33 }
34 
35 impl<'a, T: ?Sized> RwLockMappedWriteGuard<'a, T> {
skip_drop(self) -> Inner<'a, T>36     fn skip_drop(self) -> Inner<'a, T> {
37         let me = mem::ManuallyDrop::new(self);
38         // SAFETY: This duplicates the values in every field of the guard, then
39         // forgets the originals, so in the end no value is duplicated.
40         Inner {
41             #[cfg(all(tokio_unstable, feature = "tracing"))]
42             resource_span: unsafe { std::ptr::read(&me.resource_span) },
43             permits_acquired: me.permits_acquired,
44             s: me.s,
45             data: me.data,
46         }
47     }
48 
49     /// Makes a new `RwLockMappedWriteGuard` for a component of the locked data.
50     ///
51     /// This operation cannot fail as the `RwLockMappedWriteGuard` passed in already
52     /// locked the data.
53     ///
54     /// This is an associated function that needs to be used as
55     /// `RwLockMappedWriteGuard::map(..)`. A method would interfere with methods
56     /// of the same name on the contents of the locked data.
57     ///
58     /// This is an asynchronous version of [`RwLockWriteGuard::map`] from the
59     /// [`parking_lot` crate].
60     ///
61     /// [`RwLockWriteGuard::map`]: https://docs.rs/lock_api/latest/lock_api/struct.RwLockWriteGuard.html#method.map
62     /// [`parking_lot` crate]: https://crates.io/crates/parking_lot
63     ///
64     /// # Examples
65     ///
66     /// ```
67     /// use tokio::sync::{RwLock, RwLockWriteGuard};
68     ///
69     /// #[derive(Debug, Clone, Copy, PartialEq, Eq)]
70     /// struct Foo(u32);
71     ///
72     /// # #[tokio::main]
73     /// # async fn main() {
74     /// let lock = RwLock::new(Foo(1));
75     ///
76     /// {
77     ///     let mut mapped = RwLockWriteGuard::map(lock.write().await, |f| &mut f.0);
78     ///     *mapped = 2;
79     /// }
80     ///
81     /// assert_eq!(Foo(2), *lock.read().await);
82     /// # }
83     /// ```
84     #[inline]
map<F, U: ?Sized>(mut this: Self, f: F) -> RwLockMappedWriteGuard<'a, U> where F: FnOnce(&mut T) -> &mut U,85     pub fn map<F, U: ?Sized>(mut this: Self, f: F) -> RwLockMappedWriteGuard<'a, U>
86     where
87         F: FnOnce(&mut T) -> &mut U,
88     {
89         let data = f(&mut *this) as *mut U;
90         let this = this.skip_drop();
91 
92         RwLockMappedWriteGuard {
93             permits_acquired: this.permits_acquired,
94             s: this.s,
95             data,
96             marker: PhantomData,
97             #[cfg(all(tokio_unstable, feature = "tracing"))]
98             resource_span: this.resource_span,
99         }
100     }
101 
102     /// Attempts to make a new [`RwLockMappedWriteGuard`] for a component of
103     /// the locked data. The original guard is returned if the closure returns
104     /// `None`.
105     ///
106     /// This operation cannot fail as the `RwLockMappedWriteGuard` passed in already
107     /// locked the data.
108     ///
109     /// This is an associated function that needs to be
110     /// used as `RwLockMappedWriteGuard::try_map(...)`. A method would interfere
111     /// with methods of the same name on the contents of the locked data.
112     ///
113     /// This is an asynchronous version of [`RwLockWriteGuard::try_map`] from
114     /// the [`parking_lot` crate].
115     ///
116     /// [`RwLockWriteGuard::try_map`]: https://docs.rs/lock_api/latest/lock_api/struct.RwLockWriteGuard.html#method.try_map
117     /// [`parking_lot` crate]: https://crates.io/crates/parking_lot
118     ///
119     /// # Examples
120     ///
121     /// ```
122     /// use tokio::sync::{RwLock, RwLockWriteGuard};
123     ///
124     /// #[derive(Debug, Clone, Copy, PartialEq, Eq)]
125     /// struct Foo(u32);
126     ///
127     /// # #[tokio::main]
128     /// # async fn main() {
129     /// let lock = RwLock::new(Foo(1));
130     ///
131     /// {
132     ///     let guard = lock.write().await;
133     ///     let mut guard = RwLockWriteGuard::try_map(guard, |f| Some(&mut f.0)).expect("should not fail");
134     ///     *guard = 2;
135     /// }
136     ///
137     /// assert_eq!(Foo(2), *lock.read().await);
138     /// # }
139     /// ```
140     #[inline]
try_map<F, U: ?Sized>( mut this: Self, f: F, ) -> Result<RwLockMappedWriteGuard<'a, U>, Self> where F: FnOnce(&mut T) -> Option<&mut U>,141     pub fn try_map<F, U: ?Sized>(
142         mut this: Self,
143         f: F,
144     ) -> Result<RwLockMappedWriteGuard<'a, U>, Self>
145     where
146         F: FnOnce(&mut T) -> Option<&mut U>,
147     {
148         let data = match f(&mut *this) {
149             Some(data) => data as *mut U,
150             None => return Err(this),
151         };
152         let this = this.skip_drop();
153 
154         Ok(RwLockMappedWriteGuard {
155             permits_acquired: this.permits_acquired,
156             s: this.s,
157             data,
158             marker: PhantomData,
159             #[cfg(all(tokio_unstable, feature = "tracing"))]
160             resource_span: this.resource_span,
161         })
162     }
163 
164     // Note: No `downgrade`, `downgrade_map` nor `try_downgrade_map` because they would be unsound, as we're already
165     //       potentially been mapped with internal mutability.
166 }
167 
168 impl<T: ?Sized> ops::Deref for RwLockMappedWriteGuard<'_, T> {
169     type Target = T;
170 
deref(&self) -> &T171     fn deref(&self) -> &T {
172         unsafe { &*self.data }
173     }
174 }
175 
176 impl<T: ?Sized> ops::DerefMut for RwLockMappedWriteGuard<'_, T> {
deref_mut(&mut self) -> &mut T177     fn deref_mut(&mut self) -> &mut T {
178         unsafe { &mut *self.data }
179     }
180 }
181 
182 impl<'a, T: ?Sized> fmt::Debug for RwLockMappedWriteGuard<'a, T>
183 where
184     T: fmt::Debug,
185 {
fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result186     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
187         fmt::Debug::fmt(&**self, f)
188     }
189 }
190 
191 impl<'a, T: ?Sized> fmt::Display for RwLockMappedWriteGuard<'a, T>
192 where
193     T: fmt::Display,
194 {
fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result195     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
196         fmt::Display::fmt(&**self, f)
197     }
198 }
199 
200 impl<'a, T: ?Sized> Drop for RwLockMappedWriteGuard<'a, T> {
drop(&mut self)201     fn drop(&mut self) {
202         self.s.release(self.permits_acquired as usize);
203 
204         #[cfg(all(tokio_unstable, feature = "tracing"))]
205         self.resource_span.in_scope(|| {
206             tracing::trace!(
207             target: "runtime::resource::state_update",
208             write_locked = false,
209             write_locked.op = "override",
210             )
211         });
212     }
213 }
214