feat(layer): simple renames of flattened event fields (#23)

## Motivation

Closes #19.

Users sometimes want to rename the fields in events and spans. This
allows that specifically for flattened events. This should be done more
generally somehow.

## Solution

Accept a closure and a state object, calling the closure for every field
to assess whether to rename or not. This will most likely be a lookup in
a map.

Author: mladedav

src/layer/mod.rs

@@ -30,6 +30,7 @@ use uuid::Uuid;
 use crate::{
     cached::Cached,
     fields::{JsonFields, JsonFieldsInner},
+    serde::RenamedFields,
     visitor::JsonVisitor,
 };
 
@@ -790,6 +791,55 @@ where
         self
     }
 
+    /// Print all event fields, each as its own top level member of the JSON. This also allows the
+    /// fields to have different keys than the field names used in tracing.
+    ///
+    /// Only field names can be renamed this way, not other fields such as `timestamp` or `target`
+    /// which usually take the key as their parameter (such as `Self::with_target`).
+    ///
+    /// The renames are realized by a provided function that will be passed the original field name
+    /// and a user-defined context and it must produce the new name.
+    ///
+    /// For example when simple static renames are needed the following can work:
+    ///
+    /// ```rust
+    /// # use std::collections::HashMap;
+    /// # use tracing_subscriber::prelude::*;
+    ///
+    /// let mut layer = json_subscriber::JsonLayer::stdout();
+    ///
+    /// let renames = HashMap::from([
+    ///     ("message".to_owned(), "msg".to_owned()),
+    ///     ("foo".to_owned(), "bar".to_owned()),
+    /// ]);
+    /// layer.with_flattened_event_with_renames(
+    ///     move |name, map| map.get(name).map_or(name, String::as_str),
+    ///     renames,
+    /// );
+    /// # tracing_subscriber::registry().with(layer);
+    ///
+    /// // This will produce something like `{"bar":3,"msg":"x",...}`
+    /// tracing::info!(foo = 3, "x");
+    /// ```
+    ///
+    /// It is the user's responsibility to make sure that the field names will not clash with other
+    /// defined members of the output JSON. If they clash, invalid JSON with multiple fields with
+    /// the same key may be generated.
+    pub fn with_flattened_event_with_renames<F, T>(&mut self, renames: F, context: T) -> &mut Self
+    where
+        F: for<'a> Fn(&'a str, &'a T) -> &'a str + Send + Sync + 'static + Clone,
+        T: Clone + Send + Sync + 'static,
+    {
+        self.flattened_values.insert(
+            FlatSchemaKey::FlattenedEvent,
+            JsonValue::DynamicFromEvent(Box::new(move |event| {
+                serde_json::to_value(RenamedFields::new(event.event(), renames.clone(), &context))
+                    .ok()
+            })),
+        );
+        self
+    }
+
     /// Sets whether or not the log line will include the current span in formatted events.
     pub fn with_current_span(&mut self, key: impl Into<String>) -> &mut Self {
         self.keyed_values.insert(
@@ -1108,6 +1158,8 @@ fn write_escaped(writer: &mut dyn fmt::Write, value: &str) -> Result<(), fmt::Er
 
 #[cfg(test)]
 mod tests {
+    use std::collections::HashMap;
+
     use serde_json::json;
     use tracing::subscriber::with_default;
     use tracing_subscriber::{registry, Layer, Registry};
@@ -1164,4 +1216,37 @@ mod tests {
             tracing::info!(does = "not matter", "whatever");
         });
     }
+
+    #[test]
+    fn flattened_event_with_renames() {
+        let renames = HashMap::from([
+            ("message".to_owned(), "msg".to_owned()),
+            ("msg".to_owned(), "message".to_owned()),
+            ("same".to_owned(), "same".to_owned()),
+            ("different".to_owned(), "gone".to_owned()),
+        ]);
+        let mut layer = JsonLayer::stdout();
+        layer.with_flattened_event_with_renames(
+            move |name, map| map.get(name).map_or(name, String::as_str),
+            renames,
+        );
+
+        let expected = json!({
+            "message": "msg",
+            "msg": "message",
+            "same": "same",
+            "gone": "different",
+            "another": "another",
+        });
+
+        test_json(&expected, layer, || {
+            tracing::info!(
+                msg = "msg",
+                same = "same",
+                different = "different",
+                another = "another",
+                "message"
+            );
+        });
+    }
 }

