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           20 :     pub fn new(what: &'static str, data: T) -> Self {
      45           20 :         Self {
      46           20 :             what,
      47           20 :             state: State::Clean,
      48           20 :             data,
      49           20 :         }
      50           20 :     }
      51              : 
      52              :     /// Check for poisoning and return a [`Guard`] that provides access to the wrapped state.
      53           20 :     pub fn check_and_arm(&mut self) -> Result<Guard<T>, Error> {
      54           20 :         match self.state {
      55              :             State::Clean => {
      56           20 :                 self.state = State::Armed;
      57           20 :                 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           20 :     }
      66              : }
      67              : 
      68              : /// Armed pointer to a [`Poison`].
      69              : ///
      70              : /// Use [`Self::data`] and [`Self::data_mut`] to access the wrapped state.
      71              : /// Once modifications are done, use [`Self::disarm`].
      72              : /// If [`Guard`] gets dropped instead of calling [`Self::disarm`], the state is poisoned
      73              : /// and subsequent calls to [`Poison::check_and_arm`] will fail with an error.
      74              : pub struct Guard<'a, T>(&'a mut Poison<T>);
      75              : 
      76              : impl<T> Guard<'_, T> {
      77            0 :     pub fn data(&self) -> &T {
      78            0 :         &self.0.data
      79            0 :     }
      80           20 :     pub fn data_mut(&mut self) -> &mut T {
      81           20 :         &mut self.0.data
      82           20 :     }
      83              : 
      84           16 :     pub fn disarm(self) {
      85           16 :         match self.0.state {
      86            0 :             State::Clean => unreachable!("we set it to Armed in check_and_arm()"),
      87           16 :             State::Armed => {
      88           16 :                 self.0.state = State::Clean;
      89           16 :             }
      90            0 :             State::Poisoned { at } => {
      91            0 :                 unreachable!("we fail check_and_arm() if it's in that state: {at}")
      92              :             }
      93              :         }
      94           16 :     }
      95              : }
      96              : 
      97              : impl<T> Drop for Guard<'_, T> {
      98           20 :     fn drop(&mut self) {
      99           20 :         match self.0.state {
     100           16 :             State::Clean => {
     101           16 :                 // set by disarm()
     102           16 :             }
     103              :             State::Armed => {
     104              :                 // still armed => poison it
     105            4 :                 let at = chrono::Utc::now();
     106            4 :                 self.0.state = State::Poisoned { at };
     107            4 :                 warn!(at=?at, "poisoning {}", self.0.what);
     108              :             }
     109            0 :             State::Poisoned { at } => {
     110            0 :                 unreachable!("we fail check_and_arm() if it's in that state: {at}")
     111              :             }
     112              :         }
     113           20 :     }
     114              : }
     115              : 
     116            0 : #[derive(thiserror::Error, Debug)]
     117              : pub enum Error {
     118              :     #[error("poisoned at {at}: {what}")]
     119              :     Poisoned {
     120              :         what: &'static str,
     121              :         at: chrono::DateTime<chrono::Utc>,
     122              :     },
     123              : }