asupersync/time/
budget_ext.rs

1//! Budget extensions for time operations.
2
3use crate::cx::Cx;
4use crate::time::{Elapsed, Sleep, sleep_until};
5use crate::types::{Budget, Time};
6use std::future::Future;
7use std::marker::Unpin;
8use std::time::Duration;
9
10/// Extension trait for Budget deadline operations.
11pub trait BudgetTimeExt {
12    /// Get remaining time until deadline.
13    fn remaining_duration(&self, now: Time) -> Option<Duration>;
14
15    /// Create sleep that respects budget deadline.
16    fn deadline_sleep(&self) -> Option<Sleep>;
17
18    /// Check if deadline has passed.
19    fn deadline_elapsed(&self, now: Time) -> bool;
20}
21
22impl BudgetTimeExt for Budget {
23    #[inline]
24    fn remaining_duration(&self, now: Time) -> Option<Duration> {
25        self.deadline.map(|d| {
26            if now >= d {
27                Duration::ZERO
28            } else {
29                Duration::from_nanos(d.as_nanos() - now.as_nanos())
30            }
31        })
32    }
33
34    #[inline]
35    fn deadline_sleep(&self) -> Option<Sleep> {
36        self.deadline.map(sleep_until)
37    }
38
39    #[inline]
40    fn deadline_elapsed(&self, now: Time) -> bool {
41        self.deadline.is_some_and(|d| d <= now)
42    }
43}
44
45/// Sleep that integrates with the provided context's budget.
46///
47/// This sleeps for the shorter of the requested duration or the remaining budget.
48/// If the budget runs out, it returns `Err(Elapsed)`.
49pub async fn budget_sleep(cx: &Cx, duration: Duration, now: Time) -> Result<(), Elapsed> {
50    let budget = cx.budget();
51
52    // Use shorter of requested duration or remaining budget
53    // Use BudgetTimeExt::remaining_duration explicit call
54    let remaining = BudgetTimeExt::remaining_duration(&budget, now);
55
56    let effective_duration = match remaining {
57        Some(rem) if rem < duration => rem,
58        _ => duration,
59    };
60
61    if effective_duration.is_zero() && BudgetTimeExt::deadline_elapsed(&budget, now) {
62        let deadline = budget.deadline.unwrap_or(now);
63        return Err(Elapsed::new(deadline));
64    }
65
66    crate::time::sleep(now, effective_duration).await;
67
68    // Check if we were cut short by budget
69    if effective_duration < duration {
70        // We slept for 'remaining', which means deadline is hit.
71        let deadline = budget.deadline.unwrap_or(now);
72        return Err(Elapsed::new(deadline));
73    }
74
75    Ok(())
76}
77
78/// Timeout that respects budget deadline.
79pub async fn budget_timeout<F: Future + Unpin>(
80    cx: &Cx,
81    duration: Duration,
82    future: F,
83    now: Time,
84) -> Result<F::Output, Elapsed> {
85    let budget = cx.budget();
86
87    // Use shorter of requested timeout or remaining budget
88    let remaining = BudgetTimeExt::remaining_duration(&budget, now);
89    let effective_timeout = match remaining {
90        Some(rem) if rem < duration => rem,
91        _ => duration,
92    };
93
94    crate::time::timeout(now, effective_timeout, future).await
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100    use crate::cx::Cx;
101    use crate::test_utils::init_test_logging;
102    use crate::types::{Budget, RegionId, TaskId};
103    use crate::util::ArenaIndex;
104    use std::time::Duration;
105
106    fn init_test(name: &str) {
107        init_test_logging();
108        crate::test_phase!(name);
109    }
110
111    fn test_cx(budget: Budget) -> Cx {
112        Cx::new(
113            RegionId::from_arena(ArenaIndex::new(0, 0)),
114            TaskId::from_arena(ArenaIndex::new(0, 0)),
115            budget,
116        )
117    }
118
119    #[test]
120    fn test_budget_sleep() {
121        init_test("test_budget_sleep");
122        // `Sleep`'s fallback time source starts at `Time::ZERO` on first poll.
123        // Use a small deadline in the same time basis so this test remains fast.
124        let now = Time::ZERO;
125        let deadline = now.saturating_add_nanos(5_000_000); // 5ms
126        let budget = Budget::new().with_deadline(deadline);
127        let cx = test_cx(budget);
128
129        // Request longer sleep than budget allows
130        futures_lite::future::block_on(async {
131            let result = budget_sleep(&cx, Duration::from_secs(10), now).await;
132            let is_err = result.is_err();
133            crate::assert_with_log!(is_err, "budget sleep errors", true, is_err);
134        });
135        crate::test_complete!("test_budget_sleep");
136    }
137}
asupersync/time/budget_ext.rs

asupersync/time/
budget_ext.rs
