Skip to main content

kernel/resolution/
identification_cache.rs

1//! A cache over [`identify`](crate::resolution::identify) keyed by an on-disk
2//! freshness signature. Identifying a file model reads its GGUF/safetensors
3//! header and stats its directory; over a large shelf that repeats on every
4//! resolution pass. The cache skips it when a model's bytes are unchanged.
5//!
6//! The cheap source kinds (builtin/endpoint/ollama) are never cached — their
7//! identification is a fixed profile or a small manifest read, so a cache entry
8//! would cost more than it saves.
9//!
10//! The freshness signature is depth-1: it folds the model path and its immediate
11//! children, not files nested deeper (an HF-cache `snapshots/<ref>/config.json`
12//! or a diffusers `transformer/config.json`). An in-place edit *below* the top
13//! level is therefore not noticed. Adding, removing, or replacing an immediate
14//! child (the common re-download shape) is caught.
15
16use std::collections::HashMap;
17use std::path::Path;
18use std::sync::Mutex;
19use std::time::SystemTime;
20
21use crate::records::{ModelRecord, SourceKind};
22use crate::resolution::identity::{IdentifiedModel, identify as run_identify};
23
24struct Entry {
25    mtime: SystemTime,
26    size: i64,
27    identified: IdentifiedModel,
28}
29
30#[derive(Default)]
31struct CacheState {
32    entries: HashMap<String, Entry>,
33    hits: usize,
34}
35
36/// A freshness-keyed cache of identification results, safe to share across
37/// threads. A cache miss (or any lock failure) falls back to a direct
38/// [`identify`](crate::resolution::identify), so the cache is only ever an
39/// optimization. A changed top-level mtime or size (an added/removed/replaced
40/// immediate child, or a rewritten file) invalidates the entry; see the module
41/// docs for the depth-1 limitation.
42#[derive(Default)]
43pub struct IdentificationCache {
44    state: Mutex<CacheState>,
45}
46
47impl IdentificationCache {
48    /// An empty cache.
49    pub fn new() -> Self {
50        Self::default()
51    }
52
53    /// The number of cache hits so far (for diagnostics/tests).
54    pub fn hit_count(&self) -> usize {
55        self.state.lock().map_or(0, |state| state.hits)
56    }
57
58    /// Identify `record`, returning a cached result when the model's bytes are
59    /// unchanged since the last identification.
60    pub fn identify(&self, record: &ModelRecord) -> IdentifiedModel {
61        if is_uncached(&record.source.kind) {
62            return run_identify(record);
63        }
64        let key = cache_key(record);
65        let Some((mtime, size)) = freshness_signature(Path::new(&record.source.path)) else {
66            return run_identify(record);
67        };
68        if let Ok(mut state) = self.state.lock()
69            && let Some(entry) = state.entries.get(&key)
70            && entry.mtime == mtime
71            && entry.size == size
72        {
73            let identified = entry.identified.clone();
74            state.hits += 1;
75            return identified;
76        }
77        // Identify outside the lock — it does filesystem I/O — then record it.
78        let identified = run_identify(record);
79        if let Ok(mut state) = self.state.lock() {
80            state.entries.insert(
81                key,
82                Entry {
83                    mtime,
84                    size,
85                    identified: identified.clone(),
86                },
87            );
88        }
89        identified
90    }
91}
92
93/// The kinds whose identification is too cheap to be worth caching. Matched on the
94/// kind's string to avoid allocating a `SourceKind` (a `String` newtype) per check
95/// on the resolution hot path.
96fn is_uncached(kind: &SourceKind) -> bool {
97    matches!(kind.as_str(), "builtin" | "endpoint" | "ollama")
98}
99
100/// The cache key for `record`. `identify` reads not just the path but — for
101/// HF-cache/diffusers records — the source `reference` (which snapshot) and `repo`
102/// (the pipeline repo hint), so two revisions sharing a path must not collide.
103/// Keying on the path alone would collide, so these fields are included to prevent
104/// a different revision being served the wrong identification.
105fn cache_key(record: &ModelRecord) -> String {
106    format!(
107        "{}\u{1f}{}\u{1f}{}",
108        record.source.path,
109        record.source.reference.as_deref().unwrap_or(""),
110        record.source.repo.as_deref().unwrap_or(""),
111    )
112}
113
114/// A `(mtime, size)` signature that changes whenever a model's bytes change. For
115/// a file it is the file's own modification time and length; for a directory
116/// (a multi-file model) it is the newest child mtime and the total size across
117/// the directory and its immediate children — enough to notice a re-download or
118/// an added/replaced shard without hashing anything.
119fn freshness_signature(path: &Path) -> Option<(SystemTime, i64)> {
120    let metadata = std::fs::metadata(path).ok()?;
121    let mtime = metadata.modified().ok()?;
122    let size = metadata.len() as i64;
123    if !metadata.is_dir() {
124        return Some((mtime, size));
125    }
126    let mut latest = mtime;
127    let mut total = size;
128    if let Ok(entries) = std::fs::read_dir(path) {
129        for entry in entries.flatten() {
130            let Ok(child) = entry.metadata() else {
131                continue;
132            };
133            if let Ok(child_mtime) = child.modified()
134                && child_mtime > latest
135            {
136                latest = child_mtime;
137            }
138            // Wrapping add: the total is an identity signature, not a real byte
139            // count, so overflow only needs to stay consistent.
140            total = total.wrapping_add(child.len() as i64);
141        }
142    }
143    Some((latest, total))
144}