Add filtering of parent properties

README.md

@@ -65,7 +65,7 @@ pub fn a_sub_unit_of_work(sub_parameter: u64) {
 fn main() {
     let formatting_layer = BunyanFormattingLayer::new("tracing_demo".into(), std::io::stdout);
     let subscriber = Registry::default()
-        .with(JsonStorageLayer)
+        .with(JsonStorageLayer::new())
         .with(formatting_layer);
     tracing::subscriber::set_global_default(subscriber).unwrap();
 

src/lib.rs

@@ -36,7 +36,7 @@
 //! fn main() {
 //!     let formatting_layer = BunyanFormattingLayer::new("tracing_demo".into(), std::io::stdout);
 //!     let subscriber = Registry::default()
-//!         .with(JsonStorageLayer)
+//!         .with(JsonStorageLayer::new())
 //!         .with(formatting_layer);
 //!     tracing::subscriber::set_global_default(subscriber).unwrap();
 //!
@@ -90,7 +90,9 @@
 //! [`tracing`]: https://docs.rs/tracing
 //! [`tracing`]: https://docs.rs/tracing-subscriber
 mod formatting_layer;
+mod storage_filter;
 mod storage_layer;
 
 pub use formatting_layer::*;
+pub use storage_filter::*;
 pub use storage_layer::*;

src/storage_filter.rs

@@ -0,0 +1,67 @@
+use std::collections::HashSet;
+
+use tracing_subscriber::registry::{LookupSpan, SpanRef};
+
+/// `FilteringMode` can either be `Include` when allowing only some particular fields
+/// to be passed down or `Exclude` when allowing all fields except the specified ones.
+#[derive(Clone, Debug)]
+pub enum FilteringMode {
+    Include,
+    Exclude,
+}
+
+/// `JsonStorageFilter` will filter the fields passed from the parent span to the child span.
+/// It can either be applied to all spans or to a particular span using its name.
+#[derive(Clone, Debug)]
+pub struct JsonStorageFilter {
+    span_name: Option<String>,
+    fields: HashSet<String>,
+    mode: FilteringMode,
+}
+
+impl JsonStorageFilter {
+    /// Create a new `JsonStorageFilter`.
+    ///
+    /// You have to specify:
+    /// - `fields`, which are the names of fields affected by this filter;
+    /// - `mode`, which is the mode of filtering for this filter.
+    pub fn new(fields: HashSet<String>, mode: FilteringMode) -> Self {
+        Self {
+            span_name: None,
+            fields,
+            mode,
+        }
+    }
+
+    /// Create a new `JsonStorageFilter`.
+    ///
+    /// You have to specify:
+    /// - `span_name`, which is the name of the span that will be affected by this filter;
+    /// - `fields`, which are the names of fields affected by this filter;
+    /// - `mode`, which is the mode of filtering for this filter.
+    pub fn for_span(span_name: String, fields: HashSet<String>, mode: FilteringMode) -> Self {
+        Self {
+            span_name: Some(span_name),
+            fields,
+            mode,
+        }
+    }
+
+    pub(crate) fn filter_span<S>(&self, span: &SpanRef<S>) -> bool
+    where
+        S: for<'a> LookupSpan<'a>,
+    {
+        if let Some(span_name) = &self.span_name {
+            span.name() == span_name
+        } else {
+            true
+        }
+    }
+
+    pub(crate) fn filter_field(&self, name: &str) -> bool {
+        match self.mode {
+            FilteringMode::Include => self.fields.contains(name),
+            FilteringMode::Exclude => !self.fields.contains(name),
+        }
+    }
+}

src/storage_layer.rs

@@ -1,3 +1,4 @@
+use crate::storage_filter::JsonStorageFilter;
 use std::collections::HashMap;
 use std::fmt;
 use std::time::Instant;
@@ -13,7 +14,9 @@ use tracing_subscriber::Layer;
 /// for downstream layers concerned with emitting a formatted representation of
 /// spans or events.
 #[derive(Clone, Debug)]
-pub struct JsonStorageLayer;
+pub struct JsonStorageLayer {
+    filters: Vec<JsonStorageFilter>,
+}
 
 /// `JsonStorage` will collect information about a span when it's created (`new_span` handler)
 /// or when new records are attached to it (`on_record` handler) and store it in its `extensions`
@@ -35,6 +38,11 @@ impl<'a> JsonStorage<'a> {
     pub fn values(&self) -> &HashMap<&'a str, serde_json::Value> {
         &self.values
     }
+
+    /// Create a new `JsonStorage` from values
+    pub fn from_values(values: HashMap<&'a str, serde_json::Value>) -> Self {
+        Self { values }
+    }
 }
 
 /// Get a new visitor, with an empty bag of key-value pairs.
@@ -94,6 +102,27 @@ impl Visit for JsonStorage<'_> {
     }
 }
 
