budget_context/
tokio_runtime.rs1use std::future::Future;
2
3use crate::{Budget, BudgetError};
4
5impl Budget {
6 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 #[must_use]
19 pub fn is_cancelled(&self) -> bool {
20 self.node.cancellation.is_cancelled()
21 }
22
23 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}