rustrails_support/

log_subscriber.rs

use crate::notifications;
use once_cell::sync::{Lazy, OnceCell};
use parking_lot::RwLock;
use serde_json::Value;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};

/// A notification event adapted for log subscribers.
#[derive(Debug, Clone, PartialEq)]
pub struct LogEvent {
    /// The event name.
    pub name: String,
    /// When the event started.
    pub started_at: Instant,
    /// How long the event took.
    pub duration: Duration,
    /// Structured payload attached to the event.
    pub payload: HashMap<String, Value>,
}

/// Receives adapted notification events.
pub trait LogSubscriber: Send + Sync {
    /// Handles a delivered log event.
    fn call(&self, event: &LogEvent);
}

static SUBSCRIBERS: Lazy<RwLock<Vec<Arc<dyn LogSubscriber>>>> =
    Lazy::new(|| RwLock::new(Vec::new()));
static BRIDGE_SUBSCRIPTION: OnceCell<usize> = OnceCell::new();

/// Attaches a subscriber to the global notifications bridge.
pub fn attach<S>(subscriber: S)
where
    S: LogSubscriber + 'static,
{
    ensure_bridge();
    SUBSCRIBERS.write().push(Arc::new(subscriber));
}

fn ensure_bridge() {
    BRIDGE_SUBSCRIPTION.get_or_init(|| {
        notifications::subscribe(
            "*",
            Box::new(|event| {
                let log_event = LogEvent {
                    name: event.name.clone(),
                    started_at: event.time,
                    duration: event.duration,
                    payload: event.payload.clone(),
                };
                let subscribers = SUBSCRIBERS.read().clone();
                for subscriber in subscribers {
                    subscriber.call(&log_event);
                }
            }),
        )
    });
}

#[cfg(test)]
pub(crate) fn reset_subscribers() {
    SUBSCRIBERS.write().clear();
}

#[cfg(test)]
mod tests {
    use super::{LogEvent, LogSubscriber, attach, reset_subscribers};
    use crate::notifications;
    use parking_lot::Mutex;
    use serde_json::json;
    use std::collections::HashMap;
    use std::sync::{Arc, LazyLock, Mutex as StdMutex};
    use std::time::Duration;

    static TEST_LOCK: LazyLock<StdMutex<()>> = LazyLock::new(|| StdMutex::new(()));

    #[derive(Default)]
    struct RecordingSubscriber {
        events: Arc<Mutex<Vec<LogEvent>>>,
    }

    impl RecordingSubscriber {
        fn new() -> Self {
            Self::default()
        }

        fn handle(&self) -> Arc<Mutex<Vec<LogEvent>>> {
            Arc::clone(&self.events)
        }
    }

    impl LogSubscriber for RecordingSubscriber {
        fn call(&self, event: &LogEvent) {
            self.events.lock().push(event.clone());
        }
    }

    #[test]
    fn attach_receives_instrumented_notifications() {
        let _guard = TEST_LOCK.lock().expect("test lock should not be poisoned");
        reset_subscribers();
        let subscriber = RecordingSubscriber::new();
        let events = subscriber.handle();
        attach(subscriber);

        notifications::instrument(
            "log_subscriber.attach_receives_instrumented_notifications",
            HashMap::new(),
            || (),
        );

        let matching = events
            .lock()
            .iter()
            .filter(|event| {
                event.name == "log_subscriber.attach_receives_instrumented_notifications"
            })
            .count();
        assert_eq!(matching, 1);
    }

    #[test]
    fn payload_is_forwarded_to_subscribers() {
        let _guard = TEST_LOCK.lock().expect("test lock should not be poisoned");
        reset_subscribers();
        let subscriber = RecordingSubscriber::new();
        let events = subscriber.handle();
        attach(subscriber);

        notifications::instrument(
            "log_subscriber.payload_is_forwarded_to_subscribers",
            HashMap::from([(String::from("sql"), json!("select 1"))]),
            || (),
        );

        let event = events
            .lock()
            .iter()
            .find(|event| event.name == "log_subscriber.payload_is_forwarded_to_subscribers")
            .cloned()
            .expect("named event should be captured");
        assert_eq!(event.payload.get("sql"), Some(&json!("select 1")));
    }

    #[test]
    fn duration_is_forwarded_to_subscribers() {
        let _guard = TEST_LOCK.lock().expect("test lock should not be poisoned");
        reset_subscribers();
        let subscriber = RecordingSubscriber::new();
        let events = subscriber.handle();
        attach(subscriber);

        notifications::instrument(
            "log_subscriber.duration_is_forwarded_to_subscribers",
            HashMap::new(),
            || {
                std::thread::sleep(Duration::from_millis(10));
            },
        );

        let event = events
            .lock()
            .iter()
            .find(|event| event.name == "log_subscriber.duration_is_forwarded_to_subscribers")
            .cloned()
            .expect("named event should be captured");
        assert!(event.duration >= Duration::from_millis(10));
    }

    #[test]
    fn multiple_attached_subscribers_receive_the_same_event() {
        let _guard = TEST_LOCK.lock().expect("test lock should not be poisoned");
        reset_subscribers();
        let first = RecordingSubscriber::new();
        let second = RecordingSubscriber::new();
        let first_events = first.handle();
        let second_events = second.handle();
        attach(first);
        attach(second);

        notifications::instrument(
            "log_subscriber.multiple_attached_subscribers_receive_the_same_event",
            HashMap::new(),
            || (),
        );

        let first_matching = first_events
            .lock()
            .iter()
            .filter(|event| {
                event.name == "log_subscriber.multiple_attached_subscribers_receive_the_same_event"
            })
            .count();
        let second_matching = second_events
            .lock()
            .iter()
            .filter(|event| {
                event.name == "log_subscriber.multiple_attached_subscribers_receive_the_same_event"
            })
            .count();
        assert_eq!(first_matching, 1);
        assert_eq!(second_matching, 1);
    }

    #[test]
    fn started_at_is_recorded() {
        let _guard = TEST_LOCK.lock().expect("test lock should not be poisoned");
        reset_subscribers();
        let subscriber = RecordingSubscriber::new();
        let events = subscriber.handle();
        attach(subscriber);

        notifications::instrument(
            "log_subscriber.started_at_is_recorded",
            HashMap::new(),
            || (),
        );

        let event = events
            .lock()
            .iter()
            .find(|event| event.name == "log_subscriber.started_at_is_recorded")
            .cloned()
            .expect("named event should be captured");
        assert!(event.started_at <= std::time::Instant::now());
    }

    #[test]
    fn reset_subscribers_clears_attached_subscribers() {
        let _guard = TEST_LOCK.lock().expect("test lock should not be poisoned");
        reset_subscribers();
        let subscriber = RecordingSubscriber::new();
        let events = subscriber.handle();
        attach(subscriber);
        reset_subscribers();

        notifications::instrument(
            "log_subscriber.reset_subscribers_clears_attached_subscribers",
            HashMap::new(),
            || (),
        );

        let matching = events
            .lock()
            .iter()
            .filter(|event| {
                event.name == "log_subscriber.reset_subscribers_clears_attached_subscribers"
            })
            .count();
        assert_eq!(matching, 0);
    }
}