Line data    Source code

       1              : use std::time::Duration;
       2              : 
       3              : use tokio_util::sync::CancellationToken;
       4              : 
       5            0 : #[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            0 : pub async fn timeout_cancellable<F>(
      23            0 :     duration: Duration,
      24            0 :     cancel: &CancellationToken,
      25            0 :     future: F,
      26            0 : ) -> Result<F::Output, TimeoutCancellableError>
      27            0 : where
      28            0 :     F: std::future::Future,
      29            0 : {
      30            0 :     tokio::select!(
      31            0 :         r = tokio::time::timeout(duration, future) => {
      32            0 :             r.map_err(|_| TimeoutCancellableError::Timeout)
      33              : 
      34              :         },
      35            0 :         _ = cancel.cancelled() => {
      36            0 :             Err(TimeoutCancellableError::Cancelled)
      37              : 
      38              :         }
      39              :     )
      40            0 : }