1use crate::sync::batch_semaphore::Semaphore;
2use std::marker::PhantomData;
3use 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]
15pub 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.
27struct 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
35impl<'a, T: ?Sized> RwLockMappedWriteGuard<'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]
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]
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
168impl<T: ?Sized> ops::Deref for RwLockMappedWriteGuard<'_, T> {
169 type Target = T;
170
171 fn deref(&self) -> &T {
172 unsafe { &*self.data }
173 }
174}
175
176impl<T: ?Sized> ops::DerefMut for RwLockMappedWriteGuard<'_, T> {
177 fn deref_mut(&mut self) -> &mut T {
178 unsafe { &mut *self.data }
179 }
180}
181
182impl<'a, T: ?Sized> fmt::Debug for RwLockMappedWriteGuard<'a, T>
183where
184 T: fmt::Debug,
185{
186 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
187 fmt::Debug::fmt(&**self, f)
188 }
189}
190
191impl<'a, T: ?Sized> fmt::Display for RwLockMappedWriteGuard<'a, T>
192where
193 T: fmt::Display,
194{
195 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
196 fmt::Display::fmt(&**self, f)
197 }
198}
199
200impl<'a, T: ?Sized> Drop for RwLockMappedWriteGuard<'a, T> {
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