leindex 1.9.0

LeIndex MCP and semantic code search engine for AI tools and large codebases
#![allow(missing_docs)]

use serde::{Deserialize, Serialize};
use std::future::Future;
use std::sync::{
    Arc, Mutex,
    atomic::{AtomicU64, Ordering},
};
use std::time::Instant;

tokio::task_local! {
    static REQUEST_TIMINGS: Arc<Mutex<PhaseTimings>>;
}

pub static PROJECT_HYDRATIONS: AtomicU64 = AtomicU64::new(0);
pub static PDG_LOADS: AtomicU64 = AtomicU64::new(0);
pub static NEURAL_REQUESTS: AtomicU64 = AtomicU64::new(0);

pub fn reset_path_counters() {
    PROJECT_HYDRATIONS.store(0, Ordering::Relaxed);
    PDG_LOADS.store(0, Ordering::Relaxed);
    NEURAL_REQUESTS.store(0, Ordering::Relaxed);
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct WorkBudget {
    pub max_latency_ms: u64,
    pub allow_partial: bool,
}

impl WorkBudget {
    pub fn elapsed(self, started: Instant) -> bool {
        self.allow_partial && started.elapsed().as_millis() >= self.max_latency_ms as u128
    }
}

pub fn elapsed_ms(started: Instant) -> u64 {
    started.elapsed().as_millis().min(u64::MAX as u128) as u64
}

#[derive(Debug, Default, Serialize, Clone)]
pub struct PhaseTimings {
    pub lock_wait_ms: u64,
    pub hydrate_ms: u64,
    pub scan_ms: u64,
    pub parse_ms: u64,
    pub git_ms: u64,
    pub catalog_ms: u64,
    pub pdg_ms: u64,
    pub lexical_ms: u64,
    pub neural_ms: u64,
    pub cache_read_ms: u64,
    pub cache_write_ms: u64,
    pub persist_ms: u64,
    pub handler_ms: u64,
    pub transport_queue_ms: u64,
    pub total_ms: u64,
}

pub async fn collect_request_timings<F>(future: F) -> (F::Output, PhaseTimings)
where
    F: Future,
{
    REQUEST_TIMINGS
        .scope(Arc::new(Mutex::new(PhaseTimings::default())), async move {
            let output = future.await;
            let timings = REQUEST_TIMINGS.with(|timings| {
                timings
                    .lock()
                    .unwrap_or_else(|poisoned| poisoned.into_inner())
                    .clone()
            });
            (output, timings)
        })
        .await
}

pub type RequestTimingSink = Arc<Mutex<PhaseTimings>>;

pub fn current_request_timing_sink() -> Option<RequestTimingSink> {
    REQUEST_TIMINGS.try_with(Arc::clone).ok()
}

pub fn record_hydrate_ms(elapsed_ms: u64) {
    let _ = REQUEST_TIMINGS.try_with(|timings| record_hydrate_ms_to(timings, elapsed_ms));
}

pub fn record_pdg_ms(elapsed_ms: u64) {
    let _ = REQUEST_TIMINGS.try_with(|timings| record_pdg_ms_to(timings, elapsed_ms));
}

pub fn record_git_ms(elapsed_ms: u64) {
    let _ = REQUEST_TIMINGS.try_with(|timings| {
        timings
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
            .git_ms = elapsed_ms;
    });
}

pub fn record_neural_ms(elapsed_ms: u64) {
    let _ = REQUEST_TIMINGS.try_with(|timings| record_neural_ms_to(timings, elapsed_ms));
}

pub fn record_hydrate_ms_to(timings: &RequestTimingSink, elapsed_ms: u64) {
    timings
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
        .hydrate_ms = elapsed_ms;
}

pub fn record_pdg_ms_to(timings: &RequestTimingSink, elapsed_ms: u64) {
    timings
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
        .pdg_ms = elapsed_ms;
}

pub fn record_neural_ms_to(timings: &RequestTimingSink, elapsed_ms: u64) {
    timings
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
        .neural_ms = elapsed_ms;
}