feat(layer)!: make adding multiple dynamic fields cheaper (#35)

## Motivation

Adding multiple dynamic fields was expensive because it had to always
allocate the string for each key on every tracing invocation as well as
create a `serde_json::Value`.

## Solution

Now the strings may be `&'static str` if known beforehand and the values
may be left as `impl Serialize` if they're available. The fields will be
serialized into the output string as needed.

Author: mladedav

src/field_writer.rs

@@ -0,0 +1,69 @@
+use crate::cursor::Cursor;
+
+/// A writer passed to closures registered with
+/// [`add_multiple_dynamic_fields`](crate::JsonLayer::add_multiple_dynamic_fields).
+///
+/// Call [`write_field`](FieldWriter::write_field) to add key-value pairs directly to the JSON
+/// output. Keys are `&str` (no allocation required for static strings), and values accept any
+/// type that implements [`serde::Serialize`].
+pub struct FieldWriter<'a> {
+    writer: &'a mut String,
+    prefix_comma: bool,
+    wrote_anything: bool,
+}
+
+impl<'a> FieldWriter<'a> {
+    pub(crate) fn new(writer: &'a mut String, prefix_comma: bool) -> Self {
+        Self {
+            writer,
+            prefix_comma,
+            wrote_anything: false,
+        }
+    }
+
+    pub(crate) fn wrote_anything(&self) -> bool {
+        self.wrote_anything
+    }
+
+    /// Writes a single key-value pair into the JSON output.
+    ///
+    /// The key is serialized as a quoted, escaped JSON string. The value can be any type
+    /// implementing [`serde::Serialize`]: strings, numbers, booleans, `Option<T>`, structs
+    /// with `#[derive(Serialize)]`, etc.
+    ///
+    /// Returns `Err` if serialization fails. On error the writer state is rolled back so the
+    /// output remains valid JSON. Errors from this method are rare and indicate a bug in the
+    /// `Serialize` implementation of the value type.
+    ///
+    /// # Errors
+    ///
+    /// This function errors if serialization of the provided value fails. This means that either
+    /// the implementation of `Serialize` decides to fail, or if the value contains a map with
+    /// non-string keys.
+    pub fn write_field(
+        &mut self,
+        key: &str,
+        value: impl serde::Serialize,
+    ) -> serde_json::Result<()> {
+        let rollback = self.writer.len();
+
+        if self.wrote_anything || self.prefix_comma {
+            self.writer.push(',');
+        }
+
+        if let Err(error) = serde_json::to_writer(&Cursor::new(self.writer), key) {
+            self.writer.truncate(rollback);
+            return Err(error);
+        }
+
+        self.writer.push(':');
+
+        if let Err(error) = serde_json::to_writer(&Cursor::new(self.writer), &value) {
+            self.writer.truncate(rollback);
+            return Err(error);
+        }
+
+        self.wrote_anything = true;
+        Ok(())
+    }
+}

src/layer/event.rs

@@ -12,6 +12,7 @@ use tracing_subscriber::{
 use crate::{
     cached::Cached,
     cursor::Cursor,
+    field_writer::FieldWriter,
     layer::{JsonLayer, JsonValue, SchemaKey},
     serde::JsonSubscriberFormatter,
 };
@@ -176,6 +177,16 @@ where
             }
 
             for value in self.flattened_values.values() {
+                if let JsonValue::DynamicFromEventWithWriter(fun) = value {
+                    let mut inner = writer.inner_mut();
+                    let mut field_writer = FieldWriter::new(&mut inner, serialized_anything);
+                    fun(&event_ref, &mut field_writer);
+                    if field_writer.wrote_anything() {
+                        serialized_anything = true;
+                    }
+                    continue;
+                }
+
                 let Some(value) = resolve_json_value(value, &event_ref) else {
                     continue;
                 };
@@ -298,6 +309,8 @@ fn resolve_json_value<'a, S: Subscriber + for<'lookup> LookupSpan<'lookup>>(
             event.parent_span().and_then(fun).map(MaybeCached::Cached)
         },
         JsonValue::DynamicRawFromEvent(fun) => Some(MaybeCached::Raw(fun)),
+        // This cannot be used with a static key so this should never be called
+        JsonValue::DynamicFromEventWithWriter(_) => None,
     }
 }
 

src/layer/mod.rs

