feat: limit thread number to 1 when HDD is detected (#269)

Integral-Tech · KSXGitHub · web-flow · commit 18f5f2a365da · 2024-12-01T16:47:50.000+07:00
* feat: limit thread number to 1 when HDD is detected On HDD, multi-threaded du won't provide any performance benefit over the single-threaded one. Therefore, when HDD is detected on one of the files, limit the thread number to 1. This PR fixes issue #257. * Fix change requests * refactor: create a module * refactor: rename a test * refactor: unnecessary visibility scope * refactor: make `detect_hdd_in_files` testable * feat: improve warning messages * refactor: rename a function * refactor: rename an argument * refactor: extract a function * Add tests for path_is_in_hdd() & any_path_is_in_hdd() * Apply suggestions * Update src/app/sub/hdd.rs Co-authored-by: Khải <hvksmr1996@gmail.com> * Remove unused import * refactor: use mocked trait * style: remove unnecessary trailing comma * refactor: simplify the code * refactor: use constructor to reduce lines of code * test: add more cases * refactor: shorten the code * refactor: descriptive paths * refactor: rearrange * refactor: shorten the code * style: add trailing comma * refactor: early return * docs: add --------- Co-authored-by: khai96_ <hvksmr1996@gmail.com>
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -65,6 +65,7 @@ clap_complete = { version = "^4.5.36", optional = true }
 clap-utilities = { version = "^0.2.0", optional = true }
 serde = { version = "^1.0.214", optional = true }
 serde_json = { version = "^1.0.132", optional = true }
+sysinfo = "0.32.0"
 
 [dev-dependencies]
 build-fs-tree = "^0.7.1"
diff --git a/src/app/sub.rs b/src/app/sub.rs
@@ -1,3 +1,6 @@
+mod hdd;
+mod mount_point;
+
 use crate::{
     args::Fraction,
     data_tree::{DataTree, DataTreeReflection},
@@ -11,8 +14,10 @@ use crate::{
     status_board::GLOBAL_STATUS_BOARD,
     visualizer::{BarAlignment, ColumnWidthDistribution, Direction, Visualizer},
 };
+use hdd::any_path_is_in_hdd;
 use serde::Serialize;
 use std::{io::stdout, iter::once, num::NonZeroUsize, path::PathBuf};
+use sysinfo::Disks;
 
 /// The sub program of the main application.
 pub struct Sub<Size, SizeGetter, Report>
@@ -69,6 +74,17 @@ where
             no_sort,
         } = self;
 
+        // If one of the files is on HDD, set thread number to 1
+        let disks = Disks::new_with_refreshed_list();
+
+        if any_path_is_in_hdd::<hdd::RealApi>(&files, &disks) {
+            eprintln!("warning: HDD detected, the thread limit will be set to 1");
+            rayon::ThreadPoolBuilder::new()
+                .num_threads(1)
+                .build_global()
+                .unwrap_or_else(|_| eprintln!("warning: Failed to set thread limit to 1"));
+        }
+
         let mut iter = files
             .into_iter()
             .map(|root| -> DataTree<OsStringDisplay, Size> {
diff --git a/src/app/sub/hdd.rs b/src/app/sub/hdd.rs
@@ -0,0 +1,152 @@
+use super::mount_point::find_mount_point;
+use std::{
+    fs::canonicalize,
+    io,
+    path::{Path, PathBuf},
+};
+use sysinfo::{Disk, DiskKind};
+
+/// Mockable APIs to interact with the system.
+pub trait Api {
+    type Disk;
+    fn get_disk_kind(disk: &Self::Disk) -> DiskKind;
+    fn get_mount_point(disk: &Self::Disk) -> &Path;
+    fn canonicalize(path: &Path) -> io::Result<PathBuf>;
+}
+
+/// Implementation of [`Api`] that interacts with the real system.
+pub struct RealApi;
+impl Api for RealApi {
+    type Disk = Disk;
+
+    fn get_disk_kind(disk: &Self::Disk) -> DiskKind {
+        disk.kind()
+    }
+
+    fn get_mount_point(disk: &Self::Disk) -> &Path {
+        disk.mount_point()
+    }
+
+    fn canonicalize(path: &Path) -> io::Result<PathBuf> {
+        canonicalize(path)
+    }
+}
+
+/// Check if any path is in any HDD.
+pub fn any_path_is_in_hdd<Api: self::Api>(paths: &[PathBuf], disks: &[Api::Disk]) -> bool {
+    paths
+        .iter()
+        .filter_map(|file| Api::canonicalize(file).ok())
+        .any(|path| path_is_in_hdd::<Api>(&path, disks))
+}
+
+/// Check if path is in any HDD.
+fn path_is_in_hdd<Api: self::Api>(path: &Path, disks: &[Api::Disk]) -> bool {
+    let Some(mount_point) = find_mount_point(path, disks.iter().map(Api::get_mount_point)) else {
+        return false;
+    };
+    disks
+        .iter()
+        .filter(|disk| Api::get_disk_kind(disk) == DiskKind::HDD)
+        .any(|disk| Api::get_mount_point(disk) == mount_point)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{any_path_is_in_hdd, path_is_in_hdd, Api};
+    use pipe_trait::Pipe;
+    use pretty_assertions::assert_eq;
+    use std::path::{Path, PathBuf};
+    use sysinfo::DiskKind;
+
+    /// Fake disk for [`Api`].
+    struct Disk {
+        kind: DiskKind,
+        mount_point: &'static str,
+    }
+
+    impl Disk {
+        fn new(kind: DiskKind, mount_point: &'static str) -> Self {
+            Self { kind, mount_point }
+        }
+    }
+
+    /// Mocked implementation of [`Api`] for testing purposes.
+    struct MockedApi;
+    impl Api for MockedApi {
+        type Disk = Disk;
+
+        fn get_disk_kind(disk: &Self::Disk) -> DiskKind {
+            disk.kind
+        }
+
+        fn get_mount_point(disk: &Self::Disk) -> &Path {
+            Path::new(disk.mount_point)
+        }
+
+        fn canonicalize(path: &Path) -> std::io::Result<PathBuf> {
+            path.to_path_buf().pipe(Ok)
+        }
+    }
+
+    #[test]
+    fn test_any_path_in_hdd() {
+        let disks = &[
+            Disk::new(DiskKind::SSD, "/"),
+            Disk::new(DiskKind::HDD, "/home"),
+            Disk::new(DiskKind::HDD, "/mnt/hdd-data"),
+            Disk::new(DiskKind::SSD, "/mnt/ssd-data"),
+            Disk::new(DiskKind::HDD, "/mnt/hdd-data/repo"),
+        ];
+
+        let cases: &[(&[&str], bool)] = &[
+            (&[], false),
+            (&["/"], false),
+            (&["/home"], true),
+            (&["/mnt"], false),
+            (&["/mnt/ssd-data"], false),
+            (&["/mnt/hdd-data"], true),
+            (&["/mnt/hdd-data/repo"], true),
+            (&["/etc/fstab"], false),
+            (&["/home/usr/file"], true),
+            (&["/home/data/repo/test"], true),
+            (&["/usr/share"], false),
+            (&["/mnt/ssd-data/test"], false),
+            (&["/etc/fstab", "/home/user/file"], true),
+            (&["/mnt/hdd-data/file", "/mnt/hdd-data/repo/test"], true),
+            (&["/usr/share", "/mnt/ssd-data/test"], false),
+            (
+                &["/etc/fstab", "/home/user", "/mnt/hdd-data", "/usr/share"],
+                true,
+            ),
+        ];
+
+        for (paths, in_hdd) in cases {
+            let paths: Vec<_> = paths.iter().map(PathBuf::from).collect();
+            println!("CASE: {paths:?} → {in_hdd:?}");
+            assert_eq!(any_path_is_in_hdd::<MockedApi>(&paths, disks), *in_hdd);
+        }
+    }
+
+    #[test]
+    fn test_path_in_hdd() {
+        let disks = &[
+            Disk::new(DiskKind::SSD, "/"),
+            Disk::new(DiskKind::HDD, "/home"),
+            Disk::new(DiskKind::HDD, "/mnt/hdd-data"),
+            Disk::new(DiskKind::SSD, "/mnt/ssd-data"),
+            Disk::new(DiskKind::HDD, "/mnt/hdd-data/repo"),
+        ];
+
+        for (path, in_hdd) in [
+            ("/etc/fstab", false),
+            ("/mnt/", false),
+            ("/mnt/hdd-data/repo/test", true),
+            ("/mnt/hdd-data/test/test", true),
+            ("/mnt/ssd-data/test/test", false),
+        ] {
+            println!("CASE: {path} → {in_hdd:?}");
+            assert_eq!(path_is_in_hdd::<MockedApi>(Path::new(path), disks), in_hdd);
+        }
+    }
+}
diff --git a/src/app/sub/mount_point.rs b/src/app/sub/mount_point.rs
@@ -0,0 +1,39 @@
+use std::{ffi::OsStr, path::Path};
+
+/// Find a mount point that contains `absolute_path`.
+pub fn find_mount_point<'a>(
+    absolute_path: &Path,
+    all_mount_points: impl IntoIterator<Item = &'a Path>,
+) -> Option<&'a Path> {
+    all_mount_points
+        .into_iter()
+        .filter(|mnt| absolute_path.starts_with(mnt))
+        .max_by_key(|mnt| AsRef::<OsStr>::as_ref(mnt).len()) // Mount points can be nested in each other
+}
+
+#[cfg(test)]
+mod tests {
+    use super::find_mount_point;
+    use pretty_assertions::assert_eq;
+    use std::path::Path;
+
+    #[test]
+    fn test_mount_point() {
+        let all_mount_points = ["/", "/home", "/mnt/data", "/mnt/data/repo", "/mnt/repo"];
+
+        for (path, expected_mount_point) in &[
+            ("/etc/fstab", "/"),
+            ("/home/user", "/home"),
+            ("/mnt/data/repo/test", "/mnt/data/repo"),
+            ("/mnt/data/test/test", "/mnt/data/"),
+            ("/mnt/repo/test/test", "/mnt/repo/"),
+        ] {
+            println!("CASE: {path} → {expected_mount_point}");
+            let all_mount_points = all_mount_points.map(Path::new);
+            assert_eq!(
+                find_mount_point(Path::new(path), all_mount_points).unwrap(),
+                Path::new(expected_mount_point),
+            );
+        }
+    }
+}
