Line data    Source code

       1              : use std::panic::resume_unwind;
       2              : use std::sync::{Mutex, MutexGuard};
       3              : 
       4              : use tokio::task::JoinError;
       5              : 
       6              : pub(crate) trait LockExt<T> {
       7              :     fn lock_propagate_poison(&self) -> MutexGuard<'_, T>;
       8              : }
       9              : 
      10              : impl<T> LockExt<T> for Mutex<T> {
      11              :     /// Lock the mutex and panic if the mutex was poisoned.
      12              :     #[track_caller]
      13          488 :     fn lock_propagate_poison(&self) -> MutexGuard<'_, T> {
      14          488 :         match self.lock() {
      15          488 :             Ok(guard) => guard,
      16              :             // poison occurs when another thread panicked while holding the lock guard.
      17              :             // since panicking is often unrecoverable, propagating the poison panic is reasonable.
      18            0 :             Err(poison) => panic!("{poison}"),
      19              :         }
      20          488 :     }
      21              : }
      22              : 
      23              : pub(crate) trait TaskExt<T> {
      24              :     fn propagate_task_panic(self) -> T;
      25              : }
      26              : 
      27              : impl<T> TaskExt<T> for Result<T, JoinError> {
      28              :     /// Unwrap the result and panic if the inner task panicked.
      29              :     /// Also panics if the task was cancelled
      30              :     #[track_caller]
      31          107 :     fn propagate_task_panic(self) -> T {
      32            0 :         match self {
      33          107 :             Ok(t) => t,
      34              :             // Using resume_unwind prevents the panic hook being called twice.
      35              :             // Since we use this for structured concurrency, there is only
      36              :             // 1 logical panic, so this is more correct.
      37            0 :             Err(e) if e.is_panic() => resume_unwind(e.into_panic()),
      38            0 :             Err(e) => panic!("unexpected task error: {e}"),
      39              :         }
      40          107 :     }
      41              : }