@@ -1,11 +1,4 @@
-use std::{
-    borrow::Cow,
-    cell::RefCell,
-    collections::{BTreeMap, HashMap},
-    fmt,
-    io,
-    sync::Arc,
-};
+use std::{borrow::Cow, cell::RefCell, collections::BTreeMap, fmt, io, sync::Arc};
 
 use serde::Serialize;
 use tracing_core::{
@@ -29,6 +22,7 @@ use uuid::Uuid;
 
 use crate::{
     cached::Cached,
+    field_writer::FieldWriter,
     fields::{JsonFields, JsonFieldsInner},
     serde::RenamedFields,
     visitor::JsonVisitor,
@@ -94,6 +88,9 @@ pub(crate) enum JsonValue<S: for<'lookup> LookupSpan<'lookup>> {
     DynamicRawFromEvent(
         Box<dyn Fn(&EventRef<'_, '_, '_, S>, &mut dyn fmt::Write) -> fmt::Result + Send + Sync>,
     ),
+    DynamicFromEventWithWriter(
+        Box<dyn Fn(&EventRef<'_, '_, '_, S>, &mut FieldWriter<'_>) + Send + Sync>,
+    ),
 }
 
 impl<S, W> Layer<S> for JsonLayer<S, W>
@@ -451,52 +448,44 @@ where
         );
     }
 
-    /// Adds multiple new dynamic field where the keys may not be known when calling this method.
+    /// Adds multiple new dynamic fields where the keys may not be known when calling this method.
     ///
-    /// This method takes a closure argument that will be called with the event and tracing context.
-    /// Through these, the parent span can be accessed among other things. This closure returns a
-    /// value which can be iterated over to return a tuple of a `String` which will be used as a
-    /// JSON key and a `serde_json::Value` which will be used as a value. In most cases returning
-    /// `HashMap<String, serde_json::Value>` should be sufficient.
+    /// This method takes a closure argument that will be called with the event and tracing context
+    /// along with a [`FieldWriter`]. Call [`FieldWriter::write_field`] to write key-value pairs
+    /// directly into the JSON output. Values accept any type implementing [`serde::Serialize`].
     ///
     /// It is the user's responsibility to make sure that no two keys clash as that would create an
     /// invalid JSON. It's generally better to use [`add_dynamic_field`](Self::add_dynamic_field)
     /// instead if the field names are known.
     ///
     /// # Examples
     ///
-    /// Print either a question or an answer:
+    /// Print a question or an answer:
     ///
     /// ```rust
     /// # use tracing_subscriber::prelude::*;
     ///
     /// let mut layer = json_subscriber::JsonLayer::stdout();
     /// layer.add_multiple_dynamic_fields(
-    ///     |_event, _context| {
+    ///     |_event, _context, writer| {
     /// #       let condition = true;
     ///         if condition {
-    ///             [("question".to_owned(), serde_json::Value::String("What?".to_owned()))]
+    ///             _ = writer.write_field("question", "What?");
     ///         } else {
-    ///             [("answer".to_owned(), serde_json::Value::Number(42.into()))]
+    ///             _ = writer.write_field("answer", 42u64);
     ///         }
-    /// });
+    ///     }
+    /// );
     /// # tracing_subscriber::registry().with(layer);
-    /// # fn get_hostname() -> &'static str { "localhost" }
     /// ```
-    pub fn add_multiple_dynamic_fields<Fun, Res>(&mut self, mapper: Fun)
+    pub fn add_multiple_dynamic_fields<Fun>(&mut self, mapper: Fun)
     where
-        for<'a> Fun: Fn(&'a Event<'_>, &Context<'_, S>) -> Res + Send + Sync + 'a,
-        Res: IntoIterator<Item = (String, serde_json::Value)>,
+        Fun: Fn(&Event<'_>, &Context<'_, S>, &mut FieldWriter<'_>) + Send + Sync + 'static,
     {
         self.flattened_values.insert(
             FlatSchemaKey::new_uuid(),
-            JsonValue::DynamicFromEvent(Box::new(move |event| {
-                serde_json::to_value(
-                    mapper(event.event(), event.context())
-                        .into_iter()
-                        .collect::<HashMap<_, _>>(),
-                )
-                .ok()
+            JsonValue::DynamicFromEventWithWriter(Box::new(move |event, writer| {
+                mapper(event.event(), event.context(), writer);
             })),
         );
     }

src/lib.rs

@@ -105,6 +105,7 @@
 
 mod cached;
 mod cursor;
+mod field_writer;
 mod fields;
 pub mod fmt;
 mod layer;
@@ -114,5 +115,6 @@ mod visitor;
 #[cfg(test)]
 mod tests;
 
+pub use field_writer::FieldWriter;
 pub use fmt::{fmt, layer};
 pub use layer::JsonLayer;