Add the ability to skip fields (#25)

vrischmann · web-flow · commit 563b452fec65 · 2023-01-08T12:32:08.000Z
* all: run `cargo fmt`

* add the ahash dependency

* implement skipping fields in the formatter
diff --git a/Cargo.toml b/Cargo.toml
@@ -32,6 +32,7 @@ serde = "1.0.106"
 gethostname = "0.2.1"
 tracing-core = "0.1.10"
 time = { version = "0.3", default-features = false, features = ["formatting"] }
+ahash = "0.8.2"
 
 [dev-dependencies]
 claims = "0.6.0"
diff --git a/src/formatting_layer.rs b/src/formatting_layer.rs
@@ -1,9 +1,11 @@
 use crate::storage_layer::JsonStorage;
-use serde::ser::{SerializeMap, Serializer};
+use ahash::{HashSet, HashSetExt};
+use serde::ser::{Serialize, SerializeMap, Serializer};
 use serde_json::Value;
 use std::collections::HashMap;
 use std::fmt;
 use std::io::Write;
+use time::format_description::well_known::Rfc3339;
 use tracing::{Event, Id, Subscriber};
 use tracing_core::metadata::Level;
 use tracing_core::span::Attributes;
@@ -12,7 +14,6 @@ use tracing_subscriber::fmt::MakeWriter;
 use tracing_subscriber::layer::Context;
 use tracing_subscriber::registry::SpanRef;
 use tracing_subscriber::Layer;
-use time::format_description::well_known::Rfc3339;
 
 /// Keys for core fields of the Bunyan format (https://github.com/trentm/node-bunyan#core-fields)
 const BUNYAN_VERSION: &str = "v";
@@ -48,8 +49,26 @@ pub struct BunyanFormattingLayer<W: for<'a> MakeWriter<'a> + 'static> {
     bunyan_version: u8,
     name: String,
     default_fields: HashMap<String, Value>,
+    skip_fields: HashSet<String>,
+}
+
+/// This error will be returned in [`BunyanFormattingLayer::skip_fields`] if trying to skip a core field.
+#[non_exhaustive]
+#[derive(Debug)]
+pub struct InvalidFieldError(String);
+
+impl fmt::Display for InvalidFieldError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "{} is a core field in the bunyan log format, it can't be skipped",
+            &self.0
+        )
+    }
 }
 
+impl std::error::Error for InvalidFieldError {}
+
 impl<W: for<'a> MakeWriter<'a> + 'static> BunyanFormattingLayer<W> {
     /// Create a new `BunyanFormattingLayer`.
     ///
@@ -74,15 +93,36 @@ impl<W: for<'a> MakeWriter<'a> + 'static> BunyanFormattingLayer<W> {
         Self::with_default_fields(name, make_writer, HashMap::new())
     }
 
