use dusa_collection_utils::core::errors::{ErrorArrayItem, Errors};
use dusa_collection_utils::core::logger::LogLevel;
use dusa_collection_utils::core::types::pathtype::PathType;
use dusa_collection_utils::core::types::rb::RollingBuffer;
use dusa_collection_utils::core::types::rwarc::LockWithTimeout;
use dusa_collection_utils::log;
use libc::{c_int, kill, SIGKILL, SIGTERM};
use nix::sys::wait::{waitpid, WaitPidFlag, WaitStatus};
use nix::unistd::Pid;
use std::collections::{HashMap, HashSet, VecDeque};
use std::pin::Pin;
use std::process::Stdio;
use std::time::Duration;
use std::{io, thread};
use procfs::process::{all_processes, Process};
use tokio::io::{AsyncRead, AsyncReadExt};
use tokio::process::{Child, Command};
use tokio::task::JoinHandle;
use crate::aggregator::Metrics;
use crate::resource_monitor::{MonitorWatchdog, MonitorWatchdogSnapshot, ResourceMonitorLock};
use crate::state_persistence::{log_error, update_state, AppState};
const RESOURCE_MONITOR_SAMPLE_INTERVAL: Duration = Duration::from_millis(250);
const STDX_BUFFER_UPDATE_INTERVAL: Duration = Duration::from_millis(500);
pub struct ChildLock(pub LockWithTimeout<Child>);
pub struct SupervisedChild {
pub child: ChildLock,
resources: ResourceSupervisor,
monitor_std: Option<JoinHandle<()>>,
stdout_buffer: LockWithTimeout<RollingBuffer>,
stderr_buffer: LockWithTimeout<RollingBuffer>,
stdx_watchdog: MonitorWatchdog,
}
pub struct SupervisedProcess {
pid: Pid,
resources: ResourceSupervisor,
}
struct ResourceSupervisor {
monitor: ResourceMonitorLock,
handle: Option<JoinHandle<()>>,
watchdog: MonitorWatchdog,
}
impl ResourceSupervisor {
fn new(pid: i32) -> Result<Self, ErrorArrayItem> {
Ok(Self {
monitor: ResourceMonitorLock::new(pid)?,
handle: None,
watchdog: MonitorWatchdog::new(),
})
}
async fn ensure_running(&mut self, pid_hint: Option<u32>) {
if let Some(handle) = &self.handle {
if handle.is_finished() {
log!(
LogLevel::Warn,
"Resource monitor task finished unexpectedly for pid {:?}, restarting",
pid_hint
);
self.handle = None;
} else {
return;
}
}
let monitor = self.monitor.clone();
let handle: JoinHandle<()> = monitor
.monitor_with_watchdog_interval(
RESOURCE_MONITOR_SAMPLE_INTERVAL,
Some(self.watchdog.clone()),
)
.await;
self.handle = Some(handle);
}
fn terminate(&mut self) {
if let Some(handle) = &self.handle {
log!(LogLevel::Trace, "Terminating monitor");
handle.abort();
self.handle = None;
self.watchdog.mark_stopped();
}
}
fn is_running(&mut self) -> bool {
if let Some(handle) = &self.handle {
if handle.is_finished() {
self.handle = None;
self.watchdog.mark_stopped();
false
} else {
true
}
} else {
false
}
}
async fn get_metrics(&self) -> Result<Metrics, ErrorArrayItem> {
self.monitor.get_metrics().await
}
fn watchdog_snapshot(&self) -> MonitorWatchdogSnapshot {
self.watchdog.snapshot()
}
fn valid(&self, max_staleness: Duration, max_consecutive_failures: u64) -> bool {
self.watchdog
.snapshot()
.is_valid(max_staleness, max_consecutive_failures)
}
fn clone_idle(&mut self) -> Self {
self.terminate();
Self {
monitor: self.monitor.clone(),
handle: None,
watchdog: self.watchdog.clone(),
}
}
}
impl SupervisedProcess {
pub fn new(pid: Pid) -> Result<Self, ErrorArrayItem> {
if !is_pid_active(pid.as_raw()).unwrap_or(false) {
return Err(ErrorArrayItem::new(
Errors::SupervisedChild,
format!(
"Failed to create SupervisedProcess; cannot determine status of PID: {}",
pid
),
));
}
Ok(SupervisedProcess {
pid,
resources: ResourceSupervisor::new(pid.as_raw())?,
})
}
pub fn get_pid(&self) -> i32 {
self.pid.as_raw()
}
pub fn monitor(&self) -> &ResourceMonitorLock {
&self.resources.monitor
}
pub fn kill(&mut self) -> Result<(), ErrorArrayItem> {
self.resources.terminate();
let xid = self.pid.as_raw();
log!(LogLevel::Trace, "Killing supervised pid {}", xid);
kill_pgid_recursive(xid)?;
Ok(())
}
pub fn running(&self) -> bool {
is_pid_active(self.pid.as_raw()).unwrap_or(false)
}
pub fn active(&self) -> bool {
self.running()
}
pub async fn clone(&mut self) -> Self {
Self {
pid: self.pid,
resources: self.resources.clone_idle(),
}
}
pub async fn monitor_usage(&mut self) {
self.resources
.ensure_running(Some(self.pid.as_raw() as u32))
.await;
}
pub fn terminate_monitor(&mut self) {
self.resources.terminate();
}
pub fn monitoring(&mut self) -> bool {
self.resources.is_running()
}
pub async fn get_metrics(&self) -> Result<Metrics, ErrorArrayItem> {
self.resources.get_metrics().await
}
pub fn resource_watchdog_snapshot(&self) -> MonitorWatchdogSnapshot {
self.resources.watchdog_snapshot()
}
pub fn resource_monitor_valid(
&self,
max_staleness: Duration,
max_consecutive_failures: u64,
) -> bool {
self.resources.valid(max_staleness, max_consecutive_failures)
}
}
impl SupervisedChild {
pub async fn new(
command: &mut Command,
working_dir: Option<PathType>,
) -> Result<Self, ErrorArrayItem> {
spawn_complex_process(command, working_dir, false, true).await }
pub async fn get_pid(&self) -> Result<u32, ErrorArrayItem> {
let child_lock = &self.child;
let child_data = child_lock.0.try_read().await?;
match child_data.id() {
Some(xid) => Ok(xid),
None => Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid PID").into()),
}
}
pub async fn clone(&mut self) -> Self {
self.terminate_stdx();
let resources = self.resources.clone_idle();
let child_lock: ChildLock = self.child.clone();
Self {
child: child_lock,
resources,
monitor_std: None,
stdout_buffer: self.stdout_buffer.clone(),
stderr_buffer: self.stderr_buffer.clone(),
stdx_watchdog: self.stdx_watchdog.clone(),
}
}
pub async fn kill(&mut self) -> Result<(), ErrorArrayItem> {
self.resources.terminate();
self.terminate_stdx();
self.child.kill().await
}
pub async fn try_wait(&self) -> Result<Option<std::process::ExitStatus>, ErrorArrayItem> {
self.child.try_wait().await
}
pub async fn running(&self) -> bool {
self.child.running().await
}
pub fn monitor(&self) -> &ResourceMonitorLock {
&self.resources.monitor
}
pub async fn monitor_usage(&mut self) {
let pid_hint = self.get_pid().await.ok();
self.resources.ensure_running(pid_hint).await;
}
pub fn monitoring(&mut self) -> bool {
self.resources.is_running()
}
pub async fn monitor_stdx(&mut self) {
if let Some(handle) = &self.monitor_std {
if handle.is_finished() {
log!(
LogLevel::Warn,
"Stdout/stderr monitor finished unexpectedly for child pid {:?}, restarting",
self.get_pid().await.ok()
);
self.monitor_std = None;
} else {
return;
}
}
let child_lock = self.child.clone();
let stdout_buffer = self.stdout_buffer.clone();
let stderr_buffer = self.stderr_buffer.clone();
let stdx_watchdog = self.stdx_watchdog.clone();
let monitor_handle = tokio::spawn(async move {
let mut stdout_task = None;
let mut stderr_task = None;
stdx_watchdog.mark_started();
loop {
match child_lock.0.try_write().await {
Ok(mut child) => {
if let Some(stdout) = child.stdout.take() {
let reader = Box::pin(stdout) as Pin<Box<dyn AsyncRead + Send>>;
let buffer = stdout_buffer.clone();
stdout_task = Some(tokio::spawn(read_stream_to_buffer(
reader,
buffer,
STDX_BUFFER_UPDATE_INTERVAL,
)));
}
if let Some(stderr) = child.stderr.take() {
let reader = Box::pin(stderr) as Pin<Box<dyn AsyncRead + Send>>;
let buffer = stderr_buffer.clone();
stderr_task = Some(tokio::spawn(read_stream_to_buffer(
reader,
buffer,
STDX_BUFFER_UPDATE_INTERVAL,
)));
}
stdx_watchdog.record_success();
break;
}
Err(err) => {
log!(
LogLevel::Warn,
"Failed locking child for stdio monitor: {}",
err
);
stdx_watchdog.record_failure();
tokio::time::sleep(Duration::from_millis(250)).await;
}
}
}
if let Some(task) = stdout_task {
let _ = task.await;
}
if let Some(task) = stderr_task {
let _ = task.await;
}
stdx_watchdog.mark_stopped();
});
self.monitor_std = Some(monitor_handle)
}
pub fn monitoring_stdx(&mut self) -> bool {
if let Some(handle) = &self.monitor_std {
if handle.is_finished() {
self.monitor_std = None;
self.stdx_watchdog.mark_stopped();
false
} else {
true
}
} else {
false
}
}
pub async fn get_std_out(&self) -> Result<Vec<(u64, String)>, ErrorArrayItem> {
let rb = self.stdout_buffer.try_read().await?;
Ok(rb.get_latest_time())
}
pub async fn get_std_err(&self) -> Result<Vec<(u64, String)>, ErrorArrayItem> {
let rb = self.stderr_buffer.try_read().await?;
Ok(rb.get_latest_time())
}
pub fn terminate_monitor(&mut self) {
self.resources.terminate();
}
pub fn terminate_stdx(&mut self) {
if let Some(handle) = &self.monitor_std {
log!(LogLevel::Trace, "Terminating Standart X monitor");
handle.abort();
self.monitor_std = None;
self.stdx_watchdog.mark_stopped();
}
}
pub async fn get_metrics(&self) -> Result<Metrics, ErrorArrayItem> {
self.resources.get_metrics().await
}
pub fn resource_watchdog_snapshot(&self) -> MonitorWatchdogSnapshot {
self.resources.watchdog_snapshot()
}
pub fn stdx_watchdog_snapshot(&self) -> MonitorWatchdogSnapshot {
self.stdx_watchdog.snapshot()
}
pub fn resource_monitor_valid(
&self,
max_staleness: Duration,
max_consecutive_failures: u64,
) -> bool {
self.resources.valid(max_staleness, max_consecutive_failures)
}
pub fn stdx_monitor_valid(
&self,
max_staleness: Duration,
max_consecutive_failures: u64,
) -> bool {
self.stdx_watchdog
.snapshot()
.is_valid(max_staleness, max_consecutive_failures)
}
}
impl ChildLock {
pub fn new(child: Child) -> Self {
let rw_lock: LockWithTimeout<Child> = LockWithTimeout::new(child);
Self(rw_lock)
}
pub fn update(mut self, new_child: Child) -> Self {
self.0 = LockWithTimeout::new(new_child);
self
}
pub fn clone(&self) -> Self {
let child = &self.0;
let lock_clone = child.clone();
ChildLock { 0: lock_clone }
}
pub async fn kill(&self) -> Result<(), ErrorArrayItem> {
let child = self
.0
.try_read_with_timeout(Some(Duration::from_secs(5)))
.await?;
let xid = match child.id() {
Some(xid) => xid,
None => {
return Err(ErrorArrayItem::new(
dusa_collection_utils::core::errors::Errors::InputOutput,
"No PID found in child process".to_owned(),
))
}
};
log!(LogLevel::Trace, "Killing child pid {}", xid);
if let Ok(xid) = xid.try_into() {
kill_pgid_recursive(xid)?;
Ok(())
} else {
Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid PID").into())
}
}
pub async fn try_wait(&self) -> Result<Option<std::process::ExitStatus>, ErrorArrayItem> {
let mut child = self
.0
.try_write_with_timeout(Some(Duration::from_secs(1)))
.await?;
child.try_wait().map_err(ErrorArrayItem::from)
}
pub async fn running(&self) -> bool {
match self.try_wait().await {
Ok(None) => true,
Ok(Some(_)) => false,
Err(err) if err.err_type == Errors::GeneralError => true,
Err(_) => false,
}
}
}
pub async fn spawn_simple_process(
command: &mut Command,
capture_output: bool,
state: &mut AppState,
state_path: &PathType,
) -> Result<Child, io::Error> {
if capture_output {
command.stdout(Stdio::piped());
command.stderr(Stdio::piped());
} else {
command.stdout(Stdio::inherit());
command.stderr(Stdio::inherit());
}
match command.spawn() {
Ok(child_process) => {
log!(
LogLevel::Trace,
"Child process spawned successfully: {:?}",
child_process
);
state.data = String::from("Process spawned");
state.event_counter += 1;
update_state(state, state_path, None).await;
Ok(child_process)
}
Err(e) => {
log!(
LogLevel::Error,
"Failed to spawn child process: {}",
e.to_string()
);
let error_item: ErrorArrayItem = ErrorArrayItem::new(
dusa_collection_utils::core::errors::Errors::InputOutput,
e.to_string(),
);
log_error(state, error_item, state_path).await;
Err(e)
}
}
}
pub async fn spawn_complex_process(
command: &mut Command,
working_dir: Option<PathType>,
independent_process_group: bool,
capture_output: bool,
) -> Result<SupervisedChild, ErrorArrayItem> {
log!(LogLevel::Trace, "Child to spawn: {:?}", &command);
if independent_process_group {
unsafe {
command.pre_exec(|| {
if libc::setsid() == -1 {
return Err(io::Error::last_os_error());
}
Ok(())
})
};
} else {
command.kill_on_drop(true);
log!(
LogLevel::Trace,
"Complex process being spawned in the same process group"
);
}
if capture_output {
command.stdout(Stdio::piped());
command.stderr(Stdio::piped());
} else {
command.stdout(Stdio::inherit());
command.stderr(Stdio::inherit());
}
if let Some(path) = working_dir {
command.current_dir(path.canonicalize().map_err(ErrorArrayItem::from)?);
}
match command.spawn() {
Ok(mut child) => {
log!(
LogLevel::Trace,
"Child process spawned successfully: {:#?}",
child
);
let pid = match child.id() {
Some(d) => d,
None => {
return Err(ErrorArrayItem::new(
Errors::InputOutput,
"Couldn't determine if process spawned".to_owned(),
))
}
};
let monitor = match ResourceMonitorLock::new(pid as i32) {
Ok(resource_monitor) => resource_monitor,
Err(e) => {
child.kill().await?;
return Err(ErrorArrayItem::from(io::Error::new(
io::ErrorKind::InvalidData,
e.to_string(),
)));
}
};
let child = ChildLock::new(child);
Ok(SupervisedChild {
child,
resources: ResourceSupervisor {
monitor,
handle: None,
watchdog: MonitorWatchdog::new(),
},
monitor_std: None,
stdout_buffer: LockWithTimeout::new(RollingBuffer::new(500)),
stderr_buffer: LockWithTimeout::new(RollingBuffer::new(500)),
stdx_watchdog: MonitorWatchdog::new(),
})
}
Err(error) => {
log!(LogLevel::Error, "Failed to spawn child process: {}", error);
Err(ErrorArrayItem::from(error))
}
}
}
fn collect_descendants(root_pid: i32) -> Result<HashSet<i32>, ErrorArrayItem> {
let mut children_map: HashMap<i32, Vec<i32>> = HashMap::new();
let mut result: HashSet<i32> = HashSet::new();
for prc in all_processes()
.map_err(|e| ErrorArrayItem::from(io::Error::new(io::ErrorKind::Other, e.to_string())))?
{
let process: Process = match prc {
Ok(p) => p,
Err(_) => continue,
};
if let Ok(stat) = process.stat() {
children_map
.entry(stat.ppid)
.or_default()
.push(process.pid());
}
}
let mut queue: VecDeque<i32> = VecDeque::new();
queue.push_back(root_pid);
result.insert(root_pid);
while let Some(pid) = queue.pop_front() {
if let Some(children) = children_map.get(&pid) {
for child in children {
if result.insert(*child) {
queue.push_back(*child);
}
}
}
}
Ok(result)
}
fn reap_zombie_process(pid: c_int) {
match waitpid(Pid::from_raw(pid), Some(WaitPidFlag::WNOHANG)) {
Ok(WaitStatus::Exited(_, status)) => {
log!(
LogLevel::Trace,
"Reaped pid {} with exit status {}",
pid,
status
)
}
Ok(WaitStatus::Signaled(_, sig, _)) => {
log!(
LogLevel::Trace,
"Reaped pid {} terminated by signal {:?}",
pid,
sig
)
}
Ok(WaitStatus::StillAlive) => {
log!(
LogLevel::Trace,
"PID {} still alive when attempting reap",
pid
)
}
Ok(status) => {
log!(LogLevel::Trace, "PID {} wait status: {:?}", pid, status)
}
Err(e) => {
log!(LogLevel::Trace, "Failed to reap pid {}: {}", pid, e)
}
}
}
fn kill_pgid_recursive(pgid: i32) -> Result<(), ErrorArrayItem> {
log!(LogLevel::Trace, "Recursively killing pgid: {}", pgid);
let pids = collect_descendants(pgid)?;
log!(LogLevel::Trace, "Found descendant pids: {:?}", pids);
for pid in &pids {
let res = unsafe { kill(*pid, SIGTERM) };
if res == 0 {
log!(LogLevel::Trace, "Sent SIGTERM to pid: {}", pid);
} else {
let err = io::Error::last_os_error();
if err.raw_os_error() == Some(libc::ESRCH) {
log!(LogLevel::Trace, "PID {} already exited", pid);
} else {
log!(
LogLevel::Warn,
"Failed to send SIGTERM to pid {}: {}",
pid,
err
);
}
}
}
thread::sleep(Duration::from_millis(400));
for pid in &pids {
reap_zombie_process(*pid);
if is_pid_active(*pid).unwrap_or(false) {
log!(LogLevel::Warn, "PID {} still running; sending SIGKILL", pid);
let res = unsafe { kill(*pid, SIGKILL) };
if res != 0 {
let err = io::Error::last_os_error();
if err.raw_os_error() != Some(libc::ESRCH) {
return Err(ErrorArrayItem::from(err));
}
}
reap_zombie_process(*pid);
if !is_pid_active(*pid).unwrap_or(false) {
log!(LogLevel::Trace, "PID {} terminated", pid);
} else {
log!(LogLevel::Warn, "PID {} survived SIGKILL", pid);
}
} else {
log!(LogLevel::Trace, "PID {} terminated gracefully", pid);
}
}
Ok(())
}
pub fn is_pid_active(pid: i32) -> io::Result<bool> {
let ret = unsafe { libc::kill(pid, 0) };
if ret == 0 {
Ok(true)
} else {
match io::Error::last_os_error().raw_os_error() {
Some(libc::ESRCH) => Ok(false), Some(libc::EPERM) => Ok(true), Some(err) => Err(io::Error::from_raw_os_error(err)),
None => Err(io::Error::new(io::ErrorKind::Other, "Unknown error")),
}
}
}
use bytes::BytesMut;
async fn flush_lines_to_buffer(
buffer: &LockWithTimeout<RollingBuffer>,
pending_lines: &mut Vec<String>,
) {
if pending_lines.is_empty() {
return;
}
if let Ok(mut b) = buffer.try_write().await {
for line in pending_lines.drain(..) {
b.push(line);
}
}
}
async fn read_stream_to_buffer<R>(
mut reader: R,
buffer: LockWithTimeout<RollingBuffer>,
flush_interval: Duration,
) where
R: Unpin + AsyncRead,
{
let mut buf = BytesMut::with_capacity(1024);
let mut partial = String::new();
let mut pending_lines: Vec<String> = Vec::new();
let mut last_flush = std::time::Instant::now();
loop {
let remaining_until_flush = flush_interval.saturating_sub(last_flush.elapsed());
match tokio::time::timeout(remaining_until_flush, reader.read_buf(&mut buf)).await {
Err(_) => {
flush_lines_to_buffer(&buffer, &mut pending_lines).await;
last_flush = std::time::Instant::now();
continue;
}
Ok(result) => match result {
Ok(n) if n == 0 => break, Ok(_) => {}
Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
Err(e) => {
log!(LogLevel::Warn, "Read error in stdio monitor: {}", e);
break;
}
},
};
let chunk = String::from_utf8_lossy(&buf);
partial.push_str(&chunk);
while let Some(pos) = partial.find('\n') {
let line = partial[..pos].to_string();
pending_lines.push(line);
partial.drain(..=pos); }
buf.clear();
if last_flush.elapsed() >= flush_interval {
flush_lines_to_buffer(&buffer, &mut pending_lines).await;
last_flush = std::time::Instant::now();
}
}
if !partial.is_empty() {
pending_lines.push(partial);
}
flush_lines_to_buffer(&buffer, &mut pending_lines).await;
}