From 9d905be0e61b07ae82bb6ed25320792843fb622a Mon Sep 17 00:00:00 2001 From: Alexey Orlenko Date: Tue, 19 Nov 2024 18:23:30 +0100 Subject: [PATCH] Add tests for layer --- Cargo.lock | 27 ++- libs/telemetry/Cargo.toml | 3 + libs/telemetry/src/capturing/ng/collector.rs | 51 ++-- libs/telemetry/src/capturing/ng/layer.rs | 241 ++++++++++++++++++- 4 files changed, 284 insertions(+), 38 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d35a87bf358..531478fc83e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2033,18 +2033,18 @@ dependencies = [ [[package]] name = "insta" -version = "1.34.0" +version = "1.41.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5d64600be34b2fcfc267740a243fa7744441bb4947a619ac4e5bb6507f35fbfc" +checksum = "7e9ffc4d4892617c50a928c52b2961cb5174b6fc6ebf252b2fac9d21955c48b8" dependencies = [ "console", "lazy_static", "linked-hash-map", "pest", "pest_derive", + "ron", "serde", "similar", - "yaml-rust", ] [[package]] @@ -4472,6 +4472,17 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "ron" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88073939a61e5b7680558e6be56b419e208420c2adb92be54921fa6b72283f1a" +dependencies = [ + "base64 0.13.1", + "bitflags 1.3.2", + "serde", +] + [[package]] name = "rusqlite" version = "0.31.0" @@ -5462,6 +5473,7 @@ dependencies = [ "enumflags2", "futures", "indexmap 2.2.2", + "insta", "itertools 0.12.0", "lru 0.7.8", "once_cell", @@ -6779,15 +6791,6 @@ dependencies = [ "tap", ] -[[package]] -name = "yaml-rust" -version = "0.4.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "56c1936c4cc7a1c9ab21a1ebb602eb942ba868cbd44a99cb7cdc5892335e1c85" -dependencies = [ - "linked-hash-map", -] - [[package]] name = "yansi" version = "0.5.1" diff --git a/libs/telemetry/Cargo.toml b/libs/telemetry/Cargo.toml index 070b927f8f3..501f0333afa 100644 --- a/libs/telemetry/Cargo.toml +++ b/libs/telemetry/Cargo.toml @@ -30,3 +30,6 @@ crosstarget-utils = { path = "../crosstarget-utils" } lru = "0.7.7" enumflags2.workspace = true derive_more = "0.99.17" + +[dev-dependencies] +insta = { version = "1.41", features = ["redactions", "ron"] } diff --git a/libs/telemetry/src/capturing/ng/collector.rs b/libs/telemetry/src/capturing/ng/collector.rs index 6651622e67e..504ffe9ef37 100644 --- a/libs/telemetry/src/capturing/ng/collector.rs +++ b/libs/telemetry/src/capturing/ng/collector.rs @@ -1,36 +1,55 @@ -use std::{borrow::Cow, collections::HashMap}; +use std::{borrow::Cow, collections::HashMap, num::NonZeroU64}; +use serde::Serialize; use tokio::time::Instant; -use tracing::{span::Id, Level}; +use tracing::Level; use crate::models::{HrTime, SpanKind, TraceSpan}; -#[derive(Debug)] +#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize)] +pub struct SpanId(NonZeroU64); + +impl From<&tracing::span::Id> for SpanId { + fn from(id: &tracing::span::Id) -> Self { + Self(id.into_non_zero_u64()) + } +} + +impl From for SpanId { + fn from(id: tracing::span::Id) -> Self { + Self::from(&id) + } +} + +#[derive(Debug, Clone)] +#[cfg_attr(test, derive(serde::Serialize))] pub struct CollectedSpan { - id: Id, - parent_id: Option, + id: SpanId, + parent_id: Option, name: Cow<'static, str>, + #[cfg_attr(test, serde(skip_serializing))] start_time: Instant, + #[cfg_attr(test, serde(skip_serializing))] end_time: Instant, attributes: HashMap<&'static str, serde_json::Value>, kind: SpanKind, - links: Vec, + links: Vec, } pub(crate) struct SpanBuilder { - id: Id, + id: SpanId, name: Cow<'static, str>, start_time: Instant, end_time: Option, attributes: HashMap<&'static str, serde_json::Value>, kind: Option, - links: Vec, + links: Vec, } impl SpanBuilder { - pub fn new(name: &'static str, id: Id, start_time: Instant, attrs_size_hint: usize) -> Self { + pub fn new(name: &'static str, id: impl Into, start_time: Instant, attrs_size_hint: usize) -> Self { Self { - id, + id: id.into(), name: name.into(), start_time, end_time: None, @@ -52,14 +71,14 @@ impl SpanBuilder { self.attributes.insert(key, value); } - pub fn add_link(&mut self, link: Id) { + pub fn add_link(&mut self, link: SpanId) { self.links.push(link); } - pub fn end(self, parent_id: Option, end_time: Instant) -> CollectedSpan { + pub fn end(self, parent_id: Option>, end_time: Instant) -> CollectedSpan { CollectedSpan { id: self.id, - parent_id, + parent_id: parent_id.map(Into::into), name: self.name, start_time: self.start_time, end_time, @@ -70,7 +89,7 @@ impl SpanBuilder { } } -#[derive(Debug)] +#[derive(Debug, Clone)] pub(crate) struct CollectedEvent { name: &'static str, level: Level, @@ -79,7 +98,7 @@ pub(crate) struct CollectedEvent { } pub trait Collector { - fn add_span(&self, trace: Id, span: CollectedSpan); + fn add_span(&self, trace: SpanId, span: CollectedSpan); } pub struct Exporter {} @@ -97,7 +116,7 @@ impl Default for Exporter { } impl Collector for Exporter { - fn add_span(&self, _trace: Id, _span: CollectedSpan) { + fn add_span(&self, _trace: SpanId, _span: CollectedSpan) { todo!() } } diff --git a/libs/telemetry/src/capturing/ng/layer.rs b/libs/telemetry/src/capturing/ng/layer.rs index 72ae3537158..94a75ad62d0 100644 --- a/libs/telemetry/src/capturing/ng/layer.rs +++ b/libs/telemetry/src/capturing/ng/layer.rs @@ -68,7 +68,7 @@ where { fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) { let span = Self::require_span(id, &ctx); - let mut span_builder = SpanBuilder::new(span.name(), id.to_owned(), Instant::now(), attrs.fields().len()); + let mut span_builder = SpanBuilder::new(span.name(), id, Instant::now(), attrs.fields().len()); attrs.record(&mut SpanAttributeVisitor::new(&mut span_builder)); @@ -89,7 +89,7 @@ where let mut extensions = span.extensions_mut(); if let Some(span_builder) = extensions.get_mut::() { - span_builder.add_link(follows.to_owned()); + span_builder.add_link(follows.into()); } } @@ -101,17 +101,18 @@ where fn on_close(&self, id: Id, ctx: Context<'_, S>) { let span = Self::require_span(&id, &ctx); - let mut extensions = span.extensions_mut(); - if let Some(span_builder) = extensions.remove::() { - let end_time = Instant::now(); - let parent_id = span.parent().map(|parent| parent.id()); - let collected_span = span_builder.end(parent_id, end_time); + let Some(span_builder) = span.extensions_mut().remove::() else { + return; + }; - let trace_id = Self::root_span(&id, &ctx).id(); + let end_time = Instant::now(); + let parent_id = span.parent().map(|parent| parent.id()); + let collected_span = span_builder.end(parent_id, end_time); - self.collector.add_span(trace_id, collected_span); - } + let trace_id = Self::root_span(&id, &ctx).id(); + + self.collector.add_span(trace_id.into(), collected_span); } } @@ -154,3 +155,223 @@ impl<'a> field::Visit for SpanAttributeVisitor<'a> { self.record_str(field, &format!("{:?}", value)) } } + +#[cfg(test)] +mod tests { + use crate::capturing::ng::collector::{CollectedSpan, SpanId}; + + use super::*; + + use std::cell::RefCell; + use std::collections::{BTreeMap, HashMap}; + use std::sync::atomic::{AtomicU64, Ordering}; + use std::sync::{Arc, LazyLock, Mutex}; + use std::time::Duration; + + use assert_ron_snapshot; + use insta::internals::{Content, Redaction}; + use tracing::{info_span, span, Level}; + use tracing_subscriber::layer::SubscriberExt; + use tracing_subscriber::Registry; + + #[derive(Debug, Default, Clone)] + struct TestCollector { + spans: Arc>>>, + } + + impl TestCollector { + fn new() -> Self { + Self::default() + } + + fn get_spans(&self) -> BTreeMap> { + self.spans.lock().unwrap().clone() + } + } + + impl Collector for TestCollector { + fn add_span(&self, trace_id: SpanId, span: CollectedSpan) { + let mut spans = self.spans.lock().unwrap(); + spans.entry(trace_id).or_default().push(span); + } + } + + fn redact_id() -> Redaction { + thread_local! { + static SPAN_ID_TO_SEQUENTIAL_ID: RefCell> = <_>::default(); + static NEXT_ID: RefCell = const { RefCell::new(1) }; + } + + insta::dynamic_redaction(|value, _path| match value { + Content::NewtypeStruct("SpanId", ref nested) => match **nested { + Content::U64(original_id) => SPAN_ID_TO_SEQUENTIAL_ID.with_borrow_mut(|map| { + let id = map.entry(original_id).or_insert_with(|| { + NEXT_ID.with_borrow_mut(|next_id| { + let id = *next_id; + *next_id += 1; + id + }) + }); + Content::NewtypeStruct("SpanId", Box::new(Content::U64(*id))) + }), + _ => value, + }, + _ => value, + }) + } + + #[test] + fn test_basic_span_collection() { + let collector = TestCollector::new(); + let subscriber = Registry::default().with(layer(collector.clone())); + + tracing::subscriber::with_default(subscriber, || { + let span = info_span!("test_span", otel.kind = "client"); + let _guard = span.enter(); + }); + + let spans = collector.get_spans(); + + assert_ron_snapshot!( + spans, + { ".*" => redact_id(), ".*[].**" => redact_id() }, + @r#" + { + SpanId(1): [ + CollectedSpan( + id: SpanId(1), + parent_id: None, + name: "test_span", + attributes: {}, + kind: client, + links: [], + ), + ], + } + "# + ); + } + + #[test] + fn test_nested_spans() { + let collector = TestCollector::new(); + let subscriber = Registry::default().with(layer(collector.clone())); + + tracing::subscriber::with_default(subscriber, || { + let parent = info_span!("parent_span"); + let _parent_guard = parent.enter(); + + { + let child = info_span!("child_span", otel.kind = "internal"); + let _child_guard = child.enter(); + std::thread::sleep(Duration::from_millis(10)); + } + }); + + let spans = collector.get_spans(); + + assert_ron_snapshot!( + spans, + { ".*" => redact_id(), ".*[].**" => redact_id() }, + @r#" + { + SpanId(1): [ + CollectedSpan( + id: SpanId(2), + parent_id: Some(SpanId(274877906945)), + name: "child_span", + attributes: {}, + kind: internal, + links: [], + ), + CollectedSpan( + id: SpanId(1), + parent_id: None, + name: "parent_span", + attributes: {}, + kind: internal, + links: [], + ), + ], + } + "# + ); + } + + #[test] + fn test_span_attributes() { + let collector = TestCollector::new(); + let subscriber = Registry::default().with(layer(collector.clone())); + + tracing::subscriber::with_default(subscriber, || { + let span = info_span!( + "attribute_span", + otel.kind = "client", + string_attr = "value", + bool_attr = true, + int_attr = 42, + float_attr = 3.5 + ); + let _guard = span.enter(); + }); + + let spans = collector.get_spans(); + + assert_ron_snapshot!( + spans, + { ".*" => redact_id(), ".*[].**" => redact_id() }, + @r#" + { + SpanId(1): [ + CollectedSpan( + id: SpanId(1), + parent_id: None, + name: "attribute_span", + attributes: { + "string_attr": "value", + "bool_attr": true, + "int_attr": 42, + "float_attr": 3.5, + }, + kind: client, + links: [], + ), + ], + } + "# + ); + } + + #[test] + fn test_span_updates() { + let collector = TestCollector::new(); + let subscriber = Registry::default().with(layer(collector.clone())); + + tracing::subscriber::with_default(subscriber, || { + let span = info_span!("updated_span", otel.kind = "client"); + span.record("dynamic_attr", "added_later"); + let _guard = span.enter(); + }); + + let spans = collector.get_spans(); + + assert_ron_snapshot!( + spans, + { ".*" => redact_id(), ".*[].**" => redact_id() }, + @r#" + { + SpanId(1): [ + CollectedSpan( + id: SpanId(1), + parent_id: None, + name: "updated_span", + attributes: {}, + kind: client, + links: [], + ), + ], + } + "# + ); + } +}