cancel_safe_futures/sync/
poison.rs

1// This file is adapted from
2// https://github.com/rust-lang/rust/blob/475c71da0710fd1d40c046f9cee04b733b5b2b51/library/std/src/sync/poison.rs
3// and is used under the MIT and Apache 2.0 licenses.
4
5use core::sync::atomic::{AtomicU8, Ordering};
6use std::{error, fmt, thread};
7
8/// A type of error which can be returned whenever a [`RobustMutex`] is acquired.
9///
10/// [`RobustMutex`]es are poisoned whenever a thread panics or a cancellation happens while the lock
11/// is held. The precise semantics for when a lock is poisoned is documented on [`RobustMutex`], but
12/// once a lock is poisoned then all future acquisitions will return this error.
13///
14/// [`RobustMutex`]: crate::sync::RobustMutex
15pub struct PoisonError<T> {
16    guard: T,
17    flags: u8,
18}
19
20impl<T> PoisonError<T> {
21    fn new(guard: T, flags: u8) -> PoisonError<T> {
22        PoisonError { guard, flags }
23    }
24
25    /// Returns true if this error indicates that the lock was poisoned by a
26    /// panic from another task.
27    #[inline]
28    pub fn is_panic(&self) -> bool {
29        self.flags & PANIC_POISON != 0
30    }
31
32    /// Returns true if this error indicates that the lock was poisoned by an early cancellation.
33    #[inline]
34    pub fn is_cancel(&self) -> bool {
35        self.flags & CANCEL_POISON != 0
36    }
37
38    /// Consumes this error indicating that a lock is poisoned, returning the
39    /// underlying guard to allow access regardless.
40    #[inline]
41    pub fn into_inner(self) -> T {
42        self.guard
43    }
44
45    /// Reaches into this error indicating that a lock is poisoned, returning a
46    /// reference to the underlying guard to allow access regardless.
47    #[inline]
48    pub fn get_ref(&self) -> &T {
49        &self.guard
50    }
51
52    /// Reaches into this error indicating that a lock is poisoned, returning a
53    /// mutable reference to the underlying guard to allow access regardless.
54    #[inline]
55    pub fn get_mut(&mut self) -> &mut T {
56        &mut self.guard
57    }
58}
59
60impl<T> fmt::Debug for PoisonError<T> {
61    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
62        f.debug_struct("PoisonError")
63            .field("flags", &DebugFlags(self.flags))
64            .finish_non_exhaustive()
65    }
66}
67
68impl<T> fmt::Display for PoisonError<T> {
69    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
70        write!(f, "poisoned lock {:?}", DebugFlags(self.flags))
71    }
72}
73
74impl<T> error::Error for PoisonError<T> {}
75
76/// An enumeration of possible errors associated with a [`TryLockResult`] which
77/// can occur while trying to acquire a lock, from the [`try_lock`] method on a
78/// [`RobustMutex`].
79///
80/// [`try_lock`]: crate::sync::RobustMutex::try_lock
81/// [`RobustMutex`]: crate::sync::RobustMutex
82pub enum TryLockError<T> {
83    /// The lock could not be acquired because another task failed while holding
84    /// the lock, or an early cancellation occurred.
85    Poisoned(PoisonError<T>),
86
87    /// The lock could not be acquired at this time because the operation would
88    /// otherwise block.
89    WouldBlock,
90}
91
92impl<T> From<PoisonError<T>> for TryLockError<T> {
93    fn from(error: PoisonError<T>) -> TryLockError<T> {
94        TryLockError::Poisoned(error)
95    }
96}
97
98impl<T> fmt::Debug for TryLockError<T> {
99    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
100        match self {
101            TryLockError::Poisoned(error) => f.debug_tuple("Poisoned").field(&error).finish(),
102            TryLockError::WouldBlock => f.debug_tuple("WouldBlock").finish(),
103        }
104    }
105}
106
107impl<T> fmt::Display for TryLockError<T> {
108    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
109        match self {
110            TryLockError::Poisoned(error) => {
111                write!(f, "poisoned lock {:?}", DebugFlags(error.flags))
112            }
113            TryLockError::WouldBlock => {
114                f.write_str("try_lock failed because the operation would block")
115            }
116        }
117    }
118}
119
120impl<T> error::Error for TryLockError<T> {
121    fn cause(&self) -> Option<&dyn error::Error> {
122        match self {
123            TryLockError::Poisoned(error) => Some(error),
124            TryLockError::WouldBlock => None,
125        }
126    }
127}
128
129/// A type alias for the result of a lock method which can be poisoned.
130///
131/// The [`Ok`] variant of this result indicates that the primitive was not
132/// poisoned, and the `Guard` is contained within. The [`Err`] variant indicates
133/// that the primitive was poisoned. Note that the [`Err`] variant *also* carries
134/// the associated guard, and it can be acquired through the [`into_inner`]
135/// method.
136///
137/// [`into_inner`]: PoisonError::into_inner
138pub type LockResult<Guard> = Result<Guard, PoisonError<Guard>>;
139
140/// A type alias for the result of a nonblocking locking method.
141///
142/// For more information, see [`LockResult`]. A `TryLockResult` doesn't
143/// necessarily hold the associated guard in the [`Err`] type as the lock might not
144/// have been acquired for other reasons.
145pub type TryLockResult<Guard> = Result<Guard, TryLockError<Guard>>;
146
147pub const NO_POISON: u8 = 0;
148pub const PANIC_POISON: u8 = 1 << 0;
149pub const CANCEL_POISON: u8 = 1 << 1;
150
151pub struct Flag {
152    failed: AtomicU8,
153}
154
155impl fmt::Debug for Flag {
156    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
157        let flags = self.failed.load(Ordering::Relaxed);
158        DebugFlags(flags).fmt(f)
159    }
160}
161
162// Note that the Ordering uses to access the `failed` field of `Flag` below is always `Relaxed`, and
163// that's because this isn't actually protecting any data, it's just a flag whether we've panicked
164// or not.
165//
166// The actual location that this matters is when a mutex is **locked** which is where we have
167// external synchronization ensuring that we see memory reads/writes to this flag.
168//
169// As a result, if it matters, we should see the correct value for `failed` in all cases.
170
171impl Flag {
172    #[inline]
173    pub const fn new() -> Flag {
174        Flag {
175            failed: AtomicU8::new(NO_POISON),
176        }
177    }
178
179    /// Check the flag for an unguarded borrow, where we only care about existing poison.
180    #[inline]
181    pub fn borrow(&self) -> LockResult<()> {
182        let flags = self.get_flags();
183        if flags > NO_POISON {
184            Err(PoisonError::new((), flags))
185        } else {
186            Ok(())
187        }
188    }
189
190    #[inline]
191    pub fn guard_assuming_no_poison(&self) -> Guard {
192        Guard {
193            panicking: thread::panicking(),
194        }
195    }
196
197    #[inline]
198    pub fn done(&self, guard: &Guard, cancel_poison: bool) {
199        let mut new_flags = 0;
200        if !guard.panicking && thread::panicking() {
201            new_flags |= PANIC_POISON;
202        }
203        if cancel_poison {
204            new_flags |= CANCEL_POISON;
205        }
206        self.failed.fetch_or(new_flags, Ordering::Relaxed);
207    }
208
209    #[inline]
210    pub fn get_flags(&self) -> u8 {
211        self.failed.load(Ordering::Relaxed)
212    }
213}
214
215pub struct Guard {
216    panicking: bool,
217}
218
219pub fn map_result<T, U, F>(result: LockResult<T>, f: F) -> LockResult<U>
220where
221    F: FnOnce(T) -> U,
222{
223    match result {
224        Ok(t) => Ok(f(t)),
225        Err(error) => {
226            let flags = error.flags;
227            Err(PoisonError::new(f(error.into_inner()), flags))
228        }
229    }
230}
231
232struct DebugFlags(u8);
233
234impl fmt::Debug for DebugFlags {
235    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
236        if self.0 == NO_POISON {
237            write!(f, "(not poisoned)")
238        } else {
239            let mut poisoned_by = Vec::new();
240            if self.0 & PANIC_POISON != 0 {
241                poisoned_by.push("panic");
242            }
243            if self.0 & CANCEL_POISON != 0 {
244                poisoned_by.push("async cancellation");
245            }
246
247            write!(f, "(poisoned by {})", poisoned_by.join(", "))
248        }
249    }
250}