car_server_core/

inference_control.rs

//! Per-WebSocket inference lifecycle control.
//!
//! The registry is session-scoped: an ID minted for one socket is
//! indistinguishable from a never-issued ID on every other socket. Active,
//! pending-control, orphan, and terminal records are bounded. A backend task is
//! owned until it has actually returned, so an unconfirmed response cannot
//! detach work or release its admission charge.

use car_proto::InferenceControlStatus;
use std::collections::{HashMap, VecDeque};
use std::future::Future;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tokio::sync::{watch, OwnedSemaphorePermit, Semaphore};

pub const DEFAULT_MAX_ACTIVE: usize = 64;
pub const DEFAULT_MAX_TOMBSTONES: usize = 256;
pub const DEFAULT_MAX_PENDING_CONTROLS: usize = 64;
pub const DEFAULT_MAX_ORPHANS: usize = 64;
pub const DEFAULT_TOMBSTONE_TTL: Duration = Duration::from_secs(300);
pub const DEFAULT_TERMINATION_ACK_TIMEOUT: Duration = Duration::from_secs(5);
pub const MAX_DEADLINE_TIMEOUT_MS: u64 = 600_000;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ControlCause {
    Cancel,
    Deadline,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RegistryError {
    ActiveLimitReached,
    OrphanLimitReached,
    PendingControlLimitReached,
    DuplicateInferenceId,
    InferenceNotActive,
    DeadlineAlreadyScheduled,
}

#[derive(Debug, Clone, Copy)]
pub struct RegistryConfig {
    pub max_active: usize,
    pub max_tombstones: usize,
    pub max_pending_controls: usize,
    pub max_orphans: usize,
    pub tombstone_ttl: Duration,
    pub termination_ack_timeout: Duration,
}

impl Default for RegistryConfig {
    fn default() -> Self {
        Self {
            max_active: DEFAULT_MAX_ACTIVE,
            max_tombstones: DEFAULT_MAX_TOMBSTONES,
            max_pending_controls: DEFAULT_MAX_PENDING_CONTROLS,
            max_orphans: DEFAULT_MAX_ORPHANS,
            tombstone_ttl: DEFAULT_TOMBSTONE_TTL,
            termination_ack_timeout: DEFAULT_TERMINATION_ACK_TIMEOUT,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TerminalOutcome {
    Completed,
    Controlled(InferenceControlStatus),
}

struct ActiveEntry {
    run_id: Option<String>,
    request_id: Option<String>,
    phase: ActivePhase,
    terminal: watch::Sender<Option<TerminalOutcome>>,
    deadline_token: Option<u64>,
    backend_task: Option<tokio::task::JoinHandle<()>>,
    backend_terminated: bool,
}

#[derive(Clone)]
enum ActivePhase {
    Running,
    Terminating(u64),
}

struct Tombstone {
    at: Instant,
    outcome: TerminalOutcome,
    /// Retained until the backend wrapper observes the provider/runner future
    /// return. Keeping the JoinHandle makes the task abort-on-session-drop.
    backend_task: Option<tokio::task::JoinHandle<()>>,
}

#[derive(Default)]
struct RegistryState {
    active: HashMap<String, ActiveEntry>,
    tombstones: HashMap<String, Tombstone>,
    tombstone_order: VecDeque<String>,
    next_token: u64,
    orphan_count: usize,
}

enum ControlStart {
    New(u64),
    Wait(watch::Receiver<Option<TerminalOutcome>>),
    Return(InferenceControlStatus),
}

/// Owns a claimed `Running -> Terminating(token)` transition until it reaches
/// one terminal outcome. Cancellation and panic drop the control future, so a
/// trailing `finish` call is not sufficient: without this guard the entry
/// remains `Terminating` and every later control waits forever on its watch
/// channel. Drop publishes the cause-specific unconfirmed outcome under the
/// exact token; a stale guard can never overwrite another terminal.
struct ControlClaimGuard<'a> {
    registry: &'a InferenceRegistry,
    inference_id: &'a str,
    token: u64,
    abandoned_status: InferenceControlStatus,
    armed: bool,
}

impl<'a> ControlClaimGuard<'a> {
    fn new(
        registry: &'a InferenceRegistry,
        inference_id: &'a str,
        token: u64,
        cause: ControlCause,
    ) -> Self {
        let abandoned_status = match cause {
            ControlCause::Cancel => InferenceControlStatus::TerminationUnconfirmed,
            ControlCause::Deadline => InferenceControlStatus::DeadlineExceededUnconfirmed,
        };
        Self {
            registry,
            inference_id,
            token,
            abandoned_status,
            armed: true,
        }
    }

    fn finish(mut self, status: InferenceControlStatus) -> InferenceControlStatus {
        let won = self.registry.finish(
            self.inference_id,
            Some(self.token),
            TerminalOutcome::Controlled(status),
        );
        self.armed = false;
        if won {
            status
        } else {
            InferenceControlStatus::AlreadyTerminal
        }
    }
}

impl Drop for ControlClaimGuard<'_> {
    fn drop(&mut self) {
        if self.armed {
            let _ = self.registry.finish(
                self.inference_id,
                Some(self.token),
                TerminalOutcome::Controlled(self.abandoned_status),
            );
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DeadlineReservationError {
    Terminal(InferenceControlStatus),
    Registry(RegistryError),
}

pub struct DeadlineReservation {
    registry: Arc<InferenceRegistry>,
    inference_id: String,
    token: u64,
    terminal: watch::Receiver<Option<TerminalOutcome>>,
    armed: bool,
}

impl DeadlineReservation {
    /// Wait relative to receipt of the deadline RPC. Completion or another
    /// control wakes the timer early. The acknowledgement closure is not
    /// constructed or polled unless this exact reservation wins the race.
    pub async fn wait_and_control<F, Fut>(
        mut self,
        timeout: Duration,
        acknowledgement: F,
    ) -> Result<InferenceControlStatus, RegistryError>
    where
        F: FnOnce() -> Fut,
        Fut: Future<Output = bool>,
    {
        if self.terminal.borrow().is_none() {
            tokio::select! {
                _ = tokio::time::sleep(timeout) => {}
                _ = self.terminal.changed() => {
                    return Ok(InferenceControlStatus::AlreadyTerminal);
                }
            }
        } else {
            return Ok(InferenceControlStatus::AlreadyTerminal);
        }

        // A deadline reservation bounds one timer per active inference, but a
        // sleeping timer is not a termination operation. Admit it against the
        // shared control cap only once it actually expires, leaving capacity
        // for an urgent cancel while many long deadlines are armed.
        let _pending = self.registry.try_acquire_control()?;
        let start = self
            .registry
            .claim_reserved_deadline(&self.inference_id, self.token);
        self.armed = false;
        Ok(self
            .registry
            .drive_control_start(
                &self.inference_id,
                ControlCause::Deadline,
                start,
                acknowledgement,
            )
            .await)
    }
}

impl Drop for DeadlineReservation {
    fn drop(&mut self) {
        if self.armed {
            self.registry
                .release_deadline(&self.inference_id, self.token);
        }
    }
}

pub struct InferenceRegistry {
    config: RegistryConfig,
    pending_controls: Arc<Semaphore>,
    state: Mutex<RegistryState>,
}

impl Default for InferenceRegistry {
    fn default() -> Self {
        Self::with_config(RegistryConfig::default())
    }
}

impl InferenceRegistry {
    pub fn with_config(config: RegistryConfig) -> Self {
        assert!(
            config.max_active > 0,
            "active registry bound must be positive"
        );
        assert!(
            config.max_tombstones > 0,
            "tombstone bound must be positive"
        );
        assert!(
            config.max_pending_controls > 0,
            "pending-control bound must be positive"
        );
        assert!(config.max_orphans > 0, "orphan bound must be positive");
        assert!(
            config.max_orphans >= config.max_active,
            "orphan bound must cover every already-active inference"
        );
        assert!(
            config.max_orphans <= config.max_tombstones,
            "orphan bound must fit inside tombstone bound"
        );
        Self {
            pending_controls: Arc::new(Semaphore::new(config.max_pending_controls)),
            config,
            state: Mutex::new(RegistryState::default()),
        }
    }

    pub fn try_acquire_control(self: &Arc<Self>) -> Result<OwnedSemaphorePermit, RegistryError> {
        self.pending_controls
            .clone()
            .try_acquire_owned()
            .map_err(|_| RegistryError::PendingControlLimitReached)
    }

    pub fn begin(
        &self,
    ) -> Result<(String, watch::Receiver<Option<TerminalOutcome>>), RegistryError> {
        self.begin_for_run(None, None)
    }

    pub fn begin_for_run(
        &self,
        run_id: Option<String>,
        request_id: Option<String>,
    ) -> Result<(String, watch::Receiver<Option<TerminalOutcome>>), RegistryError> {
        for _ in 0..8 {
            let id = format!("inf_{}", uuid::Uuid::new_v4().simple());
            match self.begin_with_id_for_run(&id, run_id.clone(), request_id.clone()) {
                Ok(receiver) => return Ok((id, receiver)),
                Err(RegistryError::DuplicateInferenceId) => continue,
                Err(error) => return Err(error),
            }
        }
        Err(RegistryError::DuplicateInferenceId)
    }

    pub fn begin_with_id(
        &self,
        inference_id: &str,
    ) -> Result<watch::Receiver<Option<TerminalOutcome>>, RegistryError> {
        self.begin_with_id_for_run(inference_id, None, None)
    }

    fn begin_with_id_for_run(
        &self,
        inference_id: &str,
        run_id: Option<String>,
        request_id: Option<String>,
    ) -> Result<watch::Receiver<Option<TerminalOutcome>>, RegistryError> {
        let mut state = self.state.lock().expect("inference registry poisoned");
        self.prune_locked(&mut state, Instant::now());
        if state.active.contains_key(inference_id) || state.tombstones.contains_key(inference_id) {
            return Err(RegistryError::DuplicateInferenceId);
        }
        if state.active.len() >= self.config.max_active {
            return Err(RegistryError::ActiveLimitReached);
        }
        // Every admitted active inference may terminalize before its backend
        // wrapper returns, converting its owned JoinHandle into an orphaned,
        // non-evictable tombstone. Treat the active entry itself as that
        // worst-case capacity reservation. `finish` removes one active before
        // adding at most one orphan, so this sum cannot grow after admission.
        if state.active.len().saturating_add(state.orphan_count) >= self.config.max_orphans {
            return Err(RegistryError::OrphanLimitReached);
        }
        let (terminal, receiver) = watch::channel(None);
        state.active.insert(
            inference_id.to_string(),
            ActiveEntry {
                run_id,
                request_id,
                phase: ActivePhase::Running,
                terminal,
                deadline_token: None,
                backend_task: None,
                backend_terminated: false,
            },
        );
        Ok(receiver)
    }

    /// Active inference identities attributed to one run, sorted for stable
    /// receipt selection.
    pub fn active_for_run(&self, run_id: &str) -> Vec<(String, Option<String>)> {
        let mut state = self.state.lock().expect("inference registry poisoned");
        self.prune_locked(&mut state, Instant::now());
        let mut active: Vec<_> = state
            .active
            .iter()
            .filter(|(_, entry)| entry.run_id.as_deref() == Some(run_id))
            .map(|(inference_id, entry)| (inference_id.clone(), entry.request_id.clone()))
            .collect();
        active.sort();
        active
    }

    pub fn attach_backend(
        &self,
        inference_id: &str,
        task: tokio::task::JoinHandle<()>,
    ) -> Result<(), RegistryError> {
        let mut state = self.state.lock().expect("inference registry poisoned");
        let Some(entry) = state.active.get_mut(inference_id) else {
            task.abort();
            return Err(RegistryError::InferenceNotActive);
        };
        if entry.backend_terminated {
            drop(task);
        } else {
            entry.backend_task = Some(task);
        }
        Ok(())
    }

    /// Called only after the actual backend future has returned.
    pub fn backend_terminated(&self, inference_id: &str) {
        let mut state = self.state.lock().expect("inference registry poisoned");
        if let Some(entry) = state.active.get_mut(inference_id) {
            entry.backend_terminated = true;
            entry.backend_task.take();
            return;
        }
        if let Some(tombstone) = state.tombstones.get_mut(inference_id) {
            if tombstone.backend_task.take().is_some() {
                state.orphan_count = state.orphan_count.saturating_sub(1);
            }
        }
        self.prune_locked(&mut state, Instant::now());
    }

    pub fn is_active(&self, inference_id: &str) -> bool {
        let mut state = self.state.lock().expect("inference registry poisoned");
        self.prune_locked(&mut state, Instant::now());
        state.active.contains_key(inference_id)
    }

    pub fn complete(&self, inference_id: &str) -> bool {
        self.finish(inference_id, None, TerminalOutcome::Completed)
    }

    pub fn terminal_outcome(&self, inference_id: &str) -> Option<TerminalOutcome> {
        let mut state = self.state.lock().expect("inference registry poisoned");
        self.prune_locked(&mut state, Instant::now());
        state.tombstones.get(inference_id).map(|t| t.outcome)
    }

    pub fn reserve_deadline(
        self: &Arc<Self>,
        inference_id: &str,
    ) -> Result<DeadlineReservation, DeadlineReservationError> {
        let (token, terminal) = {
            let mut state = self.state.lock().expect("inference registry poisoned");
            self.prune_locked(&mut state, Instant::now());
            if state.tombstones.contains_key(inference_id) {
                return Err(DeadlineReservationError::Terminal(
                    InferenceControlStatus::AlreadyTerminal,
                ));
            }
            let Some(entry) = state.active.get(inference_id) else {
                return Err(DeadlineReservationError::Terminal(
                    InferenceControlStatus::Unknown,
                ));
            };
            if entry.deadline_token.is_some() {
                return Err(DeadlineReservationError::Registry(
                    RegistryError::DeadlineAlreadyScheduled,
                ));
            }
            state.next_token = state.next_token.wrapping_add(1);
            let token = state.next_token;
            let entry = state
                .active
                .get_mut(inference_id)
                .expect("entry inspected above");
            entry.deadline_token = Some(token);
            (token, entry.terminal.subscribe())
        };
        Ok(DeadlineReservation {
            registry: self.clone(),
            inference_id: inference_id.to_string(),
            token,
            terminal,
            armed: true,
        })
    }

    pub async fn schedule_deadline<F, Fut>(
        self: &Arc<Self>,
        inference_id: &str,
        timeout: Duration,
        acknowledgement: F,
    ) -> Result<InferenceControlStatus, DeadlineReservationError>
    where
        F: FnOnce() -> Fut,
        Fut: Future<Output = bool>,
    {
        self.reserve_deadline(inference_id)?
            .wait_and_control(timeout, acknowledgement)
            .await
            .map_err(DeadlineReservationError::Registry)
    }

    pub async fn control<F>(
        &self,
        inference_id: &str,
        cause: ControlCause,
        acknowledgement: F,
    ) -> InferenceControlStatus
    where
        F: Future<Output = bool>,
    {
        let start = self.begin_control(inference_id);
        self.drive_control_start(inference_id, cause, start, || acknowledgement)
            .await
    }

    pub fn counts(&self) -> (usize, usize) {
        let mut state = self.state.lock().expect("inference registry poisoned");
        self.prune_locked(&mut state, Instant::now());
        (state.active.len(), state.tombstones.len())
    }

    pub fn orphan_count(&self) -> usize {
        self.state
            .lock()
            .expect("inference registry poisoned")
            .orphan_count
    }

    /// Disconnect teardown: every retained active/orphan handle is aborted.
    pub fn abort_all(&self) {
        let mut state = self.state.lock().expect("inference registry poisoned");
        for (_, mut entry) in state.active.drain() {
            if let Some(task) = entry.backend_task.take() {
                task.abort();
            }
        }
        for (_, mut tombstone) in state.tombstones.drain() {
            if let Some(task) = tombstone.backend_task.take() {
                task.abort();
            }
        }
        state.tombstone_order.clear();
        state.orphan_count = 0;
    }

    /// Roll back a start whose `infer.started` frame could not be delivered.
    pub fn abandon(&self, inference_id: &str) {
        let mut state = self.state.lock().expect("inference registry poisoned");
        if let Some(mut entry) = state.active.remove(inference_id) {
            if let Some(task) = entry.backend_task.take() {
                task.abort();
            }
        }
    }

    /// Abort a backend after `infer.started` may already have entered the
    /// socket sink, retaining a terminal tombstone so a later flush can never
    /// expose an ID that control RPCs report as unknown.
    pub fn abort_after_started_admission(&self, inference_id: &str) {
        self.abort_owned_backend(inference_id);
    }

    /// Abort this handler's owned backend while preserving an existing
    /// terminal outcome. If no control has terminalized the entry yet, publish
    /// a completed tombstone first so any frame already admitted to the socket
    /// can never expose an unknown ID. If control already moved the JoinHandle
    /// into a tombstone, detach and abort it there and release its orphan charge.
    pub fn abort_owned_backend(&self, inference_id: &str) {
        let mut state = self.state.lock().expect("inference registry poisoned");
        self.prune_locked(&mut state, Instant::now());
        if let Some(mut entry) = state.active.remove(inference_id) {
            let outcome = TerminalOutcome::Completed;
            let _ = entry.terminal.send(Some(outcome));
            if let Some(task) = entry.backend_task.take() {
                task.abort();
            }
            self.insert_tombstone_locked(&mut state, inference_id.to_string(), outcome, None);
            return;
        }
        let Some(tombstone) = state.tombstones.get_mut(inference_id) else {
            return;
        };
        if let Some(task) = tombstone.backend_task.take() {
            task.abort();
            state.orphan_count = state.orphan_count.saturating_sub(1);
        }
    }

    fn begin_control(&self, inference_id: &str) -> ControlStart {
        let mut state = self.state.lock().expect("inference registry poisoned");
        self.prune_locked(&mut state, Instant::now());
        if state.tombstones.contains_key(inference_id) {
            return ControlStart::Return(InferenceControlStatus::AlreadyTerminal);
        }
        let Some(phase) = state
            .active
            .get(inference_id)
            .map(|entry| entry.phase.clone())
        else {
            return ControlStart::Return(InferenceControlStatus::Unknown);
        };
        match phase {
            ActivePhase::Running => {
                state.next_token = state.next_token.wrapping_add(1);
                let token = state.next_token;
                state
                    .active
                    .get_mut(inference_id)
                    .expect("entry inspected above")
                    .phase = ActivePhase::Terminating(token);
                ControlStart::New(token)
            }
            ActivePhase::Terminating(_) => ControlStart::Wait(
                state
                    .active
                    .get(inference_id)
                    .expect("entry inspected above")
                    .terminal
                    .subscribe(),
            ),
        }
    }

    fn claim_reserved_deadline(&self, inference_id: &str, deadline_token: u64) -> ControlStart {
        let mut state = self.state.lock().expect("inference registry poisoned");
        self.prune_locked(&mut state, Instant::now());
        if state.tombstones.contains_key(inference_id) {
            return ControlStart::Return(InferenceControlStatus::AlreadyTerminal);
        }
        let Some(entry) = state.active.get_mut(inference_id) else {
            return ControlStart::Return(InferenceControlStatus::Unknown);
        };
        if entry.deadline_token != Some(deadline_token) {
            return ControlStart::Return(InferenceControlStatus::AlreadyTerminal);
        }
        entry.deadline_token = None;
        match entry.phase.clone() {
            ActivePhase::Running => {
                state.next_token = state.next_token.wrapping_add(1);
                let token = state.next_token;
                state
                    .active
                    .get_mut(inference_id)
                    .expect("entry inspected above")
                    .phase = ActivePhase::Terminating(token);
                ControlStart::New(token)
            }
            ActivePhase::Terminating(_) => ControlStart::Wait(entry.terminal.subscribe()),
        }
    }

    async fn drive_control_start<F, Fut>(
        &self,
        inference_id: &str,
        cause: ControlCause,
        start: ControlStart,
        acknowledgement: F,
    ) -> InferenceControlStatus
    where
        F: FnOnce() -> Fut,
        Fut: Future<Output = bool>,
    {
        match start {
            ControlStart::Return(status) => status,
            ControlStart::Wait(mut receiver) => {
                if receiver.borrow().is_none() {
                    let _ = receiver.changed().await;
                }
                InferenceControlStatus::AlreadyTerminal
            }
            ControlStart::New(token) => {
                let claim = ControlClaimGuard::new(self, inference_id, token, cause);
                let confirmed =
                    tokio::time::timeout(self.config.termination_ack_timeout, acknowledgement())
                        .await
                        .unwrap_or(false);
                let status = match (cause, confirmed) {
                    (ControlCause::Cancel, true) => InferenceControlStatus::CancelledConfirmed,
                    (ControlCause::Cancel, false) => InferenceControlStatus::TerminationUnconfirmed,
                    (ControlCause::Deadline, true) => {
                        InferenceControlStatus::DeadlineExceededConfirmed
                    }
                    (ControlCause::Deadline, false) => {
                        InferenceControlStatus::DeadlineExceededUnconfirmed
                    }
                };
                claim.finish(status)
            }
        }
    }

    fn release_deadline(&self, inference_id: &str, token: u64) {
        let mut state = self.state.lock().expect("inference registry poisoned");
        if let Some(entry) = state.active.get_mut(inference_id) {
            if entry.deadline_token == Some(token) {
                entry.deadline_token = None;
            }
        }
    }

    fn finish(
        &self,
        inference_id: &str,
        expected_token: Option<u64>,
        outcome: TerminalOutcome,
    ) -> bool {
        let mut state = self.state.lock().expect("inference registry poisoned");
        self.prune_locked(&mut state, Instant::now());
        let Some(entry) = state.active.get(inference_id) else {
            return false;
        };
        if let Some(expected) = expected_token {
            if !matches!(entry.phase, ActivePhase::Terminating(actual) if actual == expected) {
                return false;
            }
        } else if !matches!(entry.phase, ActivePhase::Running) {
            return false;
        }
        let mut entry = state
            .active
            .remove(inference_id)
            .expect("entry inspected above");
        let _ = entry.terminal.send(Some(outcome));
        let backend_task = if entry.backend_terminated {
            None
        } else {
            entry.backend_task.take()
        };
        if backend_task.is_some() {
            state.orphan_count += 1;
        }
        self.insert_tombstone_locked(&mut state, inference_id.to_string(), outcome, backend_task);
        true
    }

    fn insert_tombstone_locked(
        &self,
        state: &mut RegistryState,
        inference_id: String,
        outcome: TerminalOutcome,
        backend_task: Option<tokio::task::JoinHandle<()>>,
    ) {
        state.tombstones.insert(
            inference_id.clone(),
            Tombstone {
                at: Instant::now(),
                outcome,
                backend_task,
            },
        );
        state.tombstone_order.push_back(inference_id);
        self.enforce_tombstone_bound_locked(state);
    }

    fn enforce_tombstone_bound_locked(&self, state: &mut RegistryState) {
        while state.tombstones.len() > self.config.max_tombstones {
            let Some(index) = state.tombstone_order.iter().position(|id| {
                state
                    .tombstones
                    .get(id)
                    .is_none_or(|entry| entry.backend_task.is_none())
            }) else {
                break;
            };
            let id = state
                .tombstone_order
                .remove(index)
                .expect("index inspected above");
            state.tombstones.remove(&id);
        }
    }

    fn prune_locked(&self, state: &mut RegistryState, now: Instant) {
        let mut retained = VecDeque::with_capacity(state.tombstone_order.len());
        while let Some(id) = state.tombstone_order.pop_front() {
            let remove = state.tombstones.get(&id).is_none_or(|entry| {
                entry.backend_task.is_none()
                    && now.duration_since(entry.at) >= self.config.tombstone_ttl
            });
            if remove {
                state.tombstones.remove(&id);
            } else {
                retained.push_back(id);
            }
        }
        state.tombstone_order = retained;
        self.enforce_tombstone_bound_locked(state);
    }
}