studio_worker/
thumbnail.rs

1//! Thumbnails of image jobs, so the tray UI can show what a job made.
2//!
3//! The daemon decodes an image result, scales it to fit
4//! [`THUMBNAIL_MAX_SIDE`] and keeps it as PNG in a bounded [`Thumbnails`]
5//! ring keyed by job id.  The tray UI keeps its own copy (fetched from
6//! `GET /jobs/:id/thumbnail`) in the same type.
7
8use std::collections::VecDeque;
9use std::io::Cursor;
10use std::sync::Arc;
11
12use anyhow::Context as _;
13use parking_lot::Mutex;
14
15/// Longer side of a thumbnail, in pixels.  Crisp at the UI's 96 px card
16/// size on a 2x display.
17pub const THUMBNAIL_MAX_SIDE: u32 = 192;
18
19/// Thumbnails of this many most recent image jobs are kept: the size of
20/// the studio and local job rings together.
21pub const THUMBNAILS_CAP: usize = 2 * crate::runtime::RECENT_JOBS_CAP;
22
23/// Scale an encoded image (WEBP / PNG / JPEG) to fit
24/// [`THUMBNAIL_MAX_SIDE`] and encode it as PNG.  Smaller images keep their
25/// size.
26pub fn make_thumbnail(encoded: &[u8]) -> anyhow::Result<Vec<u8>> {
27    let image = image::load_from_memory(encoded).context("decoding the job's image")?;
28    let image = if image.width() > THUMBNAIL_MAX_SIDE || image.height() > THUMBNAIL_MAX_SIDE {
29        image.thumbnail(THUMBNAIL_MAX_SIDE, THUMBNAIL_MAX_SIDE)
30    } else {
31        image
32    };
33    let mut png = Vec::new();
34    image
35        .write_to(&mut Cursor::new(&mut png), image::ImageFormat::Png)
36        .context("encoding the thumbnail")?;
37    Ok(png)
38}
39
40/// A job id and its PNG thumbnail.
41type Entry = (String, Arc<Vec<u8>>);
42
43/// Bounded ring of PNG thumbnails keyed by job id, oldest evicted first.
44/// Cheap to clone.
45#[derive(Clone, Default)]
46pub struct Thumbnails {
47    inner: Arc<Mutex<VecDeque<Entry>>>,
48}
49
50impl Thumbnails {
51    /// Keep `png` as `job_id`'s thumbnail, replacing an older one.
52    pub fn insert(&self, job_id: &str, png: Vec<u8>) {
53        let mut ring = self.inner.lock();
54        ring.retain(|(id, _)| id != job_id);
55        ring.push_back((job_id.to_string(), Arc::new(png)));
56        while ring.len() > THUMBNAILS_CAP {
57            ring.pop_front();
58        }
59    }
60
61    /// The thumbnail of `job_id`.
62    pub fn get(&self, job_id: &str) -> Option<Arc<Vec<u8>>> {
63        self.inner
64            .lock()
65            .iter()
66            .find(|(id, _)| id == job_id)
67            .map(|(_, png)| png.clone())
68    }
69
70    /// Whether `job_id` has a thumbnail.
71    pub fn contains(&self, job_id: &str) -> bool {
72        self.inner.lock().iter().any(|(id, _)| id == job_id)
73    }
74
75    /// Drop every thumbnail whose job id fails `keep`.
76    pub fn retain(&self, keep: impl Fn(&str) -> bool) {
77        self.inner.lock().retain(|(id, _)| keep(id));
78    }
79
80    /// Drop every thumbnail.
81    pub fn clear(&self) {
82        self.inner.lock().clear();
83    }
84
85    /// Number of thumbnails kept.
86    pub fn len(&self) -> usize {
87        self.inner.lock().len()
88    }
89
90    /// True when no thumbnail is kept.
91    pub fn is_empty(&self) -> bool {
92        self.len() == 0
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99
100    fn encoded(width: u32, height: u32, format: image::ImageFormat) -> Vec<u8> {
101        let image = image::RgbImage::from_pixel(width, height, image::Rgb([200, 40, 40]));
102        let mut out = Vec::new();
103        image::DynamicImage::ImageRgb8(image)
104            .write_to(&mut Cursor::new(&mut out), format)
105            .unwrap();
106        out
107    }
108
109    fn dimensions(png: &[u8]) -> (u32, u32) {
110        let image = image::load_from_memory_with_format(png, image::ImageFormat::Png).unwrap();
111        (image.width(), image.height())
112    }
113
114    #[test]
115    fn a_large_image_is_scaled_to_fit_keeping_its_aspect() {
116        let png = make_thumbnail(&encoded(1024, 512, image::ImageFormat::WebP)).unwrap();
117        assert_eq!(
118            dimensions(&png),
119            (THUMBNAIL_MAX_SIDE, THUMBNAIL_MAX_SIDE / 2)
120        );
121    }
122
123    #[test]
124    fn a_small_image_keeps_its_size() {
125        let png = make_thumbnail(&encoded(64, 48, image::ImageFormat::Png)).unwrap();
126        assert_eq!(dimensions(&png), (64, 48));
127    }
128
129    #[test]
130    fn bytes_that_are_not_an_image_are_refused() {
131        let err = make_thumbnail(b"not an image").unwrap_err();
132        assert!(err.to_string().contains("decoding"), "got {err:#}");
133    }
134
135    #[test]
136    fn the_ring_keeps_only_the_most_recent_thumbnails() {
137        let thumbnails = Thumbnails::default();
138        for i in 0..(THUMBNAILS_CAP + 2) {
139            thumbnails.insert(&format!("job-{i}"), vec![i as u8]);
140        }
141        assert_eq!(thumbnails.len(), THUMBNAILS_CAP);
142        assert!(!thumbnails.contains("job-0"));
143        assert!(!thumbnails.contains("job-1"));
144        assert_eq!(
145            thumbnails
146                .get(&format!("job-{}", THUMBNAILS_CAP + 1))
147                .unwrap()[0],
148            (THUMBNAILS_CAP + 1) as u8
149        );
150    }
151
152    #[test]
153    fn inserting_again_replaces_the_thumbnail() {
154        let thumbnails = Thumbnails::default();
155        thumbnails.insert("a", vec![1]);
156        thumbnails.insert("a", vec![2]);
157        assert_eq!(thumbnails.len(), 1);
158        assert_eq!(*thumbnails.get("a").unwrap(), vec![2]);
159    }
160
161    #[test]
162    fn retain_and_clear_drop_thumbnails() {
163        let thumbnails = Thumbnails::default();
164        thumbnails.insert("keep", vec![1]);
165        thumbnails.insert("drop", vec![2]);
166        thumbnails.retain(|id| id == "keep");
167        assert!(thumbnails.contains("keep") && !thumbnails.contains("drop"));
168        thumbnails.clear();
169        assert!(thumbnails.is_empty());
170    }
171}
studio_worker/thumbnail.rs

studio_worker/
thumbnail.rs
