aion-rs 0.31.0

//! Runtime-owned, non-consuming fan-out records for beamr process exits.

use std::collections::HashSet;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::mpsc::{self, Receiver};
use std::sync::{Arc, Mutex, MutexGuard};
use std::thread::JoinHandle;
use std::time::Duration;

use beamr::ets::OwnedTerm;
use beamr::process::ExitReason;
use beamr::scheduler::{ExitEvent, Scheduler};
use dashmap::DashMap;

use crate::runtime::drain_bound::{DrainBound, DrainProgress};
use crate::runtime::monitor::MonitorInstallation;
use crate::{EngineError, Pid};

#[path = "process_exit_callback.rs"]
mod callback;
#[path = "process_exit_drainer.rs"]
mod drainer;
#[path = "process_exit_record.rs"]
mod record;

pub(super) use record::{
    ObservedProcessExit, OwnedProcessExitOutcome, ProcessExitCallback,
    ProcessExitObservationFailure, ProcessExitRecord,
};

type ProcessExitRecords = DashMap<Pid, Arc<ProcessExitRecord>>;

struct RegistryLifecycle {
    closed: bool,
    /// Local children created by wrapped BEAM spawn BIFs and awaiting outcome release.
    unobserved_children: HashSet<Pid>,
}

/// Exclusive gate held from immediately before a scheduler spawn through pid classification.
pub(super) struct ProcessSpawnReservation<'a> {
    registry: &'a ProcessExitRegistry,
    lifecycle: MutexGuard<'a, RegistryLifecycle>,
}

impl ProcessSpawnReservation<'_> {
    pub(super) fn register(mut self, pid: Pid) -> Result<(), EngineError> {
        self.registry.register_locked(pid, &mut self.lifecycle)
    }

    pub(super) fn track_unobserved(mut self, pid: Pid) {
        self.lifecycle.unobserved_children.insert(pid);
    }
}

/// What [`ProcessExitRegistry`] knows about one pid's ending.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum ProcessEnding {
    /// A record is held and it carries the process' terminal outcome.
    Recorded,
    /// A record is held and no terminal outcome has been published on it yet.
    Pending,
    /// No record is held and the pid is at or below the registration
    /// watermark. The registry forgets a registration when it retires the
    /// record, so this is either a process whose outcome was delivered and
    /// reaped or a pid the registry never registered; the two cannot be told
    /// apart, and neither is a recorded ending.
    Forgotten,
    /// No record is held and the pid is above every pid this registry has
    /// registered, so the registry provably never registered it.
    NeverRegistered,
}

struct ExitDrainer {
    handle: Option<JoinHandle<Result<(), EngineError>>>,
    stopped: Receiver<()>,
}

/// Index of active owned exit records and owner of the one beamr event drainer.
pub(super) struct ProcessExitRegistry {
    records: Arc<ProcessExitRecords>,
    registered_through: AtomicU64,
    has_registered: AtomicBool,
    lifecycle: Mutex<RegistryLifecycle>,
    stop_drainer: AtomicBool,
    drainer: Mutex<ExitDrainer>,
    callbacks: callback::ProcessExitCallbackDispatcher,
    /// How long the drainer parks on the exit-event subscription between
    /// looks at the stop flag: [`DrainBound::poll_interval`], so a stop is
    /// observed within one percent of the operator's patience.
    park_bound: Duration,
    drain_bound: DrainBound,
    /// The drainer's own progress: one job per exit event it processes.
    drain_progress: Arc<DrainProgress>,
    /// Force the next `begin_shutdown` to raise the drain failure a degraded
    /// store produces, so the shutdown path's behaviour on a FAILING drain can
    /// be measured rather than reasoned about. The real failures on this path
    /// are lock poison and a drainer that will not stop inside its window;
    /// neither can be produced deterministically from a test, which is why the
    /// property went unpinned long enough to become a critical finding.
    #[cfg(test)]
    forced_shutdown_failure: AtomicBool,
    #[cfg(test)]
    pause_next_registration: AtomicBool,
    #[cfg(test)]
    registration_reached: AtomicU64,
    #[cfg(test)]
    registration_released: AtomicBool,
    #[cfg(test)]
    pause_next_publication: AtomicBool,
    #[cfg(test)]
    pause_drainer: AtomicBool,
    #[cfg(test)]
    drainer_paused: AtomicBool,
    #[cfg(test)]
    lag_recoveries: AtomicU64,
    #[cfg(test)]
    pause_next_callback_admission: AtomicBool,
    #[cfg(test)]
    callback_admission_reached: AtomicBool,
    #[cfg(test)]
    callback_admission_released: AtomicBool,
}

