use std::collections::HashSet;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use tokio::sync::Mutex;
use crate::error::DownloadError;
use crate::sink::{STATE_SUFFIX, TMP_SUFFIX};
#[derive(Debug, Default)]
pub struct ActiveDownloads {
protected: Mutex<HashSet<PathBuf>>,
}
impl ActiveDownloads {
pub fn new() -> Self {
ActiveDownloads::default()
}
pub async fn register(&self, path: impl Into<PathBuf>) {
self.protected.lock().await.insert(path.into());
}
pub async fn unregister(&self, path: &Path) {
self.protected.lock().await.remove(path);
}
pub async fn is_protected(&self, path: &Path) -> bool {
self.protected.lock().await.contains(path)
}
pub async fn len(&self) -> usize {
self.protected.lock().await.len()
}
pub async fn is_empty(&self) -> bool {
self.protected.lock().await.is_empty()
}
}
#[derive(Debug, Clone)]
pub struct GcConfig {
pub dir: PathBuf,
pub ttl: Duration,
pub interval: Duration,
}
impl GcConfig {
pub fn new(dir: impl Into<PathBuf>) -> Self {
GcConfig {
dir: dir.into(),
ttl: Duration::from_secs(3600),
interval: Duration::from_secs(600),
}
}
}
#[derive(Clone)]
pub struct TmpGc {
dir: PathBuf,
ttl: Duration,
registry: Arc<ActiveDownloads>,
}
impl TmpGc {
pub fn new(dir: impl Into<PathBuf>, ttl: Duration, registry: Arc<ActiveDownloads>) -> Self {
TmpGc {
dir: dir.into(),
ttl,
registry,
}
}
pub async fn sweep_at(&self, now: SystemTime) -> Result<usize, DownloadError> {
let entries = match std::fs::read_dir(&self.dir) {
Ok(e) => e,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(0),
Err(e) => return Err(DownloadError::state(e)),
};
let mut removed = 0usize;
for entry in entries {
let entry = entry.map_err(DownloadError::state)?;
let path = entry.path();
let name = match path.file_name().and_then(|n| n.to_str()) {
Some(n) => n,
None => continue,
};
if !name.ends_with(TMP_SUFFIX) || name.ends_with(STATE_SUFFIX) {
continue;
}
if self.registry.is_protected(&path).await {
continue; }
if !is_stale(&path, now, self.ttl) {
continue; }
std::fs::remove_file(&path).map_err(DownloadError::state)?;
removed += 1;
let sidecar = sidecar_state_path(&path);
match std::fs::remove_file(&sidecar) {
Ok(()) => {}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
Err(e) => return Err(DownloadError::state(e)),
}
}
Ok(removed)
}
pub async fn sweep(&self) -> Result<usize, DownloadError> {
self.sweep_at(SystemTime::now()).await
}
pub fn dir(&self) -> &Path {
&self.dir
}
}
fn sidecar_state_path(tmp_path: &Path) -> PathBuf {
let s = tmp_path.to_string_lossy();
let base = s.strip_suffix(TMP_SUFFIX).unwrap_or(&s);
PathBuf::from(format!("{base}{STATE_SUFFIX}"))
}
fn is_stale(path: &Path, now: SystemTime, ttl: Duration) -> bool {
match std::fs::metadata(path).and_then(|m| m.modified()) {
Ok(mtime) => match now.duration_since(mtime) {
Ok(age) => age >= ttl,
Err(_) => false, },
Err(_) => true,
}
}
#[cfg(test)]
mod tests {
use super::*;
fn temp_dir(tag: &str) -> PathBuf {
let d = std::env::temp_dir().join(format!(
"dig-download-gc-{tag}-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
std::fs::create_dir_all(&d).unwrap();
d
}
fn make_tmp(dir: &Path, name: &str) -> PathBuf {
let p = dir.join(format!("{name}{TMP_SUFFIX}"));
std::fs::write(&p, b"partial").unwrap();
p
}
#[tokio::test]
async fn sweeps_stale_orphan_and_its_sidecar() {
let dir = temp_dir("orphan");
let tmp = make_tmp(&dir, "resource");
let sidecar = sidecar_state_path(&tmp);
std::fs::write(&sidecar, b"{}").unwrap();
let registry = Arc::new(ActiveDownloads::new());
let gc = TmpGc::new(&dir, Duration::from_secs(60), registry);
let removed = gc
.sweep_at(SystemTime::now() + Duration::from_secs(3600))
.await
.unwrap();
assert_eq!(removed, 1);
assert!(!tmp.exists());
assert!(!sidecar.exists());
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test]
async fn keeps_fresh_orphan_within_ttl() {
let dir = temp_dir("fresh");
let tmp = make_tmp(&dir, "resource");
let gc = TmpGc::new(
&dir,
Duration::from_secs(3600),
Arc::new(ActiveDownloads::new()),
);
let removed = gc.sweep_at(SystemTime::now()).await.unwrap();
assert_eq!(removed, 0);
assert!(tmp.exists());
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test]
async fn never_reaps_a_protected_paused_download() {
let dir = temp_dir("protected");
let tmp = make_tmp(&dir, "resource");
let registry = Arc::new(ActiveDownloads::new());
registry.register(tmp.clone()).await; let gc = TmpGc::new(&dir, Duration::from_secs(60), registry.clone());
let removed = gc
.sweep_at(SystemTime::now() + Duration::from_secs(9999))
.await
.unwrap();
assert_eq!(removed, 0);
assert!(tmp.exists());
assert!(registry.is_protected(&tmp).await);
registry.unregister(&tmp).await;
let removed = gc
.sweep_at(SystemTime::now() + Duration::from_secs(9999))
.await
.unwrap();
assert_eq!(removed, 1);
assert!(!tmp.exists());
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test]
async fn sweep_missing_dir_is_ok() {
let gc = TmpGc::new(
std::env::temp_dir().join("dig-download-gc-does-not-exist-xyz"),
Duration::from_secs(1),
Arc::new(ActiveDownloads::new()),
);
assert_eq!(gc.sweep().await.unwrap(), 0);
}
#[tokio::test]
async fn registry_register_unregister() {
let r = ActiveDownloads::new();
assert!(r.is_empty().await);
r.register("/a/b.download.tmp").await;
assert_eq!(r.len().await, 1);
assert!(r.is_protected(Path::new("/a/b.download.tmp")).await);
r.unregister(Path::new("/a/b.download.tmp")).await;
assert!(r.is_empty().await);
}
#[test]
fn sidecar_path_derivation() {
let tmp = PathBuf::from("/data/x.dig.download.tmp");
assert_eq!(
sidecar_state_path(&tmp),
PathBuf::from("/data/x.dig.download.tmp.state")
);
}
}