Spaces:
Build error
Build error
File size: 4,441 Bytes
84d2a97 |
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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
use std::path::PathBuf;
use tokio::sync::Mutex;
use tokio::time::Instant;
use crate::operations::types::{CollectionError, CollectionResult};
/// Defines how often the disk usage should be checked if the disk is far from being full
const DEFAULT_FREQUENCY: usize = 128;
/// Defines how frequent the check of the disk usage should be, depending on the free space
///
/// The idea is that if we have a lot of disk space, it is unlikely that it will be out of space soon,
/// so we can check it less frequently. But if we have a little space, we should check it more often
/// and at some point, we should check it every time
const FREE_SPACE_TO_CHECK_FREQUENCY_HEURISTIC_MB: &[(usize, usize); 5] =
&[(512, 0), (1024, 8), (2048, 16), (4096, 32), (8096, 64)];
/// Even if there were no many updates, we still want to force the check of the disk usage
/// because some external process could have consumed the disk space
const MIN_DISK_CHECK_INTERVAL_MILLIS: usize = 2000;
#[derive(Default)]
struct LastCheck {
last_check_time: Option<Instant>,
next_check_count: usize,
}
pub struct DiskUsageWatcher {
disk_path: PathBuf,
disabled: bool,
min_free_disk_size_mb: usize,
last_check: Mutex<LastCheck>,
}
impl DiskUsageWatcher {
pub async fn new(disk_path: PathBuf, min_free_disk_size_mb: usize) -> Self {
let mut watcher = Self {
disk_path,
disabled: false,
min_free_disk_size_mb,
last_check: Default::default(),
};
match watcher.is_disk_full().await {
Ok(Some(_)) => {} // do nothing
Ok(None) => watcher.disabled = true,
Err(_) => {
watcher.disabled = true;
}
};
watcher
}
/// Returns true if the disk free space is less than the `disk_buffer_threshold_mb`
/// As the side effect, it updates the disk usage every `update_count_threshold` calls
pub async fn is_disk_full(&self) -> CollectionResult<Option<bool>> {
if self.disabled {
return Ok(None);
}
let mut last_check_guard = self.last_check.lock().await;
let since_last_check = last_check_guard
.last_check_time
.map(|x| x.elapsed().as_millis() as usize)
.unwrap_or(usize::MAX);
if last_check_guard.next_check_count == 0
|| since_last_check >= MIN_DISK_CHECK_INTERVAL_MILLIS
{
let free_space = self.get_free_space_bytes().await?;
last_check_guard.last_check_time = Some(Instant::now());
let is_full = match free_space {
Some(free_space) => {
let free_space = free_space as usize;
let mut next_check = DEFAULT_FREQUENCY;
for (threshold_mb, interval) in FREE_SPACE_TO_CHECK_FREQUENCY_HEURISTIC_MB {
if free_space < (*threshold_mb * 1024 * 1024) {
next_check = *interval;
break;
}
}
last_check_guard.next_check_count = next_check;
Some(free_space < self.min_free_disk_size_mb * 1024 * 1024)
}
None => {
last_check_guard.next_check_count = 0;
None
}
};
Ok(is_full)
} else {
last_check_guard.next_check_count = last_check_guard.next_check_count.saturating_sub(1);
Ok(None)
}
}
/// Return current disk usage in bytes, if available
pub async fn get_free_space_bytes(&self) -> CollectionResult<Option<u64>> {
if self.disabled {
return Ok(None);
}
let path = self.disk_path.clone();
let result = tokio::task::spawn_blocking(move || fs4::available_space(path.as_path()))
.await
.map_err(|e| {
CollectionError::service_error(format!("Failed to join async task: {e}"))
})?;
let result = match result {
Ok(result) => Some(result),
Err(err) => {
log::debug!(
"Failed to get free space for path: {} due to: {}",
self.disk_path.as_path().display(),
err
);
None
}
};
Ok(result)
}
}
|