impl ProcessExitRegistry {
    pub(super) fn new(
        scheduler: Arc<Scheduler>,
        shutdown_timeout: Duration,
        callback_queue_capacity: usize,
    ) -> Result<Arc<Self>, EngineError> {
        let subscription = scheduler
            .subscribe_exit_events()
            .ok_or(EngineError::ProcessExitSubscriptionUnavailable)?;
        let records = Arc::new(DashMap::new());
        let callbacks = callback::ProcessExitCallbackDispatcher::new(
            shutdown_timeout,
            callback_queue_capacity,
            Arc::downgrade(&records),
        )?;
        let (stopped_sender, stopped) = mpsc::sync_channel(1);
        let registry = Arc::new(Self {
            records,
            registered_through: AtomicU64::new(0),
            has_registered: AtomicBool::new(false),
            lifecycle: Mutex::new(RegistryLifecycle {
                closed: false,
                unobserved_children: HashSet::new(),
            }),
            stop_drainer: AtomicBool::new(false),
            drainer: Mutex::new(ExitDrainer {
                handle: None,
                stopped,
            }),
            callbacks,
            park_bound: DrainBound::new(shutdown_timeout).poll_interval(),
            drain_bound: DrainBound::new(shutdown_timeout),
            drain_progress: DrainProgress::new(),
            #[cfg(test)]
            forced_shutdown_failure: AtomicBool::new(false),
            #[cfg(test)]
            pause_next_registration: AtomicBool::new(false),
            #[cfg(test)]
            registration_reached: AtomicU64::new(0),
            #[cfg(test)]
            registration_released: AtomicBool::new(false),
            #[cfg(test)]
            pause_next_publication: AtomicBool::new(false),
            #[cfg(test)]
            pause_drainer: AtomicBool::new(false),
            #[cfg(test)]
            drainer_paused: AtomicBool::new(false),
            #[cfg(test)]
            lag_recoveries: AtomicU64::new(0),
            #[cfg(test)]
            pause_next_callback_admission: AtomicBool::new(false),
            #[cfg(test)]
            callback_admission_reached: AtomicBool::new(false),
            #[cfg(test)]
            callback_admission_released: AtomicBool::new(false),
        });
        let weak_registry = Arc::downgrade(&registry);
        let handle = std::thread::Builder::new()
            .name(String::from("aion-process-exit-drainer"))
            .spawn(move || {
                let result = drainer::run(&weak_registry, &scheduler, &subscription);
                let _ = stopped_sender.send(());
                result
            })
            .map_err(|error| EngineError::ProcessExitDrainerSpawn {
                reason: error.to_string(),
            })?;
        registry.lock_drainer()?.handle = Some(handle);
        Ok(registry)
    }