-    pub fn with_default_fields(name: String, make_writer: W, default_fields: HashMap<String, Value>) -> Self {
+    pub fn with_default_fields(
+        name: String,
+        make_writer: W,
+        default_fields: HashMap<String, Value>,
+    ) -> Self {
         Self {
             make_writer,
             name,
             pid: std::process::id(),
             hostname: gethostname::gethostname().to_string_lossy().into_owned(),
             bunyan_version: 0,
             default_fields,
+            skip_fields: HashSet::new(),
+        }
+    }
+
+    /// Add fields to skip when formatting with this layer.
+    pub fn skip_fields<T>(mut self, fields: &[T]) -> Result<Self, InvalidFieldError>
+    where
+        T: AsRef<str>,
+    {
+        for field in fields {
+            let field = field.as_ref();
+            if BUNYAN_RESERVED_FIELDS.contains(&field) {
+                return Err(InvalidFieldError(field.to_string()));
+            }
+            self.skip_fields.insert(field.to_string());
         }
+
+        Ok(self)
     }
 
     fn serialize_bunyan_core_fields(
@@ -103,6 +143,22 @@ impl<W: for<'a> MakeWriter<'a> + 'static> BunyanFormattingLayer<W> {
         Ok(())
     }
 
+    fn serialize_field<V>(
+        &self,
+        map_serializer: &mut impl SerializeMap<Error = serde_json::Error>,
+        key: &str,
+        value: &V,
+    ) -> Result<(), std::io::Error>
+    where
+        V: Serialize + ?Sized,
+    {
+        if !self.skip_fields.contains(key) {
+            map_serializer.serialize_entry(key, value)?;
+        }
+
+        Ok(())
+    }
+
     /// Given a span, it serialised it to a in-memory buffer (vector of bytes).
     fn serialize_span<S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>>(
         &self,
@@ -117,27 +173,27 @@ impl<W: for<'a> MakeWriter<'a> + 'static> BunyanFormattingLayer<W> {
         // Additional metadata useful for debugging
         // They should be nested under `src` (see https://github.com/trentm/node-bunyan#src )
         // but `tracing` does not support nested values yet
-        map_serializer.serialize_entry("target", span.metadata().target())?;
-        map_serializer.serialize_entry("line", &span.metadata().line())?;
-        map_serializer.serialize_entry("file", &span.metadata().file())?;
+        self.serialize_field(&mut map_serializer, "target", span.metadata().target())?;
+        self.serialize_field(&mut map_serializer, "line", &span.metadata().line())?;
+        self.serialize_field(&mut map_serializer, "file", &span.metadata().file())?;
 
         // Add all default fields
         for (key, value) in self.default_fields.iter() {
             if !BUNYAN_RESERVED_FIELDS.contains(&key.as_str()) {
-                map_serializer.serialize_entry(key, value)?;
+                self.serialize_field(&mut map_serializer, key, value)?;
             } else {
                 tracing::debug!(
-                        "{} is a reserved field in the bunyan log format. Skipping it.",
-                        key
-                    );
+                    "{} is a reserved field in the bunyan log format. Skipping it.",
+                    key
+                );
             }
         }
 
         let extensions = span.extensions();
         if let Some(visitor) = extensions.get::<JsonStorage>() {
             for (key, value) in visitor.values() {
                 if !BUNYAN_RESERVED_FIELDS.contains(key) {
-                    map_serializer.serialize_entry(key, value)?;
+                    self.serialize_field(&mut map_serializer, key, value)?;
                 } else {
                     tracing::debug!(
                         "{} is a reserved field in the bunyan log format. Skipping it.",
@@ -252,16 +308,15 @@ where
             // Additional metadata useful for debugging
             // They should be nested under `src` (see https://github.com/trentm/node-bunyan#src )
             // but `tracing` does not support nested values yet
-            map_serializer.serialize_entry("target", event.metadata().target())?;
-            map_serializer.serialize_entry("line", &event.metadata().line())?;
-            map_serializer.serialize_entry("file", &event.metadata().file())?;
+            self.serialize_field(&mut map_serializer, "target", event.metadata().target())?;
+            self.serialize_field(&mut map_serializer, "line", &event.metadata().line())?;
+            self.serialize_field(&mut map_serializer, "file", &event.metadata().file())?;
 
             // Add all default fields
-            for (key, value) in self.default_fields
-                .iter()
-                .filter(|(key, _)| key.as_str() != "message" && !BUNYAN_RESERVED_FIELDS.contains(&key.as_str()))
-            {
-                map_serializer.serialize_entry(key, value)?;
+            for (key, value) in self.default_fields.iter().filter(|(key, _)| {
+                key.as_str() != "message" && !BUNYAN_RESERVED_FIELDS.contains(&key.as_str())
+            }) {
+                self.serialize_field(&mut map_serializer, key, value)?;
             }
 
             // Add all the other fields associated with the event, expect the message we already used.
@@ -270,7 +325,7 @@ where
                 .iter()
                 .filter(|(&key, _)| key != "message" && !BUNYAN_RESERVED_FIELDS.contains(&key))
             {
-                map_serializer.serialize_entry(key, value)?;
+                self.serialize_field(&mut map_serializer, key, value)?;
             }
 
             // Add all the fields from the current span, if we have one.
@@ -279,7 +334,7 @@ where
                 if let Some(visitor) = extensions.get::<JsonStorage>() {
                     for (key, value) in visitor.values() {
                         if !BUNYAN_RESERVED_FIELDS.contains(key) {
-                            map_serializer.serialize_entry(key, value)?;
+                            self.serialize_field(&mut map_serializer, key, value)?;
                         } else {
                             tracing::debug!(
                                 "{} is a reserved field in the bunyan log format. Skipping it.",
diff --git a/tests/e2e.rs b/tests/e2e.rs
@@ -23,7 +23,13 @@ lazy_static! {
 fn run_and_get_raw_output<F: Fn()>(action: F) -> String {
     let mut default_fields = HashMap::new();
     default_fields.insert("custom_field".to_string(), json!("custom_value"));
-    let formatting_layer = BunyanFormattingLayer::with_default_fields("test".into(), || MockWriter::new(&BUFFER), default_fields);
+    let formatting_layer = BunyanFormattingLayer::with_default_fields(
+        "test".into(),
+        || MockWriter::new(&BUFFER),
+        default_fields,
+    )
+    .skip_fields(&["skipped"])
+    .unwrap();
     let subscriber = Registry::default()
         .with(JsonStorageLayer)
         .with(formatting_layer);
@@ -56,7 +62,8 @@ fn test_action() {
 
     info!("pre-shaving yaks");
     let b = 3;
-    let new_span = span!(Level::DEBUG, "inner shaving", b);
+    let skipped = false;
+    let new_span = span!(Level::DEBUG, "inner shaving", b, skipped);
     let _enter2 = new_span.enter();
 
     info!("shaving yaks");
@@ -160,3 +167,28 @@ fn elapsed_milliseconds_are_present_on_exit_span() {
         }
     }
 }
+
+#[test]
+fn skip_fields() {
+    let tracing_output = run_and_get_output(test_action);
+
+    for record in tracing_output {
+        assert!(record.get("skipped").is_none());
+    }
+}
+
+#[test]
+fn skipping_core_fields_is_not_allowed() {
+    let result = BunyanFormattingLayer::new("test".into(), || MockWriter::new(&BUFFER))
+        .skip_fields(&["level"]);
+
+    match result {
+        Err(err) => {
+            assert_eq!(
+                "level is a core field in the bunyan log format, it can't be skipped",
+                err.to_string()
+            );
+        }
+        Ok(_) => panic!("skipping core fields shouldn't work"),
+    }
+}
