Flush sinks on program exit with a special method flush_on_exit (#102)

Co-authored-by: Asuna <SpriteOvO@gmail.com>

Author: tisonkun

spdlog/src/lib.rs

@@ -747,18 +747,15 @@ pub fn log_crate_proxy() -> &'static LogCrateProxy {
     &PROXY
 }
 
-static IS_TEARING_DOWN: AtomicBool = AtomicBool::new(false);
-
 fn flush_default_logger_at_exit() {
     // Rust never calls `drop` for static variables.
     //
     // Setting up an exit handler gives us a chance to flush the default logger
     // once at the program exit, thus we don't lose the last logs.
 
     extern "C" fn handler() {
-        IS_TEARING_DOWN.store(true, Ordering::SeqCst);
         if let Some(default_logger) = DEFAULT_LOGGER.get() {
-            default_logger.load().flush()
+            default_logger.load().flush_sinks_on_exit()
         }
     }
 

spdlog/src/logger.rs

@@ -497,14 +497,22 @@ impl Logger {
         }
     }
 
-    fn flush_sinks(&self) {
+    fn flush_sinks_with(&self, with: impl Fn(&dyn Sink) -> Result<()>) {
         self.sinks.iter().for_each(|sink| {
-            if let Err(err) = sink.flush() {
+            if let Err(err) = with(&**sink) {
                 self.handle_error(err);
             }
         });
     }
 
+    pub(crate) fn flush_sinks_on_exit(&self) {
+        self.flush_sinks_with(|sink| sink.flush_on_exit());
+    }
+
+    pub(crate) fn flush_sinks(&self) {
+        self.flush_sinks_with(|sink| sink.flush());
+    }
+
     fn handle_error(&self, err: Error) {
         self.error_handler.read_expect().call_internal(
             format!(

spdlog/src/sink/async_sink/async_pool_sink.rs

@@ -129,21 +129,21 @@ impl Sink for AsyncPoolSink {
     }
 
     fn flush(&self) -> Result<()> {
-        if crate::IS_TEARING_DOWN.load(Ordering::SeqCst) {
-            // https://github.com/SpriteOvO/spdlog-rs/issues/64
-            //
-            // If the program is tearing down, this will be the final flush. `crossbeam`
-            // uses thread-local internally, which is not supported in `atexit` callback.
-            // This can be bypassed by flushing sinks directly on the current thread, but
-            // before we do that we have to destroy the thread pool to ensure that any
-            // pending log tasks are completed.
-            self.thread_pool.destroy();
-            self.backend.flush()
-        } else {
-            self.assign_task(Task::Flush {
-                backend: self.clone_backend(),
-            })
-        }
+        self.assign_task(Task::Flush {
+            backend: self.clone_backend(),
+        })
+    }
+
+    fn flush_on_exit(&self) -> Result<()> {
+        // https://github.com/SpriteOvO/spdlog-rs/issues/64
+        //
+        // If the program is tearing down, this will be the final flush. `crossbeam`
+        // uses thread-local internally, which is not supported in `atexit` callback.
+        // This can be bypassed by flushing sinks directly on the current thread, but
+        // before we do that we have to destroy the thread pool to ensure that any
+        // pending log tasks are completed.
+        self.thread_pool.destroy();
+        self.backend.flush_on_exit()
     }
 }
 
@@ -258,14 +258,22 @@ impl Backend {
         result
     }
 
-    fn flush(&self) -> Result<()> {
+    fn flush_with(&self, with: impl Fn(&dyn Sink) -> Result<()>) -> Result<()> {
         let mut result = Ok(());
         for sink in &self.sinks {
-            result = Error::push_result(result, sink.flush());
+            result = Error::push_result(result, with(&**sink));
         }
         result
     }
 
+    fn flush(&self) -> Result<()> {
+        self.flush_with(|sink| sink.flush())
+    }
+
+    fn flush_on_exit(&self) -> Result<()> {
+        self.flush_with(|sink| sink.flush_on_exit())
+    }
+
     fn handle_error(&self, err: Error) {
         self.prop.call_error_handler_internal("AsyncPoolSink", err)
     }

spdlog/src/sink/dedup_sink.rs

@@ -152,12 +152,20 @@ impl DedupSink {
         })
     }
 
-    fn flush_sinks(&self) -> Result<()> {
+    fn flush_with(&self, with: fn(&dyn Sink) -> Result<()>) -> Result<()> {
         #[allow(clippy::manual_try_fold)] // https://github.com/rust-lang/rust-clippy/issues/11554
         self.sinks.iter().fold(Ok(()), |result, sink| {
-            Error::push_result(result, sink.flush())
+            Error::push_result(result, with(sink.as_ref()))
         })
     }
+
+    fn flush_sinks(&self) -> Result<()> {
+        self.flush_with(|sink| sink.flush())
+    }
+
+    fn flush_sinks_on_exit(&self) -> Result<()> {
+        self.flush_with(|sink| sink.flush_on_exit())
+    }
 }
 
 impl GetSinkProp for DedupSink {
@@ -186,6 +194,10 @@ impl Sink for DedupSink {
     fn flush(&self) -> Result<()> {
         self.flush_sinks()
     }
+
+    fn flush_on_exit(&self) -> Result<()> {
+        self.flush_sinks_on_exit()
+    }
 }
 
 impl Drop for DedupSink {

spdlog/src/sink/mod.rs

@@ -238,6 +238,32 @@ pub trait Sink: SinkPropAccess + Sync + Send {
 
     /// Flushes any buffered records.
     fn flush(&self) -> Result<()>;
+
+    /// Flushes any buffered records at program exit.
+    ///
+    /// _spdlog-rs_ will perform a flush for sinks in the default logger when
+    /// the program exits, and the flush will be called to this method
+    /// `flush_on_exit` instead of `flush`. This is because the execution
+    /// context may be in the [`atexit`] callback or in the panic handler when
+    /// exiting. In such a context, some operations are restricted, e.g.
+    /// Thread-local Storage (TLS) may not be available in `atexit` callbacks.
+    ///
+    /// This method calls directly to `flush` method by default. When users'
+    /// `flush` method implementation is not usable in a program exit context,
+    /// users should override the implementation of this method to provide an
+    /// alternative flushing implementation. See the implementation of
+    /// [`AsyncPoolSink::flush_on_exit`] as an example.
+    ///
+    /// For [combined sink]s, this method should always be overridden to
+    /// propagate the information that "the program is exiting" to their
+    /// sub-sinks. See the implementation of [`DedupSink::flush_on_exit`] as an
+    /// example.
+    ///
+    /// [`atexit`]: https://en.cppreference.com/w/c/program/atexit
+    /// [combined sink]: index.html#combined-sink
+    fn flush_on_exit(&self) -> Result<()> {
+        self.flush()
+    }
 }
 
 /// Container type for [`Sink`]s.