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

ullun/
cache.rs
