use crate::common::dir_utils::DirUtils;
use dashmap::{DashMap, DashSet};
use parking_lot::RwLock;
use std::path::Path;
use tracing::warn;
const FAILURES_FILE: &str = "failures.json";
pub struct IconCacheService {
memory_cache: DashMap<String, Vec<u8>>,
cache_dir: RwLock<String>,
cached_file_hashes: DashSet<String>,
failures: DashSet<String>,
}
impl IconCacheService {
pub fn new(cache_dir: String) -> Self {
Self {
memory_cache: DashMap::new(),
cache_dir: RwLock::new(cache_dir),
cached_file_hashes: DashSet::new(),
failures: DashSet::new(),
}
}
pub fn init(&self) {
let cache_dir = self.cache_dir.read().clone();
let hashes = Self::scan_cached_icons(&cache_dir);
for hash in hashes {
self.cached_file_hashes.insert(hash);
}
self.load_failures();
}
pub async fn preload_l1(&self) {
let cache_dir = self.cache_dir.read().clone();
let keys: Vec<String> = self
.cached_file_hashes
.iter()
.map(|entry| entry.clone())
.collect();
let mut handles = Vec::with_capacity(keys.len());
for key in keys {
let dir = cache_dir.clone();
let key_clone = key.clone();
handles.push(tokio::spawn(async move {
let data = tokio::fs::read(Path::new(&dir).join(&key_clone)).await.ok();
(key, data)
}));
}
for handle in handles {
if let Ok((key, Some(data))) = handle.await {
self.memory_cache.insert(key, data);
}
}
let (count, bytes): (usize, usize) = self
.memory_cache
.iter()
.map(|e| (1, e.value().len()))
.fold((0, 0), |(c, b), (cc, bb)| (c + cc, b + bb));
tracing::info!(
"图标缓存预载完成: {} 个, {:.1} KB",
count,
bytes as f64 / 1024.0
);
}
pub fn get_l1(&self, key: &str) -> Option<Vec<u8>> {
self.memory_cache
.get(key)
.map(|entry| entry.value().clone())
}
pub fn set_l1(&self, key: &str, data: Vec<u8>) {
self.memory_cache.insert(key.to_string(), data);
}
pub fn clear_l1(&self) {
self.memory_cache.clear();
}
pub fn remove_l1(&self, key: &str) {
self.memory_cache.remove(key);
}
pub fn contains_l2(&self, key: &str) -> bool {
self.cached_file_hashes.contains(key)
}
fn load_failures(&self) {
self.failures.clear();
let path = Path::new(&*self.cache_dir.read()).join(FAILURES_FILE);
let bytes = match std::fs::read(&path) {
Ok(bytes) => bytes,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return,
Err(e) => {
warn!("读取图标失败记录失败: {}", e);
return;
}
};
match serde_json::from_slice::<Vec<String>>(&bytes) {
Ok(keys) => {
for key in keys {
self.failures.insert(key);
}
}
Err(e) => warn!("解析图标失败记录失败: {}", e),
}
}
fn save_failures(&self) {
let path = Path::new(&*self.cache_dir.read()).join(FAILURES_FILE);
let mut keys: Vec<String> = self.failures.iter().map(|entry| entry.clone()).collect();
keys.sort();
match serde_json::to_vec(&keys) {
Ok(bytes) => {
if let Err(e) = std::fs::write(&path, bytes) {
warn!("写入图标失败记录失败: {}", e);
}
}
Err(e) => warn!("序列化图标失败记录失败: {}", e),
}
}
pub fn should_skip(&self, key: &str) -> bool {
self.failures.contains(key)
}
pub fn mark_failed(&self, key: &str) {
if self.failures.insert(key.to_string()) {
self.save_failures();
}
}
pub fn clear_failed(&self, key: &str) {
if self.failures.remove(key).is_some() {
self.save_failures();
}
}
pub fn clear_all_failed(&self) {
if self.failures.is_empty() {
return;
}
self.failures.clear();
self.save_failures();
}
pub async fn get_l2(&self, key: &str) -> Option<Vec<u8>> {
let cache_dir = self.cache_dir.read().clone();
let icon_path = Path::new(&cache_dir).join(key);
tokio::fs::read(&icon_path).await.ok()
}
pub async fn set_l2(&self, key: &str, data: Vec<u8>) {
let cache_dir = self.cache_dir.read().clone();
let icon_path = Path::new(&cache_dir).join(key);
let key_owned = key.to_string();
if let Err(e) = tokio::fs::write(&icon_path, data).await {
warn!("写入图标缓存失败 {:?}: {}", icon_path, e);
} else {
self.cached_file_hashes.insert(key_owned);
}
}
pub fn update_cache_dir(&self, new_cache_dir: &str) {
self.memory_cache.clear();
{
let mut cache_dir = self.cache_dir.write();
*cache_dir = new_cache_dir.to_string();
}
self.cached_file_hashes.clear();
let hashes = Self::scan_cached_icons(new_cache_dir);
for hash in hashes {
self.cached_file_hashes.insert(hash);
}
self.load_failures();
}
fn scan_cached_icons(cache_dir: &str) -> DashSet<String> {
let result = DashSet::new();
match DirUtils::read_dir_or_create(cache_dir) {
Ok(entries) => {
for entry in entries.flatten() {
let file_name = entry.file_name();
let file_name = file_name.to_string_lossy();
if !file_name.ends_with(".webp") {
continue;
}
result.insert(file_name.into_owned());
}
}
Err(e) => warn!("扫描图标缓存目录失败: {}", e),
}
result
}
}