use catenary_proc::ProcessState;
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<(u64, ProcessState)>,
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((delta, state)) = sample_fn() else {
return false;
};
match state {
ProcessState::Dead => return false,
ProcessState::Blocked => {
}
ProcessState::Running | ProcessState::Sleeping => {
if state == 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;
}
}
}