Skip to main content

trillium_cache/
fs.rs

1//! Filesystem-backed [`CacheStorage`].
2
3use crate::{
4    CacheKey, CachePolicy, CacheStorage, PutHandle, StoredEntry, fs_shims, policy::PolicyRepr,
5};
6use futures_lite::{AsyncRead, AsyncWrite, AsyncWriteExt};
7use moka::{notification::RemovalCause, sync::Cache};
8use sha2::{Digest, Sha256};
9use std::{
10    fmt::{self, Debug, Formatter, Write as _},
11    io,
12    path::{Path, PathBuf},
13    pin::Pin,
14    sync::{
15        Arc,
16        atomic::{AtomicU64, Ordering},
17    },
18    task::{Context, Poll},
19    time::Duration,
20};
21use trillium_http::{Body, BodySource, Headers};
22
23const META_SUFFIX: &str = ".meta";
24const BODY_SUFFIX: &str = ".body";
25
26// Disk caches are cheap to grow relative to memory, so the default ceiling is larger than
27// `InMemoryStorage`'s.
28const DEFAULT_MAX_CAPACITY_BYTES: u64 = 1024 * 1024 * 1024;
29
30// Disambiguates concurrent temporary files under one directory. Process-local; on-disk
31// temporaries from a previous run are never read (only committed files are).
32static TEMP_COUNTER: AtomicU64 = AtomicU64::new(0);
33
34/// Filesystem-backed cache storage rooted at a directory.
35///
36/// Persists cached responses under a root directory so they survive process restarts. Each
37/// response is two files: a `<hash>.meta` sidecar holding the [`CachePolicy`] and any trailers
38/// as an rkyv-encoded binary blob, and a `<hash>.body` holding the raw body bytes and nothing
39/// else. Bodies stream in and out — [`put`] writes to a temporary file the caller feeds
40/// incrementally, and [`open`] streams the stored body back without loading it into memory. The
41/// metadata is not human-readable; it is optimized for compact, fast loading rather than
42/// inspection.
43///
44/// Defaults to a 1 GiB byte cap; override with
45/// [`with_max_capacity_bytes`][Self::with_max_capacity_bytes] or remove it with
46/// [`unbounded`][Self::unbounded]. Optional time-based eviction is available through
47/// [`with_time_to_idle`][Self::with_time_to_idle] and
48/// [`with_time_to_live`][Self::with_time_to_live] (off by default).
49///
50/// `Clone` is cheap — clones share the same root and capacity index, and see each other's
51/// writes.
52///
53/// # Layout
54///
55/// Entries live at `<root>/<key-hash>/<variant-hash>.{meta,body}`. The key hash is a SHA-256
56/// of the request method and URL; the variant hash is a SHA-256 of the `Vary` signature, so
57/// the multiple variants of one URL are sibling files in the same directory and [`get`]
58/// enumerates them by reading that directory. Writing a variant that already exists replaces
59/// it.
60///
61/// # Durability
62///
63/// Writes commit by renaming a fully-written temporary file into place, and the `.meta` is
64/// written last — a reader treats it as the commit marker, so a half-written or abandoned entry
65/// (a [`PutHandle`] dropped without [`finalize`]) is never visible to [`get`].
66///
67/// # Capacity
68///
69/// A byte cap (1 GiB by default) bounds the total stored body size. When a write would push
70/// the total past the cap, least-recently-used variants are evicted — their `.meta` and
71/// `.body` files deleted — until the cache fits. The cap counts body bytes only, per variant,
72/// matching the granularity of the on-disk layout. Reads count as use, so a frequently-served
73/// variant outlives idle ones. Override with [`with_max_capacity_bytes`] or remove the cap with
74/// [`unbounded`].
75///
76/// The cap is tracked in an in-memory index built by scanning the root at construction, so it
77/// survives restarts (recency resets to whatever order the scan encounters). A directory that
78/// grew past the current cap under an older, unbounded configuration is trimmed to fit on the
79/// next construction.
80///
81/// # Expiry
82///
83/// Beyond the size cap, entries can be evicted on a timer:
84/// [`with_time_to_idle`][Self::with_time_to_idle] drops variants not read within a duration,
85/// [`with_time_to_live`][Self::with_time_to_live] drops them a duration after they are stored.
86/// Both delete the variant's files on eviction, just like size eviction. This is best-effort
87/// space reclamation rather than a hard read gate — [`get`][CacheStorage::get] enumerates the
88/// files on disk, so a just-expired variant may still be served in the brief window before its
89/// files are deleted. It is never a correctness hazard: RFC 9111 freshness is enforced by the
90/// [`Cache`](crate::Cache) handler from the stored [`CachePolicy`], independent of this
91/// storage-level expiry. Both clocks are seeded at construction, so a reopened directory times
92/// each entry from the reopen, not from its pre-restart history.
93///
94/// # Runtime
95///
96/// Filesystem access goes through the runtime selected by the `smol`, `tokio`, or `async-std`
97/// feature. Enabling `fs` without one of those compiles but panics on use.
98///
99/// [`put`]: CacheStorage::put
100/// [`get`]: CacheStorage::get
101/// [`open`]: StoredEntry::open
102/// [`finalize`]: PutHandle::finalize
103/// [`with_max_capacity_bytes`]: FileSystemStorage::with_max_capacity_bytes
104/// [`unbounded`]: FileSystemStorage::unbounded
105#[derive(Clone)]
106pub struct FileSystemStorage {
107    root: Arc<PathBuf>,
108    index: Cache<VariantId, u64>,
109    max_capacity_bytes: Option<u64>,
110    time_to_idle: Option<Duration>,
111    time_to_live: Option<Duration>,
112}
113
114impl Debug for FileSystemStorage {
115    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
116        f.debug_struct("FileSystemStorage")
117            .field("root", &self.root)
118            .field("weighted_size", &self.index.weighted_size())
119            .field("max_capacity_bytes", &self.max_capacity_bytes)
120            .field("time_to_idle", &self.time_to_idle)
121            .field("time_to_live", &self.time_to_live)
122            .finish()
123    }
124}
125
126impl FileSystemStorage {
127    /// Construct a storage rooted at `root` with a 1 GiB byte cap. The directory is created
128    /// on demand as entries are written; it need not exist yet. If it exists, it is scanned
129    /// to seed the capacity index, so previously stored entries count against the cap.
130    pub fn new(root: impl Into<PathBuf>) -> Self {
131        let root = Arc::new(root.into());
132        let max_capacity_bytes = Some(DEFAULT_MAX_CAPACITY_BYTES);
133        let index = build_index(Arc::clone(&root), max_capacity_bytes, None, None);
134        scan_root(&root, &index);
135        Self {
136            root,
137            index,
138            max_capacity_bytes,
139            time_to_idle: None,
140            time_to_live: None,
141        }
142    }
143
144    /// Set the maximum total stored body size, in bytes. Least-recently-used variants are
145    /// evicted — their files deleted — when a write would exceed this cap. Defaults to
146    /// 1 GiB. Re-scans the root, so a directory already over the new cap is trimmed to fit.
147    pub fn with_max_capacity_bytes(mut self, bytes: u64) -> Self {
148        self.max_capacity_bytes = Some(bytes);
149        self.rebuild();
150        self
151    }
152
153    /// Remove the size cap. Stored bytes grow without bound. Useful in tests and short-lived
154    /// processes; a cache living on shared disk should prefer the default capped
155    /// configuration.
156    pub fn unbounded(mut self) -> Self {
157        self.max_capacity_bytes = None;
158        self.rebuild();
159        self
160    }
161
162    /// Evict entries that have not been read in this duration, deleting their files. Off by
163    /// default.
164    ///
165    /// This is best-effort space reclamation, not a read gate: [`get`](CacheStorage::get)
166    /// enumerates the files on disk rather than the expiry index, so a just-expired variant may
167    /// still be served in the window before the eviction is processed and its files deleted. It
168    /// never serves *stale* content — RFC 9111 freshness is enforced by the
169    /// [`Cache`](crate::Cache) handler from the stored [`CachePolicy`], independent of this
170    /// storage-level expiry. (The in-memory backend's idle eviction, by contrast, `get`
171    /// observes as a hard miss.)
172    ///
173    /// The idle clock is seeded at construction: a reopened directory counts idle time from the
174    /// reopen, not from each entry's last read before the restart.
175    pub fn with_time_to_idle(mut self, duration: Duration) -> Self {
176        self.time_to_idle = Some(duration);
177        self.rebuild();
178        self
179    }
180
181    /// Evict entries this duration after their last insert regardless of access, deleting their
182    /// files. Off by default.
183    ///
184    /// Best-effort like [`with_time_to_idle`](Self::with_time_to_idle): a just-expired variant
185    /// may be served until its files are deleted, but never past RFC 9111 freshness, which the
186    /// [`Cache`](crate::Cache) handler enforces separately. This TTL is independent of that
187    /// freshness — an entry may be evicted while still fresh, or linger briefly past it.
188    ///
189    /// The clock is seeded at construction, so a reopened directory counts each entry's TTL
190    /// from the reopen rather than its original store time.
191    pub fn with_time_to_live(mut self, duration: Duration) -> Self {
192        self.time_to_live = Some(duration);
193        self.rebuild();
194        self
195    }
196
197    /// Approximate total stored body size, in bytes, currently counted against the cap.
198    /// Eventually consistent — call [`run_pending_tasks`][Self::run_pending_tasks] first for
199    /// a settled value.
200    pub fn weighted_size(&self) -> u64 {
201        self.index.weighted_size()
202    }
203
204    /// Approximate count of stored variants. Eventually consistent — call
205    /// [`run_pending_tasks`][Self::run_pending_tasks] first for a settled value.
206    pub fn entry_count(&self) -> u64 {
207        self.index.entry_count()
208    }
209
210    /// Flush pending eviction bookkeeping, including deletion of files for evicted variants.
211    /// Call before reading [`weighted_size`][Self::weighted_size] or
212    /// [`entry_count`][Self::entry_count] when an exact value matters.
213    pub async fn run_pending_tasks(&self) {
214        self.index.run_pending_tasks();
215    }
216
217    // The capacity index has no resize API; rebuilding it and re-scanning the root applies a
218    // new cap while preserving on-disk entries (unlike the in-memory backend, disk data
219    // survives a reconfigure).
220    fn rebuild(&mut self) {
221        self.index = build_index(
222            Arc::clone(&self.root),
223            self.max_capacity_bytes,
224            self.time_to_idle,
225            self.time_to_live,
226        );
227        scan_root(&self.root, &self.index);
228    }
229}
230
231// Identity of one stored variant, sufficient to reconstruct its `.meta`/`.body` paths under
232// a known root. Keys the capacity index.
233#[derive(Clone, Hash, PartialEq, Eq)]
234struct VariantId {
235    key_hash: String,
236    variant_hash: String,
237}
238
239// Build the capacity index. The eviction listener deletes a variant's files when moka
240// evicts it for size or expiry; replacement and explicit invalidation are handled at their
241// call sites, so the listener ignores those causes.
242fn build_index(
243    root: Arc<PathBuf>,
244    max_capacity_bytes: Option<u64>,
245    time_to_idle: Option<Duration>,
246    time_to_live: Option<Duration>,
247) -> Cache<VariantId, u64> {
248    let mut builder = Cache::<VariantId, u64>::builder()
249        .weigher(|_key, &body_len| u32::try_from(body_len).unwrap_or(u32::MAX))
250        .eviction_listener(move |id: Arc<VariantId>, _body_len, cause: RemovalCause| {
251            if cause.was_evicted() {
252                let dir = root.join(&id.key_hash);
253                let _ = std::fs::remove_file(dir.join(format!("{}{META_SUFFIX}", id.variant_hash)));
254                let _ = std::fs::remove_file(dir.join(format!("{}{BODY_SUFFIX}", id.variant_hash)));
255            }
256        });
257    if let Some(cap) = max_capacity_bytes {
258        builder = builder.max_capacity(cap);
259    }
260    if let Some(tti) = time_to_idle {
261        builder = builder.time_to_idle(tti);
262    }
263    if let Some(ttl) = time_to_live {
264        builder = builder.time_to_live(ttl);
265    }
266    builder.build()
267}
268
269// Seed the index from the root, counting each committed variant's body length against the
270// cap. Runs on the calling thread with blocking IO — a one-time construction cost — and
271// forces eviction so an over-cap directory is trimmed before the storage is used.
272fn scan_root(root: &Path, index: &Cache<VariantId, u64>) {
273    let Ok(key_dirs) = std::fs::read_dir(root) else {
274        return;
275    };
276    for key_entry in key_dirs.flatten() {
277        let key_dir = key_entry.path();
278        let Some(key_hash) = file_stem_string(&key_dir) else {
279            continue;
280        };
281        let Ok(files) = std::fs::read_dir(&key_dir) else {
282            continue;
283        };
284        for file in files.flatten() {
285            let path = file.path();
286            let Some(variant_hash) = path
287                .file_name()
288                .and_then(|name| name.to_str())
289                .and_then(|name| name.strip_suffix(META_SUFFIX))
290                .map(str::to_string)
291            else {
292                continue;
293            };
294            let body = key_dir.join(format!("{variant_hash}{BODY_SUFFIX}"));
295            let Ok(metadata) = std::fs::metadata(&body) else {
296                continue;
297            };
298            index.insert(
299                VariantId {
300                    key_hash: key_hash.clone(),
301                    variant_hash,
302                },
303                metadata.len(),
304            );
305        }
306    }
307    index.run_pending_tasks();
308}
309
310fn file_stem_string(path: &Path) -> Option<String> {
311    path.file_name()
312        .and_then(|name| name.to_str())
313        .map(str::to_string)
314}
315
316// The rkyv-encoded sidecar written alongside each body. `PolicyRepr` recomputes the derived
317// cache-control fields on load, so only the directly-captured policy fields are stored.
318#[derive(rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)]
319struct StoredMeta {
320    policy: PolicyRepr,
321    trailers: Option<Headers>,
322}
323
324impl CacheStorage for FileSystemStorage {
325    type PutHandle = FsPutHandle;
326    type StoredEntry = FsStoredEntry;
327
328    async fn get(&self, key: &CacheKey) -> Vec<Self::StoredEntry> {
329        let key_hash = key_hash(key);
330        let dir = self.root.join(&key_hash);
331        let Ok(paths) = fs_shims::read_dir_paths(&dir).await else {
332            return Vec::new();
333        };
334
335        let mut entries = Vec::new();
336        for path in paths {
337            let Some(variant_hash) = path
338                .file_name()
339                .and_then(|name| name.to_str())
340                .and_then(|name| name.strip_suffix(META_SUFFIX))
341                .map(str::to_string)
342            else {
343                continue;
344            };
345            let Ok(bytes) = fs_shims::read(&path).await else {
346                continue;
347            };
348            let Ok(meta) = deserialize_meta(&bytes) else {
349                continue;
350            };
351            // Count the lookup as use so a frequently-served variant survives eviction.
352            self.index.get(&VariantId {
353                key_hash: key_hash.clone(),
354                variant_hash: variant_hash.clone(),
355            });
356            entries.push(FsStoredEntry {
357                meta_path: path,
358                body_path: dir.join(format!("{variant_hash}{BODY_SUFFIX}")),
359                policy: meta.policy.into(),
360                trailers: meta.trailers,
361            });
362        }
363        entries
364    }
365
366    async fn put(&self, key: CacheKey, policy: CachePolicy) -> io::Result<Self::PutHandle> {
367        let key_hash = key_hash(&key);
368        let dir = self.root.join(&key_hash);
369        fs_shims::create_dir_all(&dir).await?;
370
371        let variant_hash = variant_hash(&policy);
372        let n = TEMP_COUNTER.fetch_add(1, Ordering::Relaxed);
373        let body_tmp = dir.join(format!("{variant_hash}{BODY_SUFFIX}.tmp.{n}"));
374        let writer = fs_shims::create(&body_tmp).await?;
375
376        Ok(FsPutHandle {
377            writer,
378            body_tmp,
379            body_final: dir.join(format!("{variant_hash}{BODY_SUFFIX}")),
380            meta_tmp: dir.join(format!("{variant_hash}{META_SUFFIX}.tmp.{n}")),
381            meta_final: dir.join(format!("{variant_hash}{META_SUFFIX}")),
382            policy,
383            index: self.index.clone(),
384            variant_id: VariantId {
385                key_hash,
386                variant_hash,
387            },
388            written: 0,
389            committed: false,
390        })
391    }
392
393    async fn invalidate(&self, key: &CacheKey) {
394        let key_hash = key_hash(key);
395        let dir = self.root.join(&key_hash);
396        // Prune the index before removing files; the whole directory goes at once, so the
397        // per-variant eviction listener would be redundant (it skips explicit removals).
398        if let Ok(paths) = fs_shims::read_dir_paths(&dir).await {
399            for path in paths {
400                if let Some(variant_hash) = path
401                    .file_name()
402                    .and_then(|name| name.to_str())
403                    .and_then(|name| name.strip_suffix(META_SUFFIX))
404                {
405                    self.index.invalidate(&VariantId {
406                        key_hash: key_hash.clone(),
407                        variant_hash: variant_hash.to_string(),
408                    });
409                }
410            }
411        }
412        let _ = fs_shims::remove_dir_all(&dir).await;
413    }
414}
415
416/// Streaming [`PutHandle`] for [`FileSystemStorage`].
417///
418/// Body bytes are written to a temporary file as they arrive; [`finalize`][Self::finalize]
419/// renames the body into place and writes the metadata sidecar. Dropping without finalizing
420/// removes the temporary body and stores nothing.
421pub struct FsPutHandle {
422    writer: fs_shims::Writer,
423    body_tmp: PathBuf,
424    body_final: PathBuf,
425    meta_tmp: PathBuf,
426    meta_final: PathBuf,
427    policy: CachePolicy,
428    index: Cache<VariantId, u64>,
429    variant_id: VariantId,
430    written: u64,
431    committed: bool,
432}
433
434impl Debug for FsPutHandle {
435    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
436        f.debug_struct("FsPutHandle")
437            .field("body_final", &self.body_final)
438            .finish_non_exhaustive()
439    }
440}
441
442impl AsyncWrite for FsPutHandle {
443    fn poll_write(
444        self: Pin<&mut Self>,
445        cx: &mut Context<'_>,
446        buf: &[u8],
447    ) -> Poll<io::Result<usize>> {
448        let this = self.get_mut();
449        let poll = Pin::new(&mut this.writer).poll_write(cx, buf);
450        if let Poll::Ready(Ok(n)) = &poll {
451            this.written += *n as u64;
452        }
453        poll
454    }
455
456    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
457        Pin::new(&mut self.get_mut().writer).poll_flush(cx)
458    }
459
460    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
461        Pin::new(&mut self.get_mut().writer).poll_close(cx)
462    }
463}
464
465impl PutHandle for FsPutHandle {
466    async fn finalize(mut self, trailers: Option<Headers>) -> io::Result<()> {
467        self.writer.close().await?;
468        fs_shims::rename(&self.body_tmp, &self.body_final).await?;
469
470        let meta = StoredMeta {
471            policy: PolicyRepr::from(&self.policy),
472            trailers,
473        };
474        let bytes = serialize_meta(&meta)?;
475        fs_shims::write(&self.meta_tmp, &bytes).await?;
476        fs_shims::rename(&self.meta_tmp, &self.meta_final).await?;
477
478        // Account the committed body against the cap. Re-inserting the same variant replaces
479        // its prior weight; the eviction listener ignores the replacement.
480        self.index.insert(self.variant_id.clone(), self.written);
481
482        self.committed = true;
483        Ok(())
484    }
485}
486
487impl Drop for FsPutHandle {
488    fn drop(&mut self) {
489        if !self.committed {
490            let _ = std::fs::remove_file(&self.body_tmp);
491        }
492    }
493}
494
495/// One stored response returned by [`FileSystemStorage::get`].
496///
497/// Holds the metadata; the body stays on disk until [`open`][Self::open] streams it. `Clone`
498/// copies the metadata and re-opens the body file on demand.
499#[derive(Clone)]
500pub struct FsStoredEntry {
501    meta_path: PathBuf,
502    body_path: PathBuf,
503    policy: CachePolicy,
504    trailers: Option<Headers>,
505}
506
507impl Debug for FsStoredEntry {
508    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
509        f.debug_struct("FsStoredEntry")
510            .field("body_path", &self.body_path)
511            .field("has_trailers", &self.trailers.is_some())
512            .finish_non_exhaustive()
513    }
514}
515
516impl StoredEntry for FsStoredEntry {
517    fn policy(&self) -> &CachePolicy {
518        &self.policy
519    }
520
521    async fn refresh_policy(&mut self, new_policy: CachePolicy) -> io::Result<()> {
522        let meta = StoredMeta {
523            policy: PolicyRepr::from(&new_policy),
524            trailers: self.trailers.clone(),
525        };
526        let bytes = serialize_meta(&meta)?;
527        let tmp = temp_sibling(&self.meta_path);
528        fs_shims::write(&tmp, &bytes).await?;
529        fs_shims::rename(&tmp, &self.meta_path).await?;
530
531        self.policy = new_policy;
532        Ok(())
533    }
534
535    async fn open(self) -> io::Result<Body> {
536        let len = fs_shims::metadata_len(&self.body_path).await?;
537        let reader = fs_shims::open(&self.body_path).await?;
538        let source = FsBodySource {
539            reader,
540            trailers: self.trailers,
541        };
542        Ok(Body::new_with_trailers(source, Some(len)))
543    }
544}
545
546// BodySource over a stored body file. Reads stream straight from the file; trailers surface
547// after EOF.
548struct FsBodySource {
549    reader: fs_shims::Reader,
550    trailers: Option<Headers>,
551}
552
553impl AsyncRead for FsBodySource {
554    fn poll_read(
555        self: Pin<&mut Self>,
556        cx: &mut Context<'_>,
557        buf: &mut [u8],
558    ) -> Poll<io::Result<usize>> {
559        Pin::new(&mut self.get_mut().reader).poll_read(cx, buf)
560    }
561}
562
563impl BodySource for FsBodySource {
564    fn trailers(self: Pin<&mut Self>) -> Option<Headers> {
565        self.get_mut().trailers.take()
566    }
567}
568
569fn hash_hex(bytes: &[u8]) -> String {
570    let mut hasher = Sha256::new();
571    hasher.update(bytes);
572    finalize_hex(hasher)
573}
574
575fn key_hash(key: &CacheKey) -> String {
576    hash_hex(key.to_string().as_bytes())
577}
578
579fn variant_hash(policy: &CachePolicy) -> String {
580    let mut hasher = Sha256::new();
581    for (name, value) in &policy.vary_snapshot {
582        hasher.update(name.as_bytes());
583        hasher.update([0]);
584        match value {
585            Some(value) => {
586                hasher.update([1]);
587                hasher.update(value.as_bytes());
588            }
589            None => hasher.update([0]),
590        }
591        hasher.update([0]);
592    }
593    finalize_hex(hasher)
594}
595
596fn finalize_hex(hasher: Sha256) -> String {
597    let digest = hasher.finalize();
598    let mut out = String::with_capacity(digest.len() * 2);
599    for byte in digest {
600        write!(out, "{byte:02x}").expect("writing to a String cannot fail");
601    }
602    out
603}
604
605// A unique sibling temp path for atomically rewriting `path`.
606fn temp_sibling(path: &Path) -> PathBuf {
607    let n = TEMP_COUNTER.fetch_add(1, Ordering::Relaxed);
608    let mut name = path.as_os_str().to_owned();
609    name.push(format!(".tmp.{n}"));
610    PathBuf::from(name)
611}
612
613fn serialize_meta(meta: &StoredMeta) -> io::Result<rkyv::util::AlignedVec> {
614    rkyv::to_bytes::<rkyv::rancor::Error>(meta)
615        .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))
616}
617
618fn deserialize_meta(bytes: &[u8]) -> io::Result<StoredMeta> {
619    // A disk read lands in a buffer aligned only to 1, but rkyv's validated access requires
620    // the archived root to be aligned; copy into an `AlignedVec` before decoding.
621    let mut aligned = rkyv::util::AlignedVec::<16>::new();
622    aligned.extend_from_slice(bytes);
623    rkyv::from_bytes::<StoredMeta, rkyv::rancor::Error>(&aligned)
624        .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))
625}
626
627#[cfg(test)]
628mod tests {
629    use super::*;
630    use crate::test_helpers::*;
631    use futures_lite::{AsyncReadExt, AsyncWriteExt};
632    use std::time::{Duration, SystemTime};
633    use tempfile::TempDir;
634    use trillium_client::Conn;
635    use trillium_http::{KnownHeaderName::*, Method, Status};
636    use trillium_testing::{TestResult, harness, test};
637
638    fn key() -> CacheKey {
639        CacheKey::new(Method::Get, "http://example.com/".parse().unwrap())
640    }
641
642    fn new_storage() -> (TempDir, FileSystemStorage) {
643        let dir = tempfile::tempdir().unwrap();
644        let storage = FileSystemStorage::new(dir.path());
645        (dir, storage)
646    }
647
648    async fn store_at(storage: &FileSystemStorage, url: &str, body: &[u8]) {
649        let key = CacheKey::new(Method::Get, url.parse().unwrap());
650        let conn = exchange(
651            Method::Get,
652            &[],
653            Status::Ok,
654            &[(CacheControl, "max-age=600")],
655        );
656        let policy = policy_from(&conn, SystemTime::now(), private_cache());
657        let mut handle = storage.put(key, policy).await.unwrap();
658        handle.write_all(body).await.unwrap();
659        handle.finalize(None).await.unwrap();
660    }
661
662    async fn store(storage: &FileSystemStorage, key: CacheKey, conn: &Conn, body: &[u8]) {
663        let policy = policy_from(conn, SystemTime::now(), private_cache());
664        let mut handle = storage.put(key, policy).await.unwrap();
665        handle.write_all(body).await.unwrap();
666        handle.finalize(None).await.unwrap();
667    }
668
669    async fn read_body(entry: FsStoredEntry) -> Vec<u8> {
670        let mut body = entry.open().await.unwrap();
671        let mut buf = Vec::new();
672        body.read_to_end(&mut buf).await.unwrap();
673        buf
674    }
675
676    #[test(harness)]
677    async fn get_missing_key_returns_empty() -> TestResult {
678        let (_dir, storage) = new_storage();
679        assert!(storage.get(&key()).await.is_empty());
680        Ok(())
681    }
682
683    #[test(harness)]
684    async fn put_then_get_round_trips_through_disk() -> TestResult {
685        let (_dir, storage) = new_storage();
686        let conn = exchange(
687            Method::Get,
688            &[],
689            Status::Ok,
690            &[(CacheControl, "max-age=600")],
691        );
692        store(&storage, key(), &conn, b"hello").await;
693        let result = storage.get(&key()).await;
694        assert_eq!(result.len(), 1);
695        assert_eq!(read_body(result[0].clone()).await, b"hello");
696        Ok(())
697    }
698
699    #[test(harness)]
700    async fn put_with_same_vary_replaces() -> TestResult {
701        let (_dir, storage) = new_storage();
702        let conn = exchange(
703            Method::Get,
704            &[(AcceptEncoding, "gzip")],
705            Status::Ok,
706            &[(CacheControl, "max-age=600"), (Vary, "Accept-Encoding")],
707        );
708        store(&storage, key(), &conn, b"v1").await;
709        store(&storage, key(), &conn, b"v2").await;
710        let result = storage.get(&key()).await;
711        assert_eq!(result.len(), 1);
712        assert_eq!(read_body(result[0].clone()).await, b"v2");
713        Ok(())
714    }
715
716    #[test(harness)]
717    async fn put_with_different_vary_appends() -> TestResult {
718        let (_dir, storage) = new_storage();
719        let gzip = exchange(
720            Method::Get,
721            &[(AcceptEncoding, "gzip")],
722            Status::Ok,
723            &[(CacheControl, "max-age=600"), (Vary, "Accept-Encoding")],
724        );
725        let br = exchange(
726            Method::Get,
727            &[(AcceptEncoding, "br")],
728            Status::Ok,
729            &[(CacheControl, "max-age=600"), (Vary, "Accept-Encoding")],
730        );
731        store(&storage, key(), &gzip, b"gz").await;
732        store(&storage, key(), &br, b"br").await;
733        assert_eq!(storage.get(&key()).await.len(), 2);
734        Ok(())
735    }
736
737    #[test(harness)]
738    async fn invalidate_removes_all_entries_for_key() -> TestResult {
739        let (_dir, storage) = new_storage();
740        let conn = exchange(
741            Method::Get,
742            &[],
743            Status::Ok,
744            &[(CacheControl, "max-age=600")],
745        );
746        store(&storage, key(), &conn, b"x").await;
747        storage.invalidate(&key()).await;
748        assert!(storage.get(&key()).await.is_empty());
749        Ok(())
750    }
751
752    #[test(harness)]
753    async fn invalidate_does_not_touch_other_keys() -> TestResult {
754        let (_dir, storage) = new_storage();
755        let conn = exchange(
756            Method::Get,
757            &[],
758            Status::Ok,
759            &[(CacheControl, "max-age=600")],
760        );
761        let key_a = CacheKey::new(Method::Get, "http://a.example/".parse().unwrap());
762        let key_b = CacheKey::new(Method::Get, "http://b.example/".parse().unwrap());
763        store(&storage, key_a.clone(), &conn, b"a").await;
764        store(&storage, key_b.clone(), &conn, b"b").await;
765        storage.invalidate(&key_a).await;
766        assert!(storage.get(&key_a).await.is_empty());
767        assert_eq!(storage.get(&key_b).await.len(), 1);
768        Ok(())
769    }
770
771    #[test(harness)]
772    async fn drop_put_handle_without_finalize_discards() -> TestResult {
773        let (_dir, storage) = new_storage();
774        let conn = exchange(
775            Method::Get,
776            &[],
777            Status::Ok,
778            &[(CacheControl, "max-age=600")],
779        );
780        let policy = policy_from(&conn, SystemTime::now(), private_cache());
781        let mut handle = storage.put(key(), policy).await.unwrap();
782        handle.write_all(b"partial").await.unwrap();
783        drop(handle);
784        assert!(storage.get(&key()).await.is_empty());
785        Ok(())
786    }
787
788    #[test(harness)]
789    async fn refresh_policy_updates_meta_and_keeps_body() -> TestResult {
790        let (_dir, storage) = new_storage();
791        let conn = exchange(
792            Method::Get,
793            &[],
794            Status::Ok,
795            &[(CacheControl, "max-age=600")],
796        );
797        store(&storage, key(), &conn, b"body").await;
798
799        let mut entries = storage.get(&key()).await;
800        let original_time = entries[0].policy().response_time;
801        let refreshed = exchange(
802            Method::Get,
803            &[],
804            Status::Ok,
805            &[(CacheControl, "max-age=1200")],
806        );
807        let new_policy = policy_from(
808            &refreshed,
809            original_time + Duration::from_secs(100),
810            private_cache(),
811        );
812        entries[0].refresh_policy(new_policy).await.unwrap();
813
814        let fresh = storage.get(&key()).await;
815        assert_eq!(fresh.len(), 1);
816        assert_ne!(fresh[0].policy().response_time, original_time);
817        assert_eq!(read_body(fresh[0].clone()).await, b"body");
818        Ok(())
819    }
820
821    #[test(harness)]
822    async fn trailers_round_trip() -> TestResult {
823        let (_dir, storage) = new_storage();
824        let conn = exchange(
825            Method::Get,
826            &[],
827            Status::Ok,
828            &[(CacheControl, "max-age=600")],
829        );
830        let policy = policy_from(&conn, SystemTime::now(), private_cache());
831        let mut handle = storage.put(key(), policy).await.unwrap();
832        handle.write_all(b"data").await.unwrap();
833        let mut trailers = Headers::new();
834        trailers.insert("x-checksum", "abc123");
835        handle.finalize(Some(trailers)).await.unwrap();
836
837        let entry = storage.get(&key()).await.remove(0);
838        let mut body = entry.open().await.unwrap();
839        let mut buf = Vec::new();
840        body.read_to_end(&mut buf).await.unwrap();
841        assert_eq!(buf, b"data");
842        let trailers = body
843            .trailers()
844            .expect("stored trailers should surface after EOF");
845        assert_eq!(trailers.get_str("x-checksum"), Some("abc123"));
846        Ok(())
847    }
848
849    #[test(harness)]
850    async fn persists_across_new_storage_on_same_root() -> TestResult {
851        let dir = tempfile::tempdir().unwrap();
852        let conn = exchange(
853            Method::Get,
854            &[],
855            Status::Ok,
856            &[(CacheControl, "max-age=600")],
857        );
858        {
859            let storage = FileSystemStorage::new(dir.path());
860            store(&storage, key(), &conn, b"persisted").await;
861        }
862
863        // A brand-new storage over the same directory sees the prior instance's entry.
864        let reopened = FileSystemStorage::new(dir.path());
865        let result = reopened.get(&key()).await;
866        assert_eq!(result.len(), 1);
867        assert_eq!(read_body(result[0].clone()).await, b"persisted");
868        Ok(())
869    }
870
871    #[test(harness)]
872    async fn size_cap_evicts_and_deletes_files() -> TestResult {
873        // Cap at 1 KiB; write ten 600-byte bodies under distinct URLs.
874        let dir = tempfile::tempdir().unwrap();
875        let storage = FileSystemStorage::new(dir.path()).with_max_capacity_bytes(1024);
876        let body = vec![b'x'; 600];
877        for i in 0..10 {
878            store_at(&storage, &format!("http://example.com/{i}"), &body).await;
879        }
880        storage.run_pending_tasks().await;
881        assert!(
882            storage.weighted_size() <= 1024,
883            "weighted size {} should be within cap of 1024",
884            storage.weighted_size()
885        );
886
887        // A fresh unbounded scan of the same root reflects only the files still on disk, so
888        // the low total proves evicted variants' files were actually deleted, not just
889        // forgotten by the index.
890        let reopened = FileSystemStorage::new(dir.path()).unbounded();
891        assert!(
892            reopened.weighted_size() <= 1024,
893            "on-disk bytes {} should be within cap of 1024",
894            reopened.weighted_size()
895        );
896        Ok(())
897    }
898
899    #[test(harness)]
900    async fn rebuild_scan_trims_over_cap_directory() -> TestResult {
901        let dir = tempfile::tempdir().unwrap();
902        let body = vec![b'x'; 600];
903        {
904            let unbounded = FileSystemStorage::new(dir.path()).unbounded();
905            for i in 0..10 {
906                store_at(&unbounded, &format!("http://example.com/{i}"), &body).await;
907            }
908            unbounded.run_pending_tasks().await;
909            assert_eq!(unbounded.entry_count(), 10);
910        }
911
912        // Reopening with a cap trims the pre-existing directory to fit during construction.
913        let capped = FileSystemStorage::new(dir.path()).with_max_capacity_bytes(1024);
914        assert!(
915            capped.weighted_size() <= 1024,
916            "weighted size {} should be within cap of 1024",
917            capped.weighted_size()
918        );
919        Ok(())
920    }
921
922    #[test(harness)]
923    async fn unbounded_keeps_all_entries() -> TestResult {
924        let dir = tempfile::tempdir().unwrap();
925        let storage = FileSystemStorage::new(dir.path()).unbounded();
926        let body = vec![b'x'; 600];
927        for i in 0..10 {
928            store_at(&storage, &format!("http://example.com/{i}"), &body).await;
929        }
930        storage.run_pending_tasks().await;
931        assert_eq!(storage.entry_count(), 10);
932        assert_eq!(storage.weighted_size(), 6000);
933        Ok(())
934    }
935
936    #[test(harness)]
937    async fn replacing_a_variant_does_not_double_count() -> TestResult {
938        let (_dir, storage) = new_storage();
939        store_at(&storage, "http://example.com/", &vec![b'x'; 600]).await;
940        store_at(&storage, "http://example.com/", &vec![b'y'; 300]).await;
941        storage.run_pending_tasks().await;
942        assert_eq!(storage.entry_count(), 1);
943        assert_eq!(storage.weighted_size(), 300);
944        Ok(())
945    }
946
947    // Generous margin (>2x the TTL) over the real clock keeps these timing tests robust under
948    // loaded CI; blocking sleeps are fine in a test and advance moka's Instant-based expiry.
949    #[test(harness)]
950    async fn time_to_live_evicts_and_deletes_files() -> TestResult {
951        let dir = tempfile::tempdir().unwrap();
952        let storage =
953            FileSystemStorage::new(dir.path()).with_time_to_live(Duration::from_millis(50));
954        store_at(&storage, "http://example.com/", b"x").await;
955        storage.run_pending_tasks().await;
956        assert_eq!(storage.entry_count(), 1);
957
958        std::thread::sleep(Duration::from_millis(120));
959        storage.run_pending_tasks().await;
960        assert_eq!(storage.entry_count(), 0);
961
962        // A fresh scan of the same root proves the eviction listener deleted the files, rather
963        // than the index merely forgetting them.
964        let reopened = FileSystemStorage::new(dir.path()).unbounded();
965        assert_eq!(reopened.entry_count(), 0);
966        Ok(())
967    }
968
969    #[test(harness)]
970    async fn time_to_idle_evicts_unread_entries() -> TestResult {
971        let dir = tempfile::tempdir().unwrap();
972        let storage =
973            FileSystemStorage::new(dir.path()).with_time_to_idle(Duration::from_millis(50));
974        store_at(&storage, "http://example.com/", b"x").await;
975        storage.run_pending_tasks().await;
976        assert_eq!(storage.entry_count(), 1);
977
978        // No reads, so the entry sits idle past its TTI and is evicted with its files.
979        std::thread::sleep(Duration::from_millis(120));
980        storage.run_pending_tasks().await;
981        assert_eq!(storage.entry_count(), 0);
982        let reopened = FileSystemStorage::new(dir.path()).unbounded();
983        assert_eq!(reopened.entry_count(), 0);
984        Ok(())
985    }
986}