Skip to main content

nexo_driver_loop/
proactive.rs

1use std::time::{Duration, Instant};
2
3use nexo_driver_types::CancellationToken;
4use tokio::time::sleep;
5
6#[derive(Clone, Debug)]
7pub struct ScheduledWake {
8    pub duration_ms: u64,
9    pub reason: String,
10    pub sleep_started_at: Instant,
11}
12
13#[derive(Clone, Copy, Debug, PartialEq, Eq)]
14pub enum WakeResult {
15    Fired { elapsed_ms: u64 },
16    Cancelled,
17}
18
19pub fn build_tick_prompt(wake: &ScheduledWake, elapsed_ms: u64) -> String {
20    format!(
21        "<tick>\nkind: sleep_wake\nelapsed_ms: {elapsed_ms}\nreason: {}\n</tick>",
22        wake.reason
23    )
24}
25
26pub async fn wait_for_wake(wake: &ScheduledWake, cancel: &CancellationToken) -> WakeResult {
27    let duration = Duration::from_millis(wake.duration_ms);
28    tokio::select! {
29        _ = sleep(duration) => WakeResult::Fired {
30            elapsed_ms: wake.sleep_started_at.elapsed().as_millis() as u64,
31        },
32        _ = cancel.cancelled() => WakeResult::Cancelled,
33    }
34}
35
36#[cfg(test)]
37mod tests {
38    use super::*;
39
40    #[test]
41    fn sleep_wake_tick_prompt_is_stable() {
42        let wake = ScheduledWake {
43            duration_ms: 270_000,
44            reason: "waiting for new work".into(),
45            sleep_started_at: std::time::Instant::now(),
46        };
47        let prompt = build_tick_prompt(&wake, 270_000);
48        assert_eq!(
49            prompt,
50            "<tick>\nkind: sleep_wake\nelapsed_ms: 270000\nreason: waiting for new work\n</tick>"
51        );
52    }
53
54    #[tokio::test]
55    async fn wait_for_wake_stops_on_cancellation() {
56        let cancel = CancellationToken::new();
57        cancel.cancel();
58        let wake = ScheduledWake {
59            duration_ms: 60_000,
60            reason: "idle".into(),
61            sleep_started_at: std::time::Instant::now(),
62        };
63        let result = wait_for_wake(&wake, &cancel).await;
64        assert!(matches!(result, WakeResult::Cancelled));
65    }
66}