liminal-server 0.23.1

use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, Weak};
use std::thread::{self, JoinHandle};

use beamr::native::native_process::{NativeContext, NativeHandler, NativeOutcome};
use beamr::process::ExitReason;
use beamr::scheduler::{ExitEvent, ExitEventSubscription, Scheduler};

use super::ConnectionRuntime;
use crate::ServerError;

#[derive(Debug)]
struct Worker {
    thread: Option<JoinHandle<()>>,
    failure: Option<String>,
}

#[derive(Debug)]
pub(super) struct ReclaimReactor {
    stop_pid: u64,
    stopping: Arc<AtomicBool>,
    worker: Mutex<Worker>,
}

impl ReclaimReactor {
    pub(super) fn start(
        scheduler: &Arc<Scheduler>,
        runtime: &Arc<ConnectionRuntime>,
    ) -> Result<Self, ServerError> {
        let subscription = scheduler
            .subscribe_exit_events()
            .ok_or_else(|| refusal("connection scheduler exit-event subscription unavailable"))?;
        // The subscription parks on exit events. This store-free process gives
        // its owner an explicit exit event even when every connection is gone.
        let stop_pid = scheduler
            .spawn_native(Box::new(|| Box::new(StopSource)))
            .map_err(|error| {
                refusal(format!("reclamation stop source could not start: {error}"))
            })?;
        let stopping = Arc::new(AtomicBool::new(false));
        let stop = Arc::clone(&stopping);
        let reactor_scheduler = Arc::downgrade(scheduler);
        let reactor_runtime = Arc::downgrade(runtime);
        let thread = thread::Builder::new()
            .name("liminal-connection-reclaim".to_owned())
            .spawn(move || {
                run(
                    &subscription,
                    &reactor_scheduler,
                    &reactor_runtime,
                    &stop,
                    Some(stop_pid),
                );
            })
            .map_err(|error| {
                scheduler.terminate_process(stop_pid, ExitReason::Normal);
                refusal(format!(
                    "failed to start connection reclamation reactor: {error}"
                ))
            })?;
        Ok(Self {
            stop_pid,
            stopping,
            worker: Mutex::new(Worker {
                thread: Some(thread),
                failure: None,
            }),
        })
    }

    pub(super) fn finish(&self, scheduler: &Scheduler) -> Result<(), ServerError> {
        let mut worker = self
            .worker
            .lock()
            .map_err(|_| refusal("connection reclamation join lock poisoned"))?;
        if let Some(failure) = &worker.failure {
            return Err(refusal(failure.clone()));
        }
        let Some(thread) = worker.thread.take() else {
            return Ok(());
        };
        self.stopping.store(true, Ordering::Release);
        scheduler.terminate_process(self.stop_pid, ExitReason::Normal);
        if thread.join().is_err() {
            let message = "connection reclamation reactor panicked during shutdown";
            worker.failure = Some(message.to_owned());
            return Err(refusal(message));
        }
        drop(worker);
        Ok(())
    }
}

struct StopSource;

impl NativeHandler for StopSource {
    fn handle(&mut self, _: &mut NativeContext<'_>) -> NativeOutcome {
        NativeOutcome::Wait
    }
}

fn run(
    subscription: &ExitEventSubscription,
    scheduler: &Weak<Scheduler>,
    runtime: &Weak<ConnectionRuntime>,
    stopping: &AtomicBool,
    stop_pid: Option<u64>,
) {
    loop {
        let event = match subscription.recv() {
            Ok(event) => event,
            Err(error) => {
                tracing::warn!(%error, "connection reclamation reactor exit source disconnected");
                return;
            }
        };
        if stopping.load(Ordering::Acquire) {
            return;
        }
        if matches!(event, ExitEvent::Exited { pid, .. } if Some(pid) == stop_pid) {
            tracing::warn!("connection reclamation reactor stop source exited before shutdown");
            return;
        }
        match event {
            ExitEvent::Exited { pid, reason } => {
                let Some(runtime) = runtime.upgrade() else {
                    tracing::warn!("connection reclamation reactor stopped: its runtime is gone");
                    return;
                };
                runtime.deliver_reclamation(scheduler, pid, reason);
            }
            ExitEvent::Lagged => {
                let (Some(runtime), Some(scheduler)) = (runtime.upgrade(), scheduler.upgrade())
                else {
                    tracing::warn!("connection reclamation reactor stopped: its runtime is gone");
                    return;
                };
                runtime.reap_crashed(&scheduler);
            }
        }
    }
}

#[cfg(test)]
pub(super) fn run_for_test(
    subscription: &ExitEventSubscription,
    scheduler: &Weak<Scheduler>,
    runtime: &Weak<ConnectionRuntime>,
) {
    run(
        subscription,
        scheduler,
        runtime,
        &AtomicBool::new(false),
        None,
    );
}

fn refusal(message: impl Into<String>) -> ServerError {
    ServerError::ListenerAccept {
        message: message.into(),
    }
}