1
0
mirror of https://github.com/sharkdp/fd synced 2024-07-03 08:18:48 +00:00

Merge pull request #902 from tavianator/quit-senders

Quit senders more aggressively
This commit is contained in:
Tavian Barnes 2021-12-10 13:56:28 -05:00 committed by GitHub
commit 81669f4c10
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 44 additions and 34 deletions

View File

@ -1,7 +1,6 @@
use std::borrow::Cow;
use std::io::{self, Write};
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};
use lscolors::{Indicator, LsColors, Style};
@ -15,12 +14,7 @@ fn replace_path_separator(path: &str, new_path_separator: &str) -> String {
}
// TODO: this function is performance critical and can probably be optimized
pub fn print_entry<W: Write>(
stdout: &mut W,
entry: &Path,
config: &Config,
wants_to_quit: &AtomicBool,
) {
pub fn print_entry<W: Write>(stdout: &mut W, entry: &Path, config: &Config) {
let path = if config.strip_cwd_prefix {
strip_current_dir(entry)
} else {
@ -28,7 +22,7 @@ pub fn print_entry<W: Write>(
};
let r = if let Some(ref ls_colors) = config.ls_colors {
print_entry_colorized(stdout, path, config, ls_colors, wants_to_quit)
print_entry_colorized(stdout, path, config, ls_colors)
} else {
print_entry_uncolorized(stdout, path, config)
};
@ -50,7 +44,6 @@ fn print_entry_colorized<W: Write>(
path: &Path,
config: &Config,
ls_colors: &LsColors,
wants_to_quit: &AtomicBool,
) -> io::Result<()> {
// Split the path between the parent and the last component
let mut offset = 0;
@ -92,12 +85,6 @@ fn print_entry_colorized<W: Write>(
writeln!(stdout)?;
}
if wants_to_quit.load(Ordering::Relaxed) {
// Ignore any errors on flush, because we're about to exit anyway
let _ = stdout.flush();
ExitCode::KilledBySigint.exit();
}
Ok(())
}

View File

@ -134,11 +134,19 @@ pub fn scan(path_vec: &[PathBuf], pattern: Arc<Regex>, config: Arc<Config>) -> R
let parallel_walker = walker.threads(config.threads).build_parallel();
let wants_to_quit = Arc::new(AtomicBool::new(false));
// Flag for cleanly shutting down the parallel walk
let quit_flag = Arc::new(AtomicBool::new(false));
// Flag specifically for quitting due to ^C
let interrupt_flag = Arc::new(AtomicBool::new(false));
if config.ls_colors.is_some() && config.command.is_none() {
let wq = Arc::clone(&wants_to_quit);
let quit_flag = Arc::clone(&quit_flag);
let interrupt_flag = Arc::clone(&interrupt_flag);
ctrlc::set_handler(move || {
if wq.fetch_or(true, Ordering::Relaxed) {
quit_flag.store(true, Ordering::Relaxed);
if interrupt_flag.fetch_or(true, Ordering::Relaxed) {
// Ctrl-C has been pressed twice, exit NOW
ExitCode::KilledBySigint.exit();
}
@ -147,15 +155,15 @@ pub fn scan(path_vec: &[PathBuf], pattern: Arc<Regex>, config: Arc<Config>) -> R
}
// Spawn the thread that receives all results through the channel.
let receiver_thread = spawn_receiver(&config, &wants_to_quit, rx);
let receiver_thread = spawn_receiver(&config, &quit_flag, &interrupt_flag, rx);
// Spawn the sender threads.
spawn_senders(&config, &wants_to_quit, pattern, parallel_walker, tx);
spawn_senders(&config, &quit_flag, pattern, parallel_walker, tx);
// Wait for the receiver thread to print out all results.
let exit_code = receiver_thread.join().unwrap();
if wants_to_quit.load(Ordering::Relaxed) {
if interrupt_flag.load(Ordering::Relaxed) {
Ok(ExitCode::KilledBySigint)
} else {
Ok(exit_code)
@ -166,8 +174,10 @@ pub fn scan(path_vec: &[PathBuf], pattern: Arc<Regex>, config: Arc<Config>) -> R
struct ReceiverBuffer<W> {
/// The configuration.
config: Arc<Config>,
/// For shutting down the senders.
quit_flag: Arc<AtomicBool>,
/// The ^C notifier.
wants_to_quit: Arc<AtomicBool>,
interrupt_flag: Arc<AtomicBool>,
/// Receiver for worker results.
rx: Receiver<WorkerResult>,
/// Standard output.
@ -186,7 +196,8 @@ impl<W: Write> ReceiverBuffer<W> {
/// Create a new receiver buffer.
fn new(
config: Arc<Config>,
wants_to_quit: Arc<AtomicBool>,
quit_flag: Arc<AtomicBool>,
interrupt_flag: Arc<AtomicBool>,
rx: Receiver<WorkerResult>,
stdout: W,
) -> Self {
@ -195,7 +206,8 @@ impl<W: Write> ReceiverBuffer<W> {
Self {
config,
wants_to_quit,
quit_flag,
interrupt_flag,
rx,
stdout,
mode: ReceiverMode::Buffering,
@ -209,6 +221,7 @@ impl<W: Write> ReceiverBuffer<W> {
fn process(&mut self) -> ExitCode {
loop {
if let Err(ec) = self.poll() {
self.quit_flag.store(true, Ordering::Relaxed);
return ec;
}
}
@ -244,7 +257,7 @@ impl<W: Write> ReceiverBuffer<W> {
}
}
ReceiverMode::Streaming => {
self.print(&path);
self.print(&path)?;
self.flush()?;
}
}
@ -273,8 +286,16 @@ impl<W: Write> ReceiverBuffer<W> {
}
/// Output a path.
fn print(&mut self, path: &Path) {
output::print_entry(&mut self.stdout, path, &self.config, &self.wants_to_quit)
fn print(&mut self, path: &Path) -> Result<(), ExitCode> {
output::print_entry(&mut self.stdout, path, &self.config);
if self.interrupt_flag.load(Ordering::Relaxed) {
// Ignore any errors on flush, because we're about to exit anyway
let _ = self.flush();
return Err(ExitCode::KilledBySigint);
}
Ok(())
}
/// Switch ourselves into streaming mode.
@ -283,7 +304,7 @@ impl<W: Write> ReceiverBuffer<W> {
let buffer = mem::take(&mut self.buffer);
for path in buffer {
self.print(&path);
self.print(&path)?;
}
self.flush()
@ -315,11 +336,13 @@ impl<W: Write> ReceiverBuffer<W> {
fn spawn_receiver(
config: &Arc<Config>,
wants_to_quit: &Arc<AtomicBool>,
quit_flag: &Arc<AtomicBool>,
interrupt_flag: &Arc<AtomicBool>,
rx: Receiver<WorkerResult>,
) -> thread::JoinHandle<ExitCode> {
let config = Arc::clone(config);
let wants_to_quit = Arc::clone(wants_to_quit);
let quit_flag = Arc::clone(quit_flag);
let interrupt_flag = Arc::clone(interrupt_flag);
let show_filesystem_errors = config.show_filesystem_errors;
let threads = config.threads;
@ -375,7 +398,7 @@ fn spawn_receiver(
let stdout = stdout.lock();
let stdout = io::BufWriter::new(stdout);
let mut rxbuffer = ReceiverBuffer::new(config, wants_to_quit, rx, stdout);
let mut rxbuffer = ReceiverBuffer::new(config, quit_flag, interrupt_flag, rx, stdout);
rxbuffer.process()
}
})
@ -439,7 +462,7 @@ impl DirEntry {
fn spawn_senders(
config: &Arc<Config>,
wants_to_quit: &Arc<AtomicBool>,
quit_flag: &Arc<AtomicBool>,
pattern: Arc<Regex>,
parallel_walker: ignore::WalkParallel,
tx: Sender<WorkerResult>,
@ -448,10 +471,10 @@ fn spawn_senders(
let config = Arc::clone(config);
let pattern = Arc::clone(&pattern);
let tx_thread = tx.clone();
let wants_to_quit = Arc::clone(wants_to_quit);
let quit_flag = Arc::clone(quit_flag);
Box::new(move |entry_o| {
if wants_to_quit.load(Ordering::Relaxed) {
if quit_flag.load(Ordering::Relaxed) {
return ignore::WalkState::Quit;
}