#[cfg(feature = "enable")]
use core::cell::Cell;
use core::marker::PhantomData;
#[cfg(all(feature = "std", feature = "enable"))]
use std::thread_local;
use crate::SpanContext;
#[cfg(feature = "enable")]
use crate::collector::get_collector;
use crate::id::SpanId;
#[cfg(feature = "enable")]
use crate::protocol::{
SpanAddEventMessage, SpanAddLinkMessage, SpanCloseMessage, SpanCreateMessage, SpanEnterMessage,
SpanExitMessage, SpanSetAttributeMessage,
};
#[cfg(feature = "enable")]
use crate::time::now;
use crate::value::KeyValue;
#[cfg(feature = "enable")]
thread_local! {
pub(crate) static CURRENT_SPAN: Cell<Option<SpanContext>> = const { Cell::new(None) };
}
#[must_use]
#[derive(Default, Debug)]
pub struct Span {
#[cfg(feature = "enable")]
pub(crate) inner: Option<SpanInner>,
}
#[cfg(feature = "enable")]
#[derive(Debug)]
pub(crate) struct SpanInner {
pub(crate) context: SpanContext,
}
#[derive(Default, Debug)]
pub struct CurrentSpan;
impl Span {
#[inline]
pub fn noop() -> Self {
Self {
#[cfg(feature = "enable")]
inner: None,
}
}
pub fn new(name: &'static str, attributes: &'_ [KeyValue<'static>]) -> Self {
#[cfg(not(feature = "enable"))]
{
let _ = (name, attributes);
Self::noop()
}
#[cfg(feature = "enable")]
{
if let Some(parent) = CURRENT_SPAN.get() {
Self::new_inner(name, parent, attributes)
} else {
Self::noop()
}
}
}
pub fn root(
name: &'static str,
mut span_context: SpanContext,
attributes: &'_ [KeyValue<'static>],
) -> Self {
span_context.span_id = SpanId(0);
#[cfg(not(feature = "enable"))]
{
let _ = (name, attributes);
Self::noop()
}
#[cfg(feature = "enable")]
{
Self::new_inner(name, span_context, attributes)
}
}
pub fn enter(&'_ self) -> SpanGuardRef<'_> {
#[cfg(not(feature = "enable"))]
{
SpanGuardRef::noop()
}
#[cfg(feature = "enable")]
{
let Some(context) = self.inner.as_ref().map(|inner| inner.context) else {
return SpanGuardRef::noop();
};
self.do_enter();
CURRENT_SPAN
.try_with(|current| {
let parent = current.get();
current.set(Some(context));
SpanGuardRef::new(self, parent)
})
.unwrap_or(SpanGuardRef::noop())
}
}
pub fn entered(self) -> SpanGuard {
#[cfg(not(feature = "enable"))]
{
SpanGuard::noop()
}
#[cfg(feature = "enable")]
{
let Some(context) = self.inner.as_ref().map(|inner| inner.context) else {
return SpanGuard::noop();
};
self.do_enter();
CURRENT_SPAN
.try_with(|current| {
let parent = current.get();
current.set(Some(context));
SpanGuard::new(self, parent)
})
.unwrap_or(SpanGuard::noop())
}
}
pub fn add_event(&self, name: &'static str, attributes: &'_ [KeyValue<'static>]) {
#[cfg(not(feature = "enable"))]
{
let _ = (name, attributes);
}
#[cfg(feature = "enable")]
{
if let Some(inner) = &self.inner {
get_collector().span_event(SpanAddEventMessage {
trace_id: inner.context.trace_id,
span_id: inner.context.span_id,
name: name.into(),
time_unix_nano: now().as_nanos(),
attributes: attributes.into(),
});
}
}
}
pub fn add_link(&self, link: SpanContext) {
#[cfg(not(feature = "enable"))]
{
let _ = link;
}
#[cfg(feature = "enable")]
{
if let Some(inner) = self.inner.as_ref() {
get_collector().span_link(SpanAddLinkMessage {
trace_id: inner.context.trace_id,
span_id: inner.context.span_id,
link,
});
}
}
}
pub fn set_attribute(&self, attribute: KeyValue<'static>) {
#[cfg(not(feature = "enable"))]
{
let _ = attribute;
}
#[cfg(feature = "enable")]
{
if let Some(inner) = self.inner.as_ref() {
get_collector().span_attribute(SpanSetAttributeMessage {
trace_id: inner.context.trace_id,
span_id: inner.context.span_id,
attribute,
});
}
}
}
}
impl CurrentSpan {
pub fn add_event(name: &'static str, attributes: &'_ [KeyValue<'static>]) {
#[cfg(not(feature = "enable"))]
{
let _ = (name, attributes);
}
#[cfg(feature = "enable")]
{
if let Some(context) = SpanContext::current() {
get_collector().span_event(SpanAddEventMessage {
trace_id: context.trace_id,
span_id: context.span_id,
name: name.into(),
time_unix_nano: now().as_nanos(),
attributes: attributes.into(),
});
}
}
}
pub fn add_link(link: SpanContext) {
#[cfg(not(feature = "enable"))]
{
let _ = link;
}
#[cfg(feature = "enable")]
{
if let Some(context) = SpanContext::current() {
get_collector().span_link(SpanAddLinkMessage {
trace_id: context.trace_id,
span_id: context.span_id,
link,
});
}
}
}
pub fn set_attribute(attribute: KeyValue<'static>) {
#[cfg(not(feature = "enable"))]
{
let _ = attribute;
}
#[cfg(feature = "enable")]
{
if let Some(context) = SpanContext::current() {
get_collector().span_attribute(SpanSetAttributeMessage {
trace_id: context.trace_id,
span_id: context.span_id,
attribute,
});
}
}
}
}
#[cfg(feature = "enable")]
impl Span {
fn new_inner(
name: &'static str,
parent: SpanContext,
attributes: &'_ [KeyValue<'static>],
) -> Self {
let span_id = SpanId::next_id();
let context = SpanContext::new(parent.trace_id, span_id);
let parent_id = if parent.span_id == SpanId(0) {
None
} else {
Some(parent.span_id)
};
get_collector().new_span(SpanCreateMessage {
trace_id: context.trace_id,
span_id: context.span_id,
parent_span_id: parent_id,
name: name.into(),
start_time_unix_nano: now().as_nanos(),
attributes: attributes.into(),
});
Self {
inner: Some(SpanInner { context }),
}
}
fn do_enter(&self) {
#[cfg(feature = "enable")]
if let Some(inner) = self.inner.as_ref() {
let timestamp = now();
get_collector().enter_span(SpanEnterMessage {
trace_id: inner.context.trace_id,
span_id: inner.context.span_id,
time_unix_nano: timestamp.0,
});
}
}
fn do_exit(&self) {
#[cfg(feature = "enable")]
if let Some(inner) = self.inner.as_ref() {
let timestamp = now();
get_collector().exit_span(SpanExitMessage {
trace_id: inner.context.trace_id,
span_id: inner.context.span_id,
time_unix_nano: timestamp.0,
});
}
}
}
impl Drop for Span {
fn drop(&mut self) {
#[cfg(feature = "enable")]
if let Some(inner) = self.inner.take() {
let timestamp = now();
get_collector().close_span(SpanCloseMessage {
trace_id: inner.context.trace_id,
span_id: inner.context.span_id,
end_time_unix_nano: timestamp.0,
});
}
}
}
#[derive(Debug)]
pub struct SpanGuard {
#[cfg(feature = "enable")]
pub(crate) inner: Option<SpanGuardInner>,
_not_send: PhantomNotSend,
}
#[cfg(feature = "enable")]
#[derive(Debug)]
pub(crate) struct SpanGuardInner {
span: Span,
parent: Option<SpanContext>,
}
impl SpanGuard {
pub(crate) fn noop() -> Self {
Self {
#[cfg(feature = "enable")]
inner: None,
_not_send: PhantomNotSend,
}
}
#[cfg(feature = "enable")]
pub(crate) fn new(span: Span, parent: Option<SpanContext>) -> Self {
Self {
#[cfg(feature = "enable")]
inner: Some(SpanGuardInner { span, parent }),
_not_send: PhantomNotSend,
}
}
}
impl Drop for SpanGuard {
fn drop(&mut self) {
#[cfg(feature = "enable")]
if let Some(inner) = self.inner.take() {
let _ = CURRENT_SPAN.try_with(|current| current.replace(inner.parent));
inner.span.do_exit();
}
}
}
#[derive(Debug)]
pub struct SpanGuardRef<'a> {
#[cfg(feature = "enable")]
pub(crate) inner: Option<SpanGuardRefInner<'a>>,
_phantom: PhantomData<&'a ()>,
}
#[cfg(feature = "enable")]
#[derive(Debug)]
pub(crate) struct SpanGuardRefInner<'a> {
span: &'a Span,
parent: Option<SpanContext>,
}
impl<'a> SpanGuardRef<'a> {
pub(crate) fn noop() -> Self {
Self {
#[cfg(feature = "enable")]
inner: None,
_phantom: PhantomData,
}
}
#[cfg(feature = "enable")]
pub(crate) fn new(span: &'a Span, parent: Option<SpanContext>) -> Self {
Self {
#[cfg(feature = "enable")]
inner: Some(SpanGuardRefInner { span, parent }),
_phantom: PhantomData,
}
}
}
impl Drop for SpanGuardRef<'_> {
fn drop(&mut self) {
#[cfg(feature = "enable")]
if let Some(inner) = self.inner.take() {
let _ = CURRENT_SPAN.try_with(|current| current.replace(inner.parent));
inner.span.do_exit();
}
}
}
#[derive(Debug)]
struct PhantomNotSend {
ghost: PhantomData<*mut ()>,
}
#[allow(non_upper_case_globals)]
const PhantomNotSend: PhantomNotSend = PhantomNotSend { ghost: PhantomData };
unsafe impl Sync for PhantomNotSend {}
#[cfg(all(test, feature = "std"))]
mod tests {
use super::*;
use crate::{SpanContext, SpanId, TraceId};
#[test]
fn span_noop() {
let span = Span::noop();
assert!(span.inner.is_none());
}
#[test]
fn span_new_without_parent() {
CURRENT_SPAN.set(None);
let span = Span::new("test_span", &[]);
assert!(span.inner.is_none());
}
#[test]
fn span_new_with_parent() {
let parent_context = SpanContext::generate();
CURRENT_SPAN.set(Some(parent_context));
let span = Span::new("child_span", &[]);
let inner = span.inner.as_ref().unwrap();
assert_eq!(inner.context.trace_id, parent_context.trace_id);
assert_ne!(inner.context.span_id, parent_context.span_id);
CURRENT_SPAN.set(None);
}
#[test]
fn span_root() {
let root_context = SpanContext::generate();
assert_eq!(root_context.span_id, SpanId(0));
let span = Span::root("root_span", root_context, &[]);
let inner = span.inner.as_ref().unwrap();
assert_eq!(inner.context.trace_id, root_context.trace_id);
assert_ne!(inner.context.span_id, SpanId(0));
}
#[test]
fn span_context_from_span() {
let root_context = SpanContext::generate();
let span = Span::root("test_span", root_context, &[]);
let extracted_context = SpanContext::from_span(&span);
let context = extracted_context.unwrap();
assert_eq!(context.trace_id, root_context.trace_id);
}
#[test]
fn span_context_from_noop_span() {
let span = Span::noop();
let extracted_context = SpanContext::from_span(&span);
assert!(extracted_context.is_none());
}
#[test]
fn span_enter_and_current_context() {
CURRENT_SPAN.set(None);
assert!(SpanContext::current().is_none());
let root_context = SpanContext::generate();
let span = Span::root("test_span", root_context, &[]);
{
let _guard = span.enter();
let current_context = SpanContext::current();
let context = current_context.unwrap();
assert_eq!(context.trace_id, root_context.trace_id);
}
assert!(SpanContext::current().is_none());
}
#[test]
fn span_entered_guard() {
CURRENT_SPAN.set(None);
let root_context = SpanContext::generate();
let span = Span::root("test_span", root_context, &[]);
{
let _guard = span.entered();
let current_context = SpanContext::current();
assert!(current_context.is_some());
}
assert!(SpanContext::current().is_none());
}
#[test]
fn noop_span_operations() {
let noop_span = Span::noop();
{
let _guard = noop_span.enter();
assert!(SpanContext::current().is_none());
}
let _entered_guard = noop_span.entered();
assert!(SpanContext::current().is_none());
}
#[test]
fn nested_spans() {
CURRENT_SPAN.set(None);
let root_context = SpanContext::generate();
let _root_guard = Span::root("test_span", root_context, &[]).entered();
let child_span = Span::new("child", &[]);
let child_inner = child_span.inner.as_ref().unwrap();
assert_eq!(child_inner.context.trace_id, root_context.trace_id);
assert_ne!(child_inner.context.span_id, root_context.span_id);
}
#[test]
fn span_event() {
let context = SpanContext::generate();
let span = Span::root("test_span", context, &[]);
let event_attributes = [KeyValue::new("event_key", "event_value")];
span.add_event("test_event", &event_attributes);
let noop_span = Span::noop();
noop_span.add_event("noop_event", &event_attributes);
}
#[test]
fn span_link() {
let context = SpanContext::generate();
let span = Span::root("test_span", context, &[]);
let link_context = SpanContext::new(TraceId(0), SpanId(0));
span.add_link(link_context);
let noop_span = Span::noop();
noop_span.add_link(link_context);
}
#[test]
fn span_attribute() {
let context = SpanContext::generate();
let span = Span::root("test_span", context, &[]);
let attribute = KeyValue::new("test_key", "test_value");
span.set_attribute(attribute.clone());
let noop_span = Span::noop();
noop_span.set_attribute(attribute);
}
#[test]
fn span_methods_with_entered_span() {
let context = SpanContext::generate();
let span = Span::root("test_span", context, &[]);
let _guard = span.enter();
span.add_event("entered_event", &[]);
span.add_link(SpanContext::new(TraceId(0), SpanId(0)));
span.set_attribute(KeyValue::new("entered_key", true));
}
#[test]
fn current_span_event_with_active_span() {
CURRENT_SPAN.set(None);
let context = SpanContext::generate();
let _root_guard = Span::root("test_span", context, &[]).entered();
let event_attributes = [KeyValue::new("current_event_key", "current_event_value")];
CurrentSpan::add_event("current_test_event", &event_attributes);
}
#[test]
fn current_span_link_with_active_span() {
CURRENT_SPAN.set(None);
let context = SpanContext::generate();
let _root_guard = Span::root("test_span", context, &[]).entered();
let link_context = SpanContext::new(TraceId(0), SpanId(0));
CurrentSpan::add_link(link_context);
}
#[test]
fn current_span_attribute_with_active_span() {
CURRENT_SPAN.set(None);
let context = SpanContext::generate();
let span = Span::root("test_span", context, &[]);
let _guard = span.enter();
let attribute = KeyValue::new("current_attr_key", "current_attr_value");
CurrentSpan::set_attribute(attribute);
}
}