1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
use super::Error;
use async_recursion::async_recursion;
use dashmap::DashSet;
use futures::future::join_all;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::fs;
#[derive(Clone, Debug)]
pub struct UsedSpace {
max_capacity: u64,
dirs: Arc<DashSet<PathBuf>>,
}
impl UsedSpace {
pub fn new(max_capacity: u64) -> Self {
Self {
max_capacity,
dirs: Arc::new(DashSet::new()),
}
}
pub(crate) fn add_dir(&self, dir: &Path) {
let _ = self.dirs.insert(dir.to_path_buf());
}
pub(crate) fn max_capacity(&self) -> u64 {
self.max_capacity
}
pub(crate) async fn can_consume(&self, space: u64) -> bool {
self.total()
.await
.checked_add(space)
.map_or(false, |new_total| self.max_capacity >= new_total)
}
pub(crate) async fn total(&self) -> u64 {
let handles = self
.dirs
.iter()
.map(|d| d.clone())
.map(|path| tokio::spawn(async move { get_size(path).await.map_err(Error::from) }));
join_all(handles).await.iter().flatten().flatten().sum()
}
}
#[async_recursion]
async fn get_size(path: PathBuf) -> tokio::io::Result<u64> {
let metadata = fs::metadata(&path).await?;
let mut size = metadata.len();
if metadata.is_dir() {
let mut dir = fs::read_dir(&path).await?;
while let Some(entry) = dir.next_entry().await? {
size += get_size(entry.path()).await?;
}
}
Ok(size)
}