liminal-server 0.25.4

use std::collections::HashMap;
use std::sync::{Arc, Condvar, Mutex, Weak};
use std::time::Instant;

use beamr::native::native_process::{
    NativeContext, NativeHandler, NativeHandlerFactory, NativeOutcome,
};
use beamr::process::ExitReason;

use super::{SupervisorInner, recover_lock};
use crate::ServerError;

#[derive(Debug, Default)]
pub(super) struct NativeRetirements {
    active: Mutex<HashMap<u64, Arc<Retirement>>>,
    changed: Condvar,
}

#[derive(Debug, Default)]
struct Retirement {
    state: Mutex<RetirementState>,
}

#[derive(Debug, Default)]
struct RetirementState {
    dropped: bool,
    binding: Option<(u64, Weak<NativeRetirements>)>,
}

impl NativeRetirements {
    fn register(
        self: &Arc<Self>,
        pid: u64,
        retirement: &Arc<Retirement>,
    ) -> Result<(), ServerError> {
        let mut state = recover_lock(&retirement.state);
        if state.binding.is_some() {
            return Err(ServerError::ListenerAccept {
                message: format!("connection {pid} registered its native retirement twice"),
            });
        }
        if !state.dropped {
            // Keep the ticket locked until its map entry exists, so a handler
            // drop cannot remove it before registration inserts it.
            state.binding = Some((pid, Arc::downgrade(self)));
            recover_lock(&self.active).insert(pid, Arc::clone(retirement));
        }
        drop(state);
        Ok(())
    }

    fn retired(&self, pid: u64) {
        recover_lock(&self.active).remove(&pid);
        self.changed.notify_all();
    }

    pub(super) fn wait(&self, deadline: Instant) -> bool {
        let (active, result) = {
            let active = recover_lock(&self.active);
            if active.is_empty() {
                return true;
            }
            let Some(remaining) = deadline.checked_duration_since(Instant::now()) else {
                return false;
            };
            self.changed
                .wait_timeout_while(active, remaining, |active| !active.is_empty())
                .unwrap_or_else(std::sync::PoisonError::into_inner)
        };
        let drained = active.is_empty();
        drop(active);
        !result.timed_out() && drained
    }
}

struct RetiringHandler {
    handler: Option<Box<dyn NativeHandler>>,
    retirement: Arc<Retirement>,
}

impl NativeHandler for RetiringHandler {
    fn handle(&mut self, ctx: &mut NativeContext<'_>) -> NativeOutcome {
        if let Some(handler) = self.handler.as_mut() {
            handler.handle(ctx)
        } else {
            tracing::error!(
                connection_pid = ctx.self_pid(),
                "a retired native connection handler was scheduled again"
            );
            NativeOutcome::Stop(ExitReason::Error)
        }
    }
}

impl Drop for RetiringHandler {
    fn drop(&mut self) {
        // The scheduler removes its process-table entry before dropping this
        // body. Release the transport and every handler-owned resource before
        // telling the shutdown waiter that native retirement is complete.
        drop(self.handler.take());
        let binding = {
            let mut state = recover_lock(&self.retirement.state);
            state.dropped = true;
            state.binding.take()
        };
        if let Some((pid, owner)) = binding
            && let Some(owner) = owner.upgrade()
        {
            owner.retired(pid);
        }
    }
}

impl SupervisorInner {
    pub(super) fn spawn_native_connection(
        &self,
        factory: NativeHandlerFactory,
        context: &str,
    ) -> Result<u64, ServerError> {
        let retirement = Arc::new(Retirement::default());
        let handler_retirement = Arc::clone(&retirement);
        let tracked: NativeHandlerFactory = Box::new(move || {
            Box::new(RetiringHandler {
                handler: Some(factory()),
                retirement: Arc::clone(&handler_retirement),
            })
        });
        let pid =
            self.scheduler
                .spawn_native(tracked)
                .map_err(|error| ServerError::ListenerAccept {
                    message: format!("{context}: {error}"),
                })?;
        if let Err(error) = self.retirements.register(pid, &retirement) {
            self.scheduler.terminate_process(pid, ExitReason::Error);
            return Err(error);
        }
        Ok(pid)
    }
}