88 lines
2.5 KiB
Rust
88 lines
2.5 KiB
Rust
//! Command-wide execution budgets exclude the union of component compilation waits, including child spawns.
|
|
use std::{sync::Arc, time::Duration};
|
|
use tokio::{sync::watch, time::Instant};
|
|
|
|
#[derive(Clone, Copy)]
|
|
struct State {
|
|
deadline: Instant,
|
|
compiling: usize,
|
|
paused_at: Instant,
|
|
}
|
|
|
|
pub(super) struct Deadline(watch::Sender<State>);
|
|
|
|
impl Deadline {
|
|
pub fn new(duration: Duration) -> Arc<Self> {
|
|
let now = Instant::now();
|
|
Arc::new(Self(
|
|
watch::channel(State {
|
|
deadline: now + duration,
|
|
compiling: 0,
|
|
paused_at: now,
|
|
})
|
|
.0,
|
|
))
|
|
}
|
|
|
|
pub fn pause(self: &Arc<Self>) -> Compilation {
|
|
self.0.send_modify(|state| {
|
|
if state.compiling == 0 {
|
|
state.paused_at = Instant::now();
|
|
}
|
|
state.compiling += 1;
|
|
});
|
|
Compilation(self.clone())
|
|
}
|
|
|
|
pub async fn expired(&self) {
|
|
let mut state = self.0.subscribe();
|
|
loop {
|
|
let current = *state.borrow_and_update();
|
|
if current.compiling == 0 {
|
|
tokio::select! {
|
|
biased;
|
|
changed = state.changed() => { changed.unwrap(); }
|
|
_ = tokio::time::sleep_until(current.deadline) => return,
|
|
}
|
|
} else {
|
|
state.changed().await.unwrap();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(super) struct Compilation(Arc<Deadline>);
|
|
|
|
impl Drop for Compilation {
|
|
fn drop(&mut self) {
|
|
self.0.0.send_modify(|state| {
|
|
state.compiling -= 1;
|
|
if state.compiling == 0 {
|
|
state.deadline += state.paused_at.elapsed();
|
|
}
|
|
});
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn overlapping_compilations_share_one_pause_and_drop_resumes() {
|
|
let deadline = Deadline::new(Duration::from_secs(2));
|
|
let first = deadline.pause();
|
|
let paused_at = deadline.0.borrow().paused_at;
|
|
let second = deadline.pause();
|
|
assert_eq!(deadline.0.borrow().compiling, 2);
|
|
assert_eq!(deadline.0.borrow().paused_at, paused_at);
|
|
let original = deadline.0.borrow().deadline;
|
|
drop(first);
|
|
assert_eq!(deadline.0.borrow().compiling, 1);
|
|
assert_eq!(deadline.0.borrow().deadline, original);
|
|
drop(second);
|
|
assert_eq!(deadline.0.borrow().compiling, 0);
|
|
assert!(deadline.0.borrow().deadline >= original);
|
|
}
|
|
}
|