#![deny(
missing_docs,
elided_lifetimes_in_paths,
unsafe_code,
rustdoc::invalid_rust_codeblocks,
rustdoc::broken_intra_doc_links,
missing_copy_implementations,
unused_doc_comments
)]
use std::{fmt, sync::OnceLock};
use rootcause::{
Report, ReportMut,
handlers::{
AttachmentFormattingPlacement, AttachmentFormattingStyle, AttachmentHandler,
FormattingFunction,
},
hooks::report_creation::ReportCreationHook,
markers::{self, Dynamic, ObjectMarkerFor},
report_attachment::ReportAttachment,
};
use tracing::{
Span,
field::{Field, Visit},
};
use tracing_subscriber::{
Registry,
registry::{LookupSpan, SpanRef},
};
#[derive(Copy, Clone)]
pub struct SpanHandler;
struct CapturedFields(String);
impl AttachmentHandler<Span> for SpanHandler {
fn display(value: &Span, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match value
.with_subscriber(|(span_id, dispatch)| display_span_chain(span_id, dispatch, formatter))
{
Some(Ok(())) => Ok(()),
Some(Err(e)) => Err(e),
None => write!(formatter, "No tracing subscriber available"),
}
}
fn debug(value: &Span, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
std::fmt::Debug::fmt(value, formatter)
}
fn preferred_formatting_style(
span: &Span,
_report_formatting_function: FormattingFunction,
) -> AttachmentFormattingStyle {
AttachmentFormattingStyle {
placement: if span.is_none() {
AttachmentFormattingPlacement::Hidden
} else {
AttachmentFormattingPlacement::InlineWithHeader {
header: "Tracing spans:",
}
},
function: FormattingFunction::Display,
priority: 9, }
}
}
fn display_span_chain(
span_id: &tracing::span::Id,
dispatch: &tracing::Dispatch,
formatter: &mut fmt::Formatter<'_>,
) -> fmt::Result {
let Some(registry) = dispatch.downcast_ref::<Registry>() else {
write!(formatter, "No tracing registry subscriber found")?;
return Ok(());
};
let Some(span) = registry.span(span_id) else {
write!(formatter, "No span found for ID")?;
return Ok(());
};
let mut first_span = true;
for ancestor_span in span.scope() {
if first_span {
first_span = false;
} else {
writeln!(formatter)?;
}
display_span(ancestor_span, formatter)?;
}
Ok(())
}
fn display_span(
span: SpanRef<'_, Registry>,
formatter: &mut fmt::Formatter<'_>,
) -> Result<(), fmt::Error> {
write!(formatter, "{}", span.name())?;
let extensions = span.extensions();
let Some(captured_fields) = extensions.get::<CapturedFields>() else {
write!(
formatter,
"{{ Span values missing. Was the RootcauseLayer installed correctly? }}"
)?;
return Ok(());
};
if captured_fields.0.is_empty() {
Ok(())
} else {
write!(formatter, "{{{}}}", captured_fields.0)
}
}
#[derive(Copy, Clone, Debug, Default)]
pub struct RootcauseLayer;
impl<S> tracing_subscriber::Layer<S> for RootcauseLayer
where
S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
{
fn on_new_span(
&self,
attrs: &tracing::span::Attributes<'_>,
id: &tracing::span::Id,
ctx: tracing_subscriber::layer::Context<'_, S>,
) {
let span = ctx.span(id).expect("span not found");
let mut extensions = span.extensions_mut();
struct Visitor(String);
impl Visit for Visitor {
fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
use std::fmt::Write;
if self.0.is_empty() {
let _ = write!(self.0, "{}={value:?}", field.name());
} else {
let _ = write!(self.0, " {}={value:?}", field.name());
}
}
}
let mut visitor = Visitor(String::new());
attrs.record(&mut visitor);
extensions.insert(CapturedFields(visitor.0));
}
}
#[derive(Copy, Clone)]
pub struct SpanCollector {
pub capture_span_for_reports_with_children: bool,
}
#[derive(Debug)]
struct RootcauseTracingEnvOptions {
span_leafs_only: bool,
}
impl RootcauseTracingEnvOptions {
fn get() -> &'static Self {
static ROOTCAUSE_TRACING_FLAGS: OnceLock<RootcauseTracingEnvOptions> = OnceLock::new();
ROOTCAUSE_TRACING_FLAGS.get_or_init(|| {
let mut span_leafs_only = false;
if let Some(var) = std::env::var_os("ROOTCAUSE_TRACING") {
for v in var.to_string_lossy().split(',') {
if v.eq_ignore_ascii_case("leafs") {
span_leafs_only = true;
}
}
}
RootcauseTracingEnvOptions { span_leafs_only }
})
}
}
impl SpanCollector {
pub fn new() -> Self {
let env_options = RootcauseTracingEnvOptions::get();
let capture_span_for_reports_with_children = !env_options.span_leafs_only;
Self {
capture_span_for_reports_with_children,
}
}
}
impl Default for SpanCollector {
fn default() -> Self {
Self::new()
}
}
impl ReportCreationHook for SpanCollector {
fn on_local_creation(&self, mut report: ReportMut<'_, Dynamic, markers::Local>) {
let do_capture =
self.capture_span_for_reports_with_children || report.children().is_empty();
if do_capture {
let span = Span::current();
if !span.is_none() {
let attachment = ReportAttachment::new_custom::<SpanHandler>(span);
report.attachments_mut().push(attachment.into_dynamic());
}
}
}
fn on_sendsync_creation(&self, mut report: ReportMut<'_, Dynamic, markers::SendSync>) {
let do_capture =
self.capture_span_for_reports_with_children || report.children().is_empty();
if do_capture {
let span = Span::current();
if !span.is_none() {
let attachment = ReportAttachment::new_custom::<SpanHandler>(span);
report.attachments_mut().push(attachment.into_dynamic());
}
}
}
}
pub trait SpanExt: Sized {
fn attach_span(self) -> Self;
}
impl<C: ?Sized, T> SpanExt for Report<C, markers::Mutable, T>
where
Span: ObjectMarkerFor<T>,
{
fn attach_span(mut self) -> Self {
let span = Span::current();
if !span.is_disabled() {
self = self.attach_custom::<SpanHandler, _>(span);
}
self
}
}
impl<C: ?Sized, V, T> SpanExt for Result<V, Report<C, markers::Mutable, T>>
where
Span: ObjectMarkerFor<T>,
{
fn attach_span(self) -> Self {
match self {
Ok(v) => Ok(v),
Err(report) => Err(report.attach_span()),
}
}
}