use super::*;
#[cfg(unix)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum PidLiveness {
Alive,
Dead,
PermissionDenied,
}
#[cfg(unix)]
impl PidLiveness {
pub(super) fn is_running(self) -> bool {
!matches!(self, PidLiveness::Dead)
}
}
#[cfg(unix)]
pub(super) fn classify_kill_result(rc: i32, errno: i32) -> PidLiveness {
if rc == 0 {
return PidLiveness::Alive;
}
match errno {
libc::EPERM => PidLiveness::PermissionDenied,
_ => PidLiveness::Dead,
}
}
#[cfg(unix)]
pub(super) fn is_process_running(pid: u32) -> bool {
let Ok(pid) = i32::try_from(pid) else {
return false;
};
if pid <= 0 {
return false;
}
let rc = unsafe { libc::kill(pid, 0) };
let errno = std::io::Error::last_os_error().raw_os_error().unwrap_or(0);
classify_kill_result(rc, errno).is_running()
}
#[cfg(unix)]
pub(super) fn pid_can_name_incumbent(
pid: u32,
current_pid: u32,
allow_same_process_incumbent: bool,
) -> bool {
allow_same_process_incumbent || pid != current_pid
}
#[cfg(unix)]
pub(super) const DUPLICATE_PROBE_TIMEOUT: std::time::Duration =
std::time::Duration::from_millis(500);
#[cfg(unix)]
pub(super) async fn socket_speaks_khived_protocol(
sock: &std::path::Path,
expected_config_id: &str,
) -> bool {
let probe = DaemonRequestFrame {
probe_only: true,
protocol_version: PROTOCOL_VERSION,
config_id: expected_config_id.to_string(),
..Default::default()
};
let Ok(payload) = serde_json::to_vec(&probe) else {
return false;
};
let response = tokio::time::timeout(DUPLICATE_PROBE_TIMEOUT, async {
let mut stream = UnixStream::connect(sock).await.ok()?;
write_frame(&mut stream, &payload).await.ok()?;
let raw = read_frame(&mut stream).await.ok()?;
serde_json::from_slice::<DaemonResponseFrame>(&raw).ok()
})
.await
.ok()
.flatten();
let Some(resp) = response else {
return false;
};
let is_probe_ack = resp.ok
&& resp.result.is_none()
&& resp.error.is_none()
&& resp.metrics.is_none()
&& resp.request_id.is_none();
is_probe_ack
&& !resp.version_mismatch
&& !resp.namespace_mismatch
&& !resp.config_mismatch
&& resp.daemon_protocol_version == PROTOCOL_VERSION
&& resp
.served_config_id
.as_deref()
.is_some_and(|served| config_ids_compatible(expected_config_id, served))
}
#[cfg(unix)]
pub(super) async fn socket_is_unreachable(sock: &std::path::Path) -> bool {
match tokio::time::timeout(DUPLICATE_PROBE_TIMEOUT, UnixStream::connect(sock)).await {
Ok(Err(error)) => matches!(
error.kind(),
std::io::ErrorKind::NotFound | std::io::ErrorKind::ConnectionRefused
),
_ => false,
}
}
#[cfg(unix)]
pub(super) enum Incumbent {
Serving(u32),
Live(u32),
Stale,
}
#[cfg(unix)]
pub(super) async fn cleanup_stale_daemon(
sock: &std::path::Path,
pid_file: &std::path::Path,
allow_same_process_incumbent: bool,
expected_config_id: &str,
) -> Incumbent {
let mut stale_pid_file_guard = None;
if let Ok(pid_str) = std::fs::read_to_string(pid_file) {
if let Ok(pid) = pid_str.trim().parse::<u32>() {
if pid_can_name_incumbent(pid, std::process::id(), allow_same_process_incumbent)
&& is_process_running(pid)
{
if sock.exists() && socket_speaks_khived_protocol(sock, expected_config_id).await {
return Incumbent::Serving(pid);
}
if sock.exists() && !socket_is_unreachable(sock).await {
return Incumbent::Live(pid);
}
match try_acquire_pid_file_lock(pid_file) {
Ok(Some(guard)) => stale_pid_file_guard = Some(guard),
Ok(None) => return Incumbent::Live(pid),
Err(e) => {
tracing::warn!(
error = %e,
path = ?pid_file,
"cannot check daemon PID-file lock"
);
return Incumbent::Live(pid);
}
}
}
}
}
if sock.exists() {
if let Err(e) = std::fs::remove_file(sock) {
tracing::warn!(error = %e, path = ?sock, "failed to remove stale socket");
}
}
if pid_file.exists() {
if let Err(e) = std::fs::remove_file(pid_file) {
tracing::warn!(error = %e, path = ?pid_file, "failed to remove stale PID file");
}
}
drop(stale_pid_file_guard);
Incumbent::Stale
}
#[cfg(unix)]
pub(super) fn write_pid_file_exclusive(
pid_file: &std::path::Path,
) -> std::io::Result<std::fs::File> {
use std::os::unix::fs::OpenOptionsExt;
let mut opts = std::fs::OpenOptions::new();
opts.write(true).create_new(true).mode(0o600);
let mut f = opts.open(pid_file)?;
let rc = unsafe { libc::flock(f.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) };
if rc != 0 {
return Err(std::io::Error::last_os_error());
}
f.write_all(std::process::id().to_string().as_bytes())?;
Ok(f)
}
#[cfg(unix)]
pub(super) fn try_acquire_pid_file_lock(
pid_file: &std::path::Path,
) -> std::io::Result<Option<std::fs::File>> {
let file = std::fs::OpenOptions::new()
.read(true)
.write(true)
.open(pid_file)?;
let rc = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) };
if rc == 0 {
return Ok(Some(file));
}
let error = std::io::Error::last_os_error();
if error.kind() == std::io::ErrorKind::WouldBlock
|| error.raw_os_error() == Some(libc::EWOULDBLOCK)
{
Ok(None)
} else {
Err(error)
}
}
#[cfg(unix)]
pub(super) fn remove_pid_file_if_owned(pid_file: &std::path::Path, guard: &std::fs::File) {
let Ok(owned) = guard.metadata() else {
return;
};
let Ok(current) = std::fs::metadata(pid_file) else {
return;
};
if owned.dev() == current.dev() && owned.ino() == current.ino() {
if let Err(e) = std::fs::remove_file(pid_file) {
tracing::warn!(error = %e, path = ?pid_file, "failed to remove unbound PID file");
}
}
}
#[cfg(unix)]
pub(super) async fn pid_file_names_a_reachable_daemon(
pid_file: &std::path::Path,
sock: &std::path::Path,
allow_same_process_incumbent: bool,
expected_config_id: &str,
) -> bool {
let Ok(pid_str) = std::fs::read_to_string(pid_file) else {
return false;
};
let Ok(pid) = pid_str.trim().parse::<u32>() else {
return false;
};
pid_can_name_incumbent(pid, std::process::id(), allow_same_process_incumbent)
&& is_process_running(pid)
&& sock.exists()
&& socket_speaks_khived_protocol(sock, expected_config_id).await
}
#[cfg(unix)]
pub(super) async fn drain(active: &std::sync::atomic::AtomicUsize) -> bool {
drain_with_timeout(active, drain_timeout()).await
}
#[cfg(unix)]
pub(super) async fn drain_for_idle(
active: &std::sync::atomic::AtomicUsize,
timeout: std::time::Duration,
) {
let deadline = tokio::time::Instant::now() + timeout;
let mut warned = false;
while active.load(std::sync::atomic::Ordering::SeqCst) + background_task_count() != 0 {
if !warned && tokio::time::Instant::now() >= deadline {
tracing::warn!(
"idle drain interval elapsed; retaining workers and rendezvous until settled"
);
warned = true;
}
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
}
}
#[cfg(unix)]
pub(super) async fn drain_with_timeout(
active: &std::sync::atomic::AtomicUsize,
timeout: std::time::Duration,
) -> bool {
use std::sync::atomic::Ordering;
let remaining = || active.load(Ordering::SeqCst) + background_task_count();
if remaining() == 0 {
return true;
}
let deadline = tokio::time::Instant::now() + timeout;
while remaining() > 0 {
if tokio::time::Instant::now() >= deadline {
tracing::warn!(
remaining_connections = active.load(Ordering::SeqCst),
remaining_background_tasks = background_task_count(),
outstanding_background_tasks = %background_task_names().join(", "),
"drain timeout reached; forcing shutdown"
);
return false;
}
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_millis(100)) => {}
_ = tokio::time::sleep_until(deadline) => {}
}
}
true
}
#[cfg(unix)]
pub(super) async fn finish_connection_tasks(
tasks: Vec<tokio::task::JoinHandle<()>>,
drained: bool,
) {
if !drained {
for task in &tasks {
if !task.is_finished() {
task.abort();
}
}
}
for task in tasks {
let _ = task.await;
}
}
pub fn drain_timeout() -> std::time::Duration {
let secs = khive_db::env::env_parse_or("KHIVE_DRAIN_TIMEOUT_SECS", DEFAULT_DRAIN_TIMEOUT_SECS);
std::time::Duration::from_secs(secs)
}