From 40048fa47b9ec42d691c53bc14cc0dc602d392db Mon Sep 17 00:00:00 2001
From: Al Liu <scygliu1@gmail.com>
Date: Fri, 16 Jun 2023 19:26:48 +0800
Subject: [PATCH] Fix #183: support structured log for sim app. (#184)

---
 simulations/Cargo.toml                      |  2 +-
 simulations/src/bin/app/log.rs              | 47 +++++++++++++++++++++
 simulations/src/bin/{app.rs => app/main.rs} | 17 +++-----
 3 files changed, 53 insertions(+), 13 deletions(-)
 create mode 100644 simulations/src/bin/app/log.rs
 rename simulations/src/bin/{app.rs => app/main.rs} (95%)

diff --git a/simulations/Cargo.toml b/simulations/Cargo.toml
index e33ecbcc..ac7f4a8c 100644
--- a/simulations/Cargo.toml
+++ b/simulations/Cargo.toml
@@ -5,7 +5,7 @@ edition = "2021"
 
 [[bin]]
 name = "simulation"
-path = "src/bin/app.rs"
+path = "src/bin/app/main.rs"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
diff --git a/simulations/src/bin/app/log.rs b/simulations/src/bin/app/log.rs
new file mode 100644
index 00000000..6b578fc1
--- /dev/null
+++ b/simulations/src/bin/app/log.rs
@@ -0,0 +1,47 @@
+use std::str::FromStr;
+
+#[derive(Copy, Clone)]
+pub enum LogFormat {
+    Plain,
+    Json,
+}
+
+impl FromStr for LogFormat {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.trim().to_ascii_lowercase().as_str() {
+            "json" => Ok(LogFormat::Json),
+            "plain" => Ok(LogFormat::Plain),
+            _ => Err(anyhow::anyhow!("Unknown log format")),
+        }
+    }
+}
+
+pub fn config_tracing(fmt: Option<LogFormat>) {
+    let filter = std::env::var("SIMULATION_LOG").unwrap_or_else(|_| "info".to_owned());
+    if let Some(LogFormat::Json) = fmt {
+        let subscriber = tracing_subscriber::fmt::fmt()
+            .without_time()
+            .with_line_number(true)
+            .with_env_filter(filter)
+            .with_file(false)
+            .with_target(true)
+            .with_ansi(true)
+            .json()
+            .finish();
+        tracing::subscriber::set_global_default(subscriber)
+            .expect("config_tracing is only called once");
+    } else {
+        let subscriber = tracing_subscriber::fmt::fmt()
+            .without_time()
+            .with_line_number(true)
+            .with_env_filter(filter)
+            .with_file(false)
+            .with_target(true)
+            .with_ansi(true)
+            .finish();
+        tracing::subscriber::set_global_default(subscriber)
+            .expect("config_tracing is only called once");
+    }
+}
diff --git a/simulations/src/bin/app.rs b/simulations/src/bin/app/main.rs
similarity index 95%
rename from simulations/src/bin/app.rs
rename to simulations/src/bin/app/main.rs
index 27e2b430..9b3b8fed 100644
--- a/simulations/src/bin/app.rs
+++ b/simulations/src/bin/app/main.rs
@@ -33,6 +33,7 @@ use simulations::{
     node::carnot::CarnotNode, output_processors::OutData, runner::SimulationRunner,
     settings::SimulationSettings, util::node_id,
 };
+mod log;
 
 /// Main simulation wrapper
 /// Pipes together the cli arguments with the execution
@@ -43,6 +44,8 @@ pub struct SimulationApp {
     input_settings: PathBuf,
     #[clap(long)]
     stream_type: Option<StreamType>,
+    #[clap(long)]
+    log_format: Option<log::LogFormat>,
 }
 
 impl SimulationApp {
@@ -50,6 +53,7 @@ impl SimulationApp {
         let Self {
             input_settings,
             stream_type,
+            log_format: _,
         } = self;
         let simulation_settings: SimulationSettings = load_json_from_file(&input_settings)?;
 
@@ -221,19 +225,8 @@ fn generate_overlays<R: Rng>(
 }
 
 fn main() -> anyhow::Result<()> {
-    let filter = std::env::var("SIMULATION_LOG").unwrap_or_else(|_| "info".to_owned());
-    let subscriber = tracing_subscriber::fmt::fmt()
-        .without_time()
-        .with_line_number(true)
-        .with_env_filter(filter)
-        .with_file(false)
-        .with_target(true)
-        .with_ansi(true)
-        .finish();
-    tracing::subscriber::set_global_default(subscriber)
-        .expect("config_tracing is only called once");
-
     let app: SimulationApp = SimulationApp::parse();
+    log::config_tracing(app.log_format);
 
     if let Err(e) = app.run() {
         tracing::error!("error: {}", e);