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/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/vmm_config/mod.rs

@@ -278,6 +278,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`.
@@ -318,7 +320,21 @@ impl LoggerConfig {
 
         let fmt_layer = fmt_layer.with_writer(writer.with_max_level(level));
 
-        tracing_subscriber::registry().with(fmt_layer).init();
+        if let Some(profile_file) = &self.profile_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) = tracing_flame::FlameLayer::with_file(profile_file).unwrap();
+            tracing_subscriber::registry()
+                .with(fmt_layer)
+                .with(flame_layer)
+                .init();
+        } else {
+            tracing_subscriber::registry().with(fmt_layer).init();
+        }
+
         Ok(())
     }
     /// Initializes the logger with the new format.
@@ -353,7 +369,20 @@ impl LoggerConfig {
 
         let fmt_layer = fmt_layer.with_writer(writer.with_max_level(level));
 
-        tracing_subscriber::registry().with(fmt_layer).init();
+        if let Some(profile_file) = &self.profile_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) = tracing_flame::FlameLayer::with_file(profile_file).unwrap();
+            tracing_subscriber::registry()
+                .with(fmt_layer)
+                .with(flame_layer)
+                .init();
+        } else {
+            tracing_subscriber::registry().with(fmt_layer).init();
+        }
 
         Ok(())
     }