TLA  Line data    Source code

       1                 : use std::time::Duration;
       2                 : 
       3                 : use tokio_util::sync::CancellationToken;
       4                 : 
       5 GBC           1 : #[derive(thiserror::Error, Debug)]
       6                 : pub enum TimeoutCancellableError {
       7                 :     #[error("Timed out")]
       8                 :     Timeout,
       9                 :     #[error("Cancelled")]
      10                 :     Cancelled,
      11                 : }
      12                 : 
      13                 : /// Wrap [`tokio::time::timeout`] with a CancellationToken.
      14                 : ///
      15                 : /// This wrapper is appropriate for any long running operation in a task
      16                 : /// that ought to respect a CancellationToken (which means most tasks).
      17                 : ///
      18                 : /// The only time you should use a bare tokio::timeout is when the future `F`
      19                 : /// itself respects a CancellationToken: otherwise, always use this wrapper
      20                 : /// with your CancellationToken to ensure that your task does not hold up
      21                 : /// graceful shutdown.
      22 CBC       51583 : pub async fn timeout_cancellable<F>(
      23           51583 :     duration: Duration,
      24           51583 :     cancel: &CancellationToken,
      25           51583 :     future: F,
      26           51583 : ) -> Result<F::Output, TimeoutCancellableError>
      27           51583 : where
      28           51583 :     F: std::future::Future,
      29           51583 : {
      30         1219474 :     tokio::select!(
      31           51572 :         r = tokio::time::timeout(duration, future) => {
      32 UBC           0 :             r.map_err(|_| TimeoutCancellableError::Timeout)
      33                 : 
      34                 :         },
      35                 :         _ = cancel.cancelled() => {
      36                 :             Err(TimeoutCancellableError::Cancelled)
      37                 : 
      38                 :         }
      39                 :     )
      40 CBC       51576 : }