use std::time::Instant;
#[derive(Debug, Clone, Default)]
pub struct DiskMonitorData {
pub free_bytes: u64,
pub total_bytes: u64,
pub free_percent: f32,
pub last_update: Option<Instant>,
}
#[cfg(feature = "system-monitor")]
mod inner {
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::thread::JoinHandle;
use std::time::{Duration, Instant};
use parking_lot::Mutex;
use sysinfo::{Disk, DiskRefreshKind, Disks};
use super::DiskMonitorData;
pub struct DiskMonitor {
data: Arc<Mutex<DiskMonitorData>>,
cwd: Arc<Mutex<Option<String>>>,
follow_cwd: Arc<Mutex<bool>>,
launch_dir: Arc<Mutex<PathBuf>>,
running: Arc<AtomicBool>,
thread: Mutex<Option<JoinHandle<()>>>,
}
impl DiskMonitor {
pub fn new() -> Self {
let launch_dir = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("/"));
Self {
data: Arc::new(Mutex::new(DiskMonitorData::default())),
cwd: Arc::new(Mutex::new(None)),
follow_cwd: Arc::new(Mutex::new(false)),
launch_dir: Arc::new(Mutex::new(launch_dir)),
running: Arc::new(AtomicBool::new(false)),
thread: Mutex::new(None),
}
}
pub fn start(&self, poll_interval_secs: f32) {
if self
.running
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.is_err()
{
return;
}
let data = Arc::clone(&self.data);
let cwd = Arc::clone(&self.cwd);
let follow_cwd = Arc::clone(&self.follow_cwd);
let launch_dir = Arc::clone(&self.launch_dir);
let running = Arc::clone(&self.running);
let interval = Duration::from_secs_f32(poll_interval_secs.max(5.0));
let handle = std::thread::Builder::new()
.name("status-bar-disk".to_string())
.spawn(move || {
let mut disks = Disks::new_with_refreshed_list();
while running.load(Ordering::SeqCst) {
disks.refresh_specifics(true, DiskRefreshKind::nothing().with_storage());
let use_cwd = *follow_cwd.lock();
let target: PathBuf = if use_cwd {
cwd.lock()
.clone()
.filter(|c| !c.is_empty())
.map(PathBuf::from)
.unwrap_or_else(|| launch_dir.lock().clone())
} else {
launch_dir.lock().clone()
};
let fallback = use_cwd.then(|| launch_dir.lock().clone());
if let Some(snapshot) =
snapshot_for(disks.list(), &target, fallback.as_deref())
{
*data.lock() = snapshot;
}
let deadline = Instant::now() + interval;
while Instant::now() < deadline && running.load(Ordering::Relaxed) {
std::thread::sleep(Duration::from_millis(50));
}
}
});
match handle {
Ok(h) => *self.thread.lock() = Some(h),
Err(e) => {
self.running.store(false, Ordering::SeqCst);
crate::debug_error!(
"SESSION_LOGGER",
"failed to spawn disk monitor thread: {:?}",
e
);
}
}
}
pub fn signal_stop(&self) {
self.running.store(false, Ordering::SeqCst);
}
pub fn stop(&self) {
self.signal_stop();
if let Some(handle) = self.thread.lock().take() {
let _ = handle.join();
}
}
pub fn is_running(&self) -> bool {
self.running.load(Ordering::SeqCst)
}
pub fn data(&self) -> DiskMonitorData {
self.data.lock().clone()
}
pub fn set_cwd(&self, new_cwd: Option<&str>) {
*self.cwd.lock() = new_cwd.map(String::from);
}
pub fn set_follow_cwd(&self, follow: bool) {
*self.follow_cwd.lock() = follow;
}
}
impl Default for DiskMonitor {
fn default() -> Self {
Self::new()
}
}
impl Drop for DiskMonitor {
fn drop(&mut self) {
self.stop();
}
}
pub fn snapshot_for(
disks: &[Disk],
target: &Path,
fallback: Option<&Path>,
) -> Option<DiskMonitorData> {
let disk = disk_for_path(disks, target).or_else(|| disk_for_path(disks, fallback?))?;
let total = disk.total_space();
let free = disk.available_space();
let free_percent = if total > 0 {
(free as f64 / total as f64 * 100.0) as f32
} else {
0.0
};
Some(DiskMonitorData {
free_bytes: free,
total_bytes: total,
free_percent,
last_update: Some(Instant::now()),
})
}
pub fn disk_for_path<'a>(disks: &'a [Disk], target: &Path) -> Option<&'a Disk> {
disks
.iter()
.filter(|d| target.starts_with(d.mount_point()))
.max_by_key(|d| d.mount_point().components().count())
}
}
#[cfg(feature = "system-monitor")]
pub use inner::{DiskMonitor, disk_for_path, snapshot_for};
#[cfg(not(feature = "system-monitor"))]
mod inner {
use super::DiskMonitorData;
pub struct DiskMonitor;
impl DiskMonitor {
pub fn new() -> Self {
Self
}
pub fn start(&self, _poll_interval_secs: f32) {}
pub fn signal_stop(&self) {}
pub fn stop(&self) {}
pub fn is_running(&self) -> bool {
false
}
pub fn data(&self) -> DiskMonitorData {
DiskMonitorData::default()
}
pub fn set_cwd(&self, _new_cwd: Option<&str>) {}
pub fn set_follow_cwd(&self, _follow: bool) {}
}
impl Default for DiskMonitor {
fn default() -> Self {
Self::new()
}
}
}
#[cfg(not(feature = "system-monitor"))]
pub use inner::DiskMonitor;
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_disk_monitor_data_default() {
let d = DiskMonitorData::default();
assert_eq!(d.free_bytes, 0);
assert_eq!(d.total_bytes, 0);
assert_eq!(d.free_percent, 0.0);
assert!(d.last_update.is_none());
}
#[cfg(feature = "system-monitor")]
#[test]
fn test_disk_for_path_resolves() {
use sysinfo::Disks;
let disks = Disks::new_with_refreshed_list();
let tmp = std::env::temp_dir();
let d = disk_for_path(disks.list(), &tmp);
assert!(d.is_some(), "no disk matched temp dir {:?}", tmp);
let d = d.unwrap();
assert!(d.total_space() > 0);
assert!(d.available_space() <= d.total_space());
}
#[cfg(feature = "system-monitor")]
#[test]
fn test_disk_for_path_returns_mount_prefix() {
use sysinfo::Disks;
let disks = Disks::new_with_refreshed_list();
if disks.list().is_empty() {
return; }
let target = std::env::temp_dir();
let matched = disk_for_path(disks.list(), &target);
assert!(matched.is_some(), "no disk matched temp dir {target:?}");
assert!(
target.starts_with(matched.unwrap().mount_point()),
"matched mount point must be a prefix of the target path"
);
}
#[cfg(feature = "system-monitor")]
#[test]
fn test_snapshot_for_target_and_fallback() {
use sysinfo::Disks;
let disks = Disks::new_with_refreshed_list();
let tmp = std::env::temp_dir();
let s = snapshot_for(disks.list(), &tmp, None).expect("temp dir resolves to a disk");
assert!(s.total_bytes > 0 && s.free_bytes <= s.total_bytes);
assert!(s.free_percent >= 0.0 && s.free_percent <= 100.0);
assert!(s.last_update.is_some());
let unmatched = std::path::Path::new("relative/nowhere");
assert!(snapshot_for(disks.list(), unmatched, None).is_none());
assert!(snapshot_for(disks.list(), unmatched, Some(&tmp)).is_some());
}
#[cfg(feature = "system-monitor")]
#[test]
fn test_disk_monitor_start_stop() {
let monitor = DiskMonitor::new();
assert!(!monitor.is_running());
monitor.start(5.0);
assert!(monitor.is_running());
monitor.stop();
assert!(!monitor.is_running());
}
}