1use std::future::Future;
2use std::sync::Arc;
3use std::time::Duration;
4
5use moka::future::Cache;
6
7use allowthem_core::AllowThem;
8
9use crate::error::SaasError;
10use crate::tenants::{TenantId, TenantStatus};
11
12#[derive(Debug, Clone)]
13pub struct TenantMeta {
14 pub id: TenantId,
15 pub status: TenantStatus,
16 pub plan_id: Vec<u8>,
17}
18
19#[derive(Clone)]
20pub struct SlugCache(Cache<String, TenantMeta>);
21
22impl SlugCache {
23 pub fn new(max_capacity: u64, ttl_secs: u64) -> Self {
24 Self(
25 Cache::builder()
26 .max_capacity(max_capacity)
27 .time_to_live(Duration::from_secs(ttl_secs))
28 .build(),
29 )
30 }
31
32 pub async fn get_or_fetch<F, Fut>(
33 &self,
34 slug: &str,
35 init: F,
36 ) -> Result<Option<TenantMeta>, SaasError>
37 where
38 F: FnOnce() -> Fut,
39 Fut: Future<Output = Result<Option<TenantMeta>, SaasError>>,
40 {
41 if let Some(meta) = self.0.get(slug).await {
42 return Ok(Some(meta));
43 }
44 let result = init().await?;
45 if let Some(ref meta) = result {
46 self.0.insert(slug.to_owned(), meta.clone()).await;
47 }
48 Ok(result)
49 }
50
51 pub fn entry_count(&self) -> u64 {
54 self.0.entry_count()
55 }
56}
57
58#[derive(Clone)]
59pub struct HandleCache(Cache<TenantId, AllowThem>);
60
61impl HandleCache {
62 pub fn new(max_capacity: u64) -> Self {
63 let cache = Cache::builder()
64 .max_capacity(max_capacity)
65 .async_eviction_listener(|tenant_id: Arc<TenantId>, _handle, _cause| {
66 Box::pin(async move {
67 tracing::debug!(tenant_id = %tenant_id.as_uuid(), "tenant handle evicted");
68 })
69 })
70 .build();
71 Self(cache)
72 }
73
74 pub async fn get_or_init<F>(
75 &self,
76 tenant_id: TenantId,
77 init: F,
78 ) -> Result<AllowThem, Arc<SaasError>>
79 where
80 F: Future<Output = Result<AllowThem, SaasError>>,
81 {
82 self.0.try_get_with(tenant_id, init).await
83 }
84
85 pub async fn invalidate(&self, tenant_id: &TenantId) {
86 self.0.invalidate(tenant_id).await;
87 }
88
89 pub fn entry_count(&self) -> u64 {
90 self.0.entry_count()
91 }
92}
93
94#[cfg(test)]
95mod tests {
96 use std::sync::Arc;
97 use std::sync::atomic::{AtomicUsize, Ordering};
98
99 use allowthem_core::AllowThemBuilder;
100 use uuid::Uuid;
101
102 use super::*;
103
104 fn make_meta(slug_seed: u8) -> TenantMeta {
105 TenantMeta {
106 id: TenantId::from(Uuid::from_bytes([slug_seed; 16])),
107 status: TenantStatus::Active,
108 plan_id: vec![slug_seed],
109 }
110 }
111
112 #[tokio::test]
113 async fn slug_cache_miss_calls_init() {
114 let cache = SlugCache::new(10, 60);
115 let count = Arc::new(AtomicUsize::new(0));
116 let c = count.clone();
117 let meta = make_meta(1);
118
119 let result = cache
120 .get_or_fetch("acme", || {
121 let c = c.clone();
122 let m = meta.clone();
123 async move {
124 c.fetch_add(1, Ordering::SeqCst);
125 Ok(Some(m))
126 }
127 })
128 .await
129 .unwrap();
130
131 assert!(result.is_some());
132 assert_eq!(count.load(Ordering::SeqCst), 1);
133 }
134
135 #[tokio::test]
136 async fn slug_cache_hit_skips_init() {
137 let cache = SlugCache::new(10, 60);
138 let count = Arc::new(AtomicUsize::new(0));
139 let meta = make_meta(2);
140
141 for _ in 0..2 {
142 let c = count.clone();
143 let m = meta.clone();
144 cache
145 .get_or_fetch("beta", move || {
146 let c = c.clone();
147 let m = m.clone();
148 async move {
149 c.fetch_add(1, Ordering::SeqCst);
150 Ok(Some(m))
151 }
152 })
153 .await
154 .unwrap();
155 }
156
157 assert_eq!(count.load(Ordering::SeqCst), 1);
158 }
159
160 #[tokio::test]
161 async fn handle_cache_coalesces_concurrent() {
162 let cache = HandleCache::new(10);
163 let count = Arc::new(AtomicUsize::new(0));
164 let barrier = Arc::new(tokio::sync::Barrier::new(3));
165 let id = TenantId::from(Uuid::from_bytes([0xAB; 16]));
166
167 let tasks: Vec<_> = (0..3)
168 .map(|_| {
169 let cache = cache.clone();
170 let count = count.clone();
171 let b = barrier.clone();
172 tokio::spawn(async move {
173 b.wait().await;
174 let _ = cache
175 .get_or_init(id, async {
176 count.fetch_add(1, Ordering::SeqCst);
177 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
178 Err::<AllowThem, SaasError>(SaasError::TenantNotFound)
179 })
180 .await;
181 })
182 })
183 .collect();
184
185 for t in tasks {
186 t.await.unwrap();
187 }
188
189 assert_eq!(count.load(Ordering::SeqCst), 1);
190 }
191
192 #[tokio::test]
193 async fn handle_cache_error_propagates() {
194 let cache = HandleCache::new(10);
195 let id = TenantId::from(Uuid::from_bytes([0xCD; 16]));
196
197 let err = cache
198 .get_or_init(id, async {
199 Err::<AllowThem, SaasError>(SaasError::TenantNotFound)
200 })
201 .await
202 .err()
203 .expect("expected error");
204
205 assert!(matches!(*err, SaasError::TenantNotFound));
206 }
207
208 #[tokio::test]
209 async fn handle_cache_invalidate_forces_reinit() {
210 let handle = AllowThemBuilder::new("sqlite::memory:")
211 .cookie_secure(false)
212 .build()
213 .await
214 .unwrap();
215
216 let cache = HandleCache::new(10);
217 let id = TenantId::from(Uuid::from_bytes([0xEF; 16]));
218 let count = Arc::new(AtomicUsize::new(0));
219
220 let c = count.clone();
221 let h = handle.clone();
222 cache
223 .get_or_init(id, async move {
224 c.fetch_add(1, Ordering::SeqCst);
225 Ok::<_, SaasError>(h)
226 })
227 .await
228 .unwrap();
229
230 assert_eq!(count.load(Ordering::SeqCst), 1);
231
232 cache.invalidate(&id).await;
233
234 let c = count.clone();
235 cache
236 .get_or_init(id, async move {
237 c.fetch_add(1, Ordering::SeqCst);
238 Ok::<_, SaasError>(handle)
239 })
240 .await
241 .unwrap();
242
243 assert_eq!(count.load(Ordering::SeqCst), 2);
244 }
245
246 #[tokio::test]
248 async fn slug_cache_entry_count_increments_on_get_or_fetch() {
249 let cache = SlugCache::new(10, 60);
250 let meta = make_meta(42);
251 let m = meta.clone();
252
253 cache
254 .get_or_fetch("newslug", move || {
255 let m = m.clone();
256 async move { Ok(Some(m)) }
257 })
258 .await
259 .unwrap();
260
261 cache.0.run_pending_tasks().await;
264 assert!(cache.entry_count() >= 1);
265 }
266}