src/serde/mod.rs

@@ -1,4 +1,7 @@
+mod tracing_serde;
+
 use serde_json::ser::Formatter;
+pub(crate) use tracing_serde::RenamedFields;
 
 pub(crate) struct JsonSubscriberFormatter;
 

src/serde/tracing_serde.rs

@@ -0,0 +1,127 @@
+use std::{fmt, mem::transmute};
+
+use serde::{ser::SerializeMap, Serialize, Serializer};
+use tracing::{field::Visit, Event};
+use tracing_core::Field;
+
+pub(crate) struct RenamedFields<'a, F, C> {
+    event: &'a Event<'a>,
+    renames: F,
+    context: &'a C,
+}
+
+impl<'a, F, C> RenamedFields<'a, F, C> {
+    pub(crate) fn new(event: &'a Event<'a>, renames: F, context: &'a C) -> Self {
+        Self {
+            event,
+            renames,
+            context,
+        }
+    }
+}
+
+impl<F, C> Serialize for RenamedFields<'_, F, C>
+where
+    F: for<'a> Fn(&'a str, &'a C) -> &'a str + Send + Sync + 'static,
+{
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        let len = self.event.fields().count();
+        let serializer = serializer.serialize_map(Some(len))?;
+        let renames: &'static F = unsafe { transmute(&self.renames) };
+        let mut visitor = SerdeMapVisitor::new(serializer, renames, self.context);
+        self.event.record(&mut visitor);
+        visitor.finish()
+    }
+}
+
+/// Implements `tracing_core::field::Visit` for some `serde::ser::SerializeMap`.
+#[derive(Debug)]
+pub struct SerdeMapVisitor<'a, S: SerializeMap, F, C> {
+    serializer: S,
+    renames: F,
+    context: &'a C,
+    state: Result<(), S::Error>,
+}
+
+impl<'a, S, F, C> SerdeMapVisitor<'a, S, F, C>
+where
+    S: SerializeMap,
+{
+    /// Create a new map visitor.
+    pub fn new(serializer: S, renames: F, context: &'a C) -> Self {
+        Self {
+            serializer,
+            renames,
+            context,
+            state: Ok(()),
+        }
+    }
+
+    /// Completes serializing the visited object, returning `Ok(())` if all
+    /// fields were serialized correctly, or `Error(S::Error)` if a field could
+    /// not be serialized.
+    pub fn finish(self) -> Result<S::Ok, S::Error> {
+        self.state?;
+        self.serializer.end()
+    }
+}
+
+impl<S, F, C> Visit for SerdeMapVisitor<'_, S, F, C>
+where
+    S: SerializeMap,
+    F: for<'a> Fn(&'a str, &'a C) -> &'a str + Send + Sync + 'static,
+{
+    fn record_bool(&mut self, field: &Field, value: bool) {
+        // If previous fields serialized successfully, continue serializing,
+        // otherwise, short-circuit and do nothing.
+        if self.state.is_ok() {
+            self.state = self
+                .serializer
+                .serialize_entry((self.renames)(field.name(), self.context), &value);
+        }
+    }
+
+    fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
+        if self.state.is_ok() {
+            self.state = self.serializer.serialize_entry(
+                (self.renames)(field.name(), self.context),
+                &format_args!("{value:?}"),
+            );
+        }
+    }
+
+    fn record_u64(&mut self, field: &Field, value: u64) {
+        if self.state.is_ok() {
+            self.state = self
+                .serializer
+                .serialize_entry((self.renames)(field.name(), self.context), &value);
+        }
+    }
+
+    fn record_i64(&mut self, field: &Field, value: i64) {
+        if self.state.is_ok() {
+            self.state = self
+                .serializer
+                .serialize_entry((self.renames)(field.name(), self.context), &value);
+        }
+    }
+
+    fn record_f64(&mut self, field: &Field, value: f64) {
+        if self.state.is_ok() {
+            self.state = self
+                .serializer
+                .serialize_entry((self.renames)(field.name(), self.context), &value);
+        }
+    }
+
+    fn record_str(&mut self, field: &Field, value: &str) {
+        if self.state.is_ok() {
+            self.state = self
+                .serializer
+                .serialize_entry((self.renames)(field.name(), self.context), &value);
+        }
+    }
+}