use std::time::{Duration, Instant};
use crate::config::HumanDuration;
use crate::log::StageLogger;
pub const DEFAULT_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(30);
pub const HEARTBEAT_INTERVAL_ENV: &str = "ANODIZER_HEARTBEAT_INTERVAL_MS";
pub fn heartbeat_interval() -> Option<Duration> {
match std::env::var(HEARTBEAT_INTERVAL_ENV).ok().as_deref() {
Some(raw) => match raw.trim().parse::<u64>() {
Ok(0) => None,
Ok(ms) => Some(Duration::from_millis(ms)),
Err(_) => Some(DEFAULT_HEARTBEAT_INTERVAL),
},
None => Some(DEFAULT_HEARTBEAT_INTERVAL),
}
}
pub fn format_elapsed(d: Duration) -> String {
HumanDuration(d).as_humantime_string()
}
pub fn heartbeat_message(action: &str, start: Instant) -> String {
format!("still {action} ({})", format_elapsed(start.elapsed()))
}
pub fn heartbeat_period(log: &StageLogger) -> Option<Duration> {
log.heartbeats_enabled().then(heartbeat_interval).flatten()
}
pub fn run_ticker(
stop_rx: &std::sync::mpsc::Receiver<()>,
interval: Duration,
mut on_tick: impl FnMut(),
) {
while let Err(std::sync::mpsc::RecvTimeoutError::Timeout) = stop_rx.recv_timeout(interval) {
on_tick();
}
}
pub async fn with_heartbeat<F, T>(log: &StageLogger, label: &str, fut: F) -> T
where
F: std::future::Future<Output = T>,
{
tokio::pin!(fut);
let Some(period) = heartbeat_period(log) else {
return fut.await;
};
let start = Instant::now();
let mut ticker = tokio::time::interval(period);
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
ticker.tick().await;
loop {
tokio::select! {
out = &mut fut => return out,
_ = ticker.tick() => {
log.heartbeat(&heartbeat_message(label, start));
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_helpers::env::{EnvGuard, env_mutex};
#[test]
fn interval_env_zero_disables_and_malformed_degrades() {
let _lock = env_mutex().lock().unwrap_or_else(|e| e.into_inner());
let _g = EnvGuard::set(HEARTBEAT_INTERVAL_ENV, "0");
assert_eq!(heartbeat_interval(), None, "0 disables heartbeats");
let _g = EnvGuard::set(HEARTBEAT_INTERVAL_ENV, "75");
assert_eq!(heartbeat_interval(), Some(Duration::from_millis(75)));
let _g = EnvGuard::set(HEARTBEAT_INTERVAL_ENV, "not-a-number");
assert_eq!(
heartbeat_interval(),
Some(DEFAULT_HEARTBEAT_INTERVAL),
"malformed cadence degrades to the default"
);
}
#[test]
fn run_ticker_ticks_until_sender_drops() {
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};
let (stop_tx, stop_rx) = std::sync::mpsc::channel::<()>();
let ticks = Arc::new(AtomicU32::new(0));
let observed = Arc::clone(&ticks);
let handle = std::thread::spawn(move || {
let deadline = std::time::Instant::now() + Duration::from_secs(30);
while observed.load(Ordering::SeqCst) < 2 && std::time::Instant::now() < deadline {
std::thread::sleep(Duration::from_millis(5));
}
drop(stop_tx);
});
let counter = Arc::clone(&ticks);
run_ticker(&stop_rx, Duration::from_millis(10), || {
counter.fetch_add(1, Ordering::SeqCst);
});
handle.join().expect("dropper thread");
let ticks = ticks.load(Ordering::SeqCst);
assert!(
ticks >= 2,
"the ticker must keep ticking until the sender drops; got {ticks}"
);
}
#[allow(clippy::await_holding_lock)]
#[tokio::test(flavor = "current_thread")]
async fn slow_future_emits_heartbeats_fast_future_does_not() {
let _lock = env_mutex().lock().unwrap_or_else(|e| e.into_inner());
let _g = EnvGuard::set(HEARTBEAT_INTERVAL_ENV, "30");
let (fast_log, fast_cap) = StageLogger::with_capture("t", crate::log::Verbosity::Normal);
with_heartbeat(&fast_log, "uploading", async {
tokio::time::sleep(Duration::from_millis(5)).await;
})
.await;
assert_eq!(
fast_cap.heartbeat_count(),
0,
"sub-cadence future must not heartbeat"
);
let (slow_log, slow_cap) = StageLogger::with_capture("t", crate::log::Verbosity::Normal);
with_heartbeat(&slow_log, "uploading", async {
tokio::time::sleep(Duration::from_millis(200)).await;
})
.await;
assert!(
slow_cap.heartbeat_count() >= 2,
"a multi-cadence future must heartbeat repeatedly; got {}",
slow_cap.heartbeat_count()
);
}
#[allow(clippy::await_holding_lock)]
#[tokio::test(flavor = "current_thread")]
async fn verbose_suppresses_heartbeats() {
let _lock = env_mutex().lock().unwrap_or_else(|e| e.into_inner());
let _g = EnvGuard::set(HEARTBEAT_INTERVAL_ENV, "30");
let (log, cap) = StageLogger::with_capture("t", crate::log::Verbosity::Verbose);
with_heartbeat(&log, "uploading", async {
tokio::time::sleep(Duration::from_millis(150)).await;
})
.await;
assert_eq!(
cap.heartbeat_count(),
0,
"verbose tee is the progress signal; no heartbeat"
);
}
}