AppFlowy/rust-lib/flowy-log/src/layer.rs

use serde::ser::{SerializeMap, Serializer};
use serde_json::Value;
use std::{fmt, io::Write};
use tracing::{Event, Id, Subscriber};
use tracing_bunyan_formatter::JsonStorage;
use tracing_core::{metadata::Level, span::Attributes};

use tracing_subscriber::{fmt::MakeWriter, layer::Context, registry::SpanRef, Layer};
const LEVEL: &str = "level";
const TIME: &str = "time";
const MESSAGE: &str = "msg";

const FLOWY_RESERVED_FIELDS: [&str; 3] = [LEVEL, TIME, MESSAGE];

pub struct FlowyFormattingLayer<W: MakeWriter + 'static> {
    make_writer: W,
    with_target: bool,
}

impl<W: MakeWriter + 'static> FlowyFormattingLayer<W> {
    pub fn new(make_writer: W) -> Self {
        Self {
            make_writer,
            with_target: false,
        }
    }

    fn serialize_flowy_core_fields(
        &self,
        map_serializer: &mut impl SerializeMap<Error = serde_json::Error>,
        message: &str,
        level: &Level,
    ) -> Result<(), std::io::Error> {
        map_serializer.serialize_entry(MESSAGE, &message)?;
        map_serializer.serialize_entry(LEVEL, &format!("{}", level))?;
        map_serializer.serialize_entry(TIME, &chrono::Utc::now().timestamp())?;
        Ok(())
    }

    fn serialize_span<S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>>(
        &self,
        span: &SpanRef<S>,
        ty: Type,
    ) -> Result<Vec<u8>, std::io::Error> {
        let mut buffer = Vec::new();
        let mut serializer = serde_json::Serializer::new(&mut buffer);
        let mut map_serializer = serializer.serialize_map(None)?;
        let message = format_span_context(&span, ty);
        self.serialize_flowy_core_fields(&mut map_serializer, &message, span.metadata().level())?;
        if self.with_target {
            map_serializer.serialize_entry("target", &span.metadata().target())?;
        }

        map_serializer.serialize_entry("line", &span.metadata().line())?;
        map_serializer.serialize_entry("file", &span.metadata().file())?;

        let extensions = span.extensions();
        if let Some(visitor) = extensions.get::<JsonStorage>() {
            for (key, value) in visitor.values() {
                if !FLOWY_RESERVED_FIELDS.contains(key) {
                    map_serializer.serialize_entry(key, value)?;
                } else {
                    tracing::debug!("{} is a reserved field in the bunyan log format. Skipping it.", key);
                }
            }
        }
        map_serializer.end()?;
        Ok(buffer)
    }

    fn emit(&self, mut buffer: Vec<u8>) -> Result<(), std::io::Error> {
        buffer.write_all(b"\n")?;
        self.make_writer.make_writer().write_all(&buffer)
    }
}

/// The type of record we are dealing with: entering a span, exiting a span, an
/// event.
#[derive(Clone, Debug)]
pub enum Type {
    EnterSpan,
    ExitSpan,
    Event,
}

impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let repr = match self {
            Type::EnterSpan => "START",
            Type::ExitSpan => "END",
            Type::Event => "EVENT",
        };
        write!(f, "{}", repr)
    }
}

fn format_span_context<S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>>(span: &SpanRef<S>, ty: Type) -> String {
    format!("[⛳ {} - {}]", span.metadata().name().to_uppercase(), ty)
}

fn format_event_message<S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>>(
    current_span: &Option<SpanRef<S>>,
    event: &Event,
    event_visitor: &JsonStorage<'_>,
) -> String {
    // Extract the "message" field, if provided. Fallback to the target, if missing.
    let mut message = event_visitor
        .values()
        .get("message")
        .map(|v| match v {
            Value::String(s) => Some(s.as_str()),
            _ => None,
        })
        .flatten()
        .unwrap_or_else(|| event.metadata().target())
        .to_owned();

    // If the event is in the context of a span, prepend the span name to the
    // message.
    if let Some(span) = &current_span {
        message = format!("{} {}", format_span_context(span, Type::Event), message);
    }

    message
}

impl<S, W> Layer<S> for FlowyFormattingLayer<W>
where
    S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
    W: MakeWriter + 'static,
{
    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
        // Events do not necessarily happen in the context of a span, hence
        // lookup_current returns an `Option<SpanRef<_>>` instead of a
        // `SpanRef<_>`.
        let current_span = ctx.lookup_current();

        let mut event_visitor = JsonStorage::default();
        event.record(&mut event_visitor);

        // Opting for a closure to use the ? operator and get more linear code.
        let format = || {
            let mut buffer = Vec::new();

            let mut serializer = serde_json::Serializer::new(&mut buffer);
            let mut map_serializer = serializer.serialize_map(None)?;

            let message = format_event_message(&current_span, event, &event_visitor);
            self.serialize_flowy_core_fields(&mut map_serializer, &message, event.metadata().level())?;
            // Additional metadata useful for debugging
            // They should be nested under `src` (see https://github.com/trentm/node-bunyan#src )
            // but `tracing` does not support nested values yet

            if self.with_target {
                map_serializer.serialize_entry("target", event.metadata().target())?;
            }

            map_serializer.serialize_entry("line", &event.metadata().line())?;
            map_serializer.serialize_entry("file", &event.metadata().file())?;

            // Add all the other fields associated with the event, expect the message we
            // already used.
            for (key, value) in event_visitor
                .values()
                .iter()
                .filter(|(&key, _)| key != "message" && !FLOWY_RESERVED_FIELDS.contains(&key))
            {
                map_serializer.serialize_entry(key, value)?;
            }

            // Add all the fields from the current span, if we have one.
            if let Some(span) = &current_span {
                let extensions = span.extensions();
                if let Some(visitor) = extensions.get::<JsonStorage>() {
                    for (key, value) in visitor.values() {
                        if !FLOWY_RESERVED_FIELDS.contains(key) {
                            map_serializer.serialize_entry(key, value)?;
                        } else {
                            tracing::debug!("{} is a reserved field in the flowy log format. Skipping it.", key);
                        }
                    }
                }
            }
            map_serializer.end()?;
            Ok(buffer)
        };

        let result: std::io::Result<Vec<u8>> = format();
        if let Ok(formatted) = result {
            let _ = self.emit(formatted);
        }
    }

    fn new_span(&self, _attrs: &Attributes, id: &Id, ctx: Context<'_, S>) {
        let span = ctx.span(id).expect("Span not found, this is a bug");
        if let Ok(serialized) = self.serialize_span(&span, Type::EnterSpan) {
            let _ = self.emit(serialized);
        }
    }

    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
        let span = ctx.span(&id).expect("Span not found, this is a bug");
        if let Ok(serialized) = self.serialize_span(&span, Type::ExitSpan) {
            let _ = self.emit(serialized);
        }
    }
}