use catenary_proc::ProcessDelta;
use std::future::Future;
use std::time::Duration;
use tokio::sync::Notify;
const POLL_INTERVAL: Duration = Duration::from_millis(200);
const SAFETY_CAP: Duration = Duration::from_secs(300);
pub async fn load_aware_grace<S, F, Fut>(
sample_fn: &mut S,
threshold: u64,
max_wall: Option<Duration>,
notify: &Notify,
progress_active: impl Fn() -> bool,
condition: F,
) -> bool
where
S: FnMut() -> Option<ProcessDelta>,
F: Fn() -> Fut,
Fut: Future<Output = bool>,
{
let wall_deadline = tokio::time::Instant::now() + max_wall.unwrap_or(SAFETY_CAP);
let mut remaining_threshold = i64::try_from(threshold).unwrap_or(i64::MAX);
loop {
if condition().await {
return true;
}
tokio::select! {
() = notify.notified() => {
if condition().await {
return true;
}
}
() = tokio::time::sleep(POLL_INTERVAL) => {
}
}
let Some(d) = sample_fn() else {
return false;
};
match d.state {
catenary_proc::ProcessState::Dead => return false,
catenary_proc::ProcessState::Blocked => {
}
catenary_proc::ProcessState::Running | catenary_proc::ProcessState::Sleeping => {
let delta = d.delta_utime + d.delta_stime;
if d.state == catenary_proc::ProcessState::Running
&& delta > 0
&& !progress_active()
{
remaining_threshold -= i64::try_from(delta).unwrap_or(remaining_threshold);
}
}
}
if remaining_threshold <= 0 {
return false;
}
if tokio::time::Instant::now() >= wall_deadline {
return false;
}
}
}