screeps_async/
time.rs

1//! Utilities for tracking time
2
3use crate::utils::game_time;
4use crate::with_runtime;
5use std::future::Future;
6use std::pin::Pin;
7use std::task::{Context, Poll};
8
9/// Future returned by [delay_ticks]
10pub struct Delay {
11    when: u32,
12    timer_index: usize,
13}
14
15impl Delay {
16    fn new(when: u32) -> Self {
17        with_runtime(|runtime| {
18            let mut timer_map = runtime.timers.try_lock().unwrap();
19            let wakers = timer_map.entry(when).or_default();
20
21            let timer_index = wakers.len();
22            wakers.push(None); // Store an empty waker to ensure len() is incremented for the next delay
23            Delay { when, timer_index }
24        })
25    }
26}
27
28impl Future for Delay {
29    type Output = ();
30
31    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
32        if game_time() >= self.when {
33            return Poll::Ready(());
34        }
35
36        with_runtime(|runtime| {
37            let mut timers = runtime.timers.try_lock().unwrap();
38
39            // SAFETY: timers map gets populated before future is created and removed on final wake
40            let wakers = timers.get_mut(&self.when).unwrap();
41
42            if let Some(waker) = wakers.get_mut(self.timer_index).and_then(Option::as_mut) {
43                // Waker already registered, check if it needs updating
44                waker.clone_from(cx.waker());
45            } else {
46                // First time this future was polled, save the waker
47                wakers[self.timer_index] = Some(cx.waker().clone())
48            }
49        });
50
51        Poll::Pending
52    }
53}
54
55/// Sleeps for `dur` game ticks.
56///
57/// If `dur` is zero, this function completes immediately and does not yield to the scheduler.
58/// If you wish to yield execution back to the scheduler, please use [yield_now] instead
59pub fn delay_ticks(dur: u32) -> Delay {
60    let when = game_time() + dur;
61    Delay::new(when)
62}
63
64/// Sleep until [screeps::game::time()] >= `when`
65///
66/// The Future returned by this function completes immediately if [screeps::game::time()] is already
67/// >= `when` and does not yield to the scheduler.
68pub fn delay_until(when: u32) -> Delay {
69    Delay::new(when)
70}
71
72/// Delay execution until the next tick
73pub async fn yield_tick() {
74    delay_ticks(1).await
75}
76
77/// Yields execution back to the async runtime, but doesn't necessarily wait until next tick
78/// to continue execution.
79///
80/// Long-running tasks that perform a significant amount of synchronous work between `.await`s
81/// can prevent other tasks from being executed. In the worst case, too much synchronous work in a row
82/// can consume all remaining CPU time this tick since the scheduler cannot interrupt work in the middle
83/// of synchronous sections of code. To alleviate this problem, [yield_now] should be called periodically
84/// to yield control back to the scheduler and give other tasks a chance to run.
85pub async fn yield_now() {
86    struct YieldNow {
87        yielded: bool,
88    }
89
90    impl Future for YieldNow {
91        type Output = ();
92
93        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
94            if self.yielded {
95                Poll::Ready(())
96            } else {
97                self.yielded = true;
98                cx.waker().wake_by_ref();
99                Poll::Pending
100            }
101        }
102    }
103
104    YieldNow { yielded: false }.await;
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110    use crate::spawn;
111    use crate::tests::game_time;
112    use rstest::rstest;
113    use std::cell::{OnceCell, RefCell};
114    use std::rc::Rc;
115
116    #[rstest]
117    #[case(0, 0)]
118    #[case(1, 1)]
119    #[case(4, 4)]
120    fn test_delay_ticks(#[case] dur: u32, #[case] expected: u32) {
121        crate::tests::init_test();
122
123        let has_run = Rc::new(OnceCell::new());
124        {
125            let has_run = has_run.clone();
126
127            spawn(async move {
128                assert_eq!(0, game_time());
129                delay_ticks(dur).await;
130                assert_eq!(expected, game_time());
131
132                has_run.set(()).unwrap();
133            })
134            .detach();
135        }
136
137        // task hasn't run yet
138        assert!(has_run.get().is_none());
139
140        // Should complete within `dur` ticks (since we have infinite cpu time in this test)
141        while game_time() <= dur {
142            crate::tests::tick().unwrap()
143        }
144
145        // Future has been run
146        assert!(has_run.get().is_some(), "Future failed to complete");
147    }
148
149    #[test]
150    fn test_yield_now() {
151        crate::tests::init_test();
152
153        let steps = Rc::new(RefCell::new(Vec::new()));
154        {
155            let steps = steps.clone();
156            spawn(async move {
157                {
158                    steps.borrow_mut().push(1);
159                }
160                yield_now().await;
161                {
162                    // should run after second spawn if yield_now works correctly
163                    steps.borrow_mut().push(3);
164                }
165            })
166            .detach();
167        }
168        {
169            let steps = steps.clone();
170            spawn(async move {
171                steps.borrow_mut().push(2);
172            })
173            .detach();
174        }
175
176        crate::run().unwrap();
177
178        let steps = steps.take();
179
180        assert_eq!(vec![1, 2, 3], steps);
181    }
182}
screeps_async/time.rs

screeps_async/
time.rs
