searchi/crates/searchi-service/src/coordination.rs

//! Shared admission for persistent inference sessions. This bounds admitted
//! native worker threads, not all extraction/ANN/OS work. Already executing
//! inference cannot be preempted; long audio windows remain a latency limit.
use crate::EmbeddingProvider;
use anyhow::{bail, Result};
use searchi_core::metrics::{SpanToken, TraceCollector};
use searchi_semantic::TextWindow;
use std::{
    collections::BTreeMap,
    path::Path,
    sync::{Arc, Condvar, Mutex},
    time::{Duration, Instant},
};

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum InferenceRole {
    Foreground,
    Background,
}
#[derive(Default)]
struct State {
    limit: usize,
    used: usize,
    foreground: usize,
}
struct Shared {
    state: Mutex<State>,
    changed: Condvar,
}
/// Clone this handle across query, setup and indexing sessions.
#[derive(Clone)]
pub struct InferenceBudget(Arc<Shared>);
impl InferenceBudget {
    pub fn new(requested: usize) -> Self {
        Self(Arc::new(Shared {
            state: Mutex::new(State {
                limit: clamp(requested),
                ..State::default()
            }),
            changed: Condvar::new(),
        }))
    }
    pub fn threads(&self) -> usize {
        self.0.state.lock().unwrap_or_else(|e| e.into_inner()).limit
    }
    /// Leave one admitted worker thread for background indexing. On a
    /// single-thread budget, already executing inference still cannot yield.
    pub fn worker_threads(&self, role: InferenceRole) -> usize {
        match role {
            InferenceRole::Foreground => self.threads().saturating_sub(1).max(1),
            InferenceRole::Background => 1,
        }
    }
    /// Update admission immediately; already admitted work is not cancelled.
    /// Recreate sessions whose worker thread count exceeds a reduced limit.
    pub fn configure(&self, requested: usize) -> usize {
        let limit = clamp(requested);
        self.0.state.lock().unwrap_or_else(|e| e.into_inner()).limit = limit;
        self.0.changed.notify_all();
        limit
    }
    /// Register before waiting for a service mutex. This reserves priority, not
    /// threads, so the foreground does not deadlock with its indexing session.
    pub fn pending_foreground(&self) -> PendingForeground {
        self.0
            .state
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .foreground += 1;
        self.0.changed.notify_all();
        PendingForeground(self.clone())
    }
    pub fn acquire(
        &self,
        role: InferenceRole,
        actual_worker_threads: usize,
    ) -> Result<InferencePermit> {
        let weight = actual_worker_threads.max(1);
        let _pending = (role == InferenceRole::Foreground).then(|| self.pending_foreground());
        let deadline = Instant::now() + Duration::from_secs(30);
        let mut state = self.0.state.lock().unwrap_or_else(|e| e.into_inner());
        loop {
            if weight > state.limit {
                bail!("Inference session uses {weight} threads but shared budget is {}; recreate the session after changing its thread count", state.limit);
            }
            if state.used <= state.limit - weight
                && (role == InferenceRole::Foreground || state.foreground == 0)
            {
                state.used += weight;
                return Ok(InferencePermit {
                    budget: self.clone(),
                    weight,
                });
            }
            let remaining = deadline.saturating_duration_since(Instant::now());
            if remaining.is_zero() {
                bail!("Timed out waiting for shared inference budget");
            }
            state = self
                .0
                .changed
                .wait_timeout(state, remaining)
                .unwrap_or_else(|e| e.into_inner())
                .0;
        }
    }
}
fn clamp(requested: usize) -> usize {
    requested
        .clamp(1, 8)
        .min(std::thread::available_parallelism().map_or(1, usize::from))
}
pub struct PendingForeground(InferenceBudget);
impl Drop for PendingForeground {
    fn drop(&mut self) {
        let mut s = self.0 .0.state.lock().unwrap_or_else(|e| e.into_inner());
        s.foreground -= 1;
        self.0 .0.changed.notify_all();
    }
}
pub struct InferencePermit {
    budget: InferenceBudget,
    weight: usize,
}
impl Drop for InferencePermit {
    fn drop(&mut self) {
        let mut s = self
            .budget
            .0
            .state
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        s.used -= self.weight;
        self.budget.0.changed.notify_all();
    }
}

