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
160impl<T: ?Sized, U: ?Sized> ops::Deref for OwnedRwLockMappedWriteGuard<T, U> {
161 type Target = U;
162
163 fn deref(&self) -> &U {
164 unsafe { &*self.data }
165 }
166}
167
168impl<T: ?Sized, U: ?Sized> ops::DerefMut for OwnedRwLockMappedWriteGuard<T, U> {
169 fn deref_mut(&mut self) -> &mut U {
170 unsafe { &mut *self.data }
171 }
172}
173
174impl<T: ?Sized, U: ?Sized> fmt::Debug for OwnedRwLockMappedWriteGuard<T, U>
175where
176 U: fmt::Debug,
177{
178 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
179 fmt::Debug::fmt(&**self, f)
180 }
181}
182
183impl<T: ?Sized, U: ?Sized> fmt::Display for OwnedRwLockMappedWriteGuard<T, U>
184where
185 U: fmt::Display,
186{
187 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
188 fmt::Display::fmt(&**self, f)
189 }
190}
191
192impl<T: ?Sized, U: ?Sized> Drop for OwnedRwLockMappedWriteGuard<T, U> {
193 fn drop(&mut self) {
194 self.lock.s.release(self.permits_acquired as usize);
195
196 #[cfg(all(tokio_unstable, feature = "tracing"))]
197 self.resource_span.in_scope(|| {
198 tracing::trace!(
199 target: "runtime::resource::state_update",
200 write_locked = false,
201 write_locked.op = "override",
202 )
203 });
204 }
205}
206