Merge pull request #220 from solidiquis/race-to-print-bug

Fixed race to print

5 files changed, 70 insertions, 68 deletions

diff --git a/README.md b/README.md
index aef3f72..f468a85 100644
--- a/README.md
+++ b/README.md
@@ -285,7 +285,7 @@ $ rustup toolchain install nightly-2023-06-11
 Thereafter:
 
 ```
-$ cargo +nightly-2023-03-05 install erdtree
+$ cargo +nightly-2023-06-11 install erdtree
 ```
 
 ### Homebrew-core
diff --git a/src/context/mod.rs b/src/context/mod.rs
index 7764842..031bad0 100644
--- a/src/context/mod.rs
+++ b/src/context/mod.rs
@@ -257,8 +257,10 @@ impl Context {
     /// Initializes [Context], optionally reading in the configuration file to override defaults.
     /// Arguments provided will take precedence over config.
     pub fn try_init() -> Result<Self, Error> {
-        let args = Self::compute_args()?;
-        Self::from_arg_matches(&args).map_err(Error::Config)
+        Self::compute_args().and_then(|args| {
+            color::no_color_env();
+            Self::from_arg_matches(&args).map_err(Error::Config)
+        })
     }
 
     /// Determines whether or not it's appropriate to display color in output based on
diff --git a/src/main.rs b/src/main.rs
index 5f5e0b4..c9d4713 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -10,20 +10,13 @@
     clippy::style,
     clippy::suspicious
 )]
-#![allow(
-    clippy::cast_possible_truncation,
-    clippy::cast_precision_loss,
-    clippy::cast_sign_loss,
-    clippy::let_underscore_untyped,
-    clippy::struct_excessive_bools,
-    clippy::too_many_arguments
-)]
+#![allow(clippy::cast_precision_loss, clippy::struct_excessive_bools)]
 
 use clap::CommandFactory;
 use context::{layout, Context};
-use progress::Message;
+use progress::{Indicator, IndicatorHandle, Message};
 use render::{Engine, Flat, FlatInverted, Inverted, Regular};
-use std::{error::Error, io::stdout, process::ExitCode, sync::Arc};
+use std::{error::Error, io::stdout, process::ExitCode};
 use tree::Tree;
 
 /// Operations to wrangle ANSI escaped strings.
@@ -44,10 +37,7 @@ mod icons;
 /// Concerned with displaying a progress indicator when stdout is a tty.
 mod progress;
 
-/// Concerned with taking an initialized [`Tree`] and its [`Node`]s and rendering the output.
-///
-/// [`Tree`]: tree::Tree
-/// [`Node`]: tree::node::Node
+/// Concerned with taking an initialized [`tree::Tree`] and its [`tree::node::Node`]s and rendering the output.
 mod render;
 
 /// Global used throughout the program to paint the output.
@@ -80,29 +70,18 @@ fn run() -> Result<(), Box<dyn Error>> {
         return Ok(());
     }
 
-    context::color::no_color_env();
-
     styles::init(ctx.no_color());
 
-    let indicator = (ctx.stdout_is_tty && !ctx.no_progress)
-        .then(progress::Indicator::measure)
-        .map(Arc::new);
+    let indicator = Indicator::maybe_init(&ctx);
 
