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           16 :     pub fn new(what: &'static str, data: T) -> Self {
      45           16 :         Self {
      46           16 :             what,
      47           16 :             state: State::Clean,
      48           16 :             data,
      49           16 :         }
      50           16 :     }
      51              : 
      52              :     /// Check for poisoning and return a [`Guard`] that provides access to the wrapped state.
      53           16 :     pub fn check_and_arm(&mut self) -> Result<Guard<T>, Error> {
      54           16 :         match self.state {
      55              :             State::Clean => {
      56           16 :                 self.state = State::Armed;
      57           16 :                 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           16 :     }
      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           16 :     pub fn data_mut(&mut self) -> &mut T {
      79           16 :         &mut self.0.data
      80           16 :     }
      81              : 
      82           12 :     pub fn disarm(self) {
      83           12 :         match self.0.state {
      84            0 :             State::Clean => unreachable!("we set it to Armed in check_and_arm()"),
      85           12 :             State::Armed => {
      86           12 :                 self.0.state = State::Clean;
      87           12 :             }
      88            0 :             State::Poisoned { at } => {
      89            0 :                 unreachable!("we fail check_and_arm() if it's in that state: {at}")
      90              :             }
      91              :         }
      92           12 :     }
      93              : }
      94              : 
      95              : impl<'a, T> Drop for Guard<'a, T> {
      96           16 :     fn drop(&mut self) {
      97           16 :         match self.0.state {
      98           12 :             State::Clean => {
      99           12 :                 // set by disarm()
     100           12 :             }
     101              :             State::Armed => {
     102              :                 // still armed => poison it
     103            4 :                 let at = chrono::Utc::now();
     104            4 :                 self.0.state = State::Poisoned { at };
     105            4 :                 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           16 :     }
     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              : }