ogre_watchdog/

watchdog.rs

use dashmap::DashMap;
use once_cell::sync::Lazy;
use std::fmt::{Display, Formatter};
use std::{
    fmt::Debug,
    sync::{
        atomic::{AtomicBool, Ordering::Relaxed},
        Arc,
    },
    thread,
    time::{Duration, Instant},
};

static START: Lazy<Instant> = Lazy::new(Instant::now);

#[inline(always)]
fn now_ms() -> u64 {
    START.elapsed().as_millis() as u64
}

/// OS thread handle (only meaningful with demotion feature).
#[derive(Clone, Copy, Debug)]
pub enum OsThreadHandle {
    #[cfg(target_os = "linux")]
    LinuxTid(i32),
    #[cfg(target_os = "windows")]
    WindowsTid(u32),
    None,
}

#[cfg(feature = "demotion")]
impl OsThreadHandle {
    #[inline(always)]
    pub fn current() -> Self {
        #[cfg(target_os = "linux")]
        unsafe {
            OsThreadHandle::LinuxTid(libc::syscall(libc::SYS_gettid) as i32)
        }
        #[cfg(target_os = "windows")]
        {
            use windows_sys::Win32::System::Threading::GetCurrentThreadId;
            OsThreadHandle::WindowsTid(unsafe { GetCurrentThreadId() })
        }
    }
}

/// Config for the watchdog.
#[derive(Clone, Debug)]
pub struct Config {
    /// Maximum number of unhealthy (stalled) executors tolerated before meltdown.
    /// Meltdown triggers when stalled_count > max_unhealthy_executors AND there is at least 1 active work.
    pub max_unhealthy_executors: usize,
    /// Hard maximum runtime per work unit. If a work runs longer, it's considered stalled.
    pub stall_threshold: Duration,
    /// How often the watchdog thread checks health.
    pub check_interval: Duration,
    /// Enable thread priority demotion of stalled executors (only present with the demotion feature).
    #[cfg(feature = "demotion")]
    pub demote_stalled_threads: bool,
}

impl Default for Config {
    fn default() -> Self {
        Self {
            max_unhealthy_executors: 0,
            stall_threshold: Duration::from_secs(299),
            check_interval: Duration::from_secs(30),
            #[cfg(feature = "demotion")]
            demote_stalled_threads: true,
        }
    }
}

/// Hooks trait: supply your own to get compile-time inlining & zero branching.
/// IMPORTANT: All hooks must avoid taking app locks, blocking I/O, or heavy allocation.
pub trait Hooks<ExecutorIdType, EventType>: Send + Sync + 'static {
    /// Called when executor/work is admitted.
    fn on_started(&self, _key: &ExecutorIdType, _token: &EventType) {}
    /// Called exactly once when a work first transitions to stalled.
    fn on_stalled(&self, _key: &ExecutorIdType, _token: &EventType) {}
    /// Called when a previously stalled work finishes before meltdown.
    /// Use this to undo any DLQ pre-marking and (optionally) promote thread priority back.
    fn on_recovered(&self, _key: &ExecutorIdType, _token: &EventType) {}
    /// Final best-effort callback before termination/restart.
    /// You decide how to terminate (std::process::exit, abort, panic, etc).
    fn on_meltdown(&self, _stalled: &[(ExecutorIdType, EventType)]) {}
}

/// Default no-op hooks.
pub struct NoopHooks;
impl<ExecutorIdType, EventType> Hooks<ExecutorIdType, EventType> for NoopHooks {}

#[cfg(feature = "demotion")]
#[derive(Clone, Copy, Debug)]
struct DemotionInfo {
    handle: OsThreadHandle,
    #[cfg(target_os = "linux")]
    orig_nice: i32,
    #[cfg(target_os = "windows")]
    orig_prio: i32,
}

struct Entry<EventType: Clone> {
    started_ms: u64,
    is_stalled: AtomicBool,
    event_payload: EventType,
    #[cfg(feature = "demotion")]
    demotion: Option<DemotionInfo>,
}

pub struct Watchdog<ExecutorIdType, EventType, HooksType = NoopHooks>
where
    ExecutorIdType: Eq + std::hash::Hash + Clone + Debug + Send + Sync + 'static,
    EventType: Clone + Send + Sync + 'static,
    HooksType: Hooks<ExecutorIdType, EventType>,
{
    config: Config,
    accepting: AtomicBool,
    entries: DashMap<ExecutorIdType, Entry<EventType>>,
    hooks: HooksType,
    #[cfg(feature = "systemd")]
    systemd_ping: bool,
}