-    if indicator.is_some() {
-        let indicator = indicator.clone();
-
-        ctrlc::set_handler(move || {
-            let _ = progress::IndicatorHandle::terminate(indicator.clone());
-            tty::restore_tty();
-        })?;
-    }
-
-    let (tree, ctx) = match Tree::try_init(ctx, indicator.clone()) {
-        Ok(res) => res,
-        Err(err) => {
-            let _ = progress::IndicatorHandle::terminate(indicator);
-            return Err(Box::new(err));
-        },
+    let (tree, ctx) = {
+        match Tree::try_init(ctx, indicator.as_ref()) {
+            Ok(res) => res,
+            Err(err) => {
+                IndicatorHandle::terminate(indicator);
+                return Err(Box::new(err));
+            },
+        }
     };
 
     macro_rules! compute_output {
@@ -122,12 +101,11 @@ fn run() -> Result<(), Box<dyn Error>> {
     if let Some(mut progress) = indicator {
         progress.mailbox().send(Message::RenderReady)?;
 
-        if let Some(hand) = Arc::get_mut(&mut progress) {
-            hand.join_handle
-                .take()
-                .map(|h| h.join().unwrap())
-                .transpose()?;
-        }
+        progress
+            .join_handle
+            .take()
+            .map(|h| h.join().unwrap())
+            .transpose()?;
     }
 
     #[cfg(debug_assertions)]
diff --git a/src/progress.rs b/src/progress.rs
index 440f1dd..73ca59a 100644
--- a/src/progress.rs
+++ b/src/progress.rs
@@ -1,3 +1,4 @@
+use crate::{context::Context, tty};
 use crossterm::{
     cursor,
     terminal::{self, ClearType},
@@ -5,10 +6,7 @@ use crossterm::{
 };
 use std::{
     io::{self, Write},
-    sync::{
-        mpsc::{self, SendError, SyncSender},
-        Arc,
-    },
+    sync::mpsc::{self, SendError, SyncSender},
     thread::{self, JoinHandle},
 };
 
@@ -99,17 +97,26 @@ impl IndicatorHandle {
         self.mailbox.clone()
     }
 
-    /// Send a message through to the `priority_mailbox` tear down the [`Indicator`].
-    pub fn terminate(this: Option<Arc<Self>>) -> Result<(), Error> {
-        if let Some(mut handle) = this {
-            handle.mailbox().send(Message::Finish)?;
+    /// Analogous to [`Self::try_terminate`] but panics if failure.
+    pub fn terminate(this: Option<Self>) {
+        Self::try_terminate(this).expect("Failed to properly terminate the progress indicator");
+    }
 
-            if let Some(hand) = Arc::get_mut(&mut handle) {
-                hand.join_handle
-                    .take()
-                    .map(|h| h.join().unwrap())
-                    .transpose()?;
-            }
+    /// Attempts to terminate the [`Indicator`] with cleanup.
+    pub fn try_terminate(this: Option<Self>) -> Result<(), Error> {
+        if let Some(mut handle) = this {
+            // This is allowed to fail silently. If user administers interrupt then the `Indicator`
+            // will be dropped along with the receiving end of the `mailbox`.
+            //
+            // If user does not administer interrupt but file-system traversal fails for whatever
+            // reason then this will proceed as normal.
+            let _ = handle.mailbox().send(Message::Finish);
+
+            handle
+                .join_handle
+                .take()
+                .map(|h| h.join().unwrap())
+                .transpose()?;
         }
 
         Ok(())
@@ -117,6 +124,27 @@ impl IndicatorHandle {
 }
 
 impl<'a> Indicator<'a> {
+    /// Initializes an [`Indicator`] returning an atomic reference counter of an [`IndicatorHandle`] if
+    /// a progress indicator is enabled via [`Context`]. Upon initialization an interrupt handler is
+    /// also registered. Sources of panic can come from [`IndicatorHandle::terminate`] or
+    /// [`ctrlc::set_handler`].
+    pub fn maybe_init(ctx: &Context) -> Option<IndicatorHandle> {
+        (ctx.stdout_is_tty && !ctx.no_progress)
+            .then(Indicator::measure)
+            .map(|indicator| {
+                let mailbox = indicator.mailbox();
+
+                let int_handler = move || {
+                    let _ = mailbox.try_send(Message::Finish);
+                    tty::restore_tty();
+                };
+
+                ctrlc::set_handler(int_handler).expect("Failed to set interrupt handler");
+
+                indicator
+            })
+    }
+
     /// Initializes a worker thread that owns [`Indicator`] that awaits on [`Message`]s to traverse
     /// through its internal states. An [`IndicatorHandle`] is returned as a mechanism to allow the
     /// outside world to send messages to the worker thread and ultimately to the [`Indicator`].
diff --git a/src/tree/mod.rs b/src/tree/mod.rs
index fff0815..be2b7ee 100644
--- a/src/tree/mod.rs
+++ b/src/tree/mod.rs
@@ -16,10 +16,7 @@ use std::{
     fs,
     path::PathBuf,
     result::Result as StdResult,
-    sync::{
-        mpsc::{self, Sender},
-        Arc,
-    },
+    sync::mpsc::{self, Sender},
     thread,
 };
 use visitor::{BranchVisitorBuilder, TraversalState};
@@ -30,10 +27,7 @@ pub mod count;
 /// Errors related to traversal, [Tree] construction, and the like.
 pub mod error;
 
-/// Contains components of the [`Tree`] data structure that derive from [`DirEntry`].
-///
-/// [`Tree`]: Tree
-/// [`DirEntry`]: ignore::DirEntry
+/// Contains components of the [`Tree`] data structure that derive from [`ignore::DirEntry`].
 pub mod node;
 
 /// Custom visitor that operates on each thread during filesystem traversal.
@@ -57,7 +51,7 @@ impl Tree {
     /// various properties necessary to render output.
     pub fn try_init(
         mut ctx: Context,
-        indicator: Option<Arc<IndicatorHandle>>,
+        indicator: Option<&IndicatorHandle>,
     ) -> Result<(Self, Context)> {
         let mut column_properties = column::Properties::from(&ctx);
 
@@ -105,12 +99,12 @@ impl Tree {
     fn traverse(
         ctx: &Context,
         column_properties: &mut column::Properties,
-        indicator: Option<Arc<IndicatorHandle>>,
+        indicator: Option<&IndicatorHandle>,
     ) -> Result<(Arena<Node>, NodeId)> {
         let walker = WalkParallel::try_from(ctx)?;
         let (tx, rx) = mpsc::channel();
 
-        let progress_indicator_mailbox = indicator.map(|arc| arc.mailbox());
+        let progress_indicator_mailbox = indicator.map(IndicatorHandle::mailbox);
 
         thread::scope(|s| {
             let res = s.spawn(move || {