use super::protocol::{TileCoordinate, TileResponse};
use crate::error::{Result, StreamingError};
use dashmap::DashMap;
use lru::LruCache;
use std::num::NonZeroUsize;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::fs;
use tokio::sync::RwLock;
use tracing::{debug, warn};
#[derive(Debug, Clone)]
pub struct TileCacheConfig {
pub max_memory_tiles: usize,
pub max_disk_bytes: u64,
pub disk_cache_dir: Option<PathBuf>,
pub compress: bool,
pub ttl_seconds: u64,
}
impl Default for TileCacheConfig {
fn default() -> Self {
Self {
max_memory_tiles: 1000,
max_disk_bytes: 1024 * 1024 * 1024, disk_cache_dir: None,
compress: false,
ttl_seconds: 3600, }
}
}
pub struct TileCache {
config: TileCacheConfig,
memory_cache: Arc<RwLock<LruCache<TileCoordinate, CachedTile>>>,
disk_cache_map: Arc<DashMap<TileCoordinate, PathBuf>>,
}
struct CachedTile {
response: TileResponse,
cached_at: std::time::Instant,
}
impl TileCache {
pub fn new(config: TileCacheConfig) -> Result<Self> {
let max_size = NonZeroUsize::new(config.max_memory_tiles)
.ok_or_else(|| StreamingError::ConfigError("Invalid cache size".to_string()))?;
Ok(Self {
config,
memory_cache: Arc::new(RwLock::new(LruCache::new(max_size))),
disk_cache_map: Arc::new(DashMap::new()),
})
}
pub async fn get(&self, coord: &TileCoordinate) -> Option<TileResponse> {
let mut cache = self.memory_cache.write().await;
if let Some(cached) = cache.get(coord) {
if !self.is_expired(&cached.cached_at) {
debug!("Memory cache hit for tile {}", coord);
return Some(cached.response.clone());
}
}
drop(cache);
if let Some(path) = self.disk_cache_map.get(coord) {
if let Ok(response) = self.load_from_disk(coord, path.value()).await {
debug!("Disk cache hit for tile {}", coord);
self.put_memory(coord, response.clone()).await.ok();
return Some(response);
}
}
None
}
pub async fn put(&self, response: TileResponse) -> Result<()> {
let coord = response.coord;
self.put_memory(&coord, response.clone()).await?;
if self.config.disk_cache_dir.is_some() {
self.put_disk(&coord, response).await?;
}
Ok(())
}
async fn put_memory(&self, coord: &TileCoordinate, response: TileResponse) -> Result<()> {
let mut cache = self.memory_cache.write().await;
cache.put(*coord, CachedTile {
response,
cached_at: std::time::Instant::now(),
});
Ok(())
}
async fn put_disk(&self, coord: &TileCoordinate, response: TileResponse) -> Result<()> {
let cache_dir = self.config.disk_cache_dir.as_ref()
.ok_or_else(|| StreamingError::ConfigError("Disk cache not configured".to_string()))?;
let path = self.tile_path(cache_dir, coord);
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)
.await
.map_err(|e| StreamingError::Io(e))?;
}
fs::write(&path, &response.data)
.await
.map_err(|e| StreamingError::Io(e))?;
self.disk_cache_map.insert(*coord, path);
Ok(())
}
async fn load_from_disk(&self, coord: &TileCoordinate, path: &Path) -> Result<TileResponse> {
let data = fs::read(path)
.await
.map_err(|e| StreamingError::Io(e))?;
Ok(TileResponse::new(
*coord,
bytes::Bytes::from(data),
"image/png".to_string(),
))
}
fn tile_path(&self, base_dir: &Path, coord: &TileCoordinate) -> PathBuf {
base_dir.join(format!("{}/{}/{}.png", coord.z, coord.x, coord.y))
}
fn is_expired(&self, cached_at: &std::time::Instant) -> bool {
cached_at.elapsed().as_secs() > self.config.ttl_seconds
}
pub async fn clear(&self) -> Result<()> {
let mut cache = self.memory_cache.write().await;
cache.clear();
drop(cache);
self.disk_cache_map.clear();
if let Some(cache_dir) = &self.config.disk_cache_dir {
if cache_dir.exists() {
fs::remove_dir_all(cache_dir)
.await
.map_err(|e| StreamingError::Io(e))?;
}
}
Ok(())
}
pub async fn stats(&self) -> CacheStats {
let cache = self.memory_cache.read().await;
CacheStats {
memory_tiles: cache.len(),
disk_tiles: self.disk_cache_map.len(),
max_memory_tiles: self.config.max_memory_tiles,
}
}
}
#[derive(Debug, Clone)]
pub struct CacheStats {
pub memory_tiles: usize,
pub disk_tiles: usize,
pub max_memory_tiles: usize,
}
#[cfg(test)]
mod tests {
use super::*;
use bytes::Bytes;
#[tokio::test]
async fn test_memory_cache() {
let config = TileCacheConfig {
max_memory_tiles: 10,
..Default::default()
};
let cache = TileCache::new(config).ok();
assert!(cache.is_some());
if let Some(cache) = cache {
let coord = TileCoordinate::new(10, 512, 384);
let response = TileResponse::new(
coord,
Bytes::from(vec![0u8; 1024]),
"image/png".to_string(),
);
cache.put(response).await.ok();
let retrieved = cache.get(&coord).await;
assert!(retrieved.is_some());
}
}
}