diff --git a/src/commands/mod.rs b/src/commands/mod.rs index 9305e46..05e38e1 100644 --- a/src/commands/mod.rs +++ b/src/commands/mod.rs @@ -57,19 +57,17 @@ pub fn run( question_policy: QuestionPolicy, file_visibility_policy: FileVisibilityPolicy, ) -> crate::Result<()> { - let log_receiver = setup_channel(); + let (log_receiver, dropper) = setup_channel(); let synchronization_pair = Arc::new((Mutex::new(false), Condvar::new())); spawn_logger_thread(log_receiver, synchronization_pair.clone()); run_cmd(args, question_policy, file_visibility_policy)?; - // Drop our sender so when all threads are done, no clones are left. - // This is needed, otherwise the logging thread will never exit since we would be keeping a - // sender alive here. - todo!(); + // Send a message for the logger thread to shut down. + // This is needed, otherwise the logging thread will never exit. + drop(dropper); - // Prevent the main thread from exiting until the background thread handling the - // logging has set `flushed` to true. + // Hold the main thread from exiting until the background thread signals its shutdown. let (lock, cvar) = &*synchronization_pair; let guard = lock.lock().unwrap(); let _flushed = cvar.wait(guard).unwrap(); @@ -257,38 +255,36 @@ fn run_cmd( fn spawn_logger_thread(log_receiver: LogReceiver, synchronization_pair: Arc<(Mutex, Condvar)>) { rayon::spawn(move || { - const BUFFER_SIZE: usize = 10; - let mut buffer = Vec::::with_capacity(BUFFER_SIZE); + const BUFFER_CAPACITY: usize = 10; + let mut buffer = Vec::::with_capacity(BUFFER_CAPACITY); loop { - let msg = log_receiver.recv(); + let msg = log_receiver.recv().expect("Failed to receive log message"); - // Senders are still active - if let Ok(msg) = msg { - // Print messages if buffer is full otherwise append to it - if buffer.len() == BUFFER_SIZE { - let mut tmp = buffer.join("\n"); + let is_shutdown_message = msg.is_none(); - if let Some(msg) = map_message(&msg) { - tmp.push_str(&msg); - } + // Append message to buffer + if let Some(msg) = msg.as_ref().and_then(map_message) { + buffer.push(msg); + } - eprintln!("{}", tmp); - buffer.clear(); - } else if let Some(msg) = map_message(&msg) { - buffer.push(msg); - } - } else { - // All senders have been dropped - eprintln!("{}", buffer.join("\n")); + let should_flush = buffer.len() == BUFFER_CAPACITY || is_shutdown_message; - // Wake up the main thread - let (lock, cvar) = &*synchronization_pair; - let mut flushed = lock.lock().unwrap(); - *flushed = true; - cvar.notify_one(); + if should_flush { + let text = buffer.join("\n"); + eprintln!("{text}"); + buffer.clear(); + } + + if is_shutdown_message { break; } } + + // Wake up the main thread + let (lock, cvar) = &*synchronization_pair; + let mut flushed = lock.lock().unwrap(); + *flushed = true; + cvar.notify_one(); }); } diff --git a/src/utils/logger.rs b/src/utils/logger.rs index 83552d8..8f0ff8d 100644 --- a/src/utils/logger.rs +++ b/src/utils/logger.rs @@ -3,15 +3,15 @@ use std::sync::{mpsc, OnceLock}; use super::colors::{ORANGE, RESET, YELLOW}; use crate::accessible::is_running_in_accessible_mode; -pub type LogReceiver = mpsc::Receiver; -type LogSender = mpsc::Sender; +pub type LogReceiver = mpsc::Receiver>; +type LogSender = mpsc::Sender>; static SENDER: OnceLock = OnceLock::new(); -pub fn setup_channel() -> LogReceiver { +pub fn setup_channel() -> (LogReceiver, LoggerDropper) { let (tx, rx) = mpsc::channel(); SENDER.set(tx).expect("`setup_channel` should only be called once"); - rx + (rx, LoggerDropper) } #[track_caller] @@ -60,6 +60,7 @@ pub fn map_message(msg: &PrintMessage) -> Option { /// `is_running_in_accessible_mode`. /// /// Read more about accessibility mode in `accessible.rs`. +#[track_caller] pub fn info(contents: String) { info_with_accessibility(contents, false); } @@ -70,10 +71,12 @@ pub fn info(contents: String) { /// returns `true`. /// /// Read more about accessibility mode in `accessible.rs`. +#[track_caller] pub fn info_accessible(contents: String) { info_with_accessibility(contents, true); } +#[track_caller] fn info_with_accessibility(contents: String, accessible: bool) { send_log_message(PrintMessage { contents, @@ -97,6 +100,20 @@ pub enum MessageLevel { Warning, } +#[track_caller] fn send_log_message(msg: PrintMessage) { + send_message(Some(msg)); +} + +#[track_caller] +fn send_message(msg: Option) { get_sender().send(msg).expect("Failed to send internal message"); } + +pub struct LoggerDropper; + +impl Drop for LoggerDropper { + fn drop(&mut self) { + send_message(None); + } +}