v0.1.2-pre.006

crates/ksp-logging-lib/tests/tokio_span.rs

@@ -0,0 +1,118 @@
+// file: crates/ksp-logging-lib/tests/tokio_span.rs
+// version: 1
+
+//! Integration tests for KSP span instrumentation on a real Tokio executor.
+
+const TEST_TARGET: &str = "ksp-logging-lib";
+
+#[derive(Clone)]
+struct CountingSubscriber {
+    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 CountingSubscriber {
+    fn new(enters: std::sync::Arc<std::sync::atomic::AtomicU64>, exits: std::sync::Arc<std::sync::atomic::AtomicU64>) -> Self {
+        return Self { enters, exits, next_id: std::sync::Arc::new(std::sync::atomic::AtomicU64::new(1)) };
+    }
+}
+
+impl tracing::Subscriber for CountingSubscriber {
+    fn enabled(&self, _metadata: &tracing::Metadata<'_>) -> bool {
+        return true;
+    }
+
+    fn new_span(&self, _span: &tracing::span::Attributes<'_>) -> tracing::span::Id {
+        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<'_>) {
+        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 test_span(enters: std::sync::Arc<std::sync::atomic::AtomicU64>, exits: std::sync::Arc<std::sync::atomic::AtomicU64>) -> ksp_logging_lib::Span {
+    let subscriber = CountingSubscriber::new(enters, exits);
+    return tracing::subscriber::with_default(subscriber, || -> ksp_logging_lib::Span {
+        return ksp_logging_lib::trace_span!(target: TEST_TARGET, "tokio_runtime_span", domain = "logging", executor = "tokio");
+    });
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn instrumented_span_reenters_across_real_tokio_suspensions() {
+    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 span = test_span(std::sync::Arc::clone(&enters), std::sync::Arc::clone(&exits));
+    let observed_enters = std::sync::Arc::clone(&enters);
+    let future = ksp_logging_lib::instrument(span, async move {
+        assert!(observed_enters.load(std::sync::atomic::Ordering::Relaxed) >= 1);
+        tokio::task::yield_now().await;
+        assert!(observed_enters.load(std::sync::atomic::Ordering::Relaxed) >= 2);
+        tokio::task::yield_now().await;
+        assert!(observed_enters.load(std::sync::atomic::Ordering::Relaxed) >= 3);
+        return 42_u32;
+    });
+    let value = future.await;
+    assert_eq!(value, 42_u32);
+    let enter_count = enters.load(std::sync::atomic::Ordering::Relaxed);
+    let exit_count = exits.load(std::sync::atomic::Ordering::Relaxed);
+    assert!(enter_count >= 3);
+    assert_eq!(enter_count, exit_count);
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn instrumented_spans_are_usable_on_tokio_multithread_runtime() {
+    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 first_span = test_span(std::sync::Arc::clone(&enters), std::sync::Arc::clone(&exits));
+    let second_span = test_span(std::sync::Arc::clone(&enters), std::sync::Arc::clone(&exits));
+    let first_task = tokio::spawn(ksp_logging_lib::instrument(first_span, async {
+        for _iteration in 0..32 {
+            tokio::task::yield_now().await;
+        }
+        return 20_u32;
+    }));
+    let second_task = tokio::spawn(ksp_logging_lib::instrument(second_span, async {
+        for _iteration in 0..32 {
+            tokio::task::yield_now().await;
+        }
+        return 22_u32;
+    }));
+    let first_result = first_task.await;
+    assert!(first_result.is_ok(), "first Tokio task must complete successfully");
+    let first_value = match first_result {
+        std::result::Result::Ok(value) => value,
+        std::result::Result::Err(_) => return,
+    };
+    let second_result = second_task.await;
+    assert!(second_result.is_ok(), "second Tokio task must complete successfully");
+    let second_value = match second_result {
+        std::result::Result::Ok(value) => value,
+        std::result::Result::Err(_) => return,
+    };
+    assert_eq!(first_value + second_value, 42_u32);
+    let enter_count = enters.load(std::sync::atomic::Ordering::Relaxed);
+    let exit_count = exits.load(std::sync::atomic::Ordering::Relaxed);
+    assert!(enter_count >= 4);
+    assert_eq!(enter_count, exit_count);
+}