1 use crate::sync::rwlock::RwLock;
2 use std::fmt;
3 use std::marker::PhantomData;
4 use std::mem::{self, ManuallyDrop};
5 use std::ops;
6 use std::sync::Arc;
7 
8 /// Owned RAII structure used to release the exclusive write access of a lock when
9 /// dropped.
10 ///
11 /// This structure is created by [mapping] an [`OwnedRwLockWriteGuard`]. It is a
12 /// separate type from `OwnedRwLockWriteGuard` to disallow downgrading a mapped
13 /// guard, since doing so can cause undefined behavior.
14 ///
15 /// [mapping]: method@crate::sync::OwnedRwLockWriteGuard::map
16 /// [`OwnedRwLockWriteGuard`]: struct@crate::sync::OwnedRwLockWriteGuard
17 pub struct OwnedRwLockMappedWriteGuard<T: ?Sized, U: ?Sized = T> {
18     #[cfg(all(tokio_unstable, feature = "tracing"))]
19     pub(super) resource_span: tracing::Span,
20     pub(super) permits_acquired: u32,
21     // ManuallyDrop allows us to destructure into this field without running the destructor.
22     pub(super) lock: ManuallyDrop<Arc<RwLock<T>>>,
23     pub(super) data: *mut U,
24     pub(super) _p: PhantomData<T>,
25 }
26 
27 impl<T: ?Sized, U: ?Sized> OwnedRwLockMappedWriteGuard<T, U> {
28     /// Makes a new `OwnedRwLockMappedWriteGuard` for a component of the locked
29     /// data.
30     ///
31     /// This operation cannot fail as the `OwnedRwLockMappedWriteGuard` passed
32     /// in already locked the data.
33     ///
34     /// This is an associated function that needs to be used as
35     /// `OwnedRwLockWriteGuard::map(..)`. A method would interfere with methods
36     /// of the same name on the contents of the locked data.
37     ///
38     /// # Examples
39     ///
40     /// ```
41     /// use std::sync::Arc;
42     /// use tokio::sync::{RwLock, OwnedRwLockWriteGuard};
43     ///
44     /// #[derive(Debug, Clone, Copy, PartialEq, Eq)]
45     /// struct Foo(u32);
46     ///
47     /// # #[tokio::main]
48     /// # async fn main() {
49     /// let lock = Arc::new(RwLock::new(Foo(1)));
50     ///
51     /// {
52     ///     let lock = Arc::clone(&lock);
53     ///     let mut mapped = OwnedRwLockWriteGuard::map(lock.write_owned().await, |f| &mut f.0);
54     ///     *mapped = 2;
55     /// }
56     ///
57     /// assert_eq!(Foo(2), *lock.read().await);
58     /// # }
59     /// ```
60     #[inline]
map<F, V: ?Sized>(mut this: Self, f: F) -> OwnedRwLockMappedWriteGuard<T, V> where F: FnOnce(&mut U) -> &mut V,61     pub fn map<F, V: ?Sized>(mut this: Self, f: F) -> OwnedRwLockMappedWriteGuard<T, V>
62     where
63         F: FnOnce(&mut U) -> &mut V,
64     {
65         let data = f(&mut *this) as *mut V;
66         let lock = unsafe { ManuallyDrop::take(&mut this.lock) };
67         let permits_acquired = this.permits_acquired;
68         #[cfg(all(tokio_unstable, feature = "tracing"))]
69         let resource_span = this.resource_span.clone();
70         // NB: Forget to avoid drop impl from being called.
71         mem::forget(this);
72 
73         OwnedRwLockMappedWriteGuard {
74             permits_acquired,
75             lock: ManuallyDrop::new(lock),
76             data,
77             _p: PhantomData,
78             #[cfg(all(tokio_unstable, feature = "tracing"))]
79             resource_span,
80         }
81     }
82 
83     /// Attempts to make a new `OwnedRwLockMappedWriteGuard` for a component
84     /// of the locked data. The original guard is returned if the closure
85     /// returns `None`.
86     ///
87     /// This operation cannot fail as the `OwnedRwLockMappedWriteGuard` passed
88     /// in already locked the data.
89     ///
90     /// This is an associated function that needs to be
91     /// used as `OwnedRwLockMappedWriteGuard::try_map(...)`. A method would interfere with
92     /// methods of the same name on the contents of the locked data.
93     ///
94     /// # Examples
95     ///
96     /// ```
97     /// use std::sync::Arc;
98     /// use tokio::sync::{RwLock, OwnedRwLockWriteGuard};
99     ///
100     /// #[derive(Debug, Clone, Copy, PartialEq, Eq)]
101     /// struct Foo(u32);
102     ///
103     /// # #[tokio::main]
104     /// # async fn main() {
105     /// let lock = Arc::new(RwLock::new(Foo(1)));
106     ///
107     /// {
108     ///     let guard = Arc::clone(&lock).write_owned().await;
109     ///     let mut guard = OwnedRwLockWriteGuard::try_map(guard, |f| Some(&mut f.0)).expect("should not fail");
110     ///     *guard = 2;
111     /// }
112     ///
113     /// assert_eq!(Foo(2), *lock.read().await);
114     /// # }
115     /// ```
116     #[inline]
try_map<F, V: ?Sized>( mut this: Self, f: F, ) -> Result<OwnedRwLockMappedWriteGuard<T, V>, Self> where F: FnOnce(&mut U) -> Option<&mut V>,117     pub fn try_map<F, V: ?Sized>(
118         mut this: Self,
119         f: F,
120     ) -> Result<OwnedRwLockMappedWriteGuard<T, V>, Self>
121     where
122         F: FnOnce(&mut U) -> Option<&mut V>,
123     {
124         let data = match f(&mut *this) {
125             Some(data) => data as *mut V,
126             None => return Err(this),
127         };
128         let lock = unsafe { ManuallyDrop::take(&mut this.lock) };
129         let permits_acquired = this.permits_acquired;
130         #[cfg(all(tokio_unstable, feature = "tracing"))]
131         let resource_span = this.resource_span.clone();
132         // NB: Forget to avoid drop impl from being called.
133         mem::forget(this);
134 
135         Ok(OwnedRwLockMappedWriteGuard {
136             permits_acquired,
137             lock: ManuallyDrop::new(lock),
138             data,
139             _p: PhantomData,
140             #[cfg(all(tokio_unstable, feature = "tracing"))]
141             resource_span,
142         })
143     }
144 }
145 
146 impl<T: ?Sized, U: ?Sized> ops::Deref for OwnedRwLockMappedWriteGuard<T, U> {
147     type Target = U;
148 
deref(&self) -> &U149     fn deref(&self) -> &U {
150         unsafe { &*self.data }
151     }
152 }
153 
154 impl<T: ?Sized, U: ?Sized> ops::DerefMut for OwnedRwLockMappedWriteGuard<T, U> {
deref_mut(&mut self) -> &mut U155     fn deref_mut(&mut self) -> &mut U {
156         unsafe { &mut *self.data }
157     }
158 }
159 
160 impl<T: ?Sized, U: ?Sized> fmt::Debug for OwnedRwLockMappedWriteGuard<T, U>
161 where
162     U: fmt::Debug,
163 {
fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result164     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
165         fmt::Debug::fmt(&**self, f)
166     }
167 }
168 
169 impl<T: ?Sized, U: ?Sized> fmt::Display for OwnedRwLockMappedWriteGuard<T, U>
170 where
171     U: fmt::Display,
172 {
fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result173     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
174         fmt::Display::fmt(&**self, f)
175     }
176 }
177 
178 impl<T: ?Sized, U: ?Sized> Drop for OwnedRwLockMappedWriteGuard<T, U> {
drop(&mut self)179     fn drop(&mut self) {
180         self.lock.s.release(self.permits_acquired as usize);
181         #[cfg(all(tokio_unstable, feature = "tracing"))]
182         self.resource_span.in_scope(|| {
183             tracing::trace!(
184             target: "runtime::resource::state_update",
185             write_locked = false,
186             write_locked.op = "override",
187             )
188         });
189         unsafe { ManuallyDrop::drop(&mut self.lock) };
190     }
191 }
192