Use BufWriter when STDOUT is not a TTY

Cargo.toml

@@ -28,6 +28,7 @@ nameattr = ["indexmap"]
 
 [dependencies]
 ahash = "0.6"
+atty = "0.2"
 crossbeam-utils = { version = "0.8", optional = true }
 crossbeam-channel = { version = "0.5", optional = true }
 dashmap = { version = "3", optional = true }

src/bin/collapse-dtrace.rs

@@ -83,5 +83,5 @@ fn main() -> io::Result<()> {
     }
 
     let (infile, options) = opt.into_parts();
-    Folder::from(options).collapse_file(infile.as_ref(), io::stdout().lock())
+    Folder::from(options).collapse_file_to_stdout(infile.as_ref())
 }

src/bin/collapse-guess.rs

@@ -75,5 +75,5 @@ fn main() -> io::Result<()> {
     }
 
     let (infile, options) = opt.into_parts();
-    Folder::from(options).collapse_file(infile.as_ref(), io::stdout().lock())
+    Folder::from(options).collapse_file_to_stdout(infile.as_ref())
 }

src/bin/collapse-perf.rs

@@ -112,5 +112,5 @@ fn main() -> io::Result<()> {
     }
 
     let (infile, options) = opt.into_parts();
-    Folder::from(options).collapse_file(infile.as_ref(), io::stdout().lock())
+    Folder::from(options).collapse_file_to_stdout(infile.as_ref())
 }

src/bin/collapse-sample.rs

@@ -62,5 +62,5 @@ fn main() -> io::Result<()> {
     }
 
     let (infile, options) = opt.into_parts();
-    Folder::from(options).collapse_file(infile.as_ref(), io::stdout().lock())
+    Folder::from(options).collapse_file_to_stdout(infile.as_ref())
 }

src/bin/collapse-vtune.rs

@@ -64,5 +64,5 @@ fn main() -> io::Result<()> {
     }
 
     let (infile, options) = opt.into_parts();
-    Folder::from(options).collapse_file(infile.as_ref(), io::stdout().lock())
+    Folder::from(options).collapse_file_to_stdout(infile.as_ref())
 }

src/bin/diff-folded.rs

@@ -87,5 +87,15 @@ fn main() -> io::Result<()> {
     }
 
     let (folded1, folded2, options) = opt.into_parts();
-    differential::from_files(options, folded1, folded2, io::stdout().lock())
+
+    if atty::is(atty::Stream::Stdout) {
+        differential::from_files(options, folded1, folded2, io::stdout().lock())
+    } else {
+        differential::from_files(
+            options,
+            folded1,
+            folded2,
+            io::BufWriter::new(io::stdout().lock()),
+        )
+    }
 }

src/bin/flamegraph.rs

@@ -297,9 +297,19 @@ fn main() -> quick_xml::Result<()> {
     };
 
     let (infiles, mut options) = opt.into_parts();
+
     options.palette_map = palette_map.as_mut();
 
-    flamegraph::from_files(&mut options, &infiles, io::stdout().lock())?;
+    if atty::is(atty::Stream::Stdout) {
+        flamegraph::from_files(&mut options, &infiles, io::stdout().lock())?;
+    } else {
+        flamegraph::from_files(
+            &mut options,
+            &infiles,
+            io::BufWriter::new(io::stdout().lock()),
+        )?;
+    }
+
     save_consistent_palette_if_needed(&palette_map, PALETTE_MAP_FILE).map_err(quick_xml::Error::Io)
 }
 

src/collapse/common.rs

@@ -465,6 +465,7 @@ impl Occurrences {
                 }
             }
         }
+        writer.flush()?;
         Ok(())
     }
 }
@@ -510,7 +511,7 @@ pub(crate) fn fix_partially_demangled_rust_symbol(symbol: &str) -> Cow<str> {
                 demangled.push_str("::");
                 rest = &rest[2..];
             } else {
-                demangled.push_str(".");
+                demangled.push('.');
                 rest = &rest[1..];
             }
         } else if rest.starts_with('$') {

src/collapse/guess.rs

@@ -54,13 +54,17 @@ impl Collapse for Folder {
         W: io::Write,
     {
         let mut dtrace = {
-            let mut options = dtrace::Options::default();
-            options.nthreads = self.opt.nthreads;
+            let options = dtrace::Options {
+                nthreads: self.opt.nthreads,
+                ..Default::default()
+            };
             dtrace::Folder::from(options)
         };
         let mut perf = {
-            let mut options = perf::Options::default();
-            options.nthreads = self.opt.nthreads;
+            let options = perf::Options {
+                nthreads: self.opt.nthreads,
+                ..Default::default()
+            };
             perf::Folder::from(options)
         };
         let mut sample = sample::Folder::default();

src/collapse/matcher.rs

@@ -9,13 +9,9 @@
 #[inline]
 pub fn is_vmlinux(s: &str) -> bool {
     if let Some(vm) = s.rfind("vmlinux") {
-        s[vm..].chars().all(|c| {
-            c.is_ascii_alphanumeric()
-                || match c {
-                    '-' | '.' | '_' => true,
-                    _ => false,
-                }
-        })
+        s[vm..]
+            .chars()
+            .all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '.' | '_'))
     } else {
         false
     }

src/collapse/mod.rs

@@ -79,14 +79,27 @@ pub trait Collapse {
                 self.collapse(reader, writer)
             }
             None => {
-                let stdio = io::stdin();
-                let stdio_guard = stdio.lock();
-                let reader = io::BufReader::with_capacity(CAPACITY_READER, stdio_guard);
+                let stdin = io::stdin();
+                let stdin_guard = stdin.lock();
+                let reader = io::BufReader::with_capacity(CAPACITY_READER, stdin_guard);
                 self.collapse(reader, writer)
             }
         }
     }
 
