zcache/zcachedb: always flush stdout buffer data on termination (open…

…zfs#233)

cmd/zfs_object_agent/util/Cargo.toml

@@ -19,6 +19,7 @@ enum-map = { version = "1.1.1", features = ["serde"] }
 lazy_static = "1.4.0"
 log = "0.4"
 log4rs = "1.0.0"
+libc = "0.2"
 more-asserts = "0.2.1"
 num-traits = "0.2.14"
 rand = "0.8.3"

cmd/zfs_object_agent/util/src/write_stdout.rs

@@ -11,18 +11,37 @@ to stdout/stderr returns an error (like `EPIPE`) rather than cause a panic.
 These replacement macros also introduce buffering when stdout is not pointing to
 tty, significantly improving the performance of utilities that redirect their
 output to files.
-
-DISCLAIMER: When the utility macros below are used to write to stdout it is
-important to call flush_stdout!() before the program exits successfully to flush
-any leftover data to non-tty endpoints (like pipes and regular files).
 !*/
 
 use lazy_static::lazy_static;
+use libc::atexit;
 use std::{io::BufWriter, sync::Mutex};
 
 lazy_static! {
-    pub static ref BUFFERED_STDOUT_HANDLE: Mutex<BufWriter<std::io::Stdout>> =
-        Mutex::new(BufWriter::new(std::io::stdout()));
+    pub static ref BUFFERED_STDOUT_HANDLE: Mutex<BufWriter<std::io::Stdout>> = {
+        // Ensure that any leftover data in the buffer are flushed before
+        // terminating the process.
+        unsafe { atexit(flush_stdout) };
+        Mutex::new(BufWriter::new(std::io::stdout()))
+    };
+}
+
+extern "C" fn flush_stdout() {
+    use std::io::Write;
+    match BUFFERED_STDOUT_HANDLE.try_lock() {
+        Ok(mut hdl) => {
+            if let Err(e) = hdl.flush() {
+                crate::writeln_stderr!("{}", e);
+                std::process::exit(0);
+            }
+        }
+        Err(_) => {
+            crate::writeln_stderr!(
+                "CANNOT FLUSH STDOUT BUFFER BECAUSE ITS LOCK IS HELD BY ANOTHER THREAD"
+            );
+            std::process::exit(0);
+        }
+    };
 }
 
 pub extern crate atty;
@@ -63,24 +82,6 @@ macro_rules! writeln_stdout {
 	}}
 }
 
-/// Conventionally used at the end of main() so leftover buffer data are flushed
-/// to stdout. The macro terminates the process on write errors (does not panic).
-#[macro_export]
-macro_rules! flush_stdout {
-    () => {{
-        use std::io::Write;
-        let res = if !$crate::write_stdout::atty::is($crate::write_stdout::atty::Stream::Stdout) {
-            let mut hdl = $crate::write_stdout::BUFFERED_STDOUT_HANDLE.lock().unwrap();
-            hdl.flush()
-        } else {
-            Ok(())
-        };
-        if res.is_err() {
-            std::process::exit(0)
-        }
-    }};
-}
-
 /// Similar to `eprint!` macro, except it terminates the process on write errors (does not panic).
 #[macro_export]
 macro_rules! write_stderr {

cmd/zfs_object_agent/zcache/Cargo.toml

@@ -13,7 +13,6 @@ chrono = "0.4"
 async-trait = "0.1.51"
 clap = "2"
 exitcode = "1.1.2"
-libc = "0.2"
 log = "0.4"
 more-asserts = "0.2.1"
 num-traits = "0.2.14"

cmd/zfs_object_agent/zcache/src/main.rs

@@ -22,7 +22,6 @@ use log::*;
 use report_hits::ReportHits;
 use stats::Stats;
 use subcommand::ZcacheSubCommand;
-use util::flush_stdout;
 use util::writeln_stderr;
 
 fn main() -> Result<()> {
@@ -96,6 +95,5 @@ async fn async_main() -> Result<()> {
             std::process::exit(exitcode::USAGE);
         }
     }
-    flush_stdout!();
     Ok(())
 }

cmd/zfs_object_agent/zcdb/src/main.rs

@@ -2,7 +2,6 @@ use clap::AppSettings;
 use clap::Arg;
 use clap::SubCommand;
 use git_version::git_version;
-use util::flush_stdout;
 use util::writeln_stdout;
 use zettacache::DumpSlabsOptions;
 use zettacache::DumpStructuresOptions;
@@ -95,7 +94,7 @@ async fn main() -> Result<(), anyhow::Error> {
         .get_matches();
 
     let cache_paths = matches.values_of("cache-device").unwrap().collect();
-    let res = match matches.subcommand() {
+    match matches.subcommand() {
         ("dump-structures", Some(subcommand_matches)) => {
             ZettaCacheDBCommand::issue_command(
                 ZettaCacheDBCommand::DumpStructures(
@@ -135,7 +134,5 @@ async fn main() -> Result<(), anyhow::Error> {
             writeln_stdout!("{}", matches.usage());
             std::process::exit(exitcode::USAGE);
         }
-    };
-    flush_stdout!();
-    res
+    }
 }