use crate::{timing_layer::WithContext, SpanRootTiming, SpanTiming};
pub(crate) trait SpanExt {
fn take_ext<F>(&self, f: F)
where
F: FnMut(SpanTiming);
fn insert_ext(&self, item: SpanRootTiming);
}
impl SpanExt for tracing::Span {
fn take_ext<F>(&self, mut f: F)
where
F: FnMut(SpanTiming),
{
self.with_subscriber(|(id, subscriber)| {
if let Some(ctx) = subscriber.downcast_ref::<WithContext>() {
ctx.take(subscriber, id, &mut f)
}
});
}
fn insert_ext(&self, item: SpanRootTiming) {
self.with_subscriber(|(id, subscriber)| {
if let Some(ctx) = subscriber.downcast_ref::<WithContext>() {
ctx.set(subscriber, id, item)
}
});
}
}