wsi_streamer/tile/
cache.rs

1//! Tile cache for encoded JPEG tiles.
2//!
3//! This module provides an LRU cache for encoded tiles, preventing repeated
4//! decode/encode cycles for frequently accessed tiles.
5//!
6//! # Cache Key
7//!
8//! Tiles are cached by a composite key including:
9//! - Slide identifier (path or ID)
10//! - Pyramid level
11//! - Tile X coordinate
12//! - Tile Y coordinate
13//! - JPEG quality setting
14//!
15//! # Size-Based Eviction
16//!
17//! The cache tracks the total size of cached tiles in bytes and evicts
18//! least-recently-used entries when the capacity is exceeded.
19
20use std::sync::Arc;
21
22use bytes::Bytes;
23use lru::LruCache;
24use tokio::sync::RwLock;
25
26/// Default cache capacity: 100MB
27pub const DEFAULT_TILE_CACHE_CAPACITY: usize = 100 * 1024 * 1024;
28
29/// Default maximum number of entries (to bound LRU overhead)
30const DEFAULT_MAX_ENTRIES: usize = 10_000;
31
32// =============================================================================
33// Cache Key
34// =============================================================================
35
36/// Cache key for encoded tiles.
37///
38/// This key uniquely identifies a tile at a specific quality level.
39#[derive(Debug, Clone, PartialEq, Eq, Hash)]
40pub struct TileCacheKey {
41    /// Slide identifier (typically the S3 path or slide ID)
42    pub slide_id: Arc<str>,
43
44    /// Pyramid level (0 = highest resolution)
45    pub level: u32,
46
47    /// Tile X coordinate (0-indexed from left)
48    pub tile_x: u32,
49
50    /// Tile Y coordinate (0-indexed from top)
51    pub tile_y: u32,
52
53    /// JPEG quality (1-100)
54    pub quality: u8,
55}
56
57impl TileCacheKey {
58    /// Create a new cache key.
59    pub fn new(
60        slide_id: impl Into<Arc<str>>,
61        level: u32,
62        tile_x: u32,
63        tile_y: u32,
64        quality: u8,
65    ) -> Self {
66        Self {
67            slide_id: slide_id.into(),
68            level,
69            tile_x,
70            tile_y,
71            quality,
72        }
73    }
74}
75
76// =============================================================================
77// Tile Cache
78// =============================================================================
79
80/// LRU cache for encoded JPEG tiles with size-based capacity.
81///
82/// This cache stores encoded tile data and evicts least-recently-used entries
83/// when the total cached size exceeds capacity.
84///
85/// # Thread Safety
86///
87/// The cache is thread-safe and can be shared across async tasks via `Arc`.
88///
89/// # Example
90///
91/// ```
92/// use wsi_streamer::tile::{TileCache, TileCacheKey};
93/// use bytes::Bytes;
94/// use std::sync::Arc;
95///
96/// #[tokio::main]
97/// async fn main() {
98///     let cache = TileCache::new();
99///
100///     let key = TileCacheKey::new("slides/sample.svs", 0, 1, 2, 80);
101///     let tile_data = Bytes::from(vec![0xFF, 0xD8, 0xFF, 0xE0]); // JPEG header
102///
103///     // Store tile
104///     cache.put(key.clone(), tile_data.clone()).await;
105///
106///     // Retrieve tile
107///     let cached = cache.get(&key).await;
108///     assert_eq!(cached, Some(tile_data));
109/// }
110/// ```
111pub struct TileCache {
112    /// The underlying LRU cache
113    cache: RwLock<LruCache<TileCacheKey, Bytes>>,
114
115    /// Maximum total size in bytes
116    max_size: usize,
117
118    /// Current total size in bytes
119    current_size: RwLock<usize>,
120}
121
122impl TileCache {
123    /// Create a new tile cache with default capacity (100MB).
124    pub fn new() -> Self {
125        Self::with_capacity(DEFAULT_TILE_CACHE_CAPACITY)
126    }
127
128    /// Create a new tile cache with the specified capacity in bytes.
129    ///
130    /// # Arguments
131    ///
132    /// * `max_size` - Maximum total size of cached tiles in bytes
133    pub fn with_capacity(max_size: usize) -> Self {
134        Self {
135            cache: RwLock::new(LruCache::new(
136                std::num::NonZeroUsize::new(DEFAULT_MAX_ENTRIES).unwrap(),
137            )),
138            max_size,
139            current_size: RwLock::new(0),
140        }
141    }
142
143    /// Create a new tile cache with specified capacity and maximum entries.
144    ///
145    /// # Arguments
146    ///
147    /// * `max_size` - Maximum total size of cached tiles in bytes
148    /// * `max_entries` - Maximum number of entries in the cache
149    pub fn with_capacity_and_entries(max_size: usize, max_entries: usize) -> Self {
150        Self {
151            cache: RwLock::new(LruCache::new(
152                std::num::NonZeroUsize::new(max_entries).unwrap(),
153            )),
154            max_size,
155            current_size: RwLock::new(0),
156        }
157    }
158
159    /// Get a tile from the cache.
160    ///
161    /// Returns `Some(data)` if the tile is cached, `None` otherwise.
162    /// This operation marks the entry as recently used.
163    pub async fn get(&self, key: &TileCacheKey) -> Option<Bytes> {
164        let mut cache = self.cache.write().await;
165        cache.get(key).cloned()
166    }
167
168    /// Check if a tile is in the cache without updating LRU order.
169    ///
170    /// Returns `true` if the tile is cached, `false` otherwise.
171    pub async fn contains(&self, key: &TileCacheKey) -> bool {
172        let cache = self.cache.read().await;
173        cache.contains(key)
174    }
175
176    /// Store a tile in the cache.
177    ///
178    /// If the cache is over capacity after insertion, least-recently-used
179    /// entries are evicted until the cache is within capacity.
180    ///
181    /// If the tile already exists, it is updated and marked as recently used.
182    pub async fn put(&self, key: TileCacheKey, data: Bytes) {
183        let data_size = data.len();
184        let mut cache = self.cache.write().await;
185        let mut current_size = self.current_size.write().await;
186
187        // If key exists, subtract old size first
188        if let Some(old_data) = cache.peek(&key) {
189            *current_size = current_size.saturating_sub(old_data.len());
190        }
191
192        // Insert the new data
193        cache.put(key, data);
194        *current_size += data_size;
195
196        // Evict entries until we're under capacity
197        while *current_size > self.max_size {
198            if let Some((_, evicted_data)) = cache.pop_lru() {
199                *current_size = current_size.saturating_sub(evicted_data.len());
200            } else {
201                // Cache is empty, nothing more to evict
202                break;
203            }
204        }
205    }
206
207    /// Remove a tile from the cache.
208    ///
209    /// Returns the cached data if it existed, `None` otherwise.
210    pub async fn remove(&self, key: &TileCacheKey) -> Option<Bytes> {
211        let mut cache = self.cache.write().await;
212        let mut current_size = self.current_size.write().await;
213
214        if let Some(data) = cache.pop(key) {
215            *current_size = current_size.saturating_sub(data.len());
216            Some(data)
217        } else {
218            None
219        }
220    }
221
222    /// Clear all entries from the cache.
223    pub async fn clear(&self) {
224        let mut cache = self.cache.write().await;
225        let mut current_size = self.current_size.write().await;
226        cache.clear();
227        *current_size = 0;
228    }
229
230    /// Get the current number of cached tiles.
231    pub async fn len(&self) -> usize {
232        let cache = self.cache.read().await;
233        cache.len()
234    }
235
236    /// Check if the cache is empty.
237    pub async fn is_empty(&self) -> bool {
238        let cache = self.cache.read().await;
239        cache.is_empty()
240    }
241
242    /// Get the current total size of cached tiles in bytes.
243    pub async fn size(&self) -> usize {
244        let current_size = self.current_size.read().await;
245        *current_size
246    }
247
248    /// Get the maximum capacity in bytes.
249    pub fn capacity(&self) -> usize {
250        self.max_size
251    }
252}
253
254impl Default for TileCache {
255    fn default() -> Self {
256        Self::new()
257    }
258}
259
260// =============================================================================
261// Tests
262// =============================================================================
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267
268    fn make_key(slide: &str, level: u32, x: u32, y: u32, quality: u8) -> TileCacheKey {
269        TileCacheKey::new(slide, level, x, y, quality)
270    }
271
272    fn make_tile(size: usize) -> Bytes {
273        Bytes::from(vec![0u8; size])
274    }
275
276    #[tokio::test]
277    async fn test_basic_get_put() {
278        let cache = TileCache::new();
279
280        let key = make_key("slide.svs", 0, 1, 2, 80);
281        let data = make_tile(1000);
282
283        assert!(cache.get(&key).await.is_none());
284
285        cache.put(key.clone(), data.clone()).await;
286
287        let retrieved = cache.get(&key).await;
288        assert_eq!(retrieved, Some(data));
289    }
290
291    #[tokio::test]
292    async fn test_contains() {
293        let cache = TileCache::new();
294
295        let key = make_key("slide.svs", 0, 0, 0, 80);
296        assert!(!cache.contains(&key).await);
297
298        cache.put(key.clone(), make_tile(100)).await;
299        assert!(cache.contains(&key).await);
300    }
301
302    #[tokio::test]
303    async fn test_different_quality_different_key() {
304        let cache = TileCache::new();
305
306        let key_q80 = make_key("slide.svs", 0, 0, 0, 80);
307        let key_q90 = make_key("slide.svs", 0, 0, 0, 90);
308
309        let data_q80 = Bytes::from(vec![80u8; 100]);
310        let data_q90 = Bytes::from(vec![90u8; 100]);
311
312        cache.put(key_q80.clone(), data_q80.clone()).await;
313        cache.put(key_q90.clone(), data_q90.clone()).await;
314
315        assert_eq!(cache.get(&key_q80).await, Some(data_q80));
316        assert_eq!(cache.get(&key_q90).await, Some(data_q90));
317    }
318
319    #[tokio::test]
320    async fn test_size_tracking() {
321        let cache = TileCache::with_capacity(10_000);
322
323        assert_eq!(cache.size().await, 0);
324
325        cache.put(make_key("a", 0, 0, 0, 80), make_tile(1000)).await;
326        assert_eq!(cache.size().await, 1000);
327
328        cache.put(make_key("b", 0, 0, 0, 80), make_tile(2000)).await;
329        assert_eq!(cache.size().await, 3000);
330    }
331
332    #[tokio::test]
333    async fn test_size_based_eviction() {
334        // Cache with 1000 byte capacity
335        let cache = TileCache::with_capacity_and_entries(1000, 100);
336
337        // Add tiles totaling 800 bytes
338        cache.put(make_key("a", 0, 0, 0, 80), make_tile(400)).await;
339        cache.put(make_key("b", 0, 0, 0, 80), make_tile(400)).await;
340
341        assert_eq!(cache.len().await, 2);
342        assert_eq!(cache.size().await, 800);
343
344        // Add another tile that pushes us over capacity
345        cache.put(make_key("c", 0, 0, 0, 80), make_tile(400)).await;
346
347        // LRU entry ("a") should be evicted
348        assert!(cache.size().await <= 1000);
349        assert!(!cache.contains(&make_key("a", 0, 0, 0, 80)).await);
350        assert!(cache.contains(&make_key("b", 0, 0, 0, 80)).await);
351        assert!(cache.contains(&make_key("c", 0, 0, 0, 80)).await);
352    }
353
354    #[tokio::test]
355    async fn test_update_existing_entry() {
356        let cache = TileCache::with_capacity(10_000);
357
358        let key = make_key("slide.svs", 0, 0, 0, 80);
359
360        cache.put(key.clone(), make_tile(1000)).await;
361        assert_eq!(cache.size().await, 1000);
362
363        // Update with different size
364        cache.put(key.clone(), make_tile(500)).await;
365        assert_eq!(cache.size().await, 500);
366        assert_eq!(cache.len().await, 1);
367    }
368
369    #[tokio::test]
370    async fn test_remove() {
371        let cache = TileCache::with_capacity(10_000);
372
373        let key = make_key("slide.svs", 0, 0, 0, 80);
374        let data = make_tile(1000);
375
376        cache.put(key.clone(), data.clone()).await;
377        assert_eq!(cache.size().await, 1000);
378
379        let removed = cache.remove(&key).await;
380        assert_eq!(removed, Some(data));
381        assert_eq!(cache.size().await, 0);
382        assert!(cache.is_empty().await);
383    }
384
385    #[tokio::test]
386    async fn test_clear() {
387        let cache = TileCache::with_capacity(10_000);
388
389        cache.put(make_key("a", 0, 0, 0, 80), make_tile(1000)).await;
390        cache.put(make_key("b", 0, 0, 0, 80), make_tile(2000)).await;
391
392        assert_eq!(cache.len().await, 2);
393        assert_eq!(cache.size().await, 3000);
394
395        cache.clear().await;
396
397        assert!(cache.is_empty().await);
398        assert_eq!(cache.size().await, 0);
399    }
400
401    #[tokio::test]
402    async fn test_lru_order() {
403        // Small cache: 1500 bytes capacity
404        let cache = TileCache::with_capacity_and_entries(1500, 100);
405
406        // Add three tiles of 500 bytes each (total 1500)
407        cache.put(make_key("a", 0, 0, 0, 80), make_tile(500)).await;
408        cache.put(make_key("b", 0, 0, 0, 80), make_tile(500)).await;
409        cache.put(make_key("c", 0, 0, 0, 80), make_tile(500)).await;
410
411        // Access "a" to make it recently used
412        cache.get(&make_key("a", 0, 0, 0, 80)).await;
413
414        // Add new tile, should evict "b" (LRU)
415        cache.put(make_key("d", 0, 0, 0, 80), make_tile(500)).await;
416
417        assert!(cache.contains(&make_key("a", 0, 0, 0, 80)).await); // Recently accessed
418        assert!(!cache.contains(&make_key("b", 0, 0, 0, 80)).await); // Evicted (LRU)
419        assert!(cache.contains(&make_key("c", 0, 0, 0, 80)).await);
420        assert!(cache.contains(&make_key("d", 0, 0, 0, 80)).await);
421    }
422
423    #[tokio::test]
424    async fn test_different_slides_same_coords() {
425        let cache = TileCache::new();
426
427        let key1 = make_key("slide1.svs", 0, 0, 0, 80);
428        let key2 = make_key("slide2.svs", 0, 0, 0, 80);
429
430        let data1 = Bytes::from(vec![1u8; 100]);
431        let data2 = Bytes::from(vec![2u8; 100]);
432
433        cache.put(key1.clone(), data1.clone()).await;
434        cache.put(key2.clone(), data2.clone()).await;
435
436        assert_eq!(cache.get(&key1).await, Some(data1));
437        assert_eq!(cache.get(&key2).await, Some(data2));
438        assert_eq!(cache.len().await, 2);
439    }
440
441    #[tokio::test]
442    async fn test_capacity() {
443        let cache = TileCache::with_capacity(50_000);
444        assert_eq!(cache.capacity(), 50_000);
445    }
446
447    #[test]
448    fn test_cache_key_equality() {
449        let key1 = make_key("slide.svs", 0, 1, 2, 80);
450        let key2 = make_key("slide.svs", 0, 1, 2, 80);
451        let key3 = make_key("slide.svs", 0, 1, 2, 90);
452
453        assert_eq!(key1, key2);
454        assert_ne!(key1, key3);
455    }
456
457    #[test]
458    fn test_cache_key_hash() {
459        use std::collections::hash_map::DefaultHasher;
460        use std::hash::{Hash, Hasher};
461
462        fn hash<T: Hash>(t: &T) -> u64 {
463            let mut s = DefaultHasher::new();
464            t.hash(&mut s);
465            s.finish()
466        }
467
468        let key1 = make_key("slide.svs", 0, 1, 2, 80);
469        let key2 = make_key("slide.svs", 0, 1, 2, 80);
470
471        assert_eq!(hash(&key1), hash(&key2));
472    }
473}