fine_ill_do_it_myself/

lib.rs

//! This crate provides the [`DirSizeTracker`] struct, which efficiently keeps track of a directory's total size using filesystem events.
//!
//! See the list of known problems in the docs of [`notify`].

#![warn(missing_docs)]
#![warn(clippy::arithmetic_side_effects)]

mod error_data;
#[cfg(test)]
mod tests;

use error_data::ErrorData;
pub use notify;

use std::{
    collections::HashMap,
    path::{Path, PathBuf},
    sync::{Arc, Mutex},
};

use jwalk::WalkDir;
use notify::{Config, Event, EventHandler, RecommendedWatcher, RecursiveMode, Watcher};

/// Keeps track of a directory's total size.
///
/// The total size is updated in a separate thread using the specified implementation of [`Watcher`].
#[derive(Debug)]
pub struct DirSizeTracker<WatcherImpl = RecommendedWatcher> {
    state: Arc<Mutex<State>>,
    _watcher: Arc<WatcherImpl>,
}

// Implement Clone without requiring WatchingImpl to implement Clone
impl<WatcherImpl> Clone for DirSizeTracker<WatcherImpl> {
    fn clone(&self) -> Self {
        DirSizeTracker {
            state: Arc::clone(&self.state),
            _watcher: Arc::clone(&self._watcher),
        }
    }
}

impl<WatcherImpl: Watcher> DirSizeTracker<WatcherImpl> {
    /// Starts tracking the size of the given directory.
    ///
    /// For best performance, this should only be called once.
    pub fn new(path: PathBuf, mode: Mode) -> Result<Self, Error> {
        Self::try_new(path, mode).map_err(Error)
    }

    fn try_new(path: PathBuf, mode: Mode) -> Result<Self, ErrorData> {
        // The initial scan must be done after the watcher starts so that any events occurring during the scan will be handled.
        let state = Arc::new(Mutex::new(State {
            path,
            mode,
            need_scan: true,
            sizes: HashMap::new(),
            total_size: 0,
            error: None,
        }));
        let mut watcher =
            WatcherImpl::new(create_event_handler(Arc::clone(&state)), Config::default())?;
        watcher.watch(&state.lock().unwrap().path, RecursiveMode::Recursive)?;
        let dir_size_tracker = DirSizeTracker {
            state,
            _watcher: Arc::new(watcher),
        };
        Ok(dir_size_tracker)
    }

    /// Returns the current size of the directory.
    ///
    /// If an error occured when responding to a filesystem event, then the error is returned, the size tracking starts over, and the next call to this function will return the result of a re-attempt.
    pub fn get_total_size(&self) -> Result<u64, Error> {
        let mut state = self.state.lock().unwrap();
        if state.need_scan {
            state.need_scan = false;
            state.error = state.scan().err();
        }
        if let Some(error) = state.error.take() {
            state.need_scan = true;
            return Err(Error(error));
        }
        Ok(state.total_size)
    }
}

fn create_event_handler(state: Arc<Mutex<State>>) -> impl EventHandler {
    move |result| {
        let mut state = state.lock().unwrap();
        if state.error.is_none() {
            state.error = state.handle_event(result).err();
        }
    }
}

#[derive(Debug)]
struct State {
    path: PathBuf,
    mode: Mode,
    need_scan: bool,
    sizes: HashMap<PathBuf, u64>,
    total_size: u64,
    error: Option<ErrorData>,
}

impl State {
    fn scan(&mut self) -> Result<(), ErrorData> {
        self.total_size = 0;
        self.sizes = WalkDir::new(&self.path)
            .skip_hidden(false)
            .into_iter()
            .map(|result| {
                let path = result?.path();
                let size = get_size(&path, self.mode)?.unwrap_or(0);
                self.add_size(size)?;
                Ok((path, size))
            })
            .collect::<Result<_, ErrorData>>()?;
        Ok(())
    }

    fn handle_event(&mut self, result: Result<Event, notify::Error>) -> Result<(), ErrorData> {
        let event = result?;
        if event.need_rescan() || self.need_scan {
            self.need_scan = false;
            self.scan()?;
        } else if event.kind.is_access() {
            // No modification was made
        } else {
            for path in event.paths {
                let new_size = get_size(&path, self.mode)?;
                let old_size = if let Some(new_size) = new_size {
                    self.sizes.insert(path, new_size)
                } else {
                    self.sizes.remove(&path)
                };
                self.subtract_size(old_size.unwrap_or(0))?;
                self.add_size(new_size.unwrap_or(0))?;
            }
        }
        Ok(())
    }

    fn add_size(&mut self, n: u64) -> Result<(), ErrorData> {
        self.total_size = self
            .total_size
            .checked_add(n)
            .ok_or(ErrorData::IntOverflow)?;
        Ok(())
    }

    fn subtract_size(&mut self, n: u64) -> Result<(), ErrorData> {
        self.total_size = self
            .total_size
            .checked_sub(n)
            .ok_or(ErrorData::IntUnderflow)?;
        Ok(())
    }
}

/// Error that occurs in [`DirSizeTracker::new`] or when responding to filesystem events.
#[derive(Debug)]
pub struct Error(ErrorData);

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self.0 {
            ErrorData::Io(ref error) => write!(f, "{error}"),
            ErrorData::Jwalk(_) => write!(f, "failed to scan directory"),
            ErrorData::Notify(_) => write!(f, "failed to get filesystem events"),
            ErrorData::IntOverflow => write!(f, "total size is greater than `u64::MAX`"),
            ErrorData::IntUnderflow => write!(f, "calculated total size became negative"),
        }
    }
}

impl std::error::Error for Error {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self.0 {
            ErrorData::Io(ref error) => std::error::Error::source(error),
            ErrorData::Jwalk(ref error) => Some(error),
            ErrorData::Notify(ref error) => Some(error),
            ErrorData::IntOverflow => None,
            ErrorData::IntUnderflow => None,
        }
    }
}

/// Specifies the method used for calculating file sizes.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum Mode {
    /// Uses [`std::fs::Metadata::len`] for each file and folder.
    Metadata,
}

/// Returns the file's size, or `None` if it doesn't exist
fn get_size(path: &Path, mode: Mode) -> Result<Option<u64>, ErrorData> {
    match try_get_size(path, mode) {
        Ok(value) => Ok(Some(value)),
        Err(ErrorData::Io(error)) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
        Err(error) => Err(error),
    }
}

fn try_get_size(path: &Path, mode: Mode) -> Result<u64, ErrorData> {
    let metadata = std::fs::symlink_metadata(path)?;
    let size = match mode {
        Mode::Metadata => metadata.len(),
    };
    Ok(size)
}