119 lines
4.8 KiB
Rust
119 lines
4.8 KiB
Rust
// 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);
|
|
}
|