+impl JsonStorageLayer {
+    /// Create a new `JsonStorageLayer`.
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Create a new `JsonStorageLayer` with the specified filters
+    pub fn with_filters(filters: Vec<JsonStorageFilter>) -> Self {
+        Self { filters }
+    }
+}
+
+/// Get a new layer, without filter.
+impl Default for JsonStorageLayer {
+    fn default() -> Self {
+        Self {
+            filters: Vec::new(),
+        }
+    }
+}
+
 impl<S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>> Layer<S>
     for JsonStorageLayer
 {
@@ -111,7 +140,24 @@ impl<S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>> Layer
             let mut extensions = parent_span.extensions_mut();
             extensions
                 .get_mut::<JsonStorage>()
-                .map(|v| v.to_owned())
+                .map(|v| {
+                    // Only use relevant filters for the given parent span.
+                    let filters: Vec<&JsonStorageFilter> = self
+                        .filters
+                        .iter()
+                        .filter(|f| f.filter_span(&parent_span))
+                        .collect();
+
+                    // Filter the fields and clone the ones that have been retained.
+                    let values: HashMap<&'_ str, serde_json::Value> = v
+                        .values()
+                        .iter()
+                        .filter(|(key, _value)| filters.iter().all(|f| f.filter_field(key)))
+                        .map(|(key, value)| (*key, value.to_owned()))
+                        .collect();
+
+                    JsonStorage::from_values(values)
+                })
                 .unwrap_or_default()
         } else {
             JsonStorage::default()

tests/e2e.rs

@@ -2,11 +2,13 @@ use crate::mock_writer::MockWriter;
 use claim::assert_some_eq;
 use lazy_static::lazy_static;
 use serde_json::{json, Value};
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::Mutex;
 use time::format_description::well_known::Rfc3339;
 use tracing::{info, span, Level};
-use tracing_bunyan_formatter::{BunyanFormattingLayer, JsonStorageLayer};
+use tracing_bunyan_formatter::{
+    BunyanFormattingLayer, FilteringMode, JsonStorageFilter, JsonStorageLayer,
+};
 use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::Registry;
 
@@ -20,12 +22,16 @@ lazy_static! {
 }
 
 // Run a closure and collect the output emitted by the tracing instrumentation using an in-memory buffer.
-fn run_and_get_raw_output<F: Fn()>(action: F) -> String {
+fn run_and_get_raw_output<F: Fn()>(action: F, filters: Vec<JsonStorageFilter>) -> 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,
+    );
     let subscriber = Registry::default()
-        .with(JsonStorageLayer)
+        .with(JsonStorageLayer::with_filters(filters))
         .with(formatting_layer);
     tracing::subscriber::with_default(subscriber, action);
 
@@ -40,7 +46,13 @@ fn run_and_get_raw_output<F: Fn()>(action: F) -> String {
 // Run a closure and collect the output emitted by the tracing instrumentation using
 // an in-memory buffer as structured new-line-delimited JSON.
 fn run_and_get_output<F: Fn()>(action: F) -> Vec<Value> {
-    run_and_get_raw_output(action)
+    run_filtered_and_get_output(action, Vec::new())
+}
+
+// Run a closure with storage filters and collect the output emitted by the tracing instrumentation
+// using an in-memory buffer as structured new-line-delimited JSON.
+fn run_filtered_and_get_output<F: Fn()>(action: F, filters: Vec<JsonStorageFilter>) -> Vec<Value> {
+    run_and_get_raw_output(action, filters)
         .lines()
         .filter(|&l| !l.is_empty())
         .inspect(|l| println!("{}", l))
@@ -51,20 +63,21 @@ fn run_and_get_output<F: Fn()>(action: F) -> Vec<Value> {
 // Instrumented code to be run to test the behaviour of the tracing instrumentation.
 fn test_action() {
     let a = 2;
-    let span = span!(Level::DEBUG, "shaving_yaks", a);
+    let b = 3;
+    let span = span!(Level::DEBUG, "shaving_yaks", a, b);
     let _enter = span.enter();
 
     info!("pre-shaving yaks");
-    let b = 3;
-    let new_span = span!(Level::DEBUG, "inner shaving", b);
+    let c = 4;
+    let new_span = span!(Level::DEBUG, "inner shaving", c);
     let _enter2 = new_span.enter();
 
     info!("shaving yaks");
 }
 
 #[test]
 fn each_line_is_valid_json() {
-    let tracing_output = run_and_get_raw_output(test_action);
+    let tracing_output = run_and_get_raw_output(test_action, Vec::new());
 
     // Each line is valid JSON
     for line in tracing_output.lines().filter(|&l| !l.is_empty()) {
@@ -160,3 +173,64 @@ fn elapsed_milliseconds_are_present_on_exit_span() {
         }
     }
 }
+
+#[test]
+fn filter_include_parent_field() {
+    let tracing_output = run_filtered_and_get_output(
+        test_action,
+        vec![JsonStorageFilter::new(
+            HashSet::from(["a".to_string()]),
+            FilteringMode::Include,
+        )],
+    );
+
+    for record in tracing_output {
+        assert!(record.get("a").is_some());
+        if record.get("c").is_some() {
+            assert!(record.get("b").is_none())
+        } else {
+            assert!(record.get("b").is_some())
+        }
+    }
+}
+
+#[test]
+fn filter_exclude_parent_field() {
+    let tracing_output = run_filtered_and_get_output(
+        test_action,
+        vec![JsonStorageFilter::new(
+            HashSet::from(["b".to_string()]),
+            FilteringMode::Exclude,
+        )],
+    );
+
+    for record in tracing_output {
+        assert!(record.get("a").is_some());
+        if record.get("c").is_some() {
+            assert!(record.get("b").is_none())
+        } else {
+            assert!(record.get("b").is_some())
+        }
+    }
+}
+
+#[test]
+fn filter_specific_span() {
+    let tracing_output = run_filtered_and_get_output(
+        test_action,
+        vec![JsonStorageFilter::for_span(
+            "shaving_yaks".to_string(),
+            HashSet::from(["a".to_string()]),
+            FilteringMode::Include,
+        )],
+    );
+
+    for record in tracing_output {
+        assert!(record.get("a").is_some());
+        if record.get("c").is_some() {
+            assert!(record.get("b").is_none())
+        } else {
+            assert!(record.get("b").is_some())
+        }
+    }
+}