khadhroony-solana-project/crates/ksp-logging-lib/tests/callsite.rs

// file: crates/ksp-logging-lib/tests/callsite.rs
// version: 2

//! Integration tests for KSP logging callsite and async span instrumentation behavior.

const TEST_TARGET: &str = "ksp-logging-lib";

#[derive(Clone, Debug, Eq, PartialEq)]
struct CapturedMetadata {
    target: std::string::String,
    file: std::option::Option<std::string::String>,
    module_path: std::option::Option<std::string::String>,
    line: std::option::Option<u32>,
    is_event: bool,
    is_span: bool,
}

impl CapturedMetadata {
    fn from_metadata(metadata: &tracing::Metadata<'_>) -> Self {
        return Self {
            target: metadata.target().to_owned(),
            file: metadata.file().map(str::to_owned),
            module_path: metadata.module_path().map(str::to_owned),
            line: metadata.line(),
            is_event: metadata.is_event(),
            is_span: metadata.is_span(),
        };
    }
}

#[derive(Clone)]
struct CaptureSubscriber {
    captured: std::sync::Arc<std::sync::Mutex<std::vec::Vec<CapturedMetadata>>>,
    enters: std::sync::Arc<std::sync::atomic::AtomicU64>,
    exits: std::sync::Arc<std::sync::atomic::AtomicU64>,
    next_id: std::sync::Arc<std::sync::atomic::AtomicU64>,
}

impl CaptureSubscriber {
    fn new(
        captured: std::sync::Arc<std::sync::Mutex<std::vec::Vec<CapturedMetadata>>>,
        enters: std::sync::Arc<std::sync::atomic::AtomicU64>,
        exits: std::sync::Arc<std::sync::atomic::AtomicU64>,
    ) -> Self {
        return Self { captured, enters, exits, next_id: std::sync::Arc::new(std::sync::atomic::AtomicU64::new(1)) };
    }

    fn capture(&self, metadata: &tracing::Metadata<'_>) {
        let lock = self.captured.lock();
        if let std::result::Result::Ok(mut values) = lock {
            values.push(CapturedMetadata::from_metadata(metadata));
        }
    }
}

impl tracing::Subscriber for CaptureSubscriber {
    fn enabled(&self, _metadata: &tracing::Metadata<'_>) -> bool {
        return true;
    }

    fn new_span(&self, span: &tracing::span::Attributes<'_>) -> tracing::span::Id {
        self.capture(span.metadata());
        let id = self.next_id.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        return tracing::span::Id::from_u64(id);
    }

    fn record(&self, _span: &tracing::span::Id, _values: &tracing::span::Record<'_>) {
        return;
    }

    fn record_follows_from(&self, _span: &tracing::span::Id, _follows: &tracing::span::Id) {
        return;
    }

    fn event(&self, event: &tracing::Event<'_>) {
        self.capture(event.metadata());
        return;
    }

    fn enter(&self, _span: &tracing::span::Id) {
        self.enters.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        return;
    }

    fn exit(&self, _span: &tracing::span::Id) {
        self.exits.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        return;
    }
}

fn captured_values(captured: &std::sync::Arc<std::sync::Mutex<std::vec::Vec<CapturedMetadata>>>) -> std::vec::Vec<CapturedMetadata> {
    let lock = captured.lock();
    return match lock {
        std::result::Result::Ok(values) => values.clone(),
        std::result::Result::Err(error) => error.into_inner().clone(),
    };
}

#[test]
fn event_macro_preserves_consumer_callsite() {
    let captured = std::sync::Arc::new(std::sync::Mutex::new(std::vec::Vec::new()));
    let enters = std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0));
    let exits = std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0));
    let subscriber = CaptureSubscriber::new(captured.clone(), enters, exits);
    let expected_line = line!() + 2;
    tracing::subscriber::with_default(subscriber, || {
        ksp_logging_lib::info!(target: TEST_TARGET, domain = "logging", "callsite event");
        return;
    });
    let values = captured_values(&captured);
    assert_eq!(values.len(), 1);
    assert_eq!(values[0].target, TEST_TARGET);
    assert_eq!(values[0].file.as_deref(), std::option::Option::Some(file!()));
    assert_eq!(values[0].module_path.as_deref(), std::option::Option::Some(module_path!()));
    assert_eq!(values[0].line, std::option::Option::Some(expected_line));
    assert!(values[0].is_event);
    assert!(!values[0].is_span);
}

#[test]
fn span_macro_preserves_consumer_callsite() {
    let captured = std::sync::Arc::new(std::sync::Mutex::new(std::vec::Vec::new()));
    let enters = std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0));
    let exits = std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0));
    let subscriber = CaptureSubscriber::new(captured.clone(), enters, exits);
    let expected_line = line!() + 2;
    tracing::subscriber::with_default(subscriber, || {
        let _span = ksp_logging_lib::trace_span!(target: TEST_TARGET, "callsite_span", component = "test");
        return;
    });
    let values = captured_values(&captured);
    assert_eq!(values.len(), 1);
    assert_eq!(values[0].target, TEST_TARGET);
    assert_eq!(values[0].file.as_deref(), std::option::Option::Some(file!()));
    assert_eq!(values[0].module_path.as_deref(), std::option::Option::Some(module_path!()));
    assert_eq!(values[0].line, std::option::Option::Some(expected_line));
    assert!(!values[0].is_event);
    assert!(values[0].is_span);
}

#[test]
fn async_instrumentation_enters_and_exits_span_during_poll_and_drop() {
    let captured = std::sync::Arc::new(std::sync::Mutex::new(std::vec::Vec::new()));
    let enters = std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0));
    let exits = std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0));
    let subscriber = CaptureSubscriber::new(captured, std::sync::Arc::clone(&enters), std::sync::Arc::clone(&exits));
    tracing::subscriber::with_default(subscriber, || {
        let span = ksp_logging_lib::trace_span!(target: TEST_TARGET, "async_poll_span", domain = "logging");
        let future = ksp_logging_lib::instrument(span, std::future::ready(42_u32));
        let mut future = std::boxed::Box::pin(future);
        let waker = std::task::Waker::noop();
        let mut context = std::task::Context::from_waker(waker);
        let poll = std::future::Future::poll(future.as_mut(), &mut context);
        assert_eq!(poll, std::task::Poll::Ready(42_u32));
        assert_eq!(enters.load(std::sync::atomic::Ordering::Relaxed), 1);
        assert_eq!(exits.load(std::sync::atomic::Ordering::Relaxed), 1);
        std::mem::drop(future);
        assert_eq!(enters.load(std::sync::atomic::Ordering::Relaxed), 2);
        assert_eq!(exits.load(std::sync::atomic::Ordering::Relaxed), 2);
        return;
    });
    assert_eq!(enters.load(std::sync::atomic::Ordering::Relaxed), exits.load(std::sync::atomic::Ordering::Relaxed));
}