Bug 2013658 - Register for multiple tracing events

CHANGELOG.md

@@ -21,6 +21,10 @@
   We don't believe this should affect any consumers, since they're were already using the
   `tracing-logging` feature and were either also using `tracing-reporting` or not handling error
   reporting at all.
+* Reworked `register_event_sink` signature to allow it to register an event sink for muliple targets at once.
+* Reworked `unregister_event_sink`.  It now inputs the return value from `register_event_sink`.
+* Removed `register_min_level_event_sink` and `unregister_min_level_event_sink`.
+  Use the new `register_event_sink` instead.
 
 ### Logins
 * Opened count method on logins for Android. ([#7207](https://github.com/mozilla/application-services/pull/7207/))

components/support/error/android/src/main/java/mozilla/appservices/errorsupport/RustComponentsErrorTelemetry.kt

@@ -9,6 +9,8 @@ import kotlinx.serialization.Serializable
 import kotlinx.serialization.decodeFromString
 import kotlinx.serialization.json.Json
 import mozilla.appservices.tracing.EventSink
+import mozilla.appservices.tracing.EventSinkSpecification
+import mozilla.appservices.tracing.EventTarget
 import mozilla.appservices.tracing.TracingEvent
 import mozilla.appservices.tracing.TracingLevel
 import mozilla.appservices.tracing.registerEventSink
@@ -28,7 +30,10 @@ public object RustComponentsErrorTelemetry {
      */
     fun register() {
         Glean.registerPings(Pings)
-        registerEventSink("app-services-error-reporter", TracingLevel.DEBUG, ErrorEventSink())
+        val spec = EventSinkSpecification(
+            targets = listOf(EventTarget("app-services-error-reporter", TracingLevel.DEBUG)),
+        )
+        registerEventSink(spec, ErrorEventSink())
     }
 }
 

components/support/rust-log-forwarder/src/lib.rs

@@ -4,6 +4,11 @@
 
 use std::sync::{Arc, Once, OnceLock};
 
+use tracing_support::{
+    register_event_sink, unregister_event_sink, EventSinkId, EventSinkSpecification,
+};
+
+static EVENT_SINK_ID: OnceLock<EventSinkId> = OnceLock::new();
 static MAX_LEVEL: OnceLock<Level> = OnceLock::new();
 static FOREIGN_LOGGER: OnceLock<Box<dyn AppServicesLogger>> = OnceLock::new();
 static GLOBAL_SUBSCRIBER: Once = Once::new();
@@ -73,11 +78,19 @@ pub fn set_logger(logger: Option<Box<dyn AppServicesLogger>>) {
     let sink = Arc::new(ForwarderEventSink {});
     if let Some(logger) = logger {
         // Set up a tracing subscriber for crates which use tracing and forward to the foreign log forwarder.
-        tracing_support::register_min_level_event_sink((*level).into(), sink.clone());
-        // Set the `FOREIGN_LOGGER` global.  If this was called before we just ignore the error.
+        let event_sink_id = register_event_sink(
+            EventSinkSpecification {
+                targets: vec![],
+                min_level: Some((*level).into()),
+            },
+            sink.clone(),
+        );
+        // Set the `FOREIGN_LOGGER` and `EVENT_SINK_ID` globals.
+        // If either was called before we just ignore the error.
         FOREIGN_LOGGER.set(logger).ok();
-    } else {
-        tracing_support::unregister_min_level_event_sink();
+        EVENT_SINK_ID.set(event_sink_id).ok();
+    } else if let Some(event_sink_id) = EVENT_SINK_ID.get() {
+        unregister_event_sink(*event_sink_id);
     }
 }
 

components/support/tracing/src/layer.rs

@@ -2,72 +2,79 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
-use parking_lot::RwLock;
-use std::collections::{BTreeMap, HashMap};
-use std::sync::{Arc, LazyLock};
+use parking_lot::{const_rwlock, RwLock};
+use std::collections::BTreeMap;
+use std::sync::{
+    atomic::{AtomicU32, Ordering},
+    Arc,
+};
 use tracing::{subscriber::Interest, Metadata};
 use tracing_subscriber::{
     layer::{Context, Filter},
     Layer,
 };
 
-use crate::EventSink;
+use crate::{EventSink, Level};
 use tracing::field::{Field, Visit};
 
-struct LogEntry {
-    level: tracing::Level,
-    sink: Arc<dyn EventSink>,
-}
+static SINKS: RwLock<Vec<RegisteredEventSink>> = const_rwlock(Vec::new());
+static EVENT_SINK_COUNTER: AtomicU32 = AtomicU32::new(0);
 
-static SINKS_BY_TARGET: LazyLock<RwLock<HashMap<String, LogEntry>>> =
-    LazyLock::new(|| RwLock::new(HashMap::new()));
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub struct EventSinkId(u32);
 
-static MIN_LEVEL_SINK: RwLock<Option<LogEntry>> = RwLock::new(None);
+uniffi::custom_type!(EventSinkId, u32, {
+    try_lift: |raw_id| Ok(EventSinkId(raw_id)),
+    lower: |sink_id| sink_id.0,
+});
 
-pub fn register_event_sink(target: &str, level: crate::Level, sink: Arc<dyn EventSink>) {
-    SINKS_BY_TARGET.write().insert(
-        target.to_string(),
-        LogEntry {
-            level: level.into(),
-            sink,
-        },
-    );
+/// Register an event sink using an [EventSinkSpecification]
+///
+/// Returns an [EventSinkId] that can be used to unregister the sink.
+pub fn register_event_sink(spec: EventSinkSpecification, sink: Arc<dyn EventSink>) -> EventSinkId {
+    let id = EventSinkId(EVENT_SINK_COUNTER.fetch_add(1, Ordering::Relaxed));
+    SINKS.write().push(RegisteredEventSink { id, spec, sink });
+    id
 }
 
-/// Register an event sink that will receive events based on a minimum level
-///
-/// If an event's level is at least `level`, then the event will be sent to this sink.
-/// If so, sinks registered with `register_event_sink` will still be processed.
-///
-/// There can only be 1 min-level sink registered at once.
-pub fn register_min_level_event_sink(level: crate::Level, sink: Arc<dyn EventSink>) {
-    *MIN_LEVEL_SINK.write() = Some(LogEntry {
-        level: level.into(),
-        sink,
-    });
+struct RegisteredEventSink {
+    // ID that can be used to unregister this sink
+    id: EventSinkId,
+    spec: EventSinkSpecification,
+    sink: Arc<dyn EventSink>,
 }
 
-#[uniffi::export]
-pub fn unregister_event_sink(target: &str) {
-    SINKS_BY_TARGET.write().remove(target);
+#[derive(uniffi::Record)]
+/// Describes which events to an EventSink
+pub struct EventSinkSpecification {
+    // Send events that match these targets/levels
+    #[uniffi(default)]
+    pub targets: Vec<EventTarget>,
+    // Send events have a `min_level` or above.
+    #[uniffi(default)]
+    pub min_level: Option<Level>,
+}
+
+#[derive(uniffi::Record, Debug)]
+pub struct EventTarget {
+    pub target: String,
+    pub level: Level,
 }
 
-/// Remove the sink registered with [register_min_level_event_sink], if any.
 #[uniffi::export]
-pub fn unregister_min_level_event_sink() {
-    *MIN_LEVEL_SINK.write() = None;
+pub fn unregister_event_sink(id: EventSinkId) {
+    SINKS.write().retain(|info| info.id != id);
 }
 
 // UniFFI versions of the registration functions.  This input a Box to be compatible with callback
 // interfaces
-#[uniffi::export(name = "register_event_sink")]
-pub fn register_event_sink_box(target: &str, level: crate::Level, sink: Box<dyn EventSink>) {
-    register_event_sink(target, level, sink.into())
-}
 
-#[uniffi::export(name = "register_min_level_event_sink")]
-pub fn register_min_level_event_sink_box(level: crate::Level, sink: Box<dyn EventSink>) {
-    register_min_level_event_sink(level, sink.into())
+#[uniffi::export(name = "register_event_sink")]
+pub fn register_event_sink_box(
+    targets: EventSinkSpecification,
+    sink: Box<dyn EventSink>,
+) -> EventSinkId {
+    register_event_sink(targets, sink.into())
 }
 
 pub fn simple_event_layer<S>() -> impl Layer<S>
@@ -88,43 +95,59 @@ where
         event: &tracing::Event<'_>,
         _ctx: tracing_subscriber::layer::Context<'_, S>,
     ) {
-        let target = event.metadata().target();
-        let prefix = match target.find(':') {
-            Some(index) => &target[..index],
-            None => target,
-        };
-        if let Some(entry) = &*MIN_LEVEL_SINK.read() {
-            if entry.level >= *event.metadata().level() {
-                entry.send_event(event);
-            }
+        let sinks = find_sinks_for_event(event);
+        if sinks.is_empty() {
+            // return early to skip the conversion below
+            return;
         }
-
-        if let Some(entry) = SINKS_BY_TARGET.read().get(prefix) {
-            let level = *event.metadata().level();
-            if level <= entry.level {
-                entry.send_event(event);
-            }
-        }
-    }
-}
-
-impl LogEntry {
-    fn send_event(&self, event: &tracing::Event<'_>) {
         let mut fields = BTreeMap::new();
         let mut message = String::default();
         let mut visitor = JsonVisitor(&mut message, &mut fields);
         event.record(&mut visitor);
-        let event = crate::Event {
+        let tracing_event = crate::Event {
             level: (*event.metadata().level()).into(),
             target: event.metadata().target().to_string(),
             name: event.metadata().name().to_string(),
             message,
             fields: serde_json::to_value(&fields).unwrap_or_default(),
         };
-        self.sink.on_event(event);
+        for sink in sinks {
+            sink.on_event(tracing_event.clone());
+        }
     }
 }
 
+/// Find event sinks that match an event.
+fn find_sinks_for_event(event: &tracing::Event<'_>) -> Vec<Arc<dyn EventSink>> {
+    let target = event.metadata().target();
+    let prefix = match target.find(':') {
+        Some(index) => &target[..index],
+        None => target,
+    };
+    let level = Level::from(*event.metadata().level());
+
+    // This requires a iterating through the entire SINKS vec, which could have performance impacts
+    // if we have many sinks registered.  In practice, there should only be a handful of these so
+    // this should be fine.
+    SINKS
+        .read()
+        .iter()
+        .filter_map(|info| {
+            if let Some(min_level) = &info.spec.min_level {
+                if *min_level >= level {
+                    return Some(info.sink.clone());
+                }
+            }
+            for target in info.spec.targets.iter() {
+                if target.target == prefix && target.level >= level {
+                    return Some(info.sink.clone());
+                }
+            }
+            None
+        })
+        .collect()
+}
+
 struct SimpleEventFilter;
 
 impl SimpleEventFilter {

components/support/tracing/src/lib.rs

@@ -12,8 +12,8 @@ pub use filters::{
 };
 
 pub use layer::{
-    register_event_sink, register_min_level_event_sink, simple_event_layer, unregister_event_sink,
-    unregister_min_level_event_sink,
+    register_event_sink, simple_event_layer, unregister_event_sink, EventSinkId,
+    EventSinkSpecification, EventTarget,
 };
 // Re-export tracing so that our dependencies can use it.
 pub use tracing;
@@ -104,7 +104,7 @@ macro_rules! error {
 pub type Level = TracingLevel;
 
 // `tracing::Level` is a struct, we want an enum for both uniffi and `log::Level`` compat.
-#[derive(uniffi::Enum, Copy, Clone, Debug, PartialEq, Eq)]
+#[derive(uniffi::Enum, Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
 pub enum TracingLevel {
     Error,
     Warn,
@@ -146,7 +146,7 @@ impl From<Level> for tracing::Level {
 // ditto grr re swift name collisions.
 pub type Event = TracingEvent;
 
-#[derive(uniffi::Record, Debug)]
+#[derive(uniffi::Record, Clone, Debug, PartialEq, Eq)]
 pub struct TracingEvent {
     pub level: Level,
     pub target: String,
@@ -205,26 +205,37 @@ mod tests {
                 self.events.write().push(event);
             }
         }
-        let sink = Arc::new(Sink::new());
-        let level_sink = Arc::new(Sink::new());
 
-        crate::layer::register_event_sink("first_target", Level::Info, sink.clone());
-        crate::layer::register_event_sink("second_target", Level::Debug, sink.clone());
+        let sink = Arc::new(Sink::new());
 
-        // Only 1 sink can be registered with `register_min_level_event_sink`.  The first call
-        // should be ignored and only the second call should take effect.
-        crate::layer::register_min_level_event_sink(Level::Warn, sink.clone());
-        crate::layer::register_min_level_event_sink(Level::Error, level_sink.clone());
+        // Register an event sink
+        crate::layer::register_event_sink(
+            EventSinkSpecification {
+                targets: vec![
+                    EventTarget {
+                        target: "first_target".into(),
+                        level: Level::Info,
+                    },
+                    EventTarget {
+                        target: "second_target".into(),
+                        level: Level::Debug,
+                    },
+                ],
+                min_level: Some(Level::Warn),
+            },
+            sink.clone(),
+        );
 
         info!(target: "first_target", extra=-1, "event message");
         debug!(target: "first_target", extra=-2, "event message (should be filtered)");
         debug!(target: "second_target::submodule", extra=-3, "event message2");
-        info!(target: "third_target", extra=-4, "event message (should be filtered)");
-        // This should go to both the level sink and target sink, since it's an error
-        error!(target: "first_target", extra=-5, "event message");
+        warn!(target: "third_target", extra=-4, "event message3");
+        info!(target: "third_target", extra=-5, "event message (should be filtered)");
+        // Matches both for first_target and for the min-level, it should only be received once by
+        // the sink
+        error!(target: "first_target", extra=-6, "event message4");
 
-        assert_eq!(sink.events.read().len(), 3);
-        assert_eq!(level_sink.events.read().len(), 1);
+        assert_eq!(sink.events.read().len(), 4);
 
         let event = &sink.events.read()[0];
         assert_eq!(event.target, "first_target");
@@ -239,15 +250,15 @@ mod tests {
         assert_eq!(event2.fields.get("extra").unwrap().as_i64(), Some(-3));
 
         let event3 = &sink.events.read()[2];
-        assert_eq!(event3.target, "first_target");
-        assert_eq!(event3.level, Level::Error);
-        assert_eq!(event3.message, "event message");
-        assert_eq!(event3.fields.get("extra").unwrap().as_i64(), Some(-5));
-
-        let level_event = &level_sink.events.read()[0];
-        assert_eq!(level_event.target, "first_target");
-        assert_eq!(level_event.level, Level::Error);
-        assert_eq!(level_event.message, "event message");
-        assert_eq!(level_event.fields.get("extra").unwrap().as_i64(), Some(-5));
+        assert_eq!(event3.target, "third_target");
+        assert_eq!(event3.level, Level::Warn);
+        assert_eq!(event3.message, "event message3");
+        assert_eq!(event3.fields.get("extra").unwrap().as_i64(), Some(-4));
+
+        let event4 = &sink.events.read()[3];
+        assert_eq!(event4.target, "first_target");
+        assert_eq!(event4.level, Level::Error);
+        assert_eq!(event4.message, "event message4");
+        assert_eq!(event4.fields.get("extra").unwrap().as_i64(), Some(-6));
     }
 }