use std::sync::{Arc, Weak};
use std::time::{Duration, Instant};
use parking_lot::RwLock;
use tokio::task::JoinHandle;
use tokio::time::timeout;
use tokio_util::sync::CancellationToken;
use boxlite_shared::errors::BoxliteError;
use super::box_impl::BoxImpl;
use super::state::{BoxState, HealthStatus};
use crate::portal::interfaces::GuestInterface;
use crate::runtime::rt_impl::RuntimeImpl;
use crate::{BoxID, HealthCheckOptions, HealthState};
pub(crate) struct HealthProbe {
guest: GuestInterface,
interval: Duration,
check_timeout: Duration,
retries: u32,
start_period: Duration,
started_at: Instant,
last: HealthStatus,
}
enum Probe {
Skipped,
Healthy,
Failed,
}
impl HealthProbe {
pub(crate) fn new(
guest: GuestInterface,
config: HealthCheckOptions,
last: HealthStatus,
) -> Self {
Self {
guest,
interval: config.interval,
check_timeout: config.timeout,
retries: config.retries,
start_period: config.start_period,
started_at: Instant::now(),
last,
}
}
async fn check(&mut self) -> Probe {
if self.started_at.elapsed() < self.start_period {
return Probe::Skipped;
}
match timeout(self.check_timeout, self.guest.ping()).await {
Ok(Ok(_)) => Probe::Healthy,
Ok(Err(_)) | Err(_) => Probe::Failed,
}
}
}
pub(crate) struct BoxWatcher {
shim_pid: u32,
state: Arc<RwLock<BoxState>>,
runtime: Weak<RuntimeImpl>,
shutdown: CancellationToken,
box_id: BoxID,
box_name: Option<String>,
exit_file: std::path::PathBuf,
removes_on_exit: bool,
health: Option<HealthProbe>,
}
impl BoxWatcher {
pub(crate) fn new(bx: &BoxImpl, shim_pid: u32, health: Option<HealthProbe>) -> Self {
Self {
shim_pid,
state: Arc::clone(&bx.state),
runtime: Arc::downgrade(&bx.runtime),
shutdown: bx.shutdown_token.child_token(),
box_id: bx.config.id.clone(),
box_name: bx.config.name.clone(),
exit_file: bx
.layout
.container_exit_file(bx.config.container.id.as_str()),
removes_on_exit: bx.config.options.removes_on_stop(),
health,
}
}
pub(crate) fn spawn(self) -> JoinHandle<()> {
tokio::spawn(self.run())
}
async fn run(mut self) {
let shim = crate::util::ProcessMonitor::new(self.shim_pid);
let shutdown = self.shutdown.clone();
let mut interval = self.health.as_ref().map(|h| h.interval);
loop {
tokio::select! {
_ = shutdown.cancelled() => return,
_ = shim.wait_for_exit() => {
self.on_shim_exit();
return;
}
_ = tick(interval) => {
tokio::select! {
_ = shutdown.cancelled() => return,
_ = shim.wait_for_exit() => {
self.on_shim_exit();
return;
}
flow = self.on_health_tick() => {
if flow.is_break() {
interval = None;
}
}
}
}
}
}
}
fn on_shim_exit(&mut self) {
let Some(runtime) = self.runtime.upgrade() else {
return;
};
let stopped = {
let mut state = self.state.write();
if state.status.is_active() {
crate::runtime::rt_impl::record_main_command_exit(&mut state, &self.exit_file);
} else if state.exit_code.is_none()
&& let Some(record) = boxlite_shared::layout::ExitRecord::read(&self.exit_file)
{
state.exit_code = Some(record.exit_code);
} else {
return;
}
if self.health.is_some() {
state.health_status.state = HealthState::Unhealthy;
}
state.clone()
};
tracing::info!(
box_id = %self.box_id,
exit_code = ?stopped.exit_code,
"Main command exited; box stopped"
);
match runtime.box_manager.save_box(&self.box_id, &stopped) {
Ok(()) | Err(BoxliteError::NotFound(_)) => {}
Err(e) => tracing::warn!(
box_id = %self.box_id,
error = %e,
"Failed to persist the box's exit"
),
}
runtime.invalidate_box_impl(&self.box_id, self.box_name.as_deref());
if self.removes_on_exit
&& let Err(e) = runtime.remove_box(&self.box_id, false)
{
tracing::warn!(
box_id = %self.box_id,
error = %e,
"Failed to auto-remove the box after its main command exited"
);
}
}
async fn on_health_tick(&mut self) -> std::ops::ControlFlow<()> {
use std::ops::ControlFlow::{Break, Continue};
let outcome = match self.health.as_mut() {
Some(probe) => probe.check().await,
None => return Continue(()),
};
let (retries, last) = match &self.health {
Some(probe) => (probe.retries, probe.last),
None => return Continue(()),
};
match outcome {
Probe::Skipped => Continue(()),
Probe::Healthy => {
if last.state != HealthState::Healthy || last.failures != 0 {
let snapshot = {
let mut state = self.state.write();
state.mark_health_check_success();
state.clone()
};
self.persist(&snapshot);
if let Some(probe) = self.health.as_mut() {
probe.last = snapshot.health_status;
}
}
Continue(())
}
Probe::Failed => {
tracing::warn!(box_id = %self.box_id, "Health check probe failed");
let new_failures = last.failures + 1;
let new_state = if new_failures >= retries {
HealthState::Unhealthy
} else {
last.state
};
if last.state == new_state && last.failures == new_failures {
return Continue(());
}
let (snapshot, became_unhealthy) = {
let mut state = self.state.write();
let became_unhealthy = state.mark_health_check_failure(retries);
(state.clone(), became_unhealthy)
};
self.persist(&snapshot);
if let Some(probe) = self.health.as_mut() {
probe.last = snapshot.health_status;
}
if became_unhealthy {
Break(())
} else {
Continue(())
}
}
}
}
fn persist(&self, snapshot: &BoxState) {
let Some(runtime) = self.runtime.upgrade() else {
return;
};
if let Err(e) = runtime.box_manager.save_box(&self.box_id, snapshot) {
tracing::error!(
box_id = %self.box_id,
error = %e,
"Failed to persist health status to database"
);
}
}
}
async fn tick(interval: Option<Duration>) {
match interval {
Some(interval) => tokio::time::sleep(interval).await,
None => std::future::pending().await,
}
}