Blame for src/scheduler.rs

//! UltraSpeed 资源调度：本机/云端分别限流 + 元数据缓存。
//!
//! # 问题
//!
//! 本机 GPU 推理的并发度很低（一块 GPU 同时跑两个大模型会互相抢显存，
//! 甚至触发 Ollama 卸载重载）。云端 API 则有 RPM/TPM 限流。
//!
//! # 解法
//!
//! 用信号量把**本机并发压到 1**、云端控制在可控并发（默认 2），
//! 请求排队而不是互相踩踏。同时缓存 provider 的 health / models 元数据，
//! 避免每次请求都探活（探活本身也要几十毫秒）。
//!
//! 注意：调度**只影响并发度**，绝不改变模型 / prompt / 思考档 / 回复长度。

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, LazyLock, Mutex};
use std::time::{Duration, Instant};

use tokio::sync::{Mutex as AsyncMutex, OwnedSemaphorePermit, Semaphore};

/// 本机推理并发上限（GPU 单卡最优 = 1，排队比抢显存快）。
pub const LOCAL_CONCURRENCY: usize = 1;
/// 云端并发上限（按你的 API 限额调整）。
pub const CLOUD_CONCURRENCY: usize = 2;
/// 元数据缓存有效期。
pub const METADATA_TTL: Duration = Duration::from_secs(30);

static LOCAL_PERMITS: LazyLock<Arc<Semaphore>> =
    LazyLock::new(|| Arc::new(Semaphore::new(LOCAL_CONCURRENCY)));
static CLOUD_PERMITS: LazyLock<Arc<Semaphore>> =
    LazyLock::new(|| Arc::new(Semaphore::new(CLOUD_CONCURRENCY)));

/// 申请一个执行许可：`cloud = true` 走云端池，否则走本机池。
///
/// 返回的 permit 释放时自动归还名额。调用方应尽量短持有。
pub async fn acquire(cloud: bool) -> Result<OwnedSemaphorePermit, String> {
    let pool = if cloud {
        CLOUD_PERMITS.clone()
    } else {
        LOCAL_PERMITS.clone()
    };
    pool.acquire_owned()
        .await
        .map_err(|_| "资源池已关闭".to_string())
}

// ── 凭据版本号：凭据变更后让所有缓存失效 ────────────────────────────

static REVISION: AtomicU64 = AtomicU64::new(0);

/// 当前凭据版本（缓存 key 的一部分）。
pub fn revision() -> u64 {
    REVISION.load(Ordering::Relaxed)
}

/// 凭据/配置变更后调用：让所有 provider 的元数据缓存失效。
pub fn invalidate() {
    REVISION.fetch_add(1, Ordering::Relaxed);
    if let Ok(mut cache) = CACHE.lock() {
        cache.clear();
    }
}

// ── 带 TTL 的元数据缓存 ────────────────────────────────────────────

struct Metadata {
    health: AsyncMutex<Option<(Instant, bool)>>,
    models: AsyncMutex<Option<(Instant, Result<Vec<String>, String>)>>,
}

static CACHE: LazyLock<Mutex<HashMap<String, Arc<Metadata>>>> =
    LazyLock::new(|| Mutex::new(HashMap::new()));

fn entry(key: &str) -> Arc<Metadata> {
    let mut cache = CACHE.lock().unwrap_or_else(|e| e.into_inner());
    cache
        .entry(key.to_string())
        .or_insert_with(|| {
            Arc::new(Metadata {
                health: AsyncMutex::new(None),
                models: AsyncMutex::new(None),
            })
        })
        .clone()
}

/// 读取健康状态缓存（未命中/过期返回 None）。
pub async fn cached_health(key: &str) -> Option<bool> {
    let meta = entry(key);
    let guard = meta.health.lock().await;
    match guard.as_ref() {
        Some((at, value)) if at.elapsed() < METADATA_TTL => Some(*value),
        _ => None,
    }
}

/// 写入健康状态缓存。
pub async fn store_health(key: &str, healthy: bool) {
    let meta = entry(key);
    *meta.health.lock().await = Some((Instant::now(), healthy));
}

/// 读取模型列表缓存（未命中/过期返回 None）。
pub async fn cached_models(key: &str) -> Option<Result<Vec<String>, String>> {
    let meta = entry(key);
    let guard = meta.models.lock().await;
    match guard.as_ref() {
        Some((at, value)) if at.elapsed() < METADATA_TTL => Some(value.clone()),
        _ => None,
    }
}

/// 写入模型列表缓存。
pub async fn store_models(key: &str, models: Result<Vec<String>, String>) {
    let meta = entry(key);
    *meta.models.lock().await = Some((Instant::now(), models));
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn local_pool_serializes_to_one() {
        let first = acquire(false).await.unwrap();
        // 第二个申请应排队：用 try_acquire 验证名额已满
        assert!(
            LOCAL_PERMITS.clone().try_acquire_owned().is_err(),
            "本机并发必须压到 1（GPU 排队比抢显存快）"
        );
        drop(first);
        assert!(LOCAL_PERMITS.clone().try_acquire_owned().is_ok());
    }

    #[tokio::test]
    async fn metadata_cache_hits_within_ttl_and_invalidates() {
        let key = "test-provider";
        store_health(key, true).await;
        assert_eq!(cached_health(key).await, Some(true));
        store_models(key, Ok(vec!["m1".into()])).await;
        assert!(cached_models(key).await.is_some());
        invalidate();
        // invalidate 清空缓存
        assert_eq!(cached_health(key).await, None);
        assert!(cached_models(key).await.is_none());
    }

    #[test]
    fn revision_monotonically_increases() {
        let before = revision();
        invalidate();
        assert!(revision() > before);
    }
}