microvms_app/session/
keepalive.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Keeping a VM awake from outside while work runs inside it.
3//!
4//! The platform measures idleness by inbound requests through the endpoint proxy. Work
5//! inside the guest does not count: measured 2026-09-23 in us-east-1, a CPU-busy exec was
6//! suspended between 60 and 70 seconds after the last inbound request under a 60-second
7//! `maxIdleDurationSeconds`. Requests from inside the guest arrive over loopback and never
8//! reach the meter. So the only keepalive is a caller outside the VM that keeps sending
9//! requests, and [`KeepAwake`] is that caller: an unauthenticated `GET /v1/health` every
10//! `interval`, optionally ending once the daemon stops reporting `busy`.
11//!
12//! # Why the interval is capped at half the idle window
13//!
14//! One missed poll must not suspend the VM. With the interval at most half the window, a
15//! single failed or slow poll still leaves a second one inside the window. When the caller
16//! does not know the VM's window, the cap assumes the smallest window the platform accepts
17//! ([`MIN_IDLE_DURATION_SEC`]), which is the only assumption that cannot be too generous.
18//!
19//! # The probe is the caller's
20//!
21//! [`KeepAwake::run`] owns the policy (cadence, `busy`, deadline, error tolerance, stop)
22//! and takes the poll as a closure. [`super::Session::keep_awake`] passes its own health
23//! call. A binding holding a sandbox-owned session passes one that first checks, under the
24//! sandbox's lock, that the VM is still running and answers `None` when it is not: a
25//! health poll against a VM the caller just suspended would auto-resume it, which is the
26//! opposite of what a suspend asked for.
27
28use std::future::Future;
29use std::sync::{Arc, Mutex, PoisonError};
30use std::time::Duration;
31
32use protocol::health::Health;
33use tokio::sync::{oneshot, watch};
34
35use super::Session;
36use crate::constants::MIN_IDLE_DURATION_SEC;
37use crate::error::{Error, ErrorKind};
38
39/// The default cadence when the caller names none: a third of the window, at most this.
40pub const DEFAULT_KEEP_AWAKE_INTERVAL: Duration = Duration::from_secs(20);
41
42/// Consecutive retryable poll failures tolerated before the keepalive gives up.
43pub const DEFAULT_TOLERATED_ERRORS: u32 = 3;
44
45/// How soon a failed poll is retried. Shorter than any valid interval, so one transport
46/// blip does not spend half the idle window.
47const RETRY_AFTER: Duration = Duration::from_secs(1);
48
49/// Why a keepalive ended without an error.
50#[derive(Clone, Copy, Debug, Eq, PartialEq)]
51pub enum KeepAwakeEnd {
52    /// The caller stopped it.
53    Stopped,
54    /// `while_busy` was set and the daemon reported no running exec.
55    Idle,
56    /// `max_duration` passed.
57    Elapsed,
58    /// The probe reported the VM is no longer running (suspended or terminated).
59    NotRunning,
60}
61
62impl KeepAwakeEnd {
63    /// The wire spelling, shared by the bindings and the CLI envelope.
64    pub fn as_str(self) -> &'static str {
65        match self {
66            Self::Stopped => "stopped",
67            Self::Idle => "idle",
68            Self::Elapsed => "elapsed",
69            Self::NotRunning => "not-running",
70        }
71    }
72}
73
74impl std::fmt::Display for KeepAwakeEnd {
75    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
76        f.write_str(self.as_str())
77    }
78}
79
80/// What a finished keepalive did.
81#[derive(Clone, Debug, Eq, PartialEq)]
82pub struct KeepAwakeReport {
83    pub end: KeepAwakeEnd,
84    /// Health polls that answered.
85    pub polls: u64,
86    /// `busy` from the last answered poll, or `None` when none answered.
87    pub last_busy: Option<bool>,
88    pub elapsed: Duration,
89}
90
91/// A keepalive's policy. Build with [`KeepAwake::new`], check with [`KeepAwake::validate`].
92#[derive(Clone, Debug)]
93pub struct KeepAwake {
94    interval: Duration,
95    idle_window: Duration,
96    while_busy: bool,
97    max_duration: Option<Duration>,
98    tolerated_errors: u32,
99}
100
101impl KeepAwake {
102    /// A keepalive for a VM whose `maxIdleDurationSeconds` is `idle_window`, or the
103    /// platform minimum when unknown. The interval defaults to a third of the window,
104    /// at most [`DEFAULT_KEEP_AWAKE_INTERVAL`].
105    pub fn new(idle_window: Option<Duration>) -> Self {
106        let idle_window =
107            idle_window.unwrap_or(Duration::from_secs(u64::from(MIN_IDLE_DURATION_SEC)));
108        Self {
109            interval: (idle_window / 3).min(DEFAULT_KEEP_AWAKE_INTERVAL),
110            idle_window,
111            while_busy: false,
112            max_duration: None,
113            tolerated_errors: DEFAULT_TOLERATED_ERRORS,
114        }
115    }
116
117    #[must_use]
118    pub fn interval(mut self, interval: Duration) -> Self {
119        self.interval = interval;
120        self
121    }
122
123    /// End once the daemon reports no running exec.
124    #[must_use]
125    pub fn while_busy(mut self, while_busy: bool) -> Self {
126        self.while_busy = while_busy;
127        self
128    }
129
130    /// End after this long even if still busy.
131    #[must_use]
132    pub fn max_duration(mut self, max_duration: Option<Duration>) -> Self {
133        self.max_duration = max_duration;
134        self
135    }
136
137    #[must_use]
138    pub fn tolerated_errors(mut self, tolerated_errors: u32) -> Self {
139        self.tolerated_errors = tolerated_errors;
140        self
141    }
142
143    pub fn interval_value(&self) -> Duration {
144        self.interval
145    }
146
147    pub fn idle_window_value(&self) -> Duration {
148        self.idle_window
149    }
150
151    /// Refuses a policy that could let the VM suspend. See the module docs for the cap.
152    pub fn validate(&self) -> Result<(), Error> {
153        let minimum = Duration::from_secs(u64::from(MIN_IDLE_DURATION_SEC));
154        if self.idle_window < minimum {
155            return Err(Error::invalid_arg(format!(
156                "idle window {}s is below the platform minimum of {MIN_IDLE_DURATION_SEC}s \
157                 (IdlePolicy.maxIdleDurationSeconds)",
158                self.idle_window.as_secs_f64()
159            )));
160        }
161        if self.interval < Duration::from_secs(1) {
162            return Err(Error::invalid_arg(format!(
163                "keepalive interval {}s is below 1s",
164                self.interval.as_secs_f64()
165            )));
166        }
167        if self.interval > self.idle_window / 2 {
168            return Err(Error::invalid_arg(format!(
169                "keepalive interval {}s exceeds half the {}s idle window, so one missed poll \
170                 could let the VM suspend; pass a shorter interval or the VM's actual window",
171                self.interval.as_secs_f64(),
172                self.idle_window.as_secs_f64()
173            )));
174        }
175        if self.max_duration.is_some_and(|max| max.is_zero()) {
176            return Err(Error::invalid_arg(
177                "keepalive max duration must be positive",
178            ));
179        }
180        Ok(())
181    }
182
183    /// Polls `probe` until `stop` resolves, the policy ends it, or a poll fails for good.
184    ///
185    /// The first poll is immediate. `probe` answers `Ok(None)` when the VM is no longer
186    /// running. A retryable failure is retried after one second, up to `tolerated_errors`
187    /// in a row; any other failure ends the keepalive with that error.
188    pub async fn run<P, F, S>(&self, mut probe: P, stop: S) -> Result<KeepAwakeReport, Error>
189    where
190        P: FnMut() -> F,
191        F: Future<Output = Result<Option<Health>, Error>>,
192        S: Future<Output = ()>,
193    {
194        self.validate()?;
195        let started = tokio::time::Instant::now();
196        tokio::pin!(stop);
197        let mut polls = 0u64;
198        let mut errors = 0u32;
199        let mut last_busy = None;
200        let report = |end, polls, last_busy| KeepAwakeReport {
201            end,
202            polls,
203            last_busy,
204            elapsed: started.elapsed(),
205        };
206        loop {
207            let answer = tokio::select! {
208                biased;
209                () = &mut stop => return Ok(report(KeepAwakeEnd::Stopped, polls, last_busy)),
210                answer = probe() => answer,
211            };
212            let mut wait = self.interval;
213            match answer {
214                Ok(Some(health)) => {
215                    polls += 1;
216                    errors = 0;
217                    last_busy = Some(health.busy);
218                    if self.while_busy && !health.busy {
219                        return Ok(report(KeepAwakeEnd::Idle, polls, last_busy));
220                    }
221                }
222                Ok(None) => return Ok(report(KeepAwakeEnd::NotRunning, polls, last_busy)),
223                Err(error) if error.retryable() && errors < self.tolerated_errors => {
224                    errors += 1;
225                    wait = RETRY_AFTER;
226                }
227                Err(error) => return Err(error),
228            }
229            if let Some(max) = self.max_duration {
230                let elapsed = started.elapsed();
231                if elapsed >= max {
232                    return Ok(report(KeepAwakeEnd::Elapsed, polls, last_busy));
233                }
234                wait = wait.min(max - elapsed);
235            }
236            tokio::select! {
237                biased;
238                () = &mut stop => return Ok(report(KeepAwakeEnd::Stopped, polls, last_busy)),
239                () = tokio::time::sleep(wait) => {}
240            }
241            if self
242                .max_duration
243                .is_some_and(|max| started.elapsed() >= max)
244            {
245                return Ok(report(KeepAwakeEnd::Elapsed, polls, last_busy));
246            }
247        }
248    }
249}
250
251/// Whether the VM behind a keepalive is still one it may poll. See the module docs.
252pub type RunningGate = Box<dyn Fn() -> bool + Send + Sync>;
253
254type Outcome = Result<KeepAwakeReport, (ErrorKind, String)>;
255
256/// A keepalive running as a background task. Dropping it stops the task.
257pub struct KeepAwakeTask {
258    stop: Mutex<Option<oneshot::Sender<()>>>,
259    outcome: Arc<Mutex<Option<Outcome>>>,
260    done: watch::Receiver<bool>,
261}
262
263impl KeepAwake {
264    /// Validates the policy and starts polling `session` on the current tokio runtime.
265    ///
266    /// `running`, when given, is checked before every poll; once it answers `false` the
267    /// task ends with [`KeepAwakeEnd::NotRunning`] without polling again. Panics outside a
268    /// tokio runtime, like `tokio::spawn`.
269    pub fn spawn(
270        self,
271        session: Session,
272        running: Option<RunningGate>,
273    ) -> Result<KeepAwakeTask, Error> {
274        self.validate()?;
275        let (stop_tx, stop_rx) = oneshot::channel::<()>();
276        let (done_tx, done_rx) = watch::channel(false);
277        let outcome = Arc::new(Mutex::new(None));
278        let written = Arc::clone(&outcome);
279        tokio::spawn(async move {
280            let probe = || {
281                let session = running
282                    .as_ref()
283                    .is_none_or(|running| running())
284                    .then(|| session.clone());
285                async move {
286                    match session {
287                        Some(session) => session.health().await.map(Some),
288                        None => Ok(None),
289                    }
290                }
291            };
292            // A dropped sender resolves the receiver too: dropping the task stops it.
293            let stop = async move {
294                let _ = stop_rx.await;
295            };
296            let result = self
297                .run(probe, stop)
298                .await
299                .map_err(|error| (error.kind(), error.to_string()));
300            *written.lock().unwrap_or_else(PoisonError::into_inner) = Some(result);
301            done_tx.send_replace(true);
302        });
303        Ok(KeepAwakeTask {
304            stop: Mutex::new(Some(stop_tx)),
305            outcome,
306            done: done_rx,
307        })
308    }
309}
310
311impl KeepAwakeTask {
312    /// Whether the task is still polling.
313    pub fn is_running(&self) -> bool {
314        !*self.done.borrow()
315    }
316
317    /// Asks the task to stop. Idempotent; [`Self::finished`] waits for it.
318    pub fn request_stop(&self) {
319        if let Some(stop) = self
320            .stop
321            .lock()
322            .unwrap_or_else(PoisonError::into_inner)
323            .take()
324        {
325            let _ = stop.send(());
326        }
327    }
328
329    /// Waits for the task to end and returns its report, or the error that ended it.
330    ///
331    /// Callable any number of times, from any number of waiters.
332    pub async fn finished(&self) -> Result<KeepAwakeReport, Error> {
333        let mut done = self.done.clone();
334        // The sender only drops after writing the outcome, so an error here still means done.
335        let _ = done.wait_for(|finished| *finished).await;
336        self.outcome()
337            .expect("the task records its outcome before signalling done")
338    }
339
340    /// The report, once the task has ended.
341    pub fn outcome(&self) -> Option<Result<KeepAwakeReport, Error>> {
342        self.outcome
343            .lock()
344            .unwrap_or_else(PoisonError::into_inner)
345            .clone()
346            .map(|outcome| outcome.map_err(|(kind, message)| Error::new(kind, message)))
347    }
348}
349
350#[cfg(test)]
351mod tests {
352    use std::cell::RefCell;
353    use std::collections::VecDeque;
354
355    use super::*;
356    use crate::error::{ErrorKind, WireKind};
357
358    fn health(busy: bool) -> Health {
359        serde_json::from_value(serde_json::json!({
360            "version": "0.1.0",
361            "bootstrapped": true,
362            "disk": null,
363            "identity_degraded": false,
364            "identity_repaired": true,
365            "busy": busy,
366        }))
367        .expect("health parses")
368    }
369
370    type Answer = Result<Option<Health>, Error>;
371
372    /// A probe answering from a script, recording when each poll happened.
373    struct Script {
374        answers: RefCell<VecDeque<Answer>>,
375        at: RefCell<Vec<Duration>>,
376        start: tokio::time::Instant,
377    }
378
379    impl Script {
380        fn new(answers: impl IntoIterator<Item = Answer>) -> Self {
381            Self {
382                answers: RefCell::new(answers.into_iter().collect()),
383                at: RefCell::new(Vec::new()),
384                start: tokio::time::Instant::now(),
385            }
386        }
387
388        fn probe(&self) -> impl FnMut() -> std::future::Ready<Answer> + '_ {
389            move || {
390                self.at.borrow_mut().push(self.start.elapsed());
391                std::future::ready(
392                    self.answers
393                        .borrow_mut()
394                        .pop_front()
395                        .unwrap_or(Ok(Some(health(true)))),
396                )
397            }
398        }
399    }
400
401    fn secs(n: u64) -> Duration {
402        Duration::from_secs(n)
403    }
404
405    #[tokio::test(start_paused = true)]
406    async fn it_polls_at_the_interval_until_stopped() {
407        let script = Script::new([]);
408        let stop = tokio::time::sleep(secs(45));
409        let report = KeepAwake::new(None)
410            .interval(secs(10))
411            .run(script.probe(), stop)
412            .await
413            .expect("runs");
414        assert_eq!(report.end, KeepAwakeEnd::Stopped);
415        assert_eq!(*script.at.borrow(), [0, 10, 20, 30, 40].map(secs));
416        assert_eq!((report.polls, report.last_busy), (5, Some(true)));
417    }
418
419    #[tokio::test(start_paused = true)]
420    async fn while_busy_ends_on_the_first_idle_answer() {
421        let script = Script::new([
422            Ok(Some(health(true))),
423            Ok(Some(health(true))),
424            Ok(Some(health(false))),
425        ]);
426        let report = KeepAwake::new(None)
427            .while_busy(true)
428            .run(script.probe(), std::future::pending())
429            .await
430            .expect("runs");
431        assert_eq!(report.end, KeepAwakeEnd::Idle);
432        assert_eq!((report.polls, report.last_busy), (3, Some(false)));
433        assert_eq!(report.elapsed, secs(40));
434    }
435
436    #[tokio::test(start_paused = true)]
437    async fn an_idle_answer_without_while_busy_keeps_polling() {
438        let script = Script::new([Ok(Some(health(false)))]);
439        let report = KeepAwake::new(None)
440            .max_duration(Some(secs(50)))
441            .run(script.probe(), std::future::pending())
442            .await
443            .expect("runs");
444        assert_eq!(report.end, KeepAwakeEnd::Elapsed);
445        assert_eq!(*script.at.borrow(), [0, 20, 40].map(secs));
446        assert_eq!(report.elapsed, secs(50));
447    }
448
449    #[tokio::test(start_paused = true)]
450    async fn a_vm_that_stopped_running_ends_the_keepalive_without_polling_it() {
451        let script = Script::new([Ok(Some(health(true))), Ok(None)]);
452        let report = KeepAwake::new(None)
453            .run(script.probe(), std::future::pending())
454            .await
455            .expect("runs");
456        assert_eq!((report.end, report.polls), (KeepAwakeEnd::NotRunning, 1));
457    }
458
459    #[tokio::test(start_paused = true)]
460    async fn retryable_failures_are_retried_quickly_up_to_the_tolerance() {
461        let cut = || Err(Error::wire(WireKind::Transport, "cut"));
462        let script = Script::new([cut(), cut(), Ok(Some(health(true))), cut(), cut(), cut()]);
463        let report = KeepAwake::new(None)
464            .max_duration(Some(secs(30)))
465            .run(script.probe(), std::future::pending())
466            .await
467            .expect("three in a row is tolerated, and a success resets the count");
468        assert_eq!(report.end, KeepAwakeEnd::Elapsed);
469        assert_eq!(script.at.borrow()[..4], [0, 1, 2, 22].map(secs));
470
471        let script = Script::new([cut(), cut(), cut(), cut()]);
472        let error = KeepAwake::new(None)
473            .run(script.probe(), std::future::pending())
474            .await
475            .expect_err("a fourth consecutive failure ends it");
476        assert!(error.retryable());
477        assert_eq!(script.at.borrow().len(), 4);
478    }
479
480    #[tokio::test(start_paused = true)]
481    async fn a_non_retryable_failure_ends_it_at_once() {
482        let script = Script::new([Err(Error::invalid_arg("refused"))]);
483        let error = KeepAwake::new(None)
484            .run(script.probe(), std::future::pending())
485            .await
486            .expect_err("fails");
487        assert_eq!(error.kind(), ErrorKind::InvalidArg);
488        assert_eq!(script.at.borrow().len(), 1);
489    }
490
491    #[test]
492    fn the_interval_is_capped_at_half_the_window() {
493        assert_eq!(KeepAwake::new(None).interval_value(), secs(20));
494        assert_eq!(KeepAwake::new(Some(secs(600))).interval_value(), secs(20));
495        assert_eq!(KeepAwake::new(Some(secs(90))).interval_value(), secs(20));
496        assert!(KeepAwake::new(None).interval(secs(30)).validate().is_ok());
497        assert!(KeepAwake::new(None).interval(secs(31)).validate().is_err());
498        assert!(
499            KeepAwake::new(Some(secs(600)))
500                .interval(secs(300))
501                .validate()
502                .is_ok()
503        );
504        assert!(
505            KeepAwake::new(None)
506                .interval(Duration::from_millis(500))
507                .validate()
508                .is_err()
509        );
510        assert!(KeepAwake::new(Some(secs(59))).validate().is_err());
511        assert!(
512            KeepAwake::new(None)
513                .max_duration(Some(Duration::ZERO))
514                .validate()
515                .is_err()
516        );
517    }
518
519    fn busy_session(
520        replies: usize,
521    ) -> (Session, std::sync::Arc<crate::session::testing::Recorder>) {
522        use crate::session::testing::{Recorder, Reply, health_body, session_with};
523        let mut body = health_body(true);
524        body["busy"] = serde_json::json!(true);
525        let recorder = Recorder::with((0..replies).map(|_| Reply::ok(body.clone())));
526        let (session, _, _) = session_with(std::sync::Arc::clone(&recorder));
527        (session, recorder)
528    }
529
530    #[tokio::test(start_paused = true)]
531    async fn a_spawned_task_stops_on_request_and_reports() {
532        let (session, recorder) = busy_session(10);
533        let task = KeepAwake::new(None)
534            .interval(secs(10))
535            .spawn(session, None)
536            .expect("spawns");
537        tokio::time::sleep(secs(25)).await;
538        assert!(task.is_running());
539        task.request_stop();
540        let report = task.finished().await.expect("stops cleanly");
541        assert_eq!((report.end, report.polls), (KeepAwakeEnd::Stopped, 3));
542        assert!(!task.is_running());
543        assert_eq!(recorder.requests().len(), 3);
544        assert_eq!(task.finished().await.expect("again").polls, 3);
545    }
546
547    #[tokio::test(start_paused = true)]
548    async fn dropping_the_task_stops_polling() {
549        let (session, recorder) = busy_session(10);
550        let task = KeepAwake::new(None)
551            .interval(secs(10))
552            .spawn(session, None)
553            .expect("spawns");
554        tokio::time::sleep(secs(15)).await;
555        drop(task);
556        tokio::time::sleep(secs(60)).await;
557        assert_eq!(recorder.requests().len(), 2);
558    }
559
560    #[tokio::test(start_paused = true)]
561    async fn a_closed_gate_ends_the_task_without_another_poll() {
562        let (session, recorder) = busy_session(10);
563        let open = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(true));
564        let gate = std::sync::Arc::clone(&open);
565        let task = KeepAwake::new(None)
566            .interval(secs(10))
567            .spawn(
568                session,
569                Some(Box::new(move || {
570                    gate.load(std::sync::atomic::Ordering::SeqCst)
571                })),
572            )
573            .expect("spawns");
574        tokio::time::sleep(secs(15)).await;
575        open.store(false, std::sync::atomic::Ordering::SeqCst);
576        let report = task.finished().await.expect("ends");
577        assert_eq!((report.end, report.polls), (KeepAwakeEnd::NotRunning, 2));
578        assert_eq!(
579            recorder.requests().len(),
580            2,
581            "no poll after the gate closed"
582        );
583    }
584
585    #[tokio::test]
586    async fn an_invalid_policy_is_refused_before_spawning() {
587        let (session, recorder) = busy_session(0);
588        let error = KeepAwake::new(None)
589            .interval(secs(31))
590            .spawn(session, None)
591            .err()
592            .expect("refused");
593        assert_eq!(error.kind(), ErrorKind::InvalidArg);
594        assert!(recorder.requests().is_empty());
595    }
596}
microvms_app/session/keepalive.rs

microvms_app/session/
keepalive.rs
