tracing_apender: Fix race condition for logging on shutdown

tracing-appender/src/non_blocking.rs

@@ -105,6 +105,7 @@ pub const DEFAULT_BUFFERED_LINES_LIMIT: usize = 128_000;
 pub struct WorkerGuard {
     guard: Option<JoinHandle<()>>,
     sender: Sender<Msg>,
+    shutdown: Sender<()>,
 }
 
 /// A non-blocking writer.
@@ -147,8 +148,11 @@ impl NonBlocking {
     ) -> (NonBlocking, WorkerGuard) {
         let (sender, receiver) = bounded(buffered_lines_limit);
 
-        let worker = Worker::new(receiver, writer);
-        let worker_guard = WorkerGuard::new(worker.worker_thread(), sender.clone());
+        let (shutdown_sender, shutdown_receiver) = bounded(0);
+
+        let worker = Worker::new(receiver, writer, shutdown_receiver);
+        let worker_guard =
+            WorkerGuard::new(worker.worker_thread(), sender.clone(), shutdown_sender);
 
         (
             Self {
@@ -244,10 +248,11 @@ impl<'a> MakeWriter<'a> for NonBlocking {
 }
 
 impl WorkerGuard {
-    fn new(handle: JoinHandle<()>, sender: Sender<Msg>) -> Self {
+    fn new(handle: JoinHandle<()>, sender: Sender<Msg>, shutdown: Sender<()>) -> Self {
         WorkerGuard {
             guard: Some(handle),
             sender,
+            shutdown,
         }
     }
 }
@@ -258,7 +263,14 @@ impl Drop for WorkerGuard {
             .sender
             .send_timeout(Msg::Shutdown, Duration::from_millis(100))
         {
-            Ok(_) | Err(SendTimeoutError::Disconnected(_)) => (),
+            Ok(_) => {
+                // Attempt to wait for `Worker` to flush all messages before dropping. This happens
+                // when the `Worker` calls `recv()` on a zero-capacity channel. Use `send_timeout`
+                // so that drop is not blocked indefinitely.
+                // TODO: Make timeout configurable.
+                let _ = self.shutdown.send_timeout((), Duration::from_millis(1000));
+            }
+            Err(SendTimeoutError::Disconnected(_)) => (),
             Err(SendTimeoutError::Timeout(e)) => println!(
                 "Failed to send shutdown signal to logging worker. Error: {:?}",
                 e

tracing-appender/src/worker.rs

@@ -7,6 +7,7 @@ use std::{io, thread};
 pub(crate) struct Worker<T: Write + Send + Sync + 'static> {
     writer: T,
     receiver: Receiver<Msg>,
+    shutdown: Receiver<()>,
 }
 
 #[derive(Debug, Clone, Copy, Eq, PartialEq)]
@@ -18,8 +19,12 @@ pub(crate) enum WorkerState {
 }
 
 impl<T: Write + Send + Sync + 'static> Worker<T> {
-    pub(crate) fn new(receiver: Receiver<Msg>, writer: T) -> Worker<T> {
-        Self { writer, receiver }
+    pub(crate) fn new(receiver: Receiver<Msg>, writer: T, shutdown: Receiver<()>) -> Worker<T> {
+        Self {
+            writer,
+            receiver,
+            shutdown,
+        }
     }
 
     fn handle_recv(&mut self, result: &Result<Msg, RecvError>) -> io::Result<WorkerState> {
@@ -67,7 +72,10 @@ impl<T: Write + Send + Sync + 'static> Worker<T> {
             loop {
                 match self.work() {
                     Ok(WorkerState::Continue) | Ok(WorkerState::Empty) => {}
-                    Ok(WorkerState::Shutdown) | Ok(WorkerState::Disconnected) => break,
+                    Ok(WorkerState::Shutdown) | Ok(WorkerState::Disconnected) => {
+                        let _ = self.shutdown.recv();
+                        break;
+                    }
                     Err(_) => {
                         // TODO: Expose a metric for IO Errors, or print to stderr
                     }