fix: tracing-flame

Adds `tracing-flame` to generate profiles for flamegraphs.

Signed-off-by: Jonathan Woollett-Light <jcawl@amazon.co.uk>

Author: Jonathan Woollett-Light

src/api_server/src/request/logger.rs

@@ -42,6 +42,7 @@ mod tests {
             show_level: Some(false),
             show_log_origin: Some(false),
             new_format: None,
+            profile_file: None,
         };
         match vmm_action_from_request(parse_put_logger(&Body::new(body)).unwrap()) {
             VmmAction::ConfigureLogger(cfg) => assert_eq!(cfg, expected_cfg),
@@ -61,6 +62,7 @@ mod tests {
             show_level: Some(false),
             show_log_origin: Some(false),
             new_format: None,
+            profile_file: None,
         };
         match vmm_action_from_request(parse_put_logger(&Body::new(body)).unwrap()) {
             VmmAction::ConfigureLogger(cfg) => assert_eq!(cfg, expected_cfg),

src/firecracker/src/main.rs

@@ -203,6 +203,11 @@ fn main_exitable() -> FcExitCode {
                 .takes_value(false)
                 .help("Whether to use the new logging output format."),
         )
+        .arg(
+            Argument::new("profile-file")
+                .takes_value(true)
+                .help("Path to a fifo or a file used for configuring the profiler on startup."),
+        )
         .arg(
             Argument::new("metrics-path")
                 .takes_value(true)
@@ -302,6 +307,7 @@ fn main_exitable() -> FcExitCode {
             show_level: Some(arguments.flag_present("show-level")),
             show_log_origin: Some(arguments.flag_present("show-log-origin")),
             new_format: Some(arguments.flag_present("new-format")),
+            profile_file: arguments.single_value("profile-file").map(PathBuf::from),
         };
 
         if let Err(err) = logger_config.init() {

src/vmm/Cargo.toml

@@ -31,6 +31,7 @@ vm-fdt = "0.2.0"
 vm-superio = "0.7.0"
 tracing = { version = "0.1.37", features = ["attributes"] }
 tracing-subscriber = "0.3.17"
+tracing-flame = "0.2.0"
 tracing-core = "0.1.31"
 log = { version = "0.4.17", features = ["serde"] }
 

src/vmm/src/rpc_interface.rs

@@ -2078,6 +2078,7 @@ mod tests {
                 show_level: Some(false),
                 show_log_origin: Some(false),
                 new_format: Some(false),
+                profile_file: None,
             }),
             VmmActionError::OperationNotSupportedPostBoot,
         );

src/vmm/src/vmm_config/mod.rs

@@ -4,7 +4,7 @@
 use core::fmt;
 use std::convert::{From, TryInto};
 use std::fs::{File, OpenOptions};
-use std::io::{self, LineWriter, Write};
+use std::io::{self, BufWriter, LineWriter, Write};
 use std::os::unix::fs::OpenOptionsExt;
 use std::path::Path;
 use std::str::FromStr;
@@ -14,6 +14,7 @@ use libc::O_NONBLOCK;
 use rate_limiter::{BucketUpdate, RateLimiter, TokenBucket};
 use serde::{Deserialize, Serialize};
 use tracing_core::{Event, Subscriber};
+use tracing_flame::FlameLayer;
 use tracing_subscriber::fmt::format::{self, FormatEvent, FormatFields};
 use tracing_subscriber::fmt::{FmtContext, Layer};
 use tracing_subscriber::prelude::*;
@@ -278,6 +279,8 @@ pub struct LoggerConfig {
     pub show_log_origin: Option<bool>,
     /// Use the new logger format.
     pub new_format: Option<bool>,
+    /// The profile file to output.
+    pub profile_file: Option<std::path::PathBuf>,
 }
 
 /// Error with actions on the `LoggerConfig`.
@@ -294,13 +297,24 @@ pub enum LoggerConfigError {
     Write(std::io::Error),
 }
 
+macro_rules! registry {
+    ($($x:expr),*) => {
+        {
+            tracing_subscriber::registry()
+            $(
+                .with($x)
+            )*
+        }
+    }
+}
+
 impl LoggerConfig {
     const INIT_MESSAGE: &str = concat!("Running Firecracker v", env!("FIRECRACKER_VERSION"), "\n");
 
     /// Initializes the logger.
     pub fn init(&self) -> std::result::Result<(), LoggerConfigError> {
         let level = tracing::Level::from(self.level.unwrap_or_default());
-        let level_filter = tracing_subscriber::filter::LevelFilter::from_level(level);
+        let filter = tracing_subscriber::filter::LevelFilter::from_level(level);
 
         // In case we open a FIFO, in order to not block the instance if nobody is consuming the
         // message that is flushed to the two pipes, we are opening it with `O_NONBLOCK` flag.
@@ -320,17 +334,11 @@ impl LoggerConfig {
         // Wrap file to satisfy `tracing_subscriber::fmt::MakeWriter`.
         let writer = Mutex::new(LineWriter::new(file));
 
-        // Initialize the layers.
-        if self.new_format.unwrap_or_default() {
-            tracing_subscriber::registry()
-                .with(level_filter)
-                .with(new_log(self, writer))
-                .try_init()
-        } else {
-            tracing_subscriber::registry()
-                .with(level_filter)
-                .with(old_log(self, writer))
-                .try_init()
+        match (self.new_format.unwrap_or_default(), &self.profile_file) {
+            (true, Some(p)) => registry!(filter, new_log(self, writer), flame(p)).try_init(),
+            (false, Some(p)) => registry!(filter, old_log(self, writer), flame(p)).try_init(),
+            (true, None) => registry!(filter, new_log(self, writer)).try_init(),
+            (false, None) => registry!(filter, old_log(self, writer)).try_init(),
         }
         .map_err(LoggerConfigError::Init)?;
 
@@ -369,6 +377,18 @@ fn old_log<S: Subscriber + for<'span> LookupSpan<'span>>(
         .with_writer(writer)
 }
 
+fn flame<S: Subscriber + for<'span> LookupSpan<'span>>(
+    profile_file: &std::path::PathBuf,
+) -> FlameLayer<S, BufWriter<File>> {
+    // We can discard the flush guard as
+    // > This type is only needed when using
+    // > `tracing::subscriber::set_global_default`, which prevents the drop
+    // > implementation of layers from running when the program exits.
+    // See https://docs.rs/tracing-flame/0.2.0/tracing_flame/struct.FlushGuard.html
+    let (flame_layer, _guard) = FlameLayer::with_file(profile_file).unwrap();
+    flame_layer
+}
+
 // use std::sync::atomic::AtomicUsize;
 // use std::sync::atomic::Ordering;
 // static GURAD: AtomicUsize = AtomicUsize::new(0);