From 56797ca92d7eddd96a3c4d60717de3e66a6d26e1 Mon Sep 17 00:00:00 2001
From: Andy Lok <andylokandy@hotmail.com>
Date: Wed, 14 Aug 2024 17:20:39 +0800
Subject: [PATCH] fix: FastraceEvent should not hardcode layout (#46)

---
 src/append/fastrace.rs | 25 ++++++++++++++++---------
 src/layout/kv.rs       | 22 ++++++++++++++++++++++
 src/layout/mod.rs      |  1 +
 3 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/src/append/fastrace.rs b/src/append/fastrace.rs
index 617a510..f252464 100644
--- a/src/append/fastrace.rs
+++ b/src/append/fastrace.rs
@@ -12,11 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::borrow::Cow;
+
 use jiff::Zoned;
 use log::Record;
 
 use crate::append::Append;
-use crate::layout::KvDisplay;
+use crate::layout::collect_kvs;
 
 /// An appender that adds log records to fastrace as an event associated to the current span.
 #[derive(Default, Debug, Clone)]
@@ -24,14 +26,19 @@ pub struct FastraceEvent;
 
 impl Append for FastraceEvent {
     fn append(&self, record: &Record) -> anyhow::Result<()> {
-        let message = format!(
-            "{} {:>5} {}{}",
-            Zoned::now(),
-            record.level(),
-            record.args(),
-            KvDisplay::new(record.key_values()),
-        );
-        fastrace::Event::add_to_local_parent(message, || []);
+        let message = format!("{}", record.args());
+        fastrace::Event::add_to_local_parent(message, || {
+            [
+                (Cow::from("level"), Cow::from(record.level().as_str())),
+                (Cow::from("timestamp"), Cow::from(Zoned::now().to_string())),
+            ]
+            .into_iter()
+            .chain(
+                collect_kvs(record.key_values())
+                    .into_iter()
+                    .map(|(k, v)| (Cow::from(k), Cow::from(v))),
+            )
+        });
         Ok(())
     }
 
diff --git a/src/layout/kv.rs b/src/layout/kv.rs
index a619f34..8f9685b 100644
--- a/src/layout/kv.rs
+++ b/src/layout/kv.rs
@@ -45,3 +45,25 @@ impl<'a, 'kvs> log::kv::Visitor<'kvs> for KvWriter<'a, 'kvs> {
         Ok(())
     }
 }
+
+/// A helper to collect log's key-value pairs.
+pub fn collect_kvs(kv: &dyn log::kv::Source) -> Vec<(String, String)> {
+    let mut collector = KvCollector { kv: Vec::new() };
+    kv.visit(&mut collector).ok();
+    collector.kv
+}
+
+struct KvCollector {
+    kv: Vec<(String, String)>,
+}
+
+impl<'kvs> log::kv::Visitor<'kvs> for KvCollector {
+    fn visit_pair(
+        &mut self,
+        key: log::kv::Key<'kvs>,
+        value: log::kv::Value<'kvs>,
+    ) -> Result<(), log::kv::Error> {
+        self.kv.push((key.to_string(), value.to_string()));
+        Ok(())
+    }
+}
diff --git a/src/layout/mod.rs b/src/layout/mod.rs
index de9d092..5ac9a21 100644
--- a/src/layout/mod.rs
+++ b/src/layout/mod.rs
@@ -18,6 +18,7 @@ pub use custom::CustomLayout;
 pub use identical::IdenticalLayout;
 #[cfg(feature = "json")]
 pub use json::JsonLayout;
+pub use kv::collect_kvs;
 pub use kv::KvDisplay;
 pub use text::LevelColor;
 pub use text::TextLayout;