+    /// Collapses the contents of the provided file (or of STDIN if `infile` is `None`) and
+    /// writes folded stack lines to STDOUT.
+    fn collapse_file_to_stdout<P>(&mut self, infile: Option<P>) -> io::Result<()>
+    where
+        P: AsRef<Path>,
+    {
+        if atty::is(atty::Stream::Stdout) {
+            self.collapse_file(infile, io::stdout().lock())
+        } else {
+            self.collapse_file(infile, io::BufWriter::new(io::stdout().lock()))
+        }
+    }
+
     /// Returns whether this implementation is appropriate for the given input.
     ///
     /// - `None` means "not sure -- need more input"

src/collapse/perf.rs

@@ -409,12 +409,12 @@ impl Folder {
             // XXX: re-use existing memory in pname if possible
             self.pname = comm.replace(' ', "_");
             if self.opt.include_tid {
-                self.pname.push_str("-");
+                self.pname.push('-');
                 self.pname.push_str(pid);
-                self.pname.push_str("/");
+                self.pname.push('/');
                 self.pname.push_str(tid);
             } else if self.opt.include_pid {
-                self.pname.push_str("-");
+                self.pname.push('-');
                 self.pname.push_str(pid);
             }
 
@@ -548,7 +548,7 @@ impl Folder {
             stack_str.push_str(&self.pname);
             // add the other stack entries (if any)
             for e in self.stack.drain(..) {
-                stack_str.push_str(";");
+                stack_str.push(';');
                 stack_str.push_str(&e);
             }
 
@@ -587,15 +587,15 @@ fn with_module_fallback(module: &str, func: &str, pc: &str, include_addrs: bool)
     let mut res = String::with_capacity(func.len() + 12);
 
     if include_addrs {
-        res.push_str("[");
+        res.push('[');
         res.push_str(func);
         res.push_str(" <");
         res.push_str(pc);
         res.push_str(">]");
     } else {
-        res.push_str("[");
+        res.push('[');
         res.push_str(func);
-        res.push_str("]");
+        res.push(']');
     }
 
     res

src/flamegraph/color/mod.rs

@@ -347,16 +347,16 @@ pub(super) fn color(
         let mut hash: u64 = 0xcbf29ce484222325;
         // https://github.com/servo/rust-fnv/blob/4b4784ebfd3332dc61f0640764d6f1140e03a9ab/lib.rs#L118-L121
         for byte in name.as_bytes() {
-            hash = hash ^ (*byte as u64);
+            hash ^= *byte as u64;
             hash = hash.wrapping_mul(0x100000001b3);
         }
         let hash1 = (hash as f64 / std::u64::MAX as f64) as f32;
 
         // Rotate hash so we get two more distinct numbers
-        hash = hash ^ 0;
+        hash ^= 0;
         hash = hash.wrapping_mul(0x100000001b3);
         let hash2 = (hash as f64 / std::u64::MAX as f64) as f32;
-        hash = hash ^ 0;
+        hash ^= 0;
         hash = hash.wrapping_mul(0x100000001b3);
         let hash3 = (hash as f64 / std::u64::MAX as f64) as f32;
 

src/flamegraph/color/palettes.rs

@@ -31,8 +31,8 @@ pub(super) mod java {
         if name.contains("::") || name.starts_with("-[") || name.starts_with("+[") {
             // C++ or Objective C
             BasicPalette::Yellow
-        } else if java_prefix.contains("/")
-            || (java_prefix.contains(".") && !java_prefix.starts_with("["))
+        } else if java_prefix.contains('/')
+            || (java_prefix.contains('.') && !java_prefix.starts_with('['))
             || match java_prefix.chars().next() {
                 Some(c) => c.is_ascii_uppercase(),
                 _ => false,

src/flamegraph/mod.rs

@@ -693,6 +693,7 @@ where
     svg.write_event(Event::End(BytesEnd::borrowed(b"svg")))?;
     svg.write_event(Event::Eof)?;
 
+    svg.into_inner().flush()?;
     Ok(())
 }