Skip to main content

allowthem_saas/
cache.rs

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    /// Number of entries currently cached. Used by the super-admin
52    /// health page (99c.6 ยง4.5).
53    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    // Step 3 (99c.6): SlugCache::entry_count
247    #[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        // Moka's entry_count is eventually consistent; flush pending tasks
262        // so the counter is updated before we assert.
263        cache.0.run_pending_tasks().await;
264        assert!(cache.entry_count() >= 1);
265    }
266}