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

wsi_rs/core/
cache.rs
