xz_rerank/
cache.rs

1use async_trait::async_trait;
2use std::collections::HashMap;
3use std::fmt::Debug;
4use std::time::{Duration, Instant};
5use tokio::sync::RwLock;
6
7use crate::types::RerankResult;
8
9/// Rerank 结果缓存
10#[async_trait]
11pub trait RerankCache: Send + Sync + Debug {
12    /// 获取缓存结果
13    async fn get(&self, query: &str, candidate_ids: &[String]) -> Option<RerankResult>;
14
15    /// 存储缓存结果
16    async fn set(
17        &self,
18        query: &str,
19        candidate_ids: &[String],
20        result: &RerankResult,
21        ttl: Duration,
22    );
23}
24
25/// LRU 内存缓存实现
26#[derive(Debug)]
27pub struct MemoryRerankCache {
28    entries: RwLock<HashMap<String, CacheEntry>>,
29    max_entries: usize,
30}
31
32#[derive(Debug, Clone)]
33struct CacheEntry {
34    result: RerankResult,
35    expires_at: Instant,
36    last_accessed: Instant,
37}
38
39impl MemoryRerankCache {
40    /// 创建新的 LRU 内存缓存
41    ///
42    /// `max_entries` 指定缓存中最多保留的条目数，超出时会淘汰最久未使用的条目。
43    pub fn new(max_entries: usize) -> Self {
44        Self { entries: RwLock::new(HashMap::new()), max_entries }
45    }
46
47    fn make_key(query: &str, candidate_ids: &[String]) -> String {
48        let ids_hash = candidate_ids.join(",");
49        format!("rerank:{query}:{ids_hash}")
50    }
51
52    async fn evict_if_needed(&self) {
53        let mut entries = self.entries.write().await;
54        let now = Instant::now();
55        entries.retain(|_, e| e.expires_at > now);
56
57        if entries.len() >= self.max_entries {
58            let lru_key = entries
59                .iter()
60                .min_by_key(|(_, entry)| entry.last_accessed)
61                .map(|(key, _)| key.clone());
62
63            if let Some(key) = lru_key {
64                entries.remove(&key);
65            }
66        }
67    }
68}
69
70#[async_trait]
71impl RerankCache for MemoryRerankCache {
72    async fn get(&self, query: &str, candidate_ids: &[String]) -> Option<RerankResult> {
73        self.evict_if_needed().await;
74        let key = Self::make_key(query, candidate_ids);
75        let mut entries = self.entries.write().await;
76        if let Some(entry) = entries.get_mut(&key)
77            && entry.expires_at > Instant::now()
78        {
79            entry.last_accessed = Instant::now();
80            return Some(entry.result.clone());
81        }
82        None
83    }
84
85    async fn set(
86        &self,
87        query: &str,
88        candidate_ids: &[String],
89        result: &RerankResult,
90        ttl: Duration,
91    ) {
92        self.evict_if_needed().await;
93        let key = Self::make_key(query, candidate_ids);
94        let now = Instant::now();
95        let mut entries = self.entries.write().await;
96        entries.insert(
97            key,
98            CacheEntry { result: result.clone(), expires_at: now + ttl, last_accessed: now },
99        );
100    }
101}
xz_rerank/cache.rs

xz_rerank/
cache.rs
