ullun/
cache.rs

1//! Server-side response cache with request coalescing and stale-while-revalidate
2//!
3//! [`ResponseCache`] keeps successful `GET` responses in memory so an expensive
4//! handler (database aggregate, upstream call, rendered page) runs rarely:
5//!
6//! * **Coalescing**: when many requests miss at once, one runs the handler and
7//!   the others wait for its result, so a cold or expired entry cannot stampede
8//!   the backend
9//! * **Stale-while-revalidate**: for a window after expiry the old response is
10//!   served immediately while a single background request refreshes it, so
11//!   clients never wait for a refresh
12//! * **Safe by default**: requests carrying `Authorization` or `Cookie` bypass
13//!   the cache (or get their own entries with [`ResponseCache::per_user`]),
14//!   responses with `Set-Cookie`, `Cache-Control: no-store` / `private`, streams
15//!   and oversized bodies are never stored
16//!
17//! Responses get `X-Cache: HIT | MISS | STALE | COALESCED` and `Age`. The
18//! lifetime comes from the response's `Cache-Control: max-age` / `s-maxage`
19//! when present, otherwise from [`ResponseCache::ttl`].
20//!
21//! ```no_run
22//! use std::time::Duration;
23//! use ullun::cache::ResponseCache;
24//! use ullun::prelude::*;
25//!
26//! # async fn run() {
27//! let cache = ResponseCache::new()
28//!     .ttl(Duration::from_secs(30))
29//!     .stale_while_revalidate(Duration::from_secs(120))
30//!     .vary_by(["accept-language"]);
31//! App::new()
32//!     .middleware(cache.clone()) // keep a handle to invalidate after writes
33//!     .get("/report", |_req: Request| async { Ok(Response::text("expensive")) })
34//!     .post("/report/refresh", {
35//!         let cache = cache.clone();
36//!         move |_req: Request| {
37//!             let cache = cache.clone();
38//!             async move {
39//!                 cache.invalidate_prefix("/report");
40//!                 Ok(Response::text("ok"))
41//!             }
42//!         }
43//!     })
44//!     .run("127.0.0.1:3000")
45//!     .await
46//!     .unwrap();
47//! # }
48//! ```
49//!
50//! Keys longer than 2 KiB are never cached. Memory use is bounded by
51//! `max_entries` x `max_body` (1000 x 1 MiB by default): lower either for large responses. The cache lives in the memory of one process. A request that is not
52//! cacheable (see above) is not coalesced either: each waiter then runs the
53//! handler itself.
54
55use crate::error::Result;
56use crate::middleware::{Middleware, Next};
57use crate::request::Request;
58use crate::response::Response;
59use async_trait::async_trait;
60use bytes::Bytes;
61use http::StatusCode;
62use std::collections::HashMap;
63use std::hash::{BuildHasher, Hash, Hasher};
64use std::sync::{Arc, Mutex};
65use std::time::{Duration, Instant};
66use tokio::sync::watch;
67
68/// Longest cache key (path, query and varied headers) that is stored
69const MAX_KEY_LEN: usize = 2048;
70
71struct Stored {
72    status: StatusCode,
73    headers: http::HeaderMap,
74    body: Bytes,
75}
76
77impl Stored {
78    fn to_response(&self, state: &str, age: Duration) -> Response {
79        let mut res = Response::new(self.status).with_body(self.body.clone());
80        res.headers = self.headers.clone();
81        res.with_header("x-cache", state)
82            .with_header("age", age.as_secs().to_string())
83    }
84}
85
86struct Entry {
87    stored: Arc<Stored>,
88    stored_at: Instant,
89    fresh_for: Duration,
90}
91
92#[derive(Clone)]
93enum Flight {
94    Pending,
95    Done(Option<Arc<Stored>>),
96}
97
98#[derive(Default)]
99struct Slot {
100    path: String,
101    entry: Option<Entry>,
102    flight: Option<watch::Receiver<Flight>>,
103    refreshing: bool,
104}
105
106struct Inner {
107    slots: Mutex<HashMap<String, Slot>>,
108    hasher: std::collections::hash_map::RandomState,
109}
110
111/// Cache middleware, see the [module documentation](self). Cloning gives
112/// another handle to the same cache.
113#[derive(Clone)]
114pub struct ResponseCache {
115    inner: Arc<Inner>,
116    ttl: Duration,
117    swr: Duration,
118    max_entries: usize,
119    max_body: usize,
120    statuses: Vec<u16>,
121    vary: Vec<String>,
122    per_user: bool,
123}
124
125impl Default for ResponseCache {
126    fn default() -> Self {
127        Self::new()
128    }
129}
130
131impl ResponseCache {
132    pub fn new() -> Self {
133        Self {
134            inner: Arc::new(Inner {
135                slots: Mutex::new(HashMap::new()),
136                hasher: Default::default(),
137            }),
138            ttl: Duration::from_secs(60),
139            swr: Duration::ZERO,
140            max_entries: 1_000,
141            max_body: 1024 * 1024,
142            statuses: vec![200],
143            vary: Vec::new(),
144            per_user: false,
145        }
146    }
147
148    /// Lifetime of responses that do not say otherwise (default 60 seconds)
149    pub fn ttl(mut self, ttl: Duration) -> Self {
150        self.ttl = ttl;
151        self
152    }
153
154    /// Serve an expired response for this long while one request refreshes it
155    /// in the background (default off)
156    pub fn stale_while_revalidate(mut self, window: Duration) -> Self {
157        self.swr = window;
158        self
159    }
160
161    /// Most cached responses (default 1000); the oldest go first
162    pub fn max_entries(mut self, max: usize) -> Self {
163        self.max_entries = max.max(1);
164        self
165    }
166
167    /// Largest body that is cached (default 1 MiB)
168    pub fn max_body(mut self, bytes: usize) -> Self {
169        self.max_body = bytes;
170        self
171    }
172
173    /// Status codes worth caching (default only 200)
174    pub fn statuses<I: IntoIterator<Item = u16>>(mut self, statuses: I) -> Self {
175        self.statuses = statuses.into_iter().collect();
176        self
177    }
178
179    /// Request headers that change the response (`Accept-Language`, ...): each
180    /// distinct value gets its own entry, and `Vary` is added to the response
181    pub fn vary_by<I, S>(mut self, headers: I) -> Self
182    where
183        I: IntoIterator<Item = S>,
184        S: Into<String>,
185    {
186        self.vary = headers
187            .into_iter()
188            .map(|h| h.into().to_ascii_lowercase())
189            .collect();
190        self
191    }
192
193    /// Cache requests that carry `Authorization` or `Cookie` too, with a separate
194    /// entry per credential. Only for responses that depend on nothing but the
195    /// credential itself.
196    pub fn per_user(mut self) -> Self {
197        self.per_user = true;
198        self
199    }
200
201    /// Drop every cached response whose path starts with `prefix`
202    pub fn invalidate_prefix(&self, prefix: &str) {
203        let mut map = self.lock();
204        map.retain(|_, slot| !slot.path.starts_with(prefix));
205    }
206
207    /// Drop everything
208    pub fn clear(&self) {
209        self.lock().clear();
210    }
211
212    /// Number of cached responses
213    pub fn len(&self) -> usize {
214        self.lock().values().filter(|s| s.entry.is_some()).count()
215    }
216
217    pub fn is_empty(&self) -> bool {
218        self.len() == 0
219    }
220
221    fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<String, Slot>> {
222        self.inner.slots.lock().unwrap_or_else(|e| e.into_inner())
223    }
224
225    fn key(&self, req: &Request, scope: bool) -> String {
226        let mut key = req.uri.path().to_string();
227        if let Some(q) = req.uri.query() {
228            key.push('?');
229            key.push_str(q);
230        }
231        for name in &self.vary {
232            key.push('|');
233            if let Some(v) = req.headers.get(name.as_str()).and_then(|v| v.to_str().ok()) {
234                key.push_str(v);
235            }
236        }
237        if scope {
238            let mut h = self.inner.hasher.build_hasher();
239            for name in ["authorization", "cookie"] {
240                req.headers.get(name).map(|v| v.as_bytes()).hash(&mut h);
241            }
242            key.push_str(&format!("|u{:016x}", h.finish()));
243        }
244        key
245    }
246
247    /// Whether and for how long a response may be stored
248    fn lifetime(&self, res: &Response) -> Option<Duration> {
249        if !self.statuses.contains(&res.status.as_u16())
250            || res.is_streaming()
251            || res.has_cookies()
252            || res.headers.contains_key(http::header::SET_COOKIE)
253            || res.body.len() > self.max_body
254        {
255            return None;
256        }
257        let mut max_age = None;
258        if let Some(cc) = res
259            .headers
260            .get(http::header::CACHE_CONTROL)
261            .and_then(|v| v.to_str().ok())
262        {
263            for directive in cc.split(',').map(|d| d.trim().to_ascii_lowercase()) {
264                match directive.as_str() {
265                    "no-store" | "private" | "no-cache" => return None,
266                    d => {
267                        if let Some(n) = d
268                            .strip_prefix("s-maxage=")
269                            .or_else(|| d.strip_prefix("max-age="))
270                        {
271                            if let Ok(secs) = n.trim_matches('"').parse::<u64>() {
272                                max_age.get_or_insert(secs);
273                                if d.starts_with("s-maxage=") {
274                                    max_age = Some(secs);
275                                }
276                            }
277                        }
278                    }
279                }
280            }
281        }
282        match max_age {
283            Some(0) => None,
284            Some(secs) => Some(Duration::from_secs(secs)),
285            None => Some(self.ttl),
286        }
287    }
288
289    fn snapshot(&self, res: &Response) -> Stored {
290        let mut headers = res.headers.clone();
291        // freshness is tracked here; a stale Age from upstream would mislead
292        headers.remove(http::header::AGE);
293        headers.remove("x-cache");
294        Stored {
295            status: res.status,
296            headers,
297            body: res.body.clone(),
298        }
299    }
300
301    fn store(&self, key: &str, path: &str, stored: Arc<Stored>, fresh_for: Duration) {
302        let mut map = self.lock();
303        let others = |map: &HashMap<String, Slot>| map.len() - usize::from(map.contains_key(key));
304        if others(&map) >= self.max_entries {
305            let now = Instant::now();
306            let swr = self.swr;
307            map.retain(|k, s| {
308                k == key
309                    || s.flight.is_some()
310                    || s.entry
311                        .as_ref()
312                        .is_some_and(|e| now.duration_since(e.stored_at) < e.fresh_for + swr)
313            });
314            while others(&map) >= self.max_entries {
315                let oldest = map
316                    .iter()
317                    .filter(|(k, s)| *k != key && s.flight.is_none() && s.entry.is_some())
318                    .min_by_key(|(_, s)| s.entry.as_ref().map(|e| e.stored_at))
319                    .map(|(k, _)| k.clone());
320                match oldest {
321                    Some(k) => {
322                        map.remove(&k);
323                    }
324                    None => break,
325                }
326            }
327        }
328        let slot = map.entry(key.to_string()).or_default();
329        slot.path = path.to_string();
330        slot.entry = Some(Entry {
331            stored,
332            stored_at: Instant::now(),
333            fresh_for,
334        });
335    }
336
337    fn finish_flight(&self, key: &str) {
338        let mut map = self.lock();
339        if let Some(slot) = map.get_mut(key) {
340            slot.flight = None;
341            if slot.entry.is_none() {
342                map.remove(key);
343            }
344        }
345    }
346
347    fn end_refresh(&self, key: &str) {
348        let mut map = self.lock();
349        if let Some(slot) = map.get_mut(key) {
350            slot.refreshing = false;
351            if slot.entry.is_none() && slot.flight.is_none() {
352                map.remove(key);
353            }
354        }
355    }
356}
357
358/// Ends a flight for waiters even if the leader's future is dropped
359struct Leader {
360    cache: ResponseCache,
361    key: String,
362    tx: watch::Sender<Flight>,
363    done: bool,
364}
365
366impl Leader {
367    fn finish(mut self, result: Option<Arc<Stored>>) {
368        self.done = true;
369        let _ = self.tx.send(Flight::Done(result));
370        self.cache.finish_flight(&self.key);
371    }
372}
373
374impl Drop for Leader {
375    fn drop(&mut self) {
376        if !self.done {
377            let _ = self.tx.send(Flight::Done(None));
378            self.cache.finish_flight(&self.key);
379        }
380    }
381}
382
383enum Plan {
384    Hit(Arc<Stored>, Duration),
385    Stale(Arc<Stored>, Duration, bool),
386    Wait(watch::Receiver<Flight>),
387    Lead(Leader),
388}
389
390impl ResponseCache {
391    fn plan(&self, key: &str, path: &str) -> Plan {
392        let mut map = self.lock();
393        let slot = map.entry(key.to_string()).or_default();
394        if slot.path.is_empty() {
395            slot.path = path.to_string();
396        }
397        if let Some(entry) = &slot.entry {
398            let age = entry.stored_at.elapsed();
399            if age < entry.fresh_for {
400                return Plan::Hit(entry.stored.clone(), age);
401            }
402            if age < entry.fresh_for + self.swr {
403                let start_refresh = !slot.refreshing && slot.flight.is_none();
404                if start_refresh {
405                    slot.refreshing = true;
406                }
407                return Plan::Stale(entry.stored.clone(), age, start_refresh);
408            }
409        }
410        if let Some(rx) = &slot.flight {
411            return Plan::Wait(rx.clone());
412        }
413        let (tx, rx) = watch::channel(Flight::Pending);
414        slot.flight = Some(rx);
415        slot.entry = None;
416        Plan::Lead(Leader {
417            cache: self.clone(),
418            key: key.to_string(),
419            tx,
420            done: false,
421        })
422    }
423}
424
425#[async_trait]
426impl Middleware for ResponseCache {
427    async fn handle(&self, req: Request, next: Next) -> Result<Response> {
428        let credentialed = req.headers.contains_key(http::header::AUTHORIZATION)
429            || req.headers.contains_key(http::header::COOKIE);
430        if req.method != http::Method::GET || (credentialed && !self.per_user) {
431            return next.run(req).await;
432        }
433        let key = self.key(&req, credentialed);
434        if key.len() > MAX_KEY_LEN {
435            // keys are kept in memory: a request cannot make us hold a huge one
436            return next.run(req).await;
437        }
438        let path = req.uri.path().to_string();
439        let finish = |mut res: Response| {
440            for v in &self.vary {
441                crate::cors::add_vary(&mut res, v);
442            }
443            res
444        };
445
446        // A waiter whose leader produced nothing storable (error, private
447        // response) leads once itself; after that it just runs the handler.
448        let mut waited = false;
449        loop {
450            match self.plan(&key, &path) {
451                Plan::Hit(stored, age) => return Ok(finish(stored.to_response("HIT", age))),
452                Plan::Stale(stored, age, refresh) => {
453                    if refresh {
454                        self.spawn_refresh(key, path, next, req.duplicate());
455                    }
456                    return Ok(finish(stored.to_response("STALE", age)));
457                }
458                Plan::Wait(_) if waited => return next.run(req).await,
459                Plan::Wait(mut rx) => {
460                    let outcome = loop {
461                        if let Flight::Done(result) = &*rx.borrow_and_update() {
462                            break Some(result.clone());
463                        }
464                        if rx.changed().await.is_err() {
465                            break None;
466                        }
467                    };
468                    match outcome.flatten() {
469                        Some(stored) => {
470                            return Ok(finish(stored.to_response("COALESCED", Duration::ZERO)))
471                        }
472                        None => waited = true,
473                    }
474                }
475                Plan::Lead(leader) => {
476                    let res = next.run(req).await.unwrap_or_else(Response::from_error);
477                    return Ok(match self.lifetime(&res) {
478                        Some(fresh_for) => {
479                            let stored = Arc::new(self.snapshot(&res));
480                            self.store(&key, &path, stored.clone(), fresh_for);
481                            leader.finish(Some(stored));
482                            finish(res.with_header("x-cache", "MISS"))
483                        }
484                        None => {
485                            leader.finish(None);
486                            res.with_header("x-cache", "BYPASS")
487                        }
488                    });
489                }
490            }
491        }
492    }
493}
494
495impl ResponseCache {
496    /// Refresh a stale entry without making any client wait
497    fn spawn_refresh(&self, key: String, path: String, next: Next, req: Request) {
498        let cache = self.clone();
499        tokio::spawn(async move {
500            let res = next.run(req).await.unwrap_or_else(Response::from_error);
501            match cache.lifetime(&res) {
502                Some(fresh_for) => {
503                    let stored = Arc::new(cache.snapshot(&res));
504                    cache.store(&key, &path, stored, fresh_for);
505                }
506                None => {
507                    // do not keep serving something the origin no longer allows
508                    if let Some(slot) = cache.lock().get_mut(&key) {
509                        slot.entry = None;
510                    }
511                }
512            }
513            cache.end_refresh(&key);
514        });
515    }
516}
517
518#[cfg(test)]
519mod tests {
520    use super::*;
521
522    fn res(status: u16, cc: Option<&str>) -> Response {
523        let mut r = Response::new(StatusCode::from_u16(status).unwrap()).with_body("x");
524        if let Some(cc) = cc {
525            r = r.with_header("cache-control", cc);
526        }
527        r
528    }
529
530    #[test]
531    fn lifetime_follows_cache_control() {
532        let c = ResponseCache::new().ttl(Duration::from_secs(60));
533        assert_eq!(c.lifetime(&res(200, None)), Some(Duration::from_secs(60)));
534        assert_eq!(
535            c.lifetime(&res(200, Some("max-age=5"))),
536            Some(Duration::from_secs(5))
537        );
538        assert_eq!(
539            c.lifetime(&res(200, Some("public, max-age=5, s-maxage=9"))),
540            Some(Duration::from_secs(9))
541        );
542        for no in ["no-store", "private", "no-cache", "max-age=0"] {
543            assert_eq!(c.lifetime(&res(200, Some(no))), None, "{no}");
544        }
545        assert_eq!(c.lifetime(&res(404, None)), None);
546        assert_eq!(
547            c.lifetime(&res(200, None).with_header("set-cookie", "a=b")),
548            None
549        );
550    }
551
552    #[test]
553    fn big_bodies_are_not_stored() {
554        let c = ResponseCache::new().max_body(1);
555        assert_eq!(c.lifetime(&Response::ok().with_body("xx")), None);
556    }
557
558    #[test]
559    fn statuses_are_configurable() {
560        let c = ResponseCache::new().statuses([200, 404]);
561        assert!(c.lifetime(&res(404, None)).is_some());
562    }
563}
ullun/cache.rs

ullun/
cache.rs
