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