Merge pull request #116 from KSXGitHub/use-atomics-to-improve-performance

Use atomics to improve performance

Author: KSXGitHub

Cargo.toml

@@ -3,7 +3,7 @@ name = "parallel-disk-usage"
 description = "Highly parallelized, blazing fast directory tree analyzer"
 version = "0.6.1"
 authors = ["khai96_ <[email protected]>"]
-edition = "2018"
+edition = "2021"
 readme = "README.md"
 license = "Apache-2.0"
 documentation = "https://docs.rs/parallel-disk-usage"

rust-toolchain

@@ -1 +1 @@
-1.54.0
+1.58.1

src/app.rs

@@ -110,6 +110,7 @@ impl App {
         where
             Data: Size + Into<u64> + Send + Sync,
             ProgressReport<Data>: Default + 'static,
+            u64: Into<Data>,
         {
             ProgressAndErrorReporter::new(
                 ProgressReport::TEXT,

src/reporter/progress_and_error_reporter.rs

@@ -1,8 +1,10 @@
 use super::{ErrorReport, Event, ParallelReporter, ProgressReport, Reporter, Size};
-use pipe_trait::Pipe;
+use progress_report_state::ProgressReportState;
 use std::{
     any::Any,
-    sync::{Arc, RwLock},
+    marker::PhantomData,
+    ops::ControlFlow,
+    sync::{atomic::Ordering::Relaxed, Arc},
     thread::{sleep, spawn, JoinHandle},
     time::Duration,
 };
@@ -15,19 +17,23 @@ pub struct ProgressAndErrorReporter<Data, ReportError>
 where
     Data: Size + Send + Sync,
     ReportError: Fn(ErrorReport) + Sync,
+    u64: Into<Data>,
 {
     /// Progress information.
-    progress: Arc<RwLock<Option<ProgressReport<Data>>>>,
+    progress: Arc<ProgressReportState>,
     /// Report encountered error.
     report_error: ReportError,
     /// Join handle of progress reporting thread.
     progress_reporter_handle: JoinHandle<()>,
+    /// Keep generic parameters.
+    _phantom: PhantomData<Data>,
 }
 
 impl<Data, ReportError> ProgressAndErrorReporter<Data, ReportError>
 where
     Data: Size + Send + Sync,
     ReportError: Fn(ErrorReport) + Sync,
+    u64: Into<Data>,
 {
     /// Create a new [`ProgressAndErrorReporter`] from a report function.
     pub fn new<ReportProgress>(
@@ -39,41 +45,36 @@ where
         ProgressReport<Data>: Default + 'static,
         ReportProgress: Fn(ProgressReport<Data>) + Send + Sync + 'static,
     {
-        let progress = ProgressReport::default()
-            .pipe(Some)
-            .pipe(RwLock::new)
-            .pipe(Arc::new);
+        let progress = Arc::new(ProgressReportState::default());
         let progress_thread = progress.clone();
         let progress_reporter_handle = spawn(move || loop {
             sleep(progress_report_interval);
-            if let Ok(progress) = progress_thread.read().as_deref() {
-                if let Some(progress) = *progress {
-                    report_progress(progress);
-                } else {
-                    break;
-                }
-            }
+            match progress_thread.to_progress_report() {
+                ControlFlow::Continue(progress) => report_progress(progress),
+                ControlFlow::Break(()) => break,
+            };
         });
         ProgressAndErrorReporter {
             progress,
             report_error,
             progress_reporter_handle,
+            _phantom: PhantomData,
         }
     }
 
     /// Stop the thread that reports progress.
     ///
     /// This function would be automatically invoked once the value is [dropped](Drop).
     pub fn stop_progress_reporter(&self) {
-        let mut progress = self.progress.write().expect("lock progress to stop");
-        *progress = None;
+        self.progress.stopped.store(true, Relaxed);
     }
 }
 
 impl<Data, ReportError> Reporter<Data> for ProgressAndErrorReporter<Data, ReportError>
 where
-    Data: Size + Send + Sync,
+    Data: Size + Into<u64> + Send + Sync,
     ReportError: Fn(ErrorReport) + Sync,
+    u64: Into<Data>,
 {
     fn report(&self, event: Event<Data>) {
         use Event::*;
@@ -82,38 +83,35 @@ where
             report_error,
             ..
         } = self;
-        macro_rules! handle_field {
-            ($($field:ident $operator:tt $addend:expr;)+) => {
-                if let Some(progress) = progress.write().ok().as_mut().and_then(|x| x.as_mut()) {
-                    $(progress.$field $operator $addend;)+
-                }
+        macro_rules! bump {
+            ($field:ident += $delta:expr) => {
+                progress.$field.fetch_add($delta, Relaxed)
             };
         }
         match event {
             ReceiveData(data) => {
-                handle_field! {
-                    items += 1;
-                    total += data;
-                }
+                bump!(items += 1);
+                bump!(total += data.into());
             }
             EncounterError(error_report) => {
                 report_error(error_report);
-                handle_field! {
-                    errors += 1;
-                }
+                bump!(errors += 1);
             }
         }
     }
 }
 
 impl<Data, ReportError> ParallelReporter<Data> for ProgressAndErrorReporter<Data, ReportError>
 where
-    Data: Size + Send + Sync,
+    Data: Size + Into<u64> + Send + Sync,
     ReportError: Fn(ErrorReport) + Sync,
+    u64: Into<Data>,
 {
     type DestructionError = Box<dyn Any + Send + 'static>;
     fn destroy(self) -> Result<(), Self::DestructionError> {
         self.stop_progress_reporter();
         self.progress_reporter_handle.join()
     }
 }
+
+mod progress_report_state;

src/reporter/progress_and_error_reporter/progress_report_state.rs

@@ -0,0 +1,47 @@
+use crate::{reporter::ProgressReport, size::Size};
+use std::{
+    ops::ControlFlow,
+    sync::atomic::{AtomicBool, AtomicU64, Ordering::Relaxed},
+};
+
+/// Like [`ProgressReport`] but mutable.
+#[derive(Debug, Default)]
+pub struct ProgressReportState {
+    /// Whether the progress has stopped.
+    pub stopped: AtomicBool,
+    /// Number of scanned items.
+    pub items: AtomicU64,
+    /// Total size of scanned items.
+    pub total: AtomicU64,
+    /// Number of occurred errors.
+    pub errors: AtomicU64,
+}
+
+impl ProgressReportState {
+    /// Yield [`ProgressReport`] if it is running.
+    /// Return `Break` otherwise.
+    pub fn to_progress_report<Data>(&self) -> ControlFlow<(), ProgressReport<Data>>
+    where
+        Data: Size,
+        u64: Into<Data>,
+    {
+        macro_rules! load {
+            ($field:ident) => {
+                self.$field.load(Relaxed)
+            };
+        }
+
+        if load!(stopped) {
+            return ControlFlow::Break(());
+        }
+
+        let items = load!(items);
+        let total = load!(total).into();
+        let errors = load!(errors);
+        ControlFlow::Continue(ProgressReport {
+            items,
+            total,
+            errors,
+        })
+    }
+}

src/visualizer/proportion_bar.rs

@@ -76,7 +76,7 @@ impl Display for ProportionBarDisplay {
                     level2 = bar.display_level2(),
                     level3 = bar.display_level3(),
                     level4 = bar.display_level4(),
-                );
+                )
             };
         }
         match align {

tests/visualizer.rs

@@ -1,3 +1,5 @@
+#![recursion_limit = "256"]
+
 use parallel_disk_usage::{
     bytes_format::BytesFormat::*,
     data_tree::DataTree,