    pub(super) fn reserve_spawn(&self) -> Result<ProcessSpawnReservation<'_>, EngineError> {
        let lifecycle = self.lock_lifecycle()?;
        if lifecycle.closed {
            return Err(EngineError::ShuttingDown);
        }
        Ok(ProcessSpawnReservation {
            registry: self,
            lifecycle,
        })
    }

    fn register_locked(
        &self,
        pid: Pid,
        lifecycle: &mut RegistryLifecycle,
    ) -> Result<(), EngineError> {
        lifecycle.unobserved_children.remove(&pid);
        #[cfg(test)]
        self.pause_registration_if_requested(pid);
        #[cfg(test)]
        let pause_publication = self.pause_next_publication.swap(false, Ordering::AcqRel);
        let record = Arc::new(ProcessExitRecord::new(
            pid,
            #[cfg(test)]
            pause_publication,
        ));
        match self.records.entry(pid) {
            dashmap::mapref::entry::Entry::Occupied(_) => {
                return Err(EngineError::Runtime {
                    reason: format!("process {pid} already has a runtime-owned exit record"),
                });
            }
            dashmap::mapref::entry::Entry::Vacant(entry) => {
                entry.insert(record);
            }
        }
        self.registered_through.fetch_max(pid, Ordering::AcqRel);
        self.has_registered.store(true, Ordering::Release);
        Ok(())
    }

    /// Return the owned exit record for `pid`.
    ///
    /// # Errors
    ///
    /// Returns [`EngineError::ProcessExitAlreadyTerminal`] when no record is
    /// held and `pid` is at or below the registration watermark — the
    /// registry forgets a registration when it retires the record, so that is
    /// the only shape a reaped process can still present. Every other absence
    /// is a pid this registry provably never registered, which is
    /// [`EngineError::Runtime`].
    pub(super) fn get(&self, pid: Pid) -> Result<Arc<ProcessExitRecord>, EngineError> {
        self.find(pid).ok_or_else(|| {
            if self.below_registration_watermark(pid) {
                EngineError::ProcessExitAlreadyTerminal { process_id: pid }
            } else {
                EngineError::Runtime {
                    reason: format!("process {pid} has no runtime-owned exit outcome record"),
                }
            }
        })
    }

    pub(super) fn contains(&self, pid: Pid) -> bool {
        self.records.contains_key(&pid)
    }

    pub(super) fn find(&self, pid: Pid) -> Option<Arc<ProcessExitRecord>> {
        self.records
            .get(&pid)
            .map(|record| Arc::clone(record.value()))
    }

    /// Whether no record is held for `pid` while `pid` is at or below the
    /// highest pid this registry has ever registered.
    ///
    /// This is a watermark read, never a statement that the process ended.
    /// The registry forgets a registration when it retires the record, so a
    /// process that was registered and reaped and a pid that was never
    /// registered at all answer this the same way; only a pid above the
    /// watermark is provably unknown to the registry. Callers may use it to
    /// choose which absence to report, and may never use it as a verdict
    /// about an ending — [`Self::ending_knowledge`] is that verdict.
    pub(super) fn below_registration_watermark(&self, pid: Pid) -> bool {
        self.has_registered.load(Ordering::Acquire)
            && pid <= self.registered_through.load(Ordering::Acquire)
            && !self.contains(pid)
    }

    /// What this registry knows about `pid`'s ending.
    ///
    /// # Errors
    ///
    /// Returns [`EngineError::ProcessExitStatePoisoned`] when the held
    /// record's outcome lock is poisoned.
    pub(super) fn ending_knowledge(&self, pid: Pid) -> Result<ProcessEnding, EngineError> {
        if let Some(record) = self.find(pid) {
            return Ok(if record.is_terminal()? {
                ProcessEnding::Recorded
            } else {
                ProcessEnding::Pending
            });
        }
        Ok(if self.below_registration_watermark(pid) {
            ProcessEnding::Forgotten
        } else {
            ProcessEnding::NeverRegistered
        })
    }

    /// Whether this registry holds `pid`'s terminal outcome.
    ///
    /// Only a published terminal on a held record answers true. A missing
    /// record is not an ending: the registry keeps no memory of a
    /// registration once it retires the record, so a reaped process and a pid
    /// that was never registered are indistinguishable by absence alone.
    ///
    /// # Errors
    ///
    /// Returns [`EngineError::ProcessExitStatePoisoned`] when the held
    /// record's outcome lock is poisoned.
    pub(super) fn has_terminal(&self, pid: Pid) -> Result<bool, EngineError> {
        Ok(self.ending_knowledge(pid)? == ProcessEnding::Recorded)
    }

    pub(super) fn is_current(&self, pid: Pid, expected: &Arc<ProcessExitRecord>) -> bool {
        self.records
            .get(&pid)
            .is_some_and(|current| Arc::ptr_eq(current.value(), expected))
    }

    pub(super) fn retire(&self, pid: Pid, expected: &Arc<ProcessExitRecord>) {
        if let dashmap::mapref::entry::Entry::Occupied(entry) = self.records.entry(pid)
            && Arc::ptr_eq(entry.get(), expected)
        {
            entry.remove();
        }
    }

    #[cfg(test)]
    pub(super) fn len(&self) -> usize {
        self.records.len()
    }

    /// Arm the injected drain failure. See `forced_shutdown_failure`.
    #[cfg(test)]
    pub(super) fn force_shutdown_failure(&self) {
        self.forced_shutdown_failure.store(true, Ordering::Release);
    }

    /// Close spawn and callback admission, then snapshot every owned live pid.
    pub(super) fn begin_shutdown(&self) -> Result<Vec<Pid>, EngineError> {
        // 🔴 The injected variant is the one THIS function can actually raise.
        // It first returned `ProcessExitDrainerShutdownTimedOut`, which only
        // `close_and_join_all` can produce — a fault wearing a label its
        // injection point cannot issue, which is a fixture that models a
        // machine that does not exist even when the property it measures
        // survives. `lock_lifecycle` below is the real failure mode here, and
        // `ProcessExitRegistryPoisoned` is what it returns.
        #[cfg(test)]
        if self.forced_shutdown_failure.load(Ordering::Acquire) {
            return Err(EngineError::ProcessExitRegistryPoisoned);
        }
        let mut lifecycle = self.lock_lifecycle()?;
        lifecycle.closed = true;
        let records: Vec<_> = self
            .records
            .iter()
            .map(|record| Arc::clone(record.value()))
            .collect();
        for record in &records {
            record.close_without_monitor()?;
        }
        let mut pids = lifecycle.unobserved_children.clone();
        pids.extend(records.iter().map(|record| record.pid));
        Ok(pids.into_iter().collect())
    }

    pub(super) fn close_and_join_all(&self) -> Result<(), EngineError> {
        let _ = self.begin_shutdown()?;
        self.stop_drainer.store(true, Ordering::Release);
        {
            let mut drainer = self.lock_drainer()?;
            if drainer.handle.is_some() {
                self.drain_bound
                    .wait(
                        "aion-process-exit-drainer",
                        &drainer.stopped,
                        &self.drain_progress,
                    )
                    .map_err(
                        |timed_out| EngineError::ProcessExitDrainerShutdownTimedOut {
                            timeout_millis: timed_out.bound.as_millis(),
                            since_progress_millis: timed_out.since_progress.as_millis(),
                            // The drainer has no queue of its own; what is still
                            // owed is the owned processes not yet terminal.
                            queued: self.records.len(),
                        },
                    )?;
                let handle = drainer
                    .handle
                    .take()
                    .ok_or(EngineError::ProcessExitDrainerPanicked)?;
                handle
                    .join()
                    .map_err(|_| EngineError::ProcessExitDrainerPanicked)??;
            }
        }
        self.callbacks.shutdown()
    }

    pub(super) fn attach_callback(
        &self,
        record: &Arc<ProcessExitRecord>,
        installation: &Arc<MonitorInstallation>,
        callback: ProcessExitCallback,
    ) -> Result<(), EngineError> {
        let lifecycle = self.lock_lifecycle()?;
        if lifecycle.closed {
            return Err(EngineError::ShuttingDown);
        }
        let deferred = record.attach_callback(installation, callback)?;
        #[cfg(test)]
        self.pause_callback_admission_if_requested();
        if let Some((callback, terminal)) = deferred {
            self.dispatch_callback(record, callback, terminal)?;
        }
        drop(lifecycle);
        Ok(())
    }

    fn dispatch_callback(
        &self,
        record: &Arc<ProcessExitRecord>,
        callback: ProcessExitCallback,
        terminal: OwnedProcessExitOutcome,
    ) -> Result<(), EngineError> {
        match self.callbacks.dispatch(callback, terminal) {
            Ok(callback::CallbackDispatch::Submitted) => Ok(()),
            Ok(callback::CallbackDispatch::Deferred(callback)) => record.restore_callback(callback),
            Err(failure) => {
                let callback::CallbackDispatchFailure { callback, error } = *failure;
                record.restore_callback(callback)?;
                Err(error)
            }
        }
    }

    fn publish_record(
        &self,
        record: &Arc<ProcessExitRecord>,
        outcome: OwnedProcessExitOutcome,
    ) -> Result<(), EngineError> {
        if let Some(callback) = record.publish(outcome.clone())? {
            self.dispatch_callback(record, callback, outcome)?;
        }
        Ok(())
    }

    fn process_event(&self, scheduler: &Scheduler, event: ExitEvent) -> Result<(), EngineError> {
        match event {
            ExitEvent::Exited { pid, .. } => match scheduler.take_exit_outcome(pid) {
                Some((reason, result)) => self.publish_taken(scheduler, pid, reason, result),
                None if self.has_terminal(pid)? || !self.is_known(pid)? => Ok(()),
                None => Err(EngineError::ProcessExitOutcomeMissingAfterEvent { process_id: pid }),
            },
            ExitEvent::Lagged => {
                #[cfg(test)]
                self.lag_recoveries.fetch_add(1, Ordering::AcqRel);
                tracing::warn!("process exit event subscriber lagged; resynchronizing outcomes");
                self.resynchronize(scheduler)
            }
        }
    }

    /// Recover every pid classified before the reset while excluding concurrent spawns.
    ///
    /// Wrapped local BEAM spawn BIFs classify children under the same reservation gate,
    /// so lag recovery can take and discard their outcomes. Truly foreign scheduler pids
    /// are consumed on ordinary events but cannot be recovered after an overflow because
    /// beamr exposes no outcome-key enumeration API.
    fn resynchronize(&self, scheduler: &Scheduler) -> Result<(), EngineError> {
        let mut lifecycle = self.lock_lifecycle()?;
        let records: Vec<_> = self
            .records
            .iter()
            .map(|record| Arc::clone(record.value()))
            .collect();
        for record in records {
            if let Some((reason, result)) = scheduler.take_exit_outcome(record.pid) {
                self.publish_record(
                    &record,
                    Self::owned_outcome(scheduler, record.pid, reason, result),
                )?;
            }
        }
        let children: Vec<_> = lifecycle.unobserved_children.iter().copied().collect();
        for pid in children {
            if scheduler.take_exit_outcome(pid).is_some() {
                Self::discard_diagnostics(scheduler, pid);
                lifecycle.unobserved_children.remove(&pid);
            }
        }
        Ok(())
    }

    fn publish_taken(
        &self,
        scheduler: &Scheduler,
        pid: Pid,
        reason: ExitReason,
        result: OwnedTerm,
    ) -> Result<(), EngineError> {
        let mut lifecycle = self.lock_lifecycle()?;
        if let Some(record) = self.find(pid) {
            let outcome = Self::owned_outcome(scheduler, pid, reason, result);
            drop(lifecycle);
            self.publish_record(&record, outcome)
        } else {
            lifecycle.unobserved_children.remove(&pid);
            Self::discard_diagnostics(scheduler, pid);
            drop(result);
            Ok(())
        }
    }

    fn owned_outcome(
        scheduler: &Scheduler,
        pid: Pid,
        reason: ExitReason,
        result: OwnedTerm,
    ) -> OwnedProcessExitOutcome {
        OwnedProcessExitOutcome::Observed(Arc::new(ObservedProcessExit {
            reason,
            result,
            execution_error: scheduler.take_exit_error(pid),
            exception: scheduler.take_exit_exception(pid),
        }))
    }

    fn discard_diagnostics(scheduler: &Scheduler, pid: Pid) {
        drop(scheduler.take_exit_error(pid));
        drop(scheduler.take_exit_exception(pid));
    }

    fn is_known(&self, pid: Pid) -> Result<bool, EngineError> {
        Ok(self.contains(pid) || self.lock_lifecycle()?.unobserved_children.contains(&pid))
    }

    fn fail_unobserved(&self, failure: ProcessExitObservationFailure) {
        let records: Vec<_> = self
            .records
            .iter()
            .map(|record| Arc::clone(record.value()))
            .collect();
        for record in records {
            let outcome = OwnedProcessExitOutcome::ObservationFailed {
                process_id: record.pid,
                failure,
            };
            if let Err(error) = self.publish_record(&record, outcome) {
                tracing::error!(pid = record.pid, %error, "failed to publish exit observation failure");
            }
        }
    }

    fn all_owned_processes_terminal(&self) -> Result<bool, EngineError> {
        for record in self.records.iter() {
            if !record.is_terminal()? {
                return Ok(false);
            }
        }
        Ok(self.lock_lifecycle()?.unobserved_children.is_empty())
    }

    fn lock_lifecycle(&self) -> Result<MutexGuard<'_, RegistryLifecycle>, EngineError> {
        self.lifecycle
            .lock()
            .map_err(|_| EngineError::ProcessExitRegistryPoisoned)
    }

    fn lock_drainer(&self) -> Result<MutexGuard<'_, ExitDrainer>, EngineError> {
        self.drainer
            .lock()
            .map_err(|_| EngineError::ProcessExitDrainerPoisoned)
    }
}

#[cfg(test)]
#[path = "process_exit_test_support.rs"]
mod test_support;

#[cfg(test)]
#[path = "process_exit_tests.rs"]
mod tests;

#[cfg(test)]
#[path = "process_exit_round12_tests.rs"]
mod round12_tests;

#[cfg(test)]
#[path = "process_exit_round13_tests.rs"]
mod round13_tests;

#[cfg(test)]
#[path = "process_exit_round14_tests.rs"]
mod round14_tests;