yggdrasil/tasks/
image_cache_cleanup.rs

1//! 图片磁盘缓存定期清理任务。
2//!
3//! 仅在 `server` feature 启用时编译，每小时运行一次。
4//! 删除超过保留时间的文件，并在总大小超过上限时按修改时间删除最旧的文件。
5
6use std::io;
7use std::path::{Path, PathBuf};
8use std::time::{Duration, SystemTime};
9use tokio::time::interval;
10
11const CACHE_DIR: &str = "uploads/.cache";
12const BYTES_PER_MB: u64 = 1024 * 1024;
13const SECS_PER_HOUR: u64 = 3600;
14/// 保留期上限（小时）：防止误配极大值导致 `max_age_hours * SECS_PER_HOUR` 溢出，
15/// 或 `now - max_age` 在 SystemTime 上下溢 panic（release `panic="abort"` 会直接
16/// 崩溃整个进程，而非单次清理失败）。10 年对图片磁盘缓存已远超合理范围。
17const MAX_AGE_HOURS_CAP: u64 = 87_600; // 10 年
18
19/// 启动图片磁盘缓存清理循环，每小时触发一次。
20pub async fn run_cleanup() {
21    let mut ticker = interval(Duration::from_secs(SECS_PER_HOUR));
22    loop {
23        if let Err(e) = cleanup_image_cache().await {
24            tracing::error!("Image disk cache cleanup error: {:?}", e);
25        }
26        ticker.tick().await;
27    }
28}
29
30/// 读取环境变量并清理默认磁盘缓存目录。
31pub async fn cleanup_image_cache() -> io::Result<()> {
32    let base = Path::new(CACHE_DIR);
33    let max_mb = std::env::var("IMAGE_DISK_CACHE_MAX_MB")
34        .ok()
35        .and_then(|s| s.parse().ok())
36        .unwrap_or(1024);
37    let max_age_hours = std::env::var("IMAGE_DISK_CACHE_MAX_AGE_HOURS")
38        .ok()
39        .and_then(|s| s.parse().ok())
40        .unwrap_or(168);
41    let (deleted, bytes_freed) = cleanup_image_cache_at(base, max_mb, max_age_hours).await?;
42    if !deleted.is_empty() {
43        tracing::info!(
44            "Image disk cache cleanup: removed {} files, freed {} bytes",
45            deleted.len(),
46            bytes_freed
47        );
48    }
49    Ok(())
50}
51
52/// 清理指定目录下的图片磁盘缓存。
53///
54/// 返回被删除文件的路径列表以及释放的总字节数。
55pub async fn cleanup_image_cache_at(
56    base: &Path,
57    max_mb: u64,
58    max_age_hours: u64,
59) -> io::Result<(Vec<PathBuf>, u64)> {
60    if !base.exists() {
61        return Ok((Vec::new(), 0));
62    }
63
64    // C5 防护：clamp 上限防止乘法溢出；saturating_mul 双保险；checked_sub 防 SystemTime
65    // 下溢 panic（极端配置/时钟回拨时 max_age 可能超过 now 距 UNIX_EPOCH 的时长）。
66    let max_age_hours = max_age_hours.min(MAX_AGE_HOURS_CAP);
67    let max_age = Duration::from_secs(max_age_hours.saturating_mul(SECS_PER_HOUR));
68    let now = SystemTime::now();
69    let cutoff = now.checked_sub(max_age).unwrap_or(SystemTime::UNIX_EPOCH);
70
71    let mut entries: Vec<(PathBuf, u64, SystemTime)> = Vec::new();
72    collect_files(base, &mut entries).await?;
73
74    let mut deleted = Vec::new();
75    let mut bytes_freed: u64 = 0;
76
77    // 第一轮：删除超过保留期限的文件。
78    let mut remaining: Vec<(PathBuf, u64, SystemTime)> = Vec::new();
79    for (path, size, mtime) in entries {
80        if mtime < cutoff {
81            match tokio::fs::remove_file(&path).await {
82                Ok(_) => {
83                    deleted.push(path);
84                    bytes_freed += size;
85                }
86                Err(e) => {
87                    tracing::warn!("Failed to remove expired cache file {:?}: {:?}", path, e);
88                }
89            }
90        } else {
91            remaining.push((path, size, mtime));
92        }
93    }
94
95    // 第二轮：若总大小仍超过上限，按修改时间从旧到新删除。
96    let max_bytes = max_mb.saturating_mul(BYTES_PER_MB);
97    let mut total: u64 = remaining.iter().map(|(_, size, _)| size).sum();
98    if total > max_bytes {
99        remaining.sort_by_key(|a| a.2);
100        for (path, size, _) in remaining {
101            if total <= max_bytes {
102                break;
103            }
104            match tokio::fs::remove_file(&path).await {
105                Ok(_) => {
106                    total -= size;
107                    deleted.push(path);
108                    bytes_freed += size;
109                }
110                Err(e) => {
111                    tracing::warn!(
112                        "Failed to remove cache file {:?} for size cap: {:?}",
113                        path,
114                        e
115                    );
116                }
117            }
118        }
119    }
120
121    Ok((deleted, bytes_freed))
122}
123
124/// 递归收集目录下的所有常规文件，返回路径、大小与修改时间。
125///
126/// 跳过符号链接，避免 traversal 到 `uploads/.cache/` 外部。
127async fn collect_files(
128    base: &Path,
129    entries: &mut Vec<(PathBuf, u64, SystemTime)>,
130) -> io::Result<()> {
131    let mut stack = vec![base.to_path_buf()];
132    while let Some(dir) = stack.pop() {
133        let mut reader = tokio::fs::read_dir(&dir).await?;
134        while let Some(entry) = reader.next_entry().await? {
135            let file_type = entry.file_type().await?;
136            if file_type.is_symlink() {
137                continue;
138            } else if file_type.is_file() {
139                let metadata = entry.metadata().await?;
140                let mtime = metadata.modified()?;
141                entries.push((entry.path(), metadata.len(), mtime));
142            } else if file_type.is_dir() {
143                stack.push(entry.path());
144            }
145        }
146    }
147    Ok(())
148}
149
150#[cfg(all(test, feature = "server"))]
151mod tests {
152    use super::*;
153    use std::time::{Duration, UNIX_EPOCH};
154    use tokio::time::sleep;
155
156    fn temp_cache_dir() -> PathBuf {
157        let nanos = SystemTime::now()
158            .duration_since(UNIX_EPOCH)
159            .unwrap()
160            .as_nanos();
161        std::env::temp_dir().join(format!(
162            "yggdrasil_image_cache_test_{}_{}",
163            nanos,
164            std::process::id()
165        ))
166    }
167
168    #[tokio::test]
169    async fn cleanup_ignores_missing_directory() {
170        let dir = temp_cache_dir();
171        let (deleted, freed) = cleanup_image_cache_at(&dir, 1024, 168).await.unwrap();
172        assert!(deleted.is_empty());
173        assert_eq!(freed, 0);
174    }
175
176    #[tokio::test]
177    async fn cleanup_removes_expired_files_by_age() {
178        let dir = temp_cache_dir();
179        tokio::fs::create_dir_all(&dir).await.unwrap();
180
181        let old = dir.join("old.dat");
182        tokio::fs::write(&old, b"old content").await.unwrap();
183        // 确保文件的修改时间严格早于清理时计算的截止时间。
184        sleep(Duration::from_millis(1100)).await;
185
186        let (deleted, freed) = cleanup_image_cache_at(&dir, 1024, 0).await.unwrap();
187        assert_eq!(deleted.len(), 1);
188        assert!(!old.exists());
189        assert!(freed > 0);
190
191        tokio::fs::remove_dir_all(&dir).await.unwrap();
192    }
193
194    #[tokio::test]
195    async fn cleanup_keeps_recent_files() {
196        let dir = temp_cache_dir();
197        tokio::fs::create_dir_all(&dir).await.unwrap();
198
199        let recent = dir.join("recent.dat");
200        tokio::fs::write(&recent, b"recent content").await.unwrap();
201
202        let (deleted, freed) = cleanup_image_cache_at(&dir, 1024, 168).await.unwrap();
203        assert!(deleted.is_empty());
204        assert_eq!(freed, 0);
205        assert!(recent.exists());
206
207        tokio::fs::remove_dir_all(&dir).await.unwrap();
208    }
209
210    #[tokio::test]
211    async fn cleanup_enforces_size_cap_by_mtime() {
212        let dir = temp_cache_dir();
213        tokio::fs::create_dir_all(&dir).await.unwrap();
214
215        let f1 = dir.join("oldest.dat");
216        tokio::fs::write(&f1, vec![0u8; 1024 * 1024]).await.unwrap();
217        sleep(Duration::from_millis(1100)).await;
218
219        let f2 = dir.join("middle.dat");
220        tokio::fs::write(&f2, vec![0u8; 1024 * 1024]).await.unwrap();
221        sleep(Duration::from_millis(1100)).await;
222
223        let f3 = dir.join("newest.dat");
224        tokio::fs::write(&f3, vec![0u8; 1024 * 1024]).await.unwrap();
225
226        // 上限 2 MB，当前 3 MB，应删除最旧的一个文件。
227        let (deleted, freed) = cleanup_image_cache_at(&dir, 2, 1000).await.unwrap();
228        assert_eq!(deleted.len(), 1);
229        assert!(!f1.exists());
230        assert!(f2.exists());
231        assert!(f3.exists());
232        assert_eq!(freed, 1024 * 1024);
233
234        tokio::fs::remove_dir_all(&dir).await.unwrap();
235    }
236
237    #[tokio::test]
238    async fn cleanup_recurses_into_subdirectories() {
239        let dir = temp_cache_dir();
240        let sub = dir.join("nested");
241        tokio::fs::create_dir_all(&sub).await.unwrap();
242
243        let nested = sub.join("nested.dat");
244        tokio::fs::write(&nested, b"nested content").await.unwrap();
245        sleep(Duration::from_millis(1100)).await;
246
247        let (deleted, _freed) = cleanup_image_cache_at(&dir, 1024, 0).await.unwrap();
248        assert_eq!(deleted.len(), 1);
249        assert!(!nested.exists());
250
251        tokio::fs::remove_dir_all(&dir).await.unwrap();
252    }
253
254    #[tokio::test]
255    async fn cleanup_with_huge_max_age_does_not_panic() {
256        // C5 回归：极大保留期不得导致 u64 乘法溢出或 SystemTime 下溢 panic。
257        // 旧代码 `max_age_hours * SECS_PER_HOUR` 在 debug 溢出 panic、release 回绕；
258        // 回绕后的巨大 max_age 又使 `now - max_age` 下溢 panic（panic="abort" 崩进程）。
259        let dir = temp_cache_dir();
260        tokio::fs::create_dir_all(&dir).await.unwrap();
261        let recent = dir.join("recent.dat");
262        tokio::fs::write(&recent, b"recent").await.unwrap();
263
264        // u64::MAX 小时：未加防护时会溢出/下溢；加防护后被 clamp 到 10 年上限，
265        // recent 文件（刚写入）应保留。
266        let (deleted, _freed) = cleanup_image_cache_at(&dir, 1024, u64::MAX).await.unwrap();
267        assert!(
268            deleted.is_empty(),
269            "巨大 max_age 不应删除任何文件, got: {deleted:?}"
270        );
271        assert!(recent.exists());
272
273        tokio::fs::remove_dir_all(&dir).await.unwrap();
274    }
275}
yggdrasil/tasks/image_cache_cleanup.rs

yggdrasil/tasks/
image_cache_cleanup.rs
