diff options
author | Benjamin Nguyen <benjamin.van.nguyen@gmail.com> | 2023-07-01 09:58:24 +0700 |
---|---|---|
committer | Benjamin Nguyen <benjamin.van.nguyen@gmail.com> | 2023-07-01 09:58:24 +0700 |
commit | 25330d7c449917a78fcee07a364f9d0aa26794cb (patch) | |
tree | 82b2f5a6f5a3655cf604b5196ecee6a8b3f801be | |
parent | be3c86cdd74c924a6c8915bcc65860f6643d66ba (diff) |
sigint handler
-rw-r--r-- | Cargo.lock | 29 | ||||
-rw-r--r-- | Cargo.toml | 1 | ||||
-rw-r--r-- | src/main.rs | 31 | ||||
-rw-r--r-- | src/progress.rs | 85 | ||||
-rw-r--r-- | src/tree/error.rs | 5 | ||||
-rw-r--r-- | src/tree/mod.rs | 23 | ||||
-rw-r--r-- | src/tree/visitor.rs | 4 |
7 files changed, 136 insertions, 42 deletions
@@ -253,6 +253,16 @@ dependencies = [ ] [[package]] +name = "ctrlc" +version = "3.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a011bbe2c35ce9c1f143b7af6f94f29a167beb4cd1d29e6740ce836f723120e" +dependencies = [ + "nix", + "windows-sys 0.48.0", +] + +[[package]] name = "cxx" version = "1.0.94" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -342,6 +352,7 @@ dependencies = [ "clap_complete", "config", "crossterm", + "ctrlc", "dirs", "errno 0.3.1", "filesize", @@ -673,6 +684,18 @@ dependencies = [ ] [[package]] +name = "nix" +version = "0.26.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfdda3d196821d6af13126e40375cdf7da646a96114af134d5f417a9a1dc8e1a" +dependencies = [ + "bitflags", + "cfg-if", + "libc", + "static_assertions", +] + +[[package]] name = "nom" version = "7.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -1046,6 +1069,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a507befe795404456341dfab10cef66ead4c041f62b8b11bbb92bffe5d0953e0" [[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + +[[package]] name = "strip-ansi-escapes" version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -30,6 +30,7 @@ clap = { version = "4.1.1", features = ["derive"] } clap_complete = "4.1.1" config = { version = "0.13.3", features = ["toml"] } crossterm = "0.26.1" +ctrlc = "3.4.0" dirs = "5.0" errno = "0.3.1" filesize = "0.2.0" diff --git a/src/main.rs b/src/main.rs index 66d7722..5c29e85 100644 --- a/src/main.rs +++ b/src/main.rs @@ -23,7 +23,7 @@ use clap::CommandFactory; use context::{layout, Context}; use progress::Message; use render::{Engine, Flat, FlatInverted, Inverted, Regular}; -use std::{error::Error, io::stdout, process::ExitCode}; +use std::{error::Error, io::stdout, process::ExitCode, sync::Arc}; use tree::Tree; /// Operations to wrangle ANSI escaped strings. @@ -87,14 +87,23 @@ fn run() -> Result<(), Box<dyn Error>> { styles::init(ctx.no_color()); - let indicator = (ctx.stdout_is_tty && !ctx.no_progress).then(progress::Indicator::measure); + let indicator = (ctx.stdout_is_tty && !ctx.no_progress) + .then(progress::Indicator::measure) + .map(Arc::new); - let (tree, ctx) = match Tree::try_init(ctx, indicator.as_ref()) { + 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) => { - if let Some(thread) = indicator.map(|i| i.join_handle) { - thread.join().unwrap()?; - } + let _ = progress::IndicatorHandle::terminate(indicator); return Err(Box::new(err)); }, }; @@ -113,9 +122,15 @@ fn run() -> Result<(), Box<dyn Error>> { layout::Type::Regular => compute_output!(Regular), }; - if let Some(progress) = indicator { + if let Some(mut progress) = indicator { progress.mailbox().send(Message::RenderReady)?; - progress.join_handle.join().unwrap()?; + + if let Some(hand) = Arc::get_mut(&mut progress) { + hand.join_handle + .take() + .map(|h| h.join().unwrap()) + .transpose()?; + } } #[cfg(debug_assertions)] diff --git a/src/progress.rs b/src/progress.rs index 3e7a740..9acc332 100644 --- a/src/progress.rs +++ b/src/progress.rs @@ -5,10 +5,16 @@ use crossterm::{ }; use std::{ io::{self, Write}, - sync::mpsc::{self, Sender}, - thread, + sync::{ + mpsc::{self, SendError, SyncSender}, + Arc, + }, + thread::{self, JoinHandle}, + time::Duration, }; +const PRIORITY_MAIL_TIMEOUT: Duration = Duration::from_nanos(1); + /// Responsible for displying the progress indicator. This struct will be owned by a separate /// thread that is responsible for displaying the progress text whereas the [`IndicatorHandle`] /// is how the outside world will interact with it. @@ -20,10 +26,11 @@ pub struct Indicator<'a> { /// This struct is how the outside world will inform the [`Indicator`] about the progress of the /// program. The `join_handle` returns the handle to the thread that owns the [`Indicator`] and the -/// `mailbox` is the [`Sender`] channel that allows [`Message`]s to be sent to [`Indicator`]. +/// `mailbox` is the [`SyncSender`] channel that allows [`Message`]s to be sent to [`Indicator`]. pub struct IndicatorHandle { - pub join_handle: thread::JoinHandle<Result<(), Error>>, - mailbox: Sender<Message>, + pub join_handle: Option<JoinHandle<Result<(), Error>>>, + mailbox: SyncSender<Message>, + priority_mailbox: SyncSender<()>, } /// The different messages that could be sent to the thread that owns the [`Indicator`]. @@ -41,7 +48,7 @@ pub enum Message { } /// All of the different states the [`Indicator`] can be in during its life cycle. -#[derive(Default, PartialEq)] +#[derive(Debug, Default, PartialEq)] enum IndicatorState { /// We are currently reading from disk. #[default] @@ -57,7 +64,13 @@ enum IndicatorState { /// Errors associated with [`crossterm`]; #[derive(Debug, thiserror::Error)] #[error(transparent)] -pub struct Error(#[from] io::Error); +pub enum Error { + #[error("#{0}")] + Io(#[from] io::Error), + + #[error("#{0}")] + Send(#[from] SendError<()>), +} impl Default for Indicator<'_> { /// Default constructor for [`Indicator`]. @@ -73,19 +86,40 @@ impl Default for Indicator<'_> { impl IndicatorHandle { /// The constructor for an [`IndicatorHandle`]. pub fn new( - join_handle: thread::JoinHandle<Result<(), Error>>, - mailbox: Sender<Message>, + join_handle: Option<JoinHandle<Result<(), Error>>>, + mailbox: SyncSender<Message>, + priority_mailbox: SyncSender<()>, ) -> Self { Self { join_handle, mailbox, + priority_mailbox, } } /// Getter for a cloned `mailbox` wherewith to send [`Message`]s to the [`Indicator`]. - pub fn mailbox(&self) -> Sender<Message> { + pub fn mailbox(&self) -> SyncSender<Message> { self.mailbox.clone() } + + /// Getter for a cloned `priority_mailbox` wherewith to send [`Message`]s to the [`Indicator`]. + pub fn priority_mailbox(&self) -> SyncSender<()> { + self.priority_mailbox.clone() + } + + pub fn terminate(this: Option<Arc<Self>>) -> Result<(), Error> { + if let Some(mut handle) = this { + handle.priority_mailbox().send(())?; + + if let Some(hand) = Arc::get_mut(&mut handle) { + hand.join_handle + .take() + .map(|h| h.join().unwrap()) + .transpose()?; + } + } + Ok(()) + } } impl<'a> Indicator<'a> { @@ -93,7 +127,8 @@ impl<'a> Indicator<'a> { /// 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`]. pub fn measure() -> IndicatorHandle { - let (tx, rx) = mpsc::channel(); + let (tx, rx) = mpsc::sync_channel(1024); + let (ptx, prx) = mpsc::sync_channel(1); let join_handle = thread::spawn(move || { let mut indicator = Self::default(); @@ -102,27 +137,29 @@ impl<'a> Indicator<'a> { indicator.stdout.execute(cursor::Hide)?; while let Ok(msg) = rx.recv() { - if indicator.state == IndicatorState::Indexing { - match msg { - Message::Index => indicator.index()?, - Message::DoneIndexing => { - indicator.update_state(IndicatorState::Rendering)?; - }, - Message::RenderReady => {}, - } - } - - if indicator.state == IndicatorState::Rendering && msg == Message::RenderReady { + if prx.recv_timeout(PRIORITY_MAIL_TIMEOUT).is_ok() { indicator.update_state(IndicatorState::Done)?; break; } + + match msg { + Message::Index => indicator.index()?, + Message::DoneIndexing => { + indicator.update_state(IndicatorState::Rendering)?; + }, + Message::RenderReady => { + indicator.update_state(IndicatorState::Done)?; + return Ok(()); + }, + } + indicator.stdout.execute(cursor::RestorePosition)?; } Ok(()) }); - IndicatorHandle::new(join_handle, tx) + IndicatorHandle::new(Some(join_handle), tx, ptx) } /// Updates the `state` of the [`Indicator`] to `new_state`, immediately running an associated @@ -139,7 +176,7 @@ impl<'a> Indicator<'a> { self.rendering(); }, - (Rendering, Done) => { + (Rendering | Indexing, Done) => { let stdout = &mut self.stdout; stdout.execute(terminal::Clear(ClearType::CurrentLine))?; stdout.execute(cursor::RestorePosition)?; diff --git a/src/tree/error.rs b/src/tree/error.rs index f95e40d..570ee79 100644 --- a/src/tree/error.rs +++ b/src/tree/error.rs @@ -33,8 +33,11 @@ pub enum Error { #[cfg(unix)] #[error("{0}")] - Persmissions(#[from] PermissionsError), + Permissions(#[from] PermissionsError), #[error("{0}")] UninitializedTheme(#[from] StyleError<'static>), + + #[error("Terminated erdtree...")] + Terminated, } diff --git a/src/tree/mod.rs b/src/tree/mod.rs index e214569..fff0815 100644 --- a/src/tree/mod.rs +++ b/src/tree/mod.rs @@ -2,7 +2,7 @@ use crate::{ context::{column, Context}, disk_usage::file_size::FileSize, fs::inode::Inode, - progress::{self, IndicatorHandle, Message}, + progress::{IndicatorHandle, Message}, utils, }; use count::FileCount; @@ -16,7 +16,10 @@ use std::{ fs, path::PathBuf, result::Result as StdResult, - sync::mpsc::{self, Sender}, + sync::{ + mpsc::{self, Sender}, + Arc, + }, thread, }; use visitor::{BranchVisitorBuilder, TraversalState}; @@ -54,7 +57,7 @@ impl Tree { /// various properties necessary to render output. pub fn try_init( mut ctx: Context, - indicator: Option<&IndicatorHandle>, + indicator: Option<Arc<IndicatorHandle>>, ) -> Result<(Self, Context)> { let mut column_properties = column::Properties::from(&ctx); @@ -102,12 +105,12 @@ impl Tree { fn traverse( ctx: &Context, column_properties: &mut column::Properties, - indicator: Option<&IndicatorHandle>, + indicator: Option<Arc<IndicatorHandle>>, ) -> Result<(Arena<Node>, NodeId)> { let walker = WalkParallel::try_from(ctx)?; let (tx, rx) = mpsc::channel(); - let progress_indicator_mailbox = indicator.map(progress::IndicatorHandle::mailbox); + let progress_indicator_mailbox = indicator.map(|arc| arc.mailbox()); thread::scope(|s| { let res = s.spawn(move || { @@ -117,7 +120,9 @@ impl Tree { while let Ok(TraversalState::Ongoing(node)) = rx.recv() { if let Some(ref mailbox) = progress_indicator_mailbox { - let _ = mailbox.send(Message::Index); + if mailbox.send(Message::Index).is_err() { + return Err(Error::Terminated); + } } if node.is_dir() { @@ -147,7 +152,9 @@ impl Tree { } if let Some(ref mailbox) = progress_indicator_mailbox { - let _ = mailbox.send(Message::DoneIndexing); + if mailbox.send(Message::DoneIndexing).is_err() { + return Err(Error::Terminated); + } } let root_id = root_id.ok_or(Error::MissingRoot)?; @@ -179,7 +186,7 @@ impl Tree { walker.visit(&mut visitor_builder); - tx.send(TraversalState::Done).unwrap(); + let _ = tx.send(TraversalState::Done); res.join().unwrap() }) diff --git a/src/tree/visitor.rs b/src/tree/visitor.rs index f0ff955..add11a9 100644 --- a/src/tree/visitor.rs +++ b/src/tree/visitor.rs @@ -44,7 +44,9 @@ impl ParallelVisitor for Branch<'_> { match Node::try_from((dir_entry, self.ctx)) { Ok(node) => { - self.tx.send(TraversalState::from(node)).unwrap(); + if self.tx.send(TraversalState::from(node)).is_err() { + return WalkState::Quit; + } WalkState::Continue }, _ => WalkState::Skip, |