Skip to main content

nargo_ssr/
prefetch.rs

1use std::{collections::HashMap, future::Future, pin::Pin, sync::Arc};
2
3/// 数据预取上下文
4type PrefetchContext = HashMap<String, serde_json::Value>;
5
6/// 数据预取函数类型
7type PrefetchFn = Arc<dyn Fn(&HashMap<String, String>, &HashMap<String, String>) -> Pin<Box<dyn Future<Output = Result<serde_json::Value, String>>>> + Send + Sync>;
8
9/// 数据预取管理器
10pub struct PrefetchManager {
11    /// 预取函数映射
12    prefetch_functions: HashMap<String, PrefetchFn>,
13    /// 缓存的预取结果
14    prefetch_cache: HashMap<String, serde_json::Value>,
15}
16
17impl PrefetchManager {
18    /// 创建新的数据预取管理器
19    pub fn new() -> Self {
20        Self { prefetch_functions: HashMap::new(), prefetch_cache: HashMap::new() }
21    }
22
23    /// 注册数据预取函数
24    ///
25    /// # Arguments
26    /// * `key` - 预取函数的唯一标识
27    /// * `prefetch_fn` - 预取函数
28    pub fn register_prefetch<F, Fut>(&mut self, key: &str, prefetch_fn: F)
29    where
30        F: Fn(&HashMap<String, String>, &HashMap<String, String>) -> Fut + Send + Sync + 'static,
31        Fut: Future<Output = Result<serde_json::Value, String>> + Send + 'static,
32    {
33        self.prefetch_functions.insert(key.to_string(), Arc::new(move |params, query| Box::pin(prefetch_fn(params, query))));
34    }
35
36    /// 执行数据预取
37    ///
38    /// # Arguments
39    /// * `keys` - 需要预取的键列表
40    /// * `params` - 路由参数
41    /// * `query` - 查询参数
42    ///
43    /// # Returns
44    /// * `PrefetchContext` - 预取的上下文数据
45    pub async fn prefetch(&mut self, keys: &[&str], params: &HashMap<String, String>, query: &HashMap<String, String>) -> PrefetchContext {
46        let mut context = PrefetchContext::new();
47
48        for key in keys {
49            // 生成缓存键
50            let cache_key = self.generate_cache_key(key, params, query);
51
52            // 检查缓存
53            if let Some(cached) = self.prefetch_cache.get(&cache_key) {
54                context.insert(key.to_string(), cached.clone());
55                continue;
56            }
57
58            // 执行预取
59            if let Some(prefetch_fn) = self.prefetch_functions.get(*key) {
60                match prefetch_fn(params, query).await {
61                    Ok(value) => {
62                        context.insert(key.to_string(), value.clone());
63                        // 缓存结果
64                        self.prefetch_cache.insert(cache_key, value);
65                    }
66                    Err(err) => {
67                        eprintln!("Prefetch error for {}: {}", key, err);
68                    }
69                }
70            }
71        }
72
73        context
74    }
75
76    /// 生成缓存键
77    ///
78    /// # Arguments
79    /// * `key` - 预取函数的键
80    /// * `params` - 路由参数
81    /// * `query` - 查询参数
82    ///
83    /// # Returns
84    /// * `String` - 缓存键
85    fn generate_cache_key(&self, key: &str, params: &HashMap<String, String>, query: &HashMap<String, String>) -> String {
86        let mut cache_key = key.to_string();
87
88        // 添加路由参数
89        let mut param_pairs: Vec<String> = params.iter().map(|(k, v)| format!("{}={}", k, v)).collect();
90        param_pairs.sort();
91        for pair in param_pairs {
92            cache_key.push_str(&format!("&{}", pair));
93        }
94
95        // 添加查询参数
96        let mut query_pairs: Vec<String> = query.iter().map(|(k, v)| format!("{}={}", k, v)).collect();
97        query_pairs.sort();
98        for pair in query_pairs {
99            cache_key.push_str(&format!("&{}", pair));
100        }
101
102        cache_key
103    }
104
105    /// 清除预取缓存
106    pub fn clear_cache(&mut self) {
107        self.prefetch_cache.clear();
108    }
109}