1use lru::LruCache;
2use std::sync::{Arc, Mutex};
3
4use crate::core::types::{CpuTile, DatasetId};
5
6pub(crate) const DEFAULT_TILE_CACHE_SIZE: u64 = 64 * 1024 * 1024;
15const TILE_CACHE_BYTES_ENV: &str = "WSI_RS_TILE_CACHE_BYTES";
16pub(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)]
69pub 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 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 if let Some((_, existing)) = state.lru.pop_entry(&key) {
159 state.current_bytes -= existing.byte_size;
160 }
161
162 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 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 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()); 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}