Skip to main content

arama_cache/core/
image.rs

1//! `ImageCacheWriter` / `ImageCacheReader` — image-specific cache handles.
2//!
3//! Backed by `localcache` (RFC 002): a [`ConnectionPool`] serializes
4//! writes; a [`ReadPool`] of `read_conns` read-only connections serves
5//! parallel lookups. Both are `Arc`-based, so `Clone` is cheap.
6
7use std::path::{Path, PathBuf};
8use std::sync::Arc;
9
10use localcache::{CacheEntry, CacheStatus, ConnectionPool, ReadPool};
11use rayon::prelude::*;
12
13use crate::core::engine::{
14    CacheConfig, DbLocation, IMAGE_PAYLOAD_VERSION, NAMESPACE_IMAGE, Result, cache_options,
15    ensure_db_dir, ensure_schema, is_fresh, read_pool_size,
16};
17use crate::core::payload::ImagePayload;
18use crate::core::thumbnail::{generate_image_thumbnail, thumbnail_dest};
19use crate::types::{
20    CacheRead, DirCacheSummary, ImageCacheEntry, ImageFeatures, LookupResult, UpsertImageRequest,
21};
22
23// ---------------------------------------------------------------------------
24// Config
25// ---------------------------------------------------------------------------
26
27#[derive(Debug, Clone, Default)]
28pub struct ImageCacheConfig {
29    pub cache_config: CacheConfig,
30}
31
32// ---------------------------------------------------------------------------
33// ImageCacheWriter
34// ---------------------------------------------------------------------------
35
36/// Update handle for image files.
37///
38/// - Generates thumbnails with the `image` crate (224×224 JPEG).
39/// - `Clone` only bumps `Arc` counters.
40#[derive(Clone)]
41pub struct ImageCacheWriter {
42    write: ConnectionPool<ImagePayload>,
43    read: ReadPool<ImagePayload>,
44    config: Arc<ImageCacheConfig>,
45}
46
47impl ImageCacheWriter {
48    pub fn as_session(config: ImageCacheConfig) -> Result<Self> {
49        let options = cache_options(&config.cache_config, NAMESPACE_IMAGE, IMAGE_PAYLOAD_VERSION);
50        // Create the parent directory before localcache touches SQLite.
51        ensure_db_dir(&options)?;
52        // The writable engine is opened first: it creates the database
53        // file and schema, which the read-only pool cannot.
54        let write = ConnectionPool::open(options.clone())?;
55        // Entries written by an older pipeline version are dead weight.
56        write.with(|e| e.purge_stale_versions())?;
57        let read = ReadPool::open(options, read_pool_size(&config.cache_config))?;
58        Ok(Self {
59            write,
60            read,
61            config: Arc::new(config),
62        })
63    }
64
65    pub fn onetime(location: DbLocation) -> Result<Self> {
66        Self::as_session(ImageCacheConfig {
67            cache_config: CacheConfig {
68                db_location: location,
69                ..CacheConfig::default()
70            },
71        })
72    }
73
74    // -----------------------------------------------------------------------
75    // Update API
76    // -----------------------------------------------------------------------
77
78    pub fn upsert(&self, req: UpsertImageRequest) -> Result<()> {
79        let status = self.write.check_status(&req.path)?;
80        self.write_payload(&req, status)
81    }
82
83    /// Batch variant of `upsert`. Returns `(PathBuf, Result<()>)` per
84    /// request.
85    ///
86    /// ## Parallelization strategy
87    ///
88    /// - **Freshness checks and thumbnail generation** run in parallel
89    ///   (rayon over the read pool).
90    /// - **Database writes** are serialized on the write connection.
91    /// - Entries that are already fresh and carry nothing new are
92    ///   skipped entirely — the steady-state startup pass over an
93    ///   unchanged library performs no writes and no hashing.
94    ///
95    /// Individual errors are stored per element; other requests continue.
96    pub fn upsert_all(&self, reqs: Vec<UpsertImageRequest>) -> Vec<(PathBuf, Result<()>)> {
97        // Phase 1 (parallel): freshness + thumbnail generation.
98        let prepared: Vec<(UpsertImageRequest, Result<Prepared>)> = reqs
99            .into_par_iter()
100            .map(|req| {
101                let prep = self.prepare(&req);
102                (req, prep)
103            })
104            .collect();
105
106        // Phase 2 (serial): payload merge + write.
107        prepared
108            .into_iter()
109            .map(|(req, prep)| {
110                let path = req.path.clone();
111                let result = prep.and_then(|p| self.commit(&req, p));
112                (path, result)
113            })
114            .collect()
115    }
116
117    pub fn delete(&self, path: &Path) -> Result<bool> {
118        Ok(self.write.remove(path)?)
119    }
120
121    /// Remove every cached entry whose file lives directly in `dir`,
122    /// deleting the associated thumbnail file (when recorded) as well.
123    /// Non-recursive: entries in subdirectories are untouched.
124    /// Returns the number of entries removed (RFC 004).
125    pub fn delete_in_dir(&self, dir: &Path) -> Result<usize> {
126        let entries = self.read.query_run(|q| q.path_in_dir(dir, false))?;
127        let mut removed = 0;
128        for entry in entries {
129            if let Some(thumb) = &entry.payload.thumbnail_path {
130                // Best-effort: a missing thumbnail file is not an error.
131                let _ = std::fs::remove_file(thumb);
132            }
133            if self.write.remove(&entry.path)? {
134                removed += 1;
135            }
136        }
137        Ok(removed)
138    }
139
140    pub fn list_paths(&self) -> Result<Vec<String>> {
141        let keys = self.write.with(|e| e.keys(None))?;
142        Ok(keys
143            .into_iter()
144            .map(|p| p.to_string_lossy().into_owned())
145            .collect())
146    }
147
148    pub fn as_reader(&self) -> ImageCacheReader {
149        ImageCacheReader {
150            read: self.read.clone(),
151        }
152    }
153
154    pub fn lookup(&self, path: &Path) -> Result<LookupResult<ImageCacheEntry>> {
155        self.as_reader().lookup(path)
156    }
157
158    // -----------------------------------------------------------------------
159    // Internal
160    // -----------------------------------------------------------------------
161
162    /// Parallel-safe preparation: freshness check (read pool) and
163    /// thumbnail generation. No write-connection access.
164    fn prepare(&self, req: &UpsertImageRequest) -> Result<Prepared> {
165        let status = self.read.check_status(&req.path)?;
166        let thumbnail = self.ensure_thumbnail(&req.path, &status)?;
167        Ok(Prepared { status, thumbnail })
168    }
169
170    /// Serial commit: merge with the existing payload (when fresh) and
171    /// write. Skips the write entirely in the fresh, nothing-new case.
172    fn commit(&self, req: &UpsertImageRequest, prep: Prepared) -> Result<()> {
173        let existing = if is_fresh(&prep.status) {
174            self.write.get(&req.path)?.map(|e| e.payload)
175        } else {
176            None
177        };
178
179        // Steady-state skip: fresh entry, no new vector, thumbnail
180        // already recorded.
181        if is_fresh(&prep.status)
182            && req.clip_vector.is_none()
183            && let Some(p) = &existing
184            && (prep.thumbnail.is_none() || p.thumbnail_path == prep.thumbnail)
185        {
186            return Ok(());
187        }
188
189        let mut payload = existing.unwrap_or_default();
190        if let Some(t) = prep.thumbnail {
191            payload.thumbnail_path = Some(t);
192        }
193        if let Some(v) = &req.clip_vector {
194            payload.clip_vector = Some(v.clone());
195        }
196        self.write.set(&req.path, &payload)?;
197        Ok(())
198    }
199
200    /// Single-upsert path: thumbnail + merge + write in one call.
201    fn write_payload(&self, req: &UpsertImageRequest, status: CacheStatus) -> Result<()> {
202        let thumbnail = self.ensure_thumbnail(&req.path, &status)?;
203        self.commit(req, Prepared { status, thumbnail })
204    }
205
206    /// Generate (or reuse) the thumbnail for `path`, returning its
207    /// destination path string. A thumbnail is regenerated when the file
208    /// changed (`status != Fresh`), fixing the v1 behaviour of serving a
209    /// stale thumbnail after the source was modified.
210    fn ensure_thumbnail(&self, path: &Path, status: &CacheStatus) -> Result<Option<String>> {
211        let Some(thumb_dir) = &self.config.cache_config.thumbnail_dir else {
212            return Ok(None);
213        };
214        let dest = thumbnail_dest(thumb_dir, path)?;
215        if !dest.exists() || !is_fresh(status) {
216            generate_image_thumbnail(path, &dest)?;
217        }
218        Ok(Some(dest.to_string_lossy().into_owned()))
219    }
220}
221
222/// Intermediate result of the parallel preparation phase.
223struct Prepared {
224    status: CacheStatus,
225    thumbnail: Option<String>,
226}
227
228// ---------------------------------------------------------------------------
229// ImageCacheReader
230// ---------------------------------------------------------------------------
231
232/// Read-only handle for image files. `Clone` only bumps `Arc` counters;
233/// clones share the same read pool and may be used from many threads.
234#[derive(Clone)]
235pub struct ImageCacheReader {
236    read: ReadPool<ImagePayload>,
237}
238
239impl ImageCacheReader {
240    pub fn as_session(config: ImageCacheConfig) -> Result<Self> {
241        let options = cache_options(&config.cache_config, NAMESPACE_IMAGE, IMAGE_PAYLOAD_VERSION);
242        // Create the parent directory before localcache touches SQLite.
243        ensure_db_dir(&options)?;
244        // A standalone reader may be the first handle to ever touch this
245        // database; make sure the schema exists before going read-only.
246        ensure_schema::<ImagePayload>(&options)?;
247        let read = ReadPool::open(options, read_pool_size(&config.cache_config))?;
248        Ok(Self { read })
249    }
250
251    pub fn onetime(location: DbLocation) -> Result<Self> {
252        Self::as_session(ImageCacheConfig {
253            cache_config: CacheConfig {
254                db_location: location,
255                ..CacheConfig::default()
256            },
257        })
258    }
259
260    pub fn lookup(&self, path: &Path) -> Result<LookupResult<ImageCacheEntry>> {
261        let canonical = match path.canonicalize() {
262            Ok(p) => p,
263            Err(_) => return Ok(LookupResult::Miss),
264        };
265
266        match self.read.check_status(&canonical)? {
267            CacheStatus::Missing => Ok(LookupResult::Miss),
268            CacheStatus::Stale => Ok(LookupResult::Invalidated),
269            CacheStatus::Fresh => match self.read.get(&canonical)? {
270                None => Ok(LookupResult::Miss),
271                Some(entry) => Ok(LookupResult::Hit(to_image_entry(entry))),
272            },
273        }
274    }
275
276    /// Batch variant of `lookup`, parallelized with rayon over the
277    /// read pool's connections.
278    pub fn lookup_all(
279        &self,
280        paths: &[&Path],
281    ) -> Vec<(PathBuf, Result<LookupResult<ImageCacheEntry>>)> {
282        paths
283            .par_iter()
284            .map(|p| (p.to_path_buf(), self.lookup(p)))
285            .collect()
286    }
287
288    pub fn check(&self, path: &Path) -> Result<bool> {
289        Ok(is_fresh(&self.read.check_status(path)?))
290    }
291
292    pub fn list_paths(&self) -> Result<Vec<String>> {
293        Ok(self
294            .read
295            .keys(None)?
296            .into_iter()
297            .map(|p| p.to_string_lossy().into_owned())
298            .collect())
299    }
300
301    pub fn all(&self) -> Result<Vec<Result<ImageCacheEntry>>> {
302        let entries = self.read.query_run(|q| q)?;
303        Ok(entries.into_iter().map(|e| Ok(to_image_entry(e))).collect())
304    }
305
306    /// Group all cached entries by their parent directory and aggregate
307    /// count, total size, and the newest cached-at timestamp (RFC 004).
308    /// Cheap: enumerates entry metadata without decoding payloads.
309    pub fn summarize_by_dir(&self) -> Result<Vec<DirCacheSummary>> {
310        summarize_entries(self.read.list_entries()?)
311    }
312
313    /// Return all entries whose file lives directly inside `path`.
314    ///
315    /// If `path` is a file rather than a directory (the common call-site
316    /// pattern: "find all entries in the same directory as this file"),
317    /// its parent directory is used automatically.
318    pub fn all_in_dir(&self, path: &Path) -> Result<Vec<Result<ImageCacheEntry>>> {
319        let dir = dir_of(path);
320        let entries = self.read.query_run(|q| q.path_in_dir(dir, false))?;
321        Ok(entries.into_iter().map(|e| Ok(to_image_entry(e))).collect())
322    }
323
324    /// Return all entries whose file lives anywhere under `path`
325    /// (recursively).
326    ///
327    /// If `path` is a file, its parent directory is used automatically.
328    pub fn all_in_dir_and_sub_dirs(&self, path: &Path) -> Result<Vec<Result<ImageCacheEntry>>> {
329        let dir = dir_of(path);
330        let entries = self.read.query_run(|q| q.path_in_dir(dir, true))?;
331        Ok(entries.into_iter().map(|e| Ok(to_image_entry(e))).collect())
332    }
333}
334
335impl CacheRead for ImageCacheReader {
336    fn check(&self, path: &Path) -> Result<bool> {
337        ImageCacheReader::check(self, path)
338    }
339
340    fn check_all(&self, paths: &[&Path]) -> Vec<(PathBuf, Result<bool>)> {
341        paths
342            .par_iter()
343            .map(|p| (p.to_path_buf(), self.check(p)))
344            .collect()
345    }
346
347    fn list_paths(&self) -> Result<Vec<String>> {
348        ImageCacheReader::list_paths(self)
349    }
350}
351
352// ---------------------------------------------------------------------------
353// Helpers
354// ---------------------------------------------------------------------------
355
356/// Fold a flat entry listing into per-directory aggregates.
357pub(crate) fn summarize_entries(
358    entries: Vec<localcache::EntryInfo>,
359) -> Result<Vec<DirCacheSummary>> {
360    use std::collections::BTreeMap;
361
362    let mut map: BTreeMap<PathBuf, (usize, u64, i64)> = BTreeMap::new();
363    for e in entries {
364        let dir = e.path.parent().map(|p| p.to_path_buf()).unwrap_or_default();
365        let agg = map.entry(dir).or_insert((0, 0, 0));
366        agg.0 += 1;
367        agg.1 += e.metadata.file_size;
368        agg.2 = agg.2.max(e.updated_at);
369    }
370    Ok(map
371        .into_iter()
372        .map(
373            |(dir, (file_count, total_size, latest_cached_at))| DirCacheSummary {
374                dir_path: dir.to_string_lossy().into_owned(),
375                file_count,
376                total_size,
377                latest_cached_at,
378            },
379        )
380        .collect())
381}
382
383/// Return `path` itself when it is a directory; otherwise return its parent.
384///
385/// Call sites that pass a media *file* path to `all_in_dir` / `all_in_dir_and_sub_dirs`
386/// expect to query the directory that *contains* the file.  `path_in_dir`
387/// requires a directory, so we resolve automatically here.
388fn dir_of(path: &Path) -> &Path {
389    if path.is_dir() {
390        path
391    } else {
392        path.parent().unwrap_or(path)
393    }
394}
395
396// ---------------------------------------------------------------------------
397// Mapping
398// ---------------------------------------------------------------------------
399
400fn to_image_entry(entry: CacheEntry<ImagePayload>) -> ImageCacheEntry {
401    ImageCacheEntry {
402        path: entry.path.to_string_lossy().into_owned(),
403        thumbnail_path: entry.payload.thumbnail_path,
404        features: entry
405            .payload
406            .clip_vector
407            .map(|v| ImageFeatures { clip_vector: v }),
408    }
409}