use bytes::Bytes;
use lru::LruCache;
use std::fmt;
use std::hash::Hash;
use std::num::NonZeroUsize;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use thiserror::Error;
use tracing::{debug, trace};
#[derive(Debug, Error)]
pub enum CacheError {
#[error("Cache I/O error: {0}")]
Io(#[from] std::io::Error),
#[error("Invalid cache key")]
InvalidKey,
#[error("Cache is full")]
Full,
}
pub type CacheResult<T> = Result<T, CacheError>;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct CacheKey {
pub layer: String,
pub z: u8,
pub x: u32,
pub y: u32,
pub format: String,
pub style: Option<String>,
}
impl CacheKey {
pub fn new(layer: String, z: u8, x: u32, y: u32, format: String) -> Self {
Self {
layer,
z,
x,
y,
format,
style: None,
}
}
pub fn with_style(mut self, style: String) -> Self {
self.style = Some(style);
self
}
pub fn to_path(&self, base_dir: &Path) -> PathBuf {
let mut path = base_dir.to_path_buf();
path.push(&self.layer);
if let Some(ref style) = self.style {
path.push(style);
}
path.push(self.z.to_string());
path.push(self.x.to_string());
path.push(format!("{}.{}", self.y, self.format));
path
}
}
impl fmt::Display for CacheKey {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if let Some(ref style) = self.style {
write!(
f,
"{}/{}/{}/{}/{}.{}",
self.layer, style, self.z, self.x, self.y, self.format
)
} else {
write!(
f,
"{}/{}/{}/{}.{}",
self.layer, self.z, self.x, self.y, self.format
)
}
}
}
#[derive(Debug, Clone)]
struct CacheEntry {
data: Bytes,
created_at: Instant,
size: usize,
access_count: u64,
}
impl CacheEntry {
fn new(data: Bytes) -> Self {
let size = data.len();
Self {
data,
created_at: Instant::now(),
size,
access_count: 0,
}
}
fn is_expired(&self, ttl: Duration) -> bool {
self.created_at.elapsed() > ttl
}
fn record_access(&mut self) {
self.access_count += 1;
}
}
#[derive(Debug, Clone, Default)]
pub struct CacheStats {
pub hits: u64,
pub misses: u64,
pub entry_count: usize,
pub total_size: usize,
pub evictions: u64,
pub expirations: u64,
pub disk_reads: u64,
pub disk_writes: u64,
}
impl CacheStats {
pub fn hit_rate(&self) -> f64 {
let total = self.hits + self.misses;
if total == 0 {
0.0
} else {
self.hits as f64 / total as f64
}
}
pub fn avg_entry_size(&self) -> f64 {
if self.entry_count == 0 {
0.0
} else {
self.total_size as f64 / self.entry_count as f64
}
}
}
#[derive(Debug, Clone)]
pub struct TileCacheConfig {
pub max_memory_bytes: usize,
pub disk_cache_dir: Option<PathBuf>,
pub ttl: Duration,
pub enable_stats: bool,
pub compression: bool,
}
impl Default for TileCacheConfig {
fn default() -> Self {
Self {
max_memory_bytes: 256 * 1024 * 1024, disk_cache_dir: None,
ttl: Duration::from_secs(3600), enable_stats: true,
compression: false,
}
}
}
pub struct TileCache {
memory_cache: Arc<Mutex<LruCache<CacheKey, CacheEntry>>>,
memory_usage: Arc<Mutex<usize>>,
config: TileCacheConfig,
stats: Arc<Mutex<CacheStats>>,
}
const MIN_CACHE_CAPACITY: NonZeroUsize = match NonZeroUsize::new(100) {
Some(n) => n,
None => unreachable!(),
};
impl TileCache {
pub fn new(config: TileCacheConfig) -> Self {
let estimated_capacity = config.max_memory_bytes / (10 * 1024);
let capacity = NonZeroUsize::new(estimated_capacity)
.unwrap_or(MIN_CACHE_CAPACITY)
.max(MIN_CACHE_CAPACITY);
Self {
memory_cache: Arc::new(Mutex::new(LruCache::new(capacity))),
memory_usage: Arc::new(Mutex::new(0)),
config,
stats: Arc::new(Mutex::new(CacheStats::default())),
}
}
pub fn get(&self, key: &CacheKey) -> Option<Bytes> {
trace!("Cache lookup: {}", key.to_string());
if let Some(data) = self.get_from_memory(key) {
self.record_hit();
return Some(data);
}
if self.config.disk_cache_dir.is_some() {
if let Some(data) = self.get_from_disk(key) {
let _ = self.put_in_memory(key.clone(), data.clone());
self.record_hit();
return Some(data);
}
}
self.record_miss();
None
}
pub fn put(&self, key: CacheKey, data: Bytes) -> CacheResult<()> {
trace!("Caching tile: {}", key.to_string());
self.put_in_memory(key.clone(), data.clone())?;
if self.config.disk_cache_dir.is_some() {
let _ = self.put_on_disk(&key, &data);
}
Ok(())
}
fn get_from_memory(&self, key: &CacheKey) -> Option<Bytes> {
let mut cache = self.memory_cache.lock().ok()?;
let is_expired = if let Some(entry) = cache.peek(key) {
entry.is_expired(self.config.ttl)
} else {
return None;
};
if is_expired {
trace!("Entry expired: {}", key.to_string());
self.record_expiration();
let entry = cache.pop(key)?;
self.update_memory_usage(|usage| usage.saturating_sub(entry.size));
return None;
}
if let Some(entry) = cache.get_mut(key) {
entry.record_access();
Some(entry.data.clone())
} else {
None
}
}
fn put_in_memory(&self, key: CacheKey, data: Bytes) -> CacheResult<()> {
let entry = CacheEntry::new(data);
let entry_size = entry.size;
let mut cache = self.memory_cache.lock().map_err(|_| CacheError::Full)?;
while self.get_memory_usage() + entry_size > self.config.max_memory_bytes {
if let Some((_, evicted)) = cache.pop_lru() {
debug!("Evicting entry from memory cache");
self.update_memory_usage(|usage| usage.saturating_sub(evicted.size));
self.record_eviction();
} else {
break;
}
}
if let Some(old_entry) = cache.put(key, entry) {
self.update_memory_usage(|usage| usage.saturating_sub(old_entry.size));
}
self.update_memory_usage(|usage| usage + entry_size);
Ok(())
}
fn get_from_disk(&self, key: &CacheKey) -> Option<Bytes> {
let base_dir = self.config.disk_cache_dir.as_ref()?;
let path = key.to_path(base_dir);
match std::fs::read(&path) {
Ok(data) => {
trace!("Disk cache hit: {}", path.display());
self.record_disk_read();
Some(Bytes::from(data))
}
Err(_) => None,
}
}
fn put_on_disk(&self, key: &CacheKey, data: &Bytes) -> CacheResult<()> {
let base_dir =
self.config
.disk_cache_dir
.as_ref()
.ok_or(CacheError::Io(std::io::Error::new(
std::io::ErrorKind::NotFound,
"No disk cache directory",
)))?;
let path = key.to_path(base_dir);
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(&path, data)?;
self.record_disk_write();
trace!("Wrote to disk cache: {}", path.display());
Ok(())
}
pub fn clear(&self) -> CacheResult<()> {
if let Ok(mut cache) = self.memory_cache.lock() {
cache.clear();
}
self.update_memory_usage(|_| 0);
if let Some(ref dir) = self.config.disk_cache_dir {
if dir.exists() {
std::fs::remove_dir_all(dir)?;
std::fs::create_dir_all(dir)?;
}
}
if let Ok(mut stats) = self.stats.lock() {
*stats = CacheStats::default();
}
debug!("Cache cleared");
Ok(())
}
pub fn stats(&self) -> CacheStats {
self.stats.lock().map(|s| s.clone()).unwrap_or_default()
}
fn get_memory_usage(&self) -> usize {
self.memory_usage.lock().map(|u| *u).unwrap_or(0)
}
fn update_memory_usage<F>(&self, f: F)
where
F: FnOnce(usize) -> usize,
{
if let Ok(mut usage) = self.memory_usage.lock() {
*usage = f(*usage);
}
if let Ok(mut stats) = self.stats.lock() {
stats.total_size = self.get_memory_usage();
}
}
fn record_hit(&self) {
if self.config.enable_stats {
if let Ok(mut stats) = self.stats.lock() {
stats.hits += 1;
}
}
}
fn record_miss(&self) {
if self.config.enable_stats {
if let Ok(mut stats) = self.stats.lock() {
stats.misses += 1;
}
}
}
fn record_eviction(&self) {
if self.config.enable_stats {
if let Ok(mut stats) = self.stats.lock() {
stats.evictions += 1;
}
}
}
fn record_expiration(&self) {
if self.config.enable_stats {
if let Ok(mut stats) = self.stats.lock() {
stats.expirations += 1;
}
}
}
fn record_disk_read(&self) {
if self.config.enable_stats {
if let Ok(mut stats) = self.stats.lock() {
stats.disk_reads += 1;
}
}
}
fn record_disk_write(&self) {
if self.config.enable_stats {
if let Ok(mut stats) = self.stats.lock() {
stats.disk_writes += 1;
}
}
}
}
impl Clone for TileCache {
fn clone(&self) -> Self {
Self {
memory_cache: Arc::clone(&self.memory_cache),
memory_usage: Arc::clone(&self.memory_usage),
config: self.config.clone(),
stats: Arc::clone(&self.stats),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cache_key_to_string() {
let key = CacheKey::new("landsat".to_string(), 10, 512, 384, "png".to_string());
assert_eq!(key.to_string(), "landsat/10/512/384.png");
let key_with_style = key.with_style("default".to_string());
assert_eq!(key_with_style.to_string(), "landsat/default/10/512/384.png");
}
#[test]
fn test_cache_put_get() {
let config = TileCacheConfig::default();
let cache = TileCache::new(config);
let key = CacheKey::new("test".to_string(), 0, 0, 0, "png".to_string());
let data = Bytes::from(vec![1, 2, 3, 4, 5]);
cache.put(key.clone(), data.clone()).expect("put failed");
let retrieved = cache.get(&key).expect("get failed");
assert_eq!(retrieved, data);
}
#[test]
fn test_cache_miss() {
let config = TileCacheConfig::default();
let cache = TileCache::new(config);
let key = CacheKey::new("test".to_string(), 0, 0, 0, "png".to_string());
assert!(cache.get(&key).is_none());
let stats = cache.stats();
assert_eq!(stats.misses, 1);
assert_eq!(stats.hits, 0);
}
#[test]
fn test_cache_stats() {
let config = TileCacheConfig::default();
let cache = TileCache::new(config);
let key1 = CacheKey::new("test".to_string(), 0, 0, 0, "png".to_string());
let key2 = CacheKey::new("test".to_string(), 0, 0, 1, "png".to_string());
let data = Bytes::from(vec![1, 2, 3]);
cache.put(key1.clone(), data.clone()).expect("put failed");
cache.put(key2.clone(), data.clone()).expect("put failed");
cache.get(&key1);
cache.get(&key2);
cache.get(&CacheKey::new(
"nonexistent".to_string(),
0,
0,
0,
"png".to_string(),
));
let stats = cache.stats();
assert_eq!(stats.hits, 2);
assert_eq!(stats.misses, 1);
assert!(stats.hit_rate() > 0.6);
}
#[test]
fn test_cache_clear() {
let config = TileCacheConfig::default();
let cache = TileCache::new(config);
let key = CacheKey::new("test".to_string(), 0, 0, 0, "png".to_string());
let data = Bytes::from(vec![1, 2, 3]);
cache.put(key.clone(), data).expect("put failed");
assert!(cache.get(&key).is_some());
cache.clear().expect("clear failed");
assert!(cache.get(&key).is_none());
}
}