Revert "Update tracing.rs"

This reverts commit 036296875e.

swap/src/cli/tracing.rs

@@ -1,31 +1,74 @@
 use anyhow::Result;
 use std::fmt::Debug;
 use std::fs::OpenOptions;
-use std::io;
 use std::io::Write;
 use std::option::Option::Some;
 use std::path::{Path, PathBuf};
-use tracing::{Event, Level, Subscriber};
-use tracing_subscriber::fmt::format::{DefaultFields, Format, Json, JsonFields};
+use tracing::field::Field;
+use tracing::span::Attributes;
+use tracing::{Id, Level, Subscriber};
+use tracing_subscriber::fmt::format::{Format, Json, JsonFields};
 use tracing_subscriber::fmt::time::UtcTime;
 use tracing_subscriber::layer::{Context, SubscriberExt};
 use tracing_subscriber::util::SubscriberInitExt;
-use tracing_subscriber::{EnvFilter, fmt, FmtSubscriber, Layer, Registry, registry};
-use tracing_subscriber::fmt::{FormatEvent, MakeWriter};
+use tracing_subscriber::{EnvFilter, FmtSubscriber, Layer, Registry};
+use tracing_subscriber::fmt::{format, FormatEvent, MakeWriter};
 use serde_json::{json, Value};
-use tracing::subscriber::set_global_default;
+use serde::Serialize;
+
+#[derive(Debug, Serialize)]
+struct LogEvent<'a> {
+    message: String,
+    level: &'a str,
+    target: &'a str,
+    module: &'a str,
+    file: Option<&'a str>,
+    line: Option<u32>,
+    fields: Value,
+}
 
+/// Transforms a `tracing::Event` into a JSON-formatted string.
+fn format_event_as_json<'a, S>(
+    _: &Context<'a, S>,
+    event: &'a tracing::Event<'a>,
+) -> serde_json::Result<String> {
+    // Extracting fields from the event into a serde_json::Value
+    let mut fields = json!({});
+    let mut message = String::new();  // For capturing the main message
+
+    event.record(&mut |field: &Field, value: &dyn Debug| {
+        if field.name() == "message" {
+            message = format!("{:?}", value);
+        } else {
+            fields[field.name()] = json!(format!("{:?}", value));
+        }
+    });
+
+    let log = LogEvent {
+        message, // Use the captured message here
+        level: event.metadata().level().as_str(),
+        target: event.metadata().target(),
+        module: event.metadata().module_path().unwrap_or_default(),
+        file: event.metadata().file(),
+        line: event.metadata().line(),
+        fields,
+    };
+
+    serde_json::to_string(&log)
+}
 
