Skip to main content

wsi_rs/core/
cache.rs

1use lru::LruCache;
2use std::sync::{Arc, Mutex};
3
4use crate::core::types::{CpuTile, DatasetId};
5
6// ── TileCache (axis-aware) ────────────────────────────────────────
7
8/// Default shared decoded tile cache.
9///
10/// Standard Aperio SVS JPEG tiles are commonly 240x240 RGB, or about 170 KiB
11/// per decoded tile. A 64 MiB budget keeps a few hundred such source tiles
12/// resident, which is enough for normal viewport overlap during quick zooms
13/// without forcing users to tune cache options before the viewer is usable.
14pub(crate) const DEFAULT_TILE_CACHE_SIZE: u64 = 64 * 1024 * 1024;
15const TILE_CACHE_BYTES_ENV: &str = "WSI_RS_TILE_CACHE_BYTES";
16/// Default display-tile cache.
17///
18/// Display-tile reads on regular tiled slides cache the decoded source tiles
19/// used for composition. Keep enough room for at least a dense viewport plus
20/// adjacent zoom/pan overlap; 1 MiB only held a handful of SVS tiles and caused
21/// immediate churn during zoom-out bursts.
22pub(crate) const DEFAULT_DISPLAY_TILE_CACHE_SIZE: u64 = 32 * 1024 * 1024;
23const DISPLAY_TILE_CACHE_BYTES_ENV: &str = "WSI_RS_DISPLAY_TILE_CACHE_BYTES";
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26#[non_exhaustive]
27pub struct CacheConfig {
28    pub shared_tile_bytes: Option<u64>,
29    pub display_tile_bytes: Option<u64>,
30}
31
32impl CacheConfig {
33    pub const fn deterministic() -> Self {
34        Self {
35            shared_tile_bytes: None,
36            display_tile_bytes: None,
37        }
38    }
39
40    pub const fn with_shared_tile_bytes(mut self, bytes: u64) -> Self {
41        self.shared_tile_bytes = Some(bytes);
42        self
43    }
44
45    pub const fn with_display_tile_bytes(mut self, bytes: u64) -> Self {
46        self.display_tile_bytes = Some(bytes);
47        self
48    }
49
50    pub(crate) fn shared_tile_budget(self, source_hint: Option<u64>) -> u64 {
51        self.shared_tile_bytes
52            .or(source_hint)
53            .unwrap_or(DEFAULT_TILE_CACHE_SIZE)
54    }
55
56    pub(crate) fn display_tile_budget(self) -> u64 {
57        self.display_tile_bytes
58            .unwrap_or(DEFAULT_DISPLAY_TILE_CACHE_SIZE)
59    }
60}
61
62impl Default for CacheConfig {
63    fn default() -> Self {
64        Self::deterministic()
65    }
66}
67
68#[derive(Hash, Eq, PartialEq, Clone, Debug)]
69/// Note: scene/series are u32 here (not usize) to keep CacheKey compact and
70/// Hash-friendly. TileRequest/RegionRequest use usize for ergonomic indexing.
71/// Slide converts usize → u32 via `as u32` when constructing cache keys.
72/// Overflow is not a practical concern (>4B scenes/series is impossible).
73pub struct CacheKey {
74    pub(crate) dataset_id: DatasetId,
75    pub(crate) scene: u32,
76    pub(crate) series: u32,
77    pub(crate) level: u32,
78    pub(crate) z: u32,
79    pub(crate) c: u32,
80    pub(crate) t: u32,
81    pub(crate) tile_col: i64,
82    pub(crate) tile_row: i64,
83}
84
85pub struct TileCache {
86    inner: Mutex<TileCacheState>,
87}
88
89#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
90pub(crate) struct CacheStats {
91    pub(crate) hits: u64,
92    pub(crate) misses: u64,
93    pub(crate) puts: u64,
94    pub(crate) evictions: u64,
95    pub(crate) rejected_oversize: u64,
96    pub(crate) capacity_bytes: u64,
97    pub(crate) current_bytes: u64,
98    pub(crate) entries: usize,
99}
100
101impl std::fmt::Debug for TileCache {
102    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
103        let state = self.inner.lock().unwrap_or_else(|e| e.into_inner());
104        f.debug_struct("TileCache")
105            .field("capacity_bytes", &state.capacity_bytes)
106            .field("current_bytes", &state.current_bytes)
107            .field("entries", &state.lru.len())
108            .field("hits", &state.hits)
109            .field("misses", &state.misses)
110            .finish()
111    }
112}
113
114struct TileCacheState {
115    lru: LruCache<CacheKey, CachedTile>,
116    capacity_bytes: u64,
117    current_bytes: u64,
118    hits: u64,
119    misses: u64,
120    puts: u64,
121    evictions: u64,
122    rejected_oversize: u64,
123}
124
125struct CachedTile {
126    data: Arc<CpuTile>,
127    byte_size: u64,
128}
129
130impl TileCache {
131    pub(crate) fn new(capacity_bytes: u64) -> Self {
132        Self {
133            inner: Mutex::new(TileCacheState {
134                // The cache is byte-budgeted only. The backing LRU stays unbounded
135                // and eviction is driven by `capacity_bytes`.
136                lru: LruCache::unbounded(),
137                capacity_bytes,
138                current_bytes: 0,
139                hits: 0,
140                misses: 0,
141                puts: 0,
142                evictions: 0,
143                rejected_oversize: 0,
144            }),
145        }
146    }
147
148    pub(crate) fn put(&self, key: CacheKey, data: Arc<CpuTile>) {
149        let byte_size = data.data.byte_size() as u64;
150        let mut state = self.inner.lock().unwrap_or_else(|e| e.into_inner());
151
152        if byte_size > state.capacity_bytes {
153            state.rejected_oversize += 1;
154            return;
155        }
156
157        // Remove existing entry if present
158        if let Some((_, existing)) = state.lru.pop_entry(&key) {
159            state.current_bytes -= existing.byte_size;
160        }
161
162        // Evict LRU entries until there's room
163        while state.current_bytes + byte_size > state.capacity_bytes {
164            if let Some((_, evicted)) = state.lru.pop_lru() {
165                state.current_bytes -= evicted.byte_size;
166                state.evictions += 1;
167            } else {
168                break;
169            }
170        }
171
172        state.lru.put(key, CachedTile { data, byte_size });
173        state.current_bytes += byte_size;
174        state.puts += 1;
175    }
176
177    pub(crate) fn get(&self, key: &CacheKey) -> Option<Arc<CpuTile>> {
178        let mut state = self.inner.lock().unwrap_or_else(|e| e.into_inner());
179        let cached = state.lru.get(key).map(|entry| entry.data.clone());
180        if cached.is_some() {
181            state.hits += 1;
182        } else {
183            state.misses += 1;
184        }
185        cached
186    }
187
188    pub(crate) fn stats(&self) -> CacheStats {
189        let state = self.inner.lock().unwrap_or_else(|e| e.into_inner());
190        CacheStats {
191            hits: state.hits,
192            misses: state.misses,
193            puts: state.puts,
194            evictions: state.evictions,
195            rejected_oversize: state.rejected_oversize,
196            capacity_bytes: state.capacity_bytes,
197            current_bytes: state.current_bytes,
198            entries: state.lru.len(),
199        }
200    }
201
202    pub(crate) fn display_default() -> Self {
203        Self::new(capacity_from_env(
204            DISPLAY_TILE_CACHE_BYTES_ENV,
205            DEFAULT_DISPLAY_TILE_CACHE_SIZE,
206        ))
207    }
208
209    pub(crate) fn display_with_config(config: CacheConfig) -> Self {
210        Self::new(config.display_tile_budget())
211    }
212
213    pub(crate) fn shared_default_with_hint(default_bytes: u64) -> Self {
214        Self::new(capacity_from_env(TILE_CACHE_BYTES_ENV, default_bytes))
215    }
216
217    pub(crate) fn shared_with_config(config: CacheConfig, source_hint: Option<u64>) -> Self {
218        Self::new(config.shared_tile_budget(source_hint))
219    }
220}
221
222impl Default for TileCache {
223    fn default() -> Self {
224        Self::shared_default_with_hint(DEFAULT_TILE_CACHE_SIZE)
225    }
226}
227
228fn capacity_from_env(env_name: &str, default_bytes: u64) -> u64 {
229    std::env::var(env_name)
230        .ok()
231        .and_then(|value| value.parse::<u64>().ok())
232        .filter(|bytes| *bytes > 0)
233        .unwrap_or(default_bytes)
234}
235
236#[cfg(test)]
237mod tile_cache_tests {
238    use super::*;
239    use crate::core::types::*;
240
241    const SVS_RGB_240_TILE_BYTES: usize = 240 * 240 * 3;
242    const COMMON_ZOOM_VIEWPORT_TILE_COUNT: i64 = 96;
243
244    fn make_sample_buffer(size: usize) -> CpuTile {
245        CpuTile {
246            width: 256,
247            height: 256,
248            channels: 3,
249            color_space: ColorSpace::Rgb,
250            layout: CpuTileLayout::Interleaved,
251            data: CpuTileData::u8(vec![0u8; size]),
252        }
253    }
254
255    fn make_key(dataset_id: u128, level: u32, col: i64, row: i64) -> CacheKey {
256        CacheKey {
257            dataset_id: DatasetId::new(dataset_id),
258            scene: 0,
259            series: 0,
260            level,
261            z: 0,
262            c: 0,
263            t: 0,
264            tile_col: col,
265            tile_row: row,
266        }
267    }
268
269    #[test]
270    fn put_and_get() {
271        let cache = TileCache::new(1024 * 1024);
272        let buf = Arc::new(make_sample_buffer(100));
273        let key = make_key(1, 0, 0, 0);
274        cache.put(key.clone(), buf.clone());
275        let result = cache.get(&key).unwrap();
276        assert_eq!(result.width, 256);
277    }
278
279    #[test]
280    fn miss_returns_none() {
281        let cache = TileCache::new(1024);
282        let key = make_key(1, 0, 0, 0);
283        assert!(cache.get(&key).is_none());
284    }
285
286    #[test]
287    fn eviction_by_byte_size() {
288        let cache = TileCache::new(250);
289        cache.put(make_key(1, 0, 0, 0), Arc::new(make_sample_buffer(100)));
290        cache.put(make_key(1, 0, 1, 0), Arc::new(make_sample_buffer(100)));
291        // Both fit: 200 bytes
292        assert!(cache.get(&make_key(1, 0, 0, 0)).is_some());
293        assert!(cache.get(&make_key(1, 0, 1, 0)).is_some());
294
295        // Third pushes over 250
296        cache.put(make_key(1, 0, 2, 0), Arc::new(make_sample_buffer(100)));
297        assert!(cache.get(&make_key(1, 0, 0, 0)).is_none()); // evicted
298        assert!(cache.get(&make_key(1, 0, 1, 0)).is_some());
299        assert!(cache.get(&make_key(1, 0, 2, 0)).is_some());
300    }
301
302    #[test]
303    fn different_datasets_are_independent() {
304        let cache = TileCache::new(1024);
305        cache.put(make_key(1, 0, 0, 0), Arc::new(make_sample_buffer(10)));
306        cache.put(make_key(2, 0, 0, 0), Arc::new(make_sample_buffer(10)));
307        assert!(cache.get(&make_key(1, 0, 0, 0)).is_some());
308        assert!(cache.get(&make_key(2, 0, 0, 0)).is_some());
309    }
310
311    #[test]
312    fn axis_aware_keys() {
313        let cache = TileCache::new(1024);
314        let mut key_z0 = make_key(1, 0, 0, 0);
315        key_z0.z = 0;
316        let mut key_z1 = make_key(1, 0, 0, 0);
317        key_z1.z = 1;
318        cache.put(key_z0.clone(), Arc::new(make_sample_buffer(10)));
319        cache.put(key_z1.clone(), Arc::new(make_sample_buffer(10)));
320        assert!(cache.get(&key_z0).is_some());
321        assert!(cache.get(&key_z1).is_some());
322    }
323
324    #[test]
325    fn oversize_entry_rejected() {
326        let cache = TileCache::new(50);
327        cache.put(make_key(1, 0, 0, 0), Arc::new(make_sample_buffer(100)));
328        assert!(cache.get(&make_key(1, 0, 0, 0)).is_none());
329    }
330
331    #[test]
332    fn shared_across_threads() {
333        let cache = Arc::new(TileCache::new(4096));
334        let cache_clone = cache.clone();
335        let handle = std::thread::spawn(move || {
336            cache_clone.put(make_key(1, 0, 5, 5), Arc::new(make_sample_buffer(10)));
337        });
338        handle.join().unwrap();
339        assert!(cache.get(&make_key(1, 0, 5, 5)).is_some());
340    }
341
342    #[test]
343    fn display_default_holds_common_svs_zoom_viewport_working_set() {
344        let cache = TileCache::new(DEFAULT_DISPLAY_TILE_CACHE_SIZE);
345        for col in 0..COMMON_ZOOM_VIEWPORT_TILE_COUNT {
346            cache.put(
347                make_key(1, 0, col, 0),
348                Arc::new(make_sample_buffer(SVS_RGB_240_TILE_BYTES)),
349            );
350        }
351
352        let stats = cache.stats();
353        assert_eq!(stats.entries, COMMON_ZOOM_VIEWPORT_TILE_COUNT as usize);
354        assert_eq!(stats.evictions, 0);
355        assert_eq!(stats.rejected_oversize, 0);
356    }
357
358    #[test]
359    fn stats_count_hits_misses_puts_evictions_and_oversize_rejections() {
360        let cache = TileCache::new(150);
361        let missing = make_key(1, 0, 9, 9);
362        assert!(cache.get(&missing).is_none());
363
364        cache.put(make_key(1, 0, 0, 0), Arc::new(make_sample_buffer(100)));
365        assert!(cache.get(&make_key(1, 0, 0, 0)).is_some());
366
367        cache.put(make_key(1, 0, 1, 0), Arc::new(make_sample_buffer(100)));
368        cache.put(make_key(1, 0, 2, 0), Arc::new(make_sample_buffer(200)));
369
370        let stats = cache.stats();
371        assert_eq!(stats.hits, 1);
372        assert_eq!(stats.misses, 1);
373        assert_eq!(stats.puts, 2);
374        assert_eq!(stats.evictions, 1);
375        assert_eq!(stats.rejected_oversize, 1);
376        assert_eq!(stats.capacity_bytes, 150);
377        assert_eq!(stats.current_bytes, 100);
378        assert_eq!(stats.entries, 1);
379    }
380}