Line data Source code
1 : //! Protect a piece of state from reuse after it is left in an inconsistent state.
2 : //!
3 : //! # Example
4 : //!
5 : //! ```
6 : //! # tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap().block_on(async {
7 : //! use utils::poison::Poison;
8 : //! use std::time::Duration;
9 : //!
10 : //! struct State {
11 : //! clean: bool,
12 : //! }
13 : //! let state = tokio::sync::Mutex::new(Poison::new("mystate", State { clean: true }));
14 : //!
15 : //! let mut mutex_guard = state.lock().await;
16 : //! let mut poison_guard = mutex_guard.check_and_arm()?;
17 : //! let state = poison_guard.data_mut();
18 : //! state.clean = false;
19 : //! // If we get cancelled at this await point, subsequent check_and_arm() calls will fail.
20 : //! tokio::time::sleep(Duration::from_secs(10)).await;
21 : //! state.clean = true;
22 : //! poison_guard.disarm();
23 : //! # Ok::<(), utils::poison::Error>(())
24 : //! # });
25 : //! ```
26 :
27 : use tracing::warn;
28 :
29 : pub struct Poison<T> {
30 : what: &'static str,
31 : state: State,
32 : data: T,
33 : }
34 :
35 : #[derive(Clone, Copy)]
36 : enum State {
37 : Clean,
38 : Armed,
39 : Poisoned { at: chrono::DateTime<chrono::Utc> },
40 : }
41 :
42 : impl<T> Poison<T> {
43 : /// We log `what` `warning!` level if the [`Guard`] gets dropped without being [`Guard::disarm`]ed.
44 0 : pub fn new(what: &'static str, data: T) -> Self {
45 0 : Self {
46 0 : what,
47 0 : state: State::Clean,
48 0 : data,
49 0 : }
50 0 : }
51 :
52 : /// Check for poisoning and return a [`Guard`] that provides access to the wrapped state.
53 0 : pub fn check_and_arm(&mut self) -> Result<Guard<T>, Error> {
54 0 : match self.state {
55 : State::Clean => {
56 0 : self.state = State::Armed;
57 0 : Ok(Guard(self))
58 : }
59 0 : State::Armed => unreachable!("transient state"),
60 0 : State::Poisoned { at } => Err(Error::Poisoned {
61 0 : what: self.what,
62 0 : at,
63 0 : }),
64 : }
65 0 : }
66 : }
67 :
68 : /// Use [`Self::data`] and [`Self::data_mut`] to access the wrapped state.
69 : /// Once modifications are done, use [`Self::disarm`].
70 : /// If [`Guard`] gets dropped instead of calling [`Self::disarm`], the state is poisoned
71 : /// and subsequent calls to [`Poison::check_and_arm`] will fail with an error.
72 : pub struct Guard<'a, T>(&'a mut Poison<T>);
73 :
74 : impl<'a, T> Guard<'a, T> {
75 0 : pub fn data(&self) -> &T {
76 0 : &self.0.data
77 0 : }
78 0 : pub fn data_mut(&mut self) -> &mut T {
79 0 : &mut self.0.data
80 0 : }
81 :
82 0 : pub fn disarm(self) {
83 0 : match self.0.state {
84 0 : State::Clean => unreachable!("we set it to Armed in check_and_arm()"),
85 0 : State::Armed => {
86 0 : self.0.state = State::Clean;
87 0 : }
88 0 : State::Poisoned { at } => {
89 0 : unreachable!("we fail check_and_arm() if it's in that state: {at}")
90 : }
91 : }
92 0 : }
93 : }
94 :
95 : impl<'a, T> Drop for Guard<'a, T> {
96 0 : fn drop(&mut self) {
97 0 : match self.0.state {
98 0 : State::Clean => {
99 0 : // set by disarm()
100 0 : }
101 : State::Armed => {
102 : // still armed => poison it
103 0 : let at = chrono::Utc::now();
104 0 : self.0.state = State::Poisoned { at };
105 0 : warn!(at=?at, "poisoning {}", self.0.what);
106 : }
107 0 : State::Poisoned { at } => {
108 0 : unreachable!("we fail check_and_arm() if it's in that state: {at}")
109 : }
110 : }
111 0 : }
112 : }
113 :
114 0 : #[derive(thiserror::Error, Debug)]
115 : pub enum Error {
116 : #[error("poisoned at {at}: {what}")]
117 : Poisoned {
118 : what: &'static str,
119 : at: chrono::DateTime<chrono::Utc>,
120 : },
121 : }
|