-pub struct MyMakeWriter {
+struct SwapIdVisitor {
+    swap_id: Option<String>,
+}
+
+pub struct FileLayer {
     dir: PathBuf,
-    // other members like stdout and stderr if needed
 }
 
-impl MyMakeWriter {
+impl FileLayer {
     pub fn new(dir: impl AsRef<Path>) -> Self {
         Self {
             dir: dir.as_ref().to_path_buf(),
-            // initialize other members
         }
     }
 
@@ -33,80 +76,69 @@ impl MyMakeWriter {
         self.dir.join(format!("swap-{}.log", swap_id))
     }
 
-    fn get_file_writer(&self, swap_id: &str) -> io::Result<impl Write> {
-        OpenOptions::new().append(true).create(true).open(self.get_log_path(swap_id))
+    fn append_to_file(&self, swap_id: &str, message: &str) -> std::io::Result<()> {
+        let path = self.get_log_path(swap_id);
+        let mut file = OpenOptions::new().append(true).create(true).open(path)?;
+        file.write_all(message.as_bytes())
     }
 }
 
-impl<'a> MakeWriter<'a> for MyMakeWriter {
-    type Writer = Box<dyn Write + 'a>;
 
-    fn make_writer(&'a self) -> Self::Writer {
-        unreachable!();
+impl<S> Layer<S> for FileLayer
+    where
+        S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
+{
+    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
+        let mut visitor = SwapIdVisitor { swap_id: None };
+        attrs.record(&mut visitor);
+
+        if let Some(swap_id) = visitor.swap_id {
+            if let Some(span) = ctx.span(id) {
+                span.extensions_mut().insert(swap_id);
+            }
+        }
     }
 
-    fn make_writer_for(&'a self, meta: &tracing::Metadata<'_>) -> Self::Writer {
-        // Print all attributes of the event
-        println!("Event attributes: {:?}", meta);
-        let swap_id = "dummy-swap-id";
-
-        Box::new(self.get_file_writer(swap_id).expect("Failed to open log file"))
+    fn on_event(&self, event: &tracing::Event<'_>, ctx: Context<'_, S>) {
+        if let Some(current_span_id) = ctx.current_span().id() {
+            if let Some(span) = ctx.span(current_span_id) {
+                if let Some(swap_id) = span.extensions().get::<String>() {
+                    // TODO: This is a hack, I need to figure out how to get the JSON formatter to work
+
+                    if let Ok(json_log) = format_event_as_json(&ctx, event) {
+                        self.append_to_file(swap_id, &format!("{}\n", json_log))
+                            .expect("Failed to write to file");
+                    }
+                }
+            }
+        }
     }
 }
 
+impl tracing::field::Visit for SwapIdVisitor {
+    fn record_debug(&mut self, field: &Field, value: &dyn Debug) {
+        if field.name() == "swap_id" {
+            self.swap_id = Some(format!("{:?}", value));
+        }
+    }
+}
 
 pub fn init(debug: bool, json: bool, dir: impl AsRef<Path>) -> Result<()> {
-    let level_filter = EnvFilter::try_new("swap=debug")?;
-    let registry = Registry::default().with(level_filter);
-
+    let level = if debug { Level::DEBUG } else { Level::INFO };
+    let is_terminal = atty::is(atty::Stream::Stderr);
 
-    let file_logger = registry.with(
-        fmt::layer()
-            .with_ansi(false)
-            .with_target(false)
-            .json()
-            .with_writer(MyMakeWriter::new(dir))
-    );
+    let file_layer = FileLayer::new(dir.as_ref());
 
-    set_global_default(file_logger.with(info_terminal_printer()))?;
+    FmtSubscriber::builder()
+        .with_env_filter(format!("swap={}", level))
+        .with_writer(std::io::stderr)
+        .with_ansi(is_terminal)
+        .with_timer(UtcTime::rfc_3339())
+        .with_target(false)
+        .finish()
+        .with(file_layer)
+        .init();
 
-    //tracing::info!("Logging initialized to {}", dir.as_ref().display());
+    tracing::info!("Logging initialized to {}", dir.as_ref().display());
     Ok(())
 }
-
-pub struct StdErrPrinter<L> {
-    inner: L,
-    level: Level,
-}
-
-type StdErrLayer<S, T> = fmt::Layer<
-    S,
-    DefaultFields,
-    Format<fmt::format::Full, T>,
-    fn() -> io::Stderr,
->;
-
-fn info_terminal_printer<S>() -> StdErrPrinter<StdErrLayer<S, ()>> {
-    let is_terminal = atty::is(atty::Stream::Stderr);
-    StdErrPrinter {
-        inner: fmt::layer()
-            .with_ansi(is_terminal)
-            .with_target(false)
-            .with_level(false)
-            .without_time()
-            .with_writer(std::io::stderr),
-        level: Level::INFO,
-    }
-}
-
-impl<L, S> Layer<S> for StdErrPrinter<L>
-    where
-        L: 'static + Layer<S>,
-        S: Subscriber + for<'a> registry::LookupSpan<'a>,
-{
-    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
-        if self.level.ge(event.metadata().level()) {
-            self.inner.on_event(event, ctx);
-        }
-    }
-}