tracing-example.rs

use futures::future::join_all;
use std::{collections::HashMap, error::Error};
use tracing::{debug, field::Visit, info};
use tracing_attributes::instrument;
use tracing_core::Subscriber;
use tracing_subscriber::{
    layer::Context, layer::SubscriberExt, registry::LookupSpan, Layer, Registry,
};

#[derive(Clone, Debug)]
struct FieldMap {
    fields: HashMap<&'static str, String>,
}

impl FieldMap {
    fn new() -> Self {
        let fields = HashMap::new();
        Self { fields }
    }
    fn fold(mut self, other: Self) -> Self {
        self.fields.extend(other.fields);
        self
    }
}

impl Default for FieldMap {
    fn default() -> Self {
        Self::new()
    }
}

impl Visit for FieldMap {
    fn record_debug(&mut self, field: &tracing_core::Field, value: &dyn std::fmt::Debug) {
        let key = field.name();
        let val = format!("{:?}", value);
        self.fields.insert(key, val);
    }
    fn record_str(&mut self, field: &tracing_core::Field, value: &str) {
        self.fields.insert(field.name(), value.to_owned());
    }
}

struct TestLog {}

impl<S> Layer<S> for TestLog
where
    S: Subscriber,
    S: for<'a> LookupSpan<'a>,
{
    fn on_event(&self, event: &tracing::Event<'_>, ctx: Context<'_, S>) {
        let parent = event.parent();
        let span = parent
            .and_then(|id| ctx.span(id))
            .or_else(|| ctx.lookup_current());

        let mut fields = span
            .into_iter()
            .flat_map(|span| span.scope().from_root())
            .map(|span| {
                span.extensions()
                    .get::<FieldMap>()
                    .map(Clone::clone)
                    .unwrap_or_default()
            })
            .fold(FieldMap::new(), FieldMap::fold);
        event.record(&mut fields);
        println!("{:?}", fields.fields);
    }

    fn new_span(
        &self,
        attrs: &tracing_core::span::Attributes<'_>,
        id: &tracing_core::span::Id,
        ctx: Context<'_, S>,
    ) {
        let span = ctx.span(id).expect("unknown span");
        let mut fieldmap = FieldMap::new();
        attrs.record(&mut fieldmap);
        span.extensions_mut().insert(fieldmap);
    }

    fn on_record(
        &self,
        id: &tracing_core::span::Id,
        values: &tracing_core::span::Record<'_>,
        ctx: Context<'_, S>,
    ) {
        let span = ctx.span(id).expect("unknown span");
        let mut exts = span.extensions_mut();
        let fieldmap = exts.get_mut::<FieldMap>().expect("Missing FieldMap");
        values.record(fieldmap);
    }
}

#[instrument]
async fn parent_task(subtasks: usize) {
    info!("spawning subtasks...");
    let subtasks = (1..=subtasks)
        .map(|number| {
            debug!(message = "creating subtask;", number);
            subtask(number)
        })
        .collect::<Vec<_>>();

    let result = join_all(subtasks).await;

    debug!("all subtasks completed");
    let sum: usize = result.into_iter().sum();
    info!(sum);
}

#[instrument]
async fn subtask(number: usize) -> usize {
    info!("polling subtask...");
    number
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error + Send + Sync + 'static>> {
    let tl = TestLog {};
    let subscriber = Registry::default().with(tl);
    tracing::subscriber::set_global_default(subscriber)?;
    // tracing_subscriber::fmt()
    //     .with_max_level(tracing::Level::DEBUG)
    //     .try_init()?;
    parent_task(10).await;
    Ok(())
}