sa_token_core/service/
grant_cache.rs

1// Author: 金书记 | Author: Jin Shuji
2//
3//! Sharded Grant Cache | 分片授权缓存
4//!
5//! 权限/角色读缓存。默认**关闭**（`SaTokenConfig::grant_cache_ttl <= 0` 时
6//! 连结构都不分配），因为多实例部署下缓存意味着授权变更存在滞后窗口，
7//! 必须由使用者显式权衡后开启。
8//!
9//! Read cache for permissions and roles. Disabled by default — when
10//! `grant_cache_ttl <= 0` no structure is even allocated — because in a
11//! multi-instance deployment a cache introduces a staleness window for
12//! authorization decisions, which must be an explicit opt-in.
13
14use std::collections::HashMap;
15use std::future::Future;
16use std::sync::{Arc, RwLock};
17use std::time::{Duration, Instant};
18
19use crate::config::SaTokenConfig;
20use crate::error::SaTokenResult;
21
22/// 分片数量：必须是 2 的幂，以便用位与代替取模。
23/// Shard count; must be a power of two so the index can use a bitmask.
24const SHARD_COUNT: usize = 8;
25
26/// 分片掩码 | Shard bitmask
27const SHARD_MASK: usize = SHARD_COUNT - 1;
28
29/// 缓存键内部分隔符
30/// Internal key separator.
31const KEY_SEP: char = '\u{1}';
32
33/// 缓存数据类别 | Cached data kind
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub enum GrantKind {
36    /// 权限列表 | Permission list
37    Permission,
38    /// 角色列表 | Role list
39    Role,
40}
41
42impl GrantKind {
43    /// 键前缀标记 | Key tag
44    #[inline]
45    fn tag(self) -> char {
46        match self {
47            Self::Permission => 'p',
48            Self::Role => 'r',
49        }
50    }
51}
52
53/// 缓存条目 | Cache entry
54struct CacheEntry {
55    value: Arc<[String]>,
56    expires_at: Instant,
57}
58
59type FlightMap = HashMap<String, Arc<tokio::sync::Mutex<()>>>;
60type FlightShards = Box<[tokio::sync::Mutex<FlightMap>]>;
61
62/// 分片授权缓存 | Sharded grant cache
63pub struct GrantCache {
64    shards: Box<[RwLock<HashMap<String, CacheEntry>>]>,
65    /// 单飞门闩也按分片，避免全局一把 Mutex 串行。
66    /// Single-flight gates are sharded too, avoiding one global Mutex.
67    flights: FlightShards,
68    ttl: Duration,
69    max_per_shard: usize,
70    single_flight: bool,
71}
72
73impl std::fmt::Debug for GrantCache {
74    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75        f.debug_struct("GrantCache")
76            .field("ttl", &self.ttl)
77            .field("max_per_shard", &self.max_per_shard)
78            .field("single_flight", &self.single_flight)
79            .finish()
80    }
81}
82
83impl GrantCache {
84    /// 构造缓存 | Construct the cache
85    pub fn new(ttl: Duration, max_entries: usize, single_flight: bool) -> Self {
86        let max_per_shard = (max_entries / SHARD_COUNT).max(1);
87        let shards = (0..SHARD_COUNT)
88            .map(|_| RwLock::new(HashMap::new()))
89            .collect::<Vec<_>>()
90            .into_boxed_slice();
91        let flights = (0..SHARD_COUNT)
92            .map(|_| tokio::sync::Mutex::new(HashMap::new()))
93            .collect::<Vec<_>>()
94            .into_boxed_slice();
95        Self {
96            shards,
97            flights,
98            ttl,
99            max_per_shard,
100            single_flight,
101        }
102    }
103
104    /// 依配置构造；返回 `None` 表示**缓存关闭**。
105    /// Build from config; `None` means the cache is off.
106    pub fn from_config(config: &SaTokenConfig) -> Option<Arc<Self>> {
107        config.grant_cache_duration().map(|ttl| {
108            Arc::new(Self::new(
109                ttl,
110                config.grant_cache_max_entries,
111                config.grant_cache_single_flight,
112            ))
113        })
114    }
115
116    /// 构造缓存键 | Build a cache key
117    pub fn cache_key(kind: GrantKind, login_type: &str, login_id: &str) -> String {
118        let mut key = String::with_capacity(2 + login_type.len() + 1 + login_id.len());
119        key.push(kind.tag());
120        key.push(KEY_SEP);
121        key.push_str(login_type);
122        key.push(KEY_SEP);
123        key.push_str(login_id);
124        key
125    }
126
127    /// 内部 cache key 由 login_type + login_id 组成，非攻击者可控 HTTP 输入；SipHash 无安全收益。
128    /// Cache keys are composed of login_type + login_id, not attacker-controlled HTTP input; SipHash buys no safety here.
129    #[inline]
130    fn fnv1a(key: &str) -> usize {
131        let mut hash = 0xcbf29ce484222325u64;
132        for byte in key.as_bytes() {
133            hash ^= u64::from(*byte);
134            hash = hash.wrapping_mul(0x100000001b3);
135        }
136        hash as usize
137    }
138
139    /// 计算 key 所属分片 | Resolve the shard owning a key
140    #[inline]
141    fn shard_of(&self, key: &str) -> &RwLock<HashMap<String, CacheEntry>> {
142        match self.shards.get(Self::fnv1a(key) & SHARD_MASK) {
143            Some(s) => s,
144            None => unreachable!("fnv1a masked to SHARD_COUNT"),
145        }
146    }
147
148    #[inline]
149    fn flight_shard(&self, key: &str) -> &tokio::sync::Mutex<FlightMap> {
150        match self.flights.get(Self::fnv1a(key) & SHARD_MASK) {
151            Some(s) => s,
152            None => unreachable!("fnv1a masked to SHARD_COUNT"),
153        }
154    }
155
156    /// 只读探测 | Read-only probe
157    fn peek(&self, key: &str) -> Option<Arc<[String]>> {
158        let guard = self.shard_of(key).read().ok()?;
159        let entry = guard.get(key)?;
160        if entry.expires_at > Instant::now() {
161            Some(Arc::clone(&entry.value))
162        } else {
163            None
164        }
165    }
166
167    /// 写入并保证容量有界 | Insert while keeping capacity bounded
168    fn put(&self, key: String, value: Arc<[String]>) {
169        let Ok(mut guard) = self.shard_of(&key).write() else {
170            return;
171        };
172        let now = Instant::now();
173
174        if guard.len() >= self.max_per_shard && !guard.contains_key(&key) {
175            guard.retain(|_, entry| entry.expires_at > now);
176
177            if guard.len() >= self.max_per_shard {
178                let victim = guard
179                    .iter()
180                    .min_by_key(|(_, entry)| entry.expires_at)
181                    .map(|(k, _)| k.clone());
182                if let Some(victim) = victim {
183                    guard.remove(&victim);
184                }
185            }
186        }
187
188        guard.insert(
189            key,
190            CacheEntry {
191                value,
192                expires_at: now + self.ttl,
193            },
194        );
195    }
196
197    /// 读取或加载 | Get or load
198    pub async fn get_or_load<F, Fut>(&self, key: String, loader: F) -> SaTokenResult<Arc<[String]>>
199    where
200        F: FnOnce() -> Fut,
201        Fut: Future<Output = SaTokenResult<Vec<String>>>,
202    {
203        if let Some(hit) = self.peek(&key) {
204            return Ok(hit);
205        }
206
207        if !self.single_flight {
208            let value: Arc<[String]> = loader().await?.into();
209            self.put(key, Arc::clone(&value));
210            return Ok(value);
211        }
212
213        let gate = {
214            let mut flights = self.flight_shard(&key).lock().await;
215            Arc::clone(
216                flights
217                    .entry(key.clone())
218                    .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))),
219            )
220        };
221        let _permit = gate.lock().await;
222
223        if let Some(hit) = self.peek(&key) {
224            self.release_flight(&key).await;
225            return Ok(hit);
226        }
227
228        match loader().await {
229            Ok(list) => {
230                let value: Arc<[String]> = list.into();
231                self.put(key.clone(), Arc::clone(&value));
232                self.release_flight(&key).await;
233                Ok(value)
234            }
235            Err(err) => {
236                self.release_flight(&key).await;
237                Err(err)
238            }
239        }
240    }
241
242    /// 移除门闩表中的条目 | Drop the gate entry
243    async fn release_flight(&self, key: &str) {
244        self.flight_shard(key).lock().await.remove(key);
245    }
246
247    /// 失效单个键 | Invalidate one key
248    pub fn invalidate(&self, key: &str) {
249        if let Ok(mut guard) = self.shard_of(key).write() {
250            guard.remove(key);
251        }
252    }
253
254    /// 失效某账号在某体系下的权限与角色两条缓存。
255    /// Invalidate both the permission and role entries of an account.
256    pub fn invalidate_account(&self, login_type: &str, login_id: &str) {
257        self.invalidate(&Self::cache_key(
258            GrantKind::Permission,
259            login_type,
260            login_id,
261        ));
262        self.invalidate(&Self::cache_key(GrantKind::Role, login_type, login_id));
263    }
264
265    /// 清空全部分片 | Clear every shard
266    pub fn clear(&self) {
267        for shard in self.shards.iter() {
268            if let Ok(mut guard) = shard.write() {
269                guard.clear();
270            }
271        }
272    }
273
274    /// 当前条目总数（诊断用）| Total entry count for diagnostics
275    pub fn len(&self) -> usize {
276        self.shards
277            .iter()
278            .filter_map(|shard| shard.read().ok().map(|g| g.len()))
279            .sum()
280    }
281
282    /// 是否为空 | Whether the cache is empty
283    pub fn is_empty(&self) -> bool {
284        self.len() == 0
285    }
286}
sa_token_core/service/grant_cache.rs

sa_token_core/service/
grant_cache.rs
