Skip to main content

budget_context/
tokio_runtime.rs

1use std::future::Future;
2
3use crate::{Budget, BudgetError};
4
5impl Budget {
6    /// Requests cancellation of this budget and all existing or future children.
7    pub fn cancel(&self) {
8        self.node.cancellation.cancel();
9        #[cfg(feature = "tracing")]
10        tracing::event!(
11            tracing::Level::DEBUG,
12            budget.id = self.id().get(),
13            "budget.cancelled"
14        );
15    }
16
17    /// Returns whether cancellation has been requested for this node.
18    #[must_use]
19    pub fn is_cancelled(&self) -> bool {
20        self.node.cancellation.is_cancelled()
21    }
22
23    /// Runs a future until it completes, cancellation is requested, or the
24    /// effective deadline expires.
25    ///
26    /// If several branches are ready simultaneously, cancellation wins over
27    /// the deadline, which wins over future completion. Dropping the future is
28    /// only safe when the future itself is cancellation-safe.
29    ///
30    /// # Errors
31    ///
32    /// Returns [`BudgetError::Cancelled`] or [`BudgetError::DeadlineExceeded`]
33    /// when the corresponding terminal condition wins.
34    pub async fn run<F>(&self, future: F) -> Result<F::Output, BudgetError>
35    where
36        F: Future,
37    {
38        self.ensure_active()?;
39        tokio::pin!(future);
40
41        if let Some(deadline) = self.node.deadline {
42            tokio::select! {
43                biased;
44                () = self.node.cancellation.cancelled() => Err(BudgetError::Cancelled),
45                () = tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)) => {
46                    Err(BudgetError::DeadlineExceeded)
47                }
48                output = &mut future => Ok(output),
49            }
50        } else {
51            tokio::select! {
52                biased;
53                () = self.node.cancellation.cancelled() => Err(BudgetError::Cancelled),
54                output = &mut future => Ok(output),
55            }
56        }
57    }
58}