generalize the throttle implementation to allow usagein UI

2 files changed, 49 insertions, 66 deletions

diff --git a/src/aggregate.rs b/src/aggregate.rs
index 6e1cc2f..e897dbd 100644
--- a/src/aggregate.rs
+++ b/src/aggregate.rs
@@ -1,73 +1,16 @@
-use crate::{crossdev, InodeFilter, WalkOptions, WalkResult};
+use crate::{crossdev, InodeFilter, Throttle, WalkOptions, WalkResult};
 use anyhow::Result;
 use filesize::PathExt;
 use owo_colors::{AnsiColors as Color, OwoColorize};
+use std::time::Duration;
 use std::{io, path::Path};
-use std::{
-    sync::{
-        atomic::{AtomicBool, Ordering},
-        Arc,
-    },
-    thread,
-    time::Duration,
-};
-
-/// Throttle access to an optional `io::Write` to the specified `Duration`
-#[derive(Debug)]
-struct ThrottleWriter<W> {
-    out: Option<W>,
-    trigger: Arc<AtomicBool>,
-}
-
-impl<W> ThrottleWriter<W>
-where
-    W: io::Write,
-{
-    fn new(out: Option<W>, duration: Duration) -> Self {
-        let writer = Self {
-            out,
-            trigger: Default::default(),
-        };
-
-        if writer.out.is_some() {
-            let trigger = Arc::downgrade(&writer.trigger);
-            thread::spawn(move || {
-                thread::sleep(Duration::from_secs(1));
-                while let Some(t) = trigger.upgrade() {
-                    t.store(true, Ordering::Relaxed);
-                    thread::sleep(duration);
-                }
-            });
-        }
-
-        writer
-    }
-
-    fn throttled<F>(&mut self, f: F)
-    where
-        F: FnOnce(&mut W),
-    {
-        if self.trigger.swap(false, Ordering::Relaxed) {
-            self.unthrottled(f);
-        }
-    }
-
-    fn unthrottled<F>(&mut self, f: F)
-    where
-        F: FnOnce(&mut W),
-    {
-        if let Some(ref mut out) = self.out {
-            f(out);
-        }
-    }
-}
 
 /// Aggregate the given `paths` and write information about them to `out` in a human-readable format.
 /// If `compute_total` is set, it will write an additional line with the total size across all given `paths`.
 /// If `sort_by_size_in_bytes` is set, we will sort all sizes (ascending) before outputting them.
 pub fn aggregate(
     mut out: impl io::Write,
-    err: Option<impl io::Write>,
+    mut err: Option<impl io::Write>,
     walk_options: WalkOptions,
     compute_total: bool,
     sort_by_size_in_bytes: bool,
@@ -82,7 +25,7 @@ pub fn aggregate(
     let mut num_roots = 0;
     let mut aggregates = Vec::new();
     let mut inodes = InodeFilter::default();
-    let mut progress = ThrottleWriter::new(err, Duration::from_millis(100));
+    let progress = Throttle::new(Duration::from_millis(100));
 
     for path in paths.into_iter() {
         num_roots += 1;
@@ -99,8 +42,10 @@ pub fn aggregate(
         };
         for entry in walk_options.iter_from_path(path.as_ref(), device_id) {
             stats.entries_traversed += 1;
-            progress.throttled(|out| {
-                write!(out, "Enumerating {} entries\r", stats.entries_traversed).ok();
+            progress.throttled(|| {
+                if let Some(err) = err.as_mut() {
+                    write!(err, "Enumerating {} entries\r", stats.entries_traversed).ok();
+                }
             });
             match entry {
                 Ok(entry) => {
@@ -134,9 +79,10 @@ pub fn aggregate(
                 Err(_) => num_errors += 1,
             }
         }
-        progress.unthrottled(|out| {
-            write!(out, "\x1b[2K\r").ok();
-        });
+
+        if let Some(err) = err.as_mut() {
+            write!(err, "\x1b[2K\r").ok();
+        }
 
         if sort_by_size_in_bytes {
             aggregates.push((path.as_ref().to_owned(), num_bytes, num_errors));
diff --git a/src/common.rs b/src/common.rs
index 5eb094e..a277c1f 100644
--- a/src/common.rs
+++ b/src/common.rs
@@ -2,6 +2,9 @@ use crate::crossdev;
 use crate::traverse::{EntryData, Tree, TreeIndex};
 use byte_unit::{n_gb_bytes, n_gib_bytes, n_mb_bytes, n_mib_bytes, ByteUnit};
 use std::path::PathBuf;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+use std::time::Duration;
 use std::{fmt, path::Path};
 
 pub fn get_entry_or_panic(tree: &Tree, node_idx: TreeIndex) -> &EntryData {
@@ -114,6 +117,40 @@ pub enum TraversalSorting {
     AlphabeticalByFileName,
 }
 
+/// Throttle access to an optional `io::Write` to the specified `Duration`
+#[derive(Debug)]
+pub struct Throttle {
+    trigger: Arc<AtomicBool>,
+}
+
+impl Throttle {
+    pub fn new(duration: Duration) -> Self {
+        let instance = Self {
+            trigger: Default::default(),
+        };
+
+        let trigger = Arc::downgrade(&instance.trigger);
+        std::thread::spawn(move || {
+            std::thread::sleep(Duration::from_secs(1));
+            while let Some(t) = trigger.upgrade() {
+                t.store(true, Ordering::Relaxed);
+                std::thread::sleep(duration);
+            }
+        });
+
+        instance
+    }
+
+    pub fn throttled<F>(&self, f: F)
+    where
+        F: FnOnce(),
+    {
+        if self.trigger.swap(false, Ordering::Relaxed) {
+            f()
+        }
+    }
+}
+
 /// Configures a filesystem walk, including output and formatting options.
 #[derive(Clone)]
 pub struct WalkOptions {