/// Delegates genuine inference unchanged. Background documents yield between
/// individual encodes even when the caller supplies a larger batch.
pub struct BudgetedProvider {
    inner: Box<dyn EmbeddingProvider>,
    budget: InferenceBudget,
    role: InferenceRole,
    threads: usize,
    query_started: bool,
}
impl BudgetedProvider {
    pub fn new(
        inner: Box<dyn EmbeddingProvider>,
        budget: &InferenceBudget,
        role: InferenceRole,
        actual_worker_threads: usize,
    ) -> Self {
        Self {
            inner,
            budget: budget.clone(),
            role,
            threads: actual_worker_threads.max(1),
            query_started: false,
        }
    }
    fn permit(&self) -> Result<InferencePermit> {
        self.budget.acquire(self.role, self.threads)
    }
}
impl EmbeddingProvider for BudgetedProvider {
    fn document_batch_limit(&self) -> usize {
        if self.role == InferenceRole::Background {
            1
        } else {
            self.inner.document_batch_limit()
        }
    }
    fn fingerprint(&self) -> &str {
        self.inner.fingerprint()
    }
    fn dimensions(&self) -> usize {
        self.inner.dimensions()
    }
    fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>> {
        let _p = self.permit()?;
        self.inner.split_windows(text, title)
    }
    fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result<Vec<Vec<f32>>> {
        if self.role == InferenceRole::Background {
            let mut output = Vec::with_capacity(texts.len());
            for text in texts {
                let _p = self.permit()?;
                output.extend(self.inner.embed_documents(&[*text], title)?);
            }
            Ok(output)
        } else {
            let _p = self.permit()?;
            self.inner.embed_documents(texts, title)
        }
    }
    fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
        self.query_started = false;
        let _p = self.permit()?;
        self.query_started = true;
        self.inner.embed_query(query)
    }
    fn embed_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
        self.query_started = false;
        let _p = self.permit()?;
        self.query_started = true;
        self.inner.embed_query_debug(query)
    }
    fn embed_code_query(&mut self, query: &str) -> Result<Vec<f32>> {
        self.query_started = false;
        let _p = self.permit()?;
        self.query_started = true;
        self.inner.embed_code_query(query)
    }
    fn embed_code_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
        self.query_started = false;
        let _p = self.permit()?;
        self.query_started = true;
        self.inner.embed_code_query_debug(query)
    }
    fn query_diagnostics(&self) -> BTreeMap<String, u64> {
        if self.query_started {
            self.inner.query_diagnostics()
        } else {
            Default::default()
        }
    }
    fn embed_query_traced(
        &mut self,
        query: &str,
        debug: bool,
        trace: &mut TraceCollector,
        parent: Option<SpanToken>,
    ) -> Result<Vec<f32>> {
        self.query_started = false;
        let span = trace.begin("semantic.inference.admission_wait", parent);
        let permit = self.permit();
        trace.end(
            span,
            if trace.is_debug() {
                BTreeMap::from([
                    ("worker_threads".into(), self.threads as u64),
                    ("failed".into(), u64::from(permit.is_err())),
                ])
            } else {
                Default::default()
            },
        );
        let _permit = permit?;
        self.query_started = true;
        self.inner.embed_query_traced(query, debug, trace, parent)
    }
    fn embed_code_query_traced(
        &mut self,
        query: &str,
        debug: bool,
        trace: &mut TraceCollector,
        parent: Option<SpanToken>,
    ) -> Result<Vec<f32>> {
        self.query_started = false;
        let span = trace.begin("semantic.inference.admission_wait", parent);
        let permit = self.permit();
        trace.end(
            span,
            if trace.is_debug() {
                BTreeMap::from([
                    ("worker_threads".into(), self.threads as u64),
                    ("failed".into(), u64::from(permit.is_err())),
                ])
            } else {
                Default::default()
            },
        );
        let _permit = permit?;
        self.query_started = true;
        self.inner
            .embed_code_query_traced(query, debug, trace, parent)
    }
    fn embed_images(&mut self, paths: &[&Path], budget: u32) -> Result<Vec<Vec<f32>>> {
        if self.role == InferenceRole::Background {
            let mut output = Vec::with_capacity(paths.len());
            for path in paths {
                let _p = self.permit()?;
                output.extend(self.inner.embed_images(&[*path], budget)?);
            }
            return Ok(output);
        }
        let _p = self.permit()?;
        self.inner.embed_images(paths, budget)
    }
    fn embed_audio(&mut self, pcm: &[f32]) -> Result<Vec<f32>> {
        let _p = self.permit()?;
        self.inner.embed_audio(pcm)
    }
    fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
        self.inner.worker_info()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn failed_admission_clears_prior_worker_diagnostics_without_calling_worker() {
        use searchi_core::metrics::MetricsMode;
        use std::sync::atomic::{AtomicUsize, Ordering};
        struct ObservedInvocation(Arc<AtomicUsize>);
        impl EmbeddingProvider for ObservedInvocation {
            fn fingerprint(&self) -> &str {
                "scheduler-test-only"
            }
            fn dimensions(&self) -> usize {
                0
            }
            fn split_windows(&mut self, _: &str, _: &str) -> Result<Vec<TextWindow>> {
                unreachable!()
            }
            fn embed_documents(&mut self, _: &[&str], _: &str) -> Result<Vec<Vec<f32>>> {
                unreachable!()
            }
            fn embed_query(&mut self, _: &str) -> Result<Vec<f32>> {
                self.0.fetch_add(1, Ordering::SeqCst);
                Ok(Vec::new())
            }
            fn embed_code_query(&mut self, query: &str) -> Result<Vec<f32>> {
                self.embed_query(query)
            }
            fn query_diagnostics(&self) -> BTreeMap<String, u64> {
                BTreeMap::from([("prior_worker_call".into(), 1)])
            }
        }
        let budget = InferenceBudget::new(1);
        // Test worker weight two without requiring particular test host hardware.
        budget.0.state.lock().unwrap().limit = 2;
        let calls = Arc::new(AtomicUsize::new(0));
        let mut provider = BudgetedProvider::new(
            Box::new(ObservedInvocation(calls.clone())),
            &budget,
            InferenceRole::Foreground,
            2,
        );
        provider
            .embed_query_debug("first admitted invocation")
            .unwrap();
        assert_eq!(calls.load(Ordering::SeqCst), 1);
        assert_eq!(provider.query_diagnostics()["prior_worker_call"], 1);
        assert_eq!(budget.configure(1), 1);
        let mut trace = TraceCollector::new(MetricsMode::Debug);
        assert!(provider
            .embed_query_traced("rejected invocation", true, &mut trace, None)
            .is_err());
        assert_eq!(calls.load(Ordering::SeqCst), 1);
        assert!(provider.query_diagnostics().is_empty());
        let metrics = trace.finish(0);
        let stages = metrics.trace.unwrap();
        assert_eq!(stages.len(), 1);
        assert_eq!(stages[0].name, "semantic.inference.admission_wait");
        assert_eq!(stages[0].counters["failed"], 1);
        assert_eq!(stages[0].counters["worker_threads"], 2);
        assert!(stages[0].duration_ms >= 0.0);
        assert!(provider.embed_code_query_debug("also rejected").is_err());
        assert!(provider.query_diagnostics().is_empty());
        assert_eq!(calls.load(Ordering::SeqCst), 1);
        assert_eq!(budget.0.state.lock().unwrap().foreground, 0);
        assert_eq!(budget.0.state.lock().unwrap().used, 0);
    }
    #[test]
    fn reserved_query_capacity_admits_foreground_during_background() {
        let budget = InferenceBudget::new(1);
        // Deliberately exercise a two-thread scheduler independent of CI hardware.
        budget.0.state.lock().unwrap().limit = 2;
        assert_eq!(budget.worker_threads(InferenceRole::Foreground), 1);
        assert_eq!(budget.worker_threads(InferenceRole::Background), 1);
        let background = budget.acquire(InferenceRole::Background, 1).unwrap();
        let foreground = budget.acquire(InferenceRole::Foreground, 1).unwrap();
        assert_eq!(budget.0.state.lock().unwrap().used, 2);
        drop(foreground);
        drop(background);
    }
    #[test]
    fn pending_query_blocks_background_until_released() {
        let budget = InferenceBudget::new(1);
        let pending = budget.pending_foreground();
        let other = budget.clone();
        let (tx, rx) = std::sync::mpsc::channel();
        let worker = std::thread::spawn(move || {
            let _permit = other.acquire(InferenceRole::Background, 1).unwrap();
            tx.send(()).unwrap();
        });
        assert!(rx.recv_timeout(Duration::from_millis(50)).is_err());
        let foreground = budget.acquire(InferenceRole::Foreground, 1).unwrap();
        drop(pending);
        assert!(rx.recv_timeout(Duration::from_millis(50)).is_err());
        drop(foreground);
        rx.recv_timeout(Duration::from_secs(2)).unwrap();
        worker.join().unwrap();
    }
    #[test]
    fn panic_releases_threads_and_priority() {
        let budget = InferenceBudget::new(1);
        let shared = budget.clone();
        assert!(std::panic::catch_unwind(move || {
            let _pending = shared.pending_foreground();
            let _permit = shared.acquire(InferenceRole::Foreground, 1).unwrap();
            panic!("owned inference failure");
        })
        .is_err());
        let _permit = budget.acquire(InferenceRole::Background, 1).unwrap();
    }
    #[test]
    fn invalid_session_weight_fails_instead_of_oversubscribing() {
        let budget = InferenceBudget::new(1);
        assert!(budget.acquire(InferenceRole::Foreground, 2).is_err());
        assert_eq!(budget.configure(0), 1);
        assert!(InferenceBudget::new(usize::MAX).threads() <= 8);
    }
}