impl<ExecutorIdType, EventType> Watchdog<ExecutorIdType, EventType, NoopHooks>
where
    ExecutorIdType: Eq + std::hash::Hash + Clone + Debug + Send + Sync + 'static,
    EventType: Clone + Send + Sync + 'static,
{
    pub fn new(cfg: Config) -> Self {
        Self {
            config: cfg,
            accepting: AtomicBool::new(true),
            entries: DashMap::new(),
            hooks: NoopHooks,
            #[cfg(feature = "systemd")]
            systemd_ping: false,
        }
    }
}

impl<ExecutorIdType, EventType, HooksType> Watchdog<ExecutorIdType, EventType, HooksType>
where
    ExecutorIdType: Eq + std::hash::Hash + Clone + Debug + Send + Sync + 'static,
    EventType: Clone + Send + Sync + 'static,
    HooksType: Hooks<ExecutorIdType, EventType>,
{
    /// Swap in custom hooks (compile-time monomorphized)
    pub fn with_hooks<NewHooksType: Hooks<ExecutorIdType, EventType>>(
        self,
        hooks: NewHooksType,
    ) -> Watchdog<ExecutorIdType, EventType, NewHooksType> {
        Watchdog {
            config: self.config,
            accepting: self.accepting,
            entries: self.entries,
            hooks,
            #[cfg(feature = "systemd")]
            systemd_ping: self.systemd_ping,
        }
    }
    /// Enable systemd Watchdog pings (Linux) in the watchdog loop (optional).
    #[cfg(feature = "systemd")]
    pub fn with_systemd_ping(mut self, enabled: bool) -> Self {
        self.systemd_ping = enabled;
        self
    }

    /// Set the admission gate
    pub fn set_accepting(&self, v: bool) {
        self.accepting.store(v, Relaxed);
    }

    /// Called when a unit of work is admitted. Returns Err if draining (ACCEPTING=false)
    pub fn on_work_start(
        &self,
        key: &ExecutorIdType,
        event_payload: &EventType,
    ) -> Result<(), MeltdownInProgress> {
        if !self.accepting.load(Relaxed) {
            return Err(MeltdownInProgress);
        }

        #[cfg(feature = "demotion")]
        let demotion = if self.config.demote_stalled_threads {
            // Capture OS handle & current priority so we can restore on recovery.
            let handle = OsThreadHandle::current();
            #[cfg(target_os = "linux")]
            unsafe {
                let orig = libc::getpriority(libc::PRIO_PROCESS, 0) as i32;
                Some(DemotionInfo {
                    handle,
                    orig_nice: orig,
                })
            }
            #[cfg(target_os = "windows")]
            {
                use windows_sys::Win32::System::Threading::{GetCurrentThread, GetThreadPriority};
                unsafe {
                    let h = GetCurrentThread();
                    let p = GetThreadPriority(h);
                    Some(DemotionInfo {
                        handle,
                        orig_prio: p,
                    })
                }
            }
        } else {
            None
        };

        self.entries.insert(
            key.clone(),
            Entry {
                started_ms: now_ms(),
                is_stalled: AtomicBool::new(false),
                event_payload: event_payload.clone(),
                #[cfg(feature = "demotion")]
                demotion,
            },
        );
        self.hooks.on_started(key, event_payload);
        Ok(())
    }

    /// Finish and clear the entry; missing keys are ignored (idempotent)
    pub fn on_work_finish(&self, key: &ExecutorIdType) {
        if let Some((executor_id, entry)) = self.entries.remove(key) {
            if entry.is_stalled.load(Relaxed) {
                // If it had been stalled, signal recovery and restore priority.
                self.hooks.on_recovered(&executor_id, &entry.event_payload);
                #[cfg(feature = "demotion")]
                if let Some(info) = entry.demotion {
                    let _ = promote_thread(info);
                }
            }
        }
    }

    /// Spawn the watchdog thread
    pub fn start(self) -> Arc<Self> {
        let this = Arc::new(self);
        let this_clone = this.clone();
        thread::spawn(move || {
            let stall_ms = this.config.stall_threshold.as_millis() as u64;

            loop {
                #[cfg(feature = "systemd")]
                if this.systemd_ping {
                    let _ = sd_notify::notify(false, &[sd_notify::NotifyState::Watchdog]);
                }

                thread::sleep(this.config.check_interval);

                let now = now_ms();

                let mut stalled_count = 0;

                // Snapshot health. Mark stalled exactly once.
                for kv in this.entries.iter() {
                    let started = kv.value().started_ms;
                    if now.saturating_sub(started) > stall_ms {
                        if !kv.value().is_stalled.swap(true, Relaxed) {
                            #[cfg(feature = "demotion")]
                            if this.config.demote_stalled_threads {
                                if let Some(info) = kv.value().demotion {
                                    let _ = demote_thread(info);
                                }
                            }
                            this.hooks.on_stalled(kv.key(), &kv.value().event_payload);
                        }
                        stalled_count += 1;
                    }
                }

                // should the watchdog's meltdown kick-in?
                if stalled_count > this.config.max_unhealthy_executors {
                    break;
                }
            }

            // stop admission, give a bounded grace equal to `stall_threshold`,
            // then invoke user-provided meltdown hook (which, typically, forcibly terminate the process).
            this.accepting.store(false, Relaxed);
            thread::sleep(this.config.stall_threshold);

            if !this.entries.is_empty() {
                let pairs: Vec<(ExecutorIdType, EventType)> = this
                    .entries
                    .iter()
                    .map(|kv| (kv.key().clone(), kv.value().event_payload.clone()))
                    .collect();
                this.hooks.on_meltdown(&pairs);
            }
        });
        this_clone
    }
}

/// Returned when admission is closed (draining).
#[derive(Debug)]
pub struct Draining;

#[cfg(feature = "demotion")]
fn demote_thread(info: DemotionInfo) -> std::io::Result<()> {
    match info.handle {
        #[cfg(target_os = "linux")]
        OsThreadHandle::LinuxTid(tid) => unsafe {
            let r = libc::setpriority(libc::PRIO_PROCESS, tid as u32, 19);
            if r == 0 {
                Ok(())
            } else {
                Err(std::io::Error::last_os_error())
            }
        },
        #[cfg(target_os = "windows")]
        OsThreadHandle::WindowsTid(tid) => {
            use windows_sys::Win32::Foundation::CloseHandle;
            use windows_sys::Win32::System::Threading::{
                OpenThread, SetThreadPriority, THREAD_PRIORITY_IDLE, THREAD_QUERY_INFORMATION,
                THREAD_SET_INFORMATION,
            };
            unsafe {
                let h = OpenThread(THREAD_SET_INFORMATION | THREAD_QUERY_INFORMATION, 0, tid);
                if h == std::ptr::null_mut() {
                    return Err(std::io::Error::last_os_error());
                }
                let ok = SetThreadPriority(h, THREAD_PRIORITY_IDLE);
                let last = std::io::Error::last_os_error();
                CloseHandle(h);
                if ok != 0 {
                    Ok(())
                } else {
                    Err(last)
                }
            }
        }
        _ => Ok(()),
    }
}

#[cfg(feature = "demotion")]
fn promote_thread(info: DemotionInfo) -> std::io::Result<()> {
    match info.handle {
        #[cfg(target_os = "linux")]
        OsThreadHandle::LinuxTid(tid) => unsafe {
            let r = libc::setpriority(libc::PRIO_PROCESS, tid as u32, info.orig_nice);
            if r == 0 {
                Ok(())
            } else {
                Err(std::io::Error::last_os_error())
            }
        },
        #[cfg(target_os = "windows")]
        OsThreadHandle::WindowsTid(tid) => {
            use windows_sys::Win32::Foundation::CloseHandle;
            use windows_sys::Win32::System::Threading::{
                OpenThread, SetThreadPriority, THREAD_QUERY_INFORMATION, THREAD_SET_INFORMATION,
            };
            unsafe {
                let h = OpenThread(THREAD_SET_INFORMATION | THREAD_QUERY_INFORMATION, 0, tid);
                if h == std::ptr::null_mut() {
                    return Err(std::io::Error::last_os_error());
                }
                let ok = SetThreadPriority(h, info.orig_prio);
                let last = std::io::Error::last_os_error();
                CloseHandle(h);
                if ok != 0 {
                    Ok(())
                } else {
                    Err(last)
                }
            }
        }
        _ => Ok(()),
    }
}

#[derive(Debug)]
pub struct MeltdownInProgress;
impl Display for MeltdownInProgress {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        <Self as Debug>::fmt(self, f)
    }
}
impl std::error::Error for MeltdownInProgress {}