gitdataai/libs/observability/src/msg_json_fmt.rs

//! Custom JSON formatter that injects `_msg` as the first field.
//!
//! VictoriaLogs requires every log line to have a `_msg` field as the message
//! subject for full-text search. This formatter ensures `_msg` is always the
//! first key in the serialized JSON object.
//!
//! `_msg` derivation rules:
//! - If the event has a `message` field → use it
//! - If a parent span stored a `_msg` in the thread-local buffer (set by the
//!   HTTP middleware) → use it
//! - Fallback → use the event target (module path)

use serde::Serialize;
use std::cell::RefCell;
use std::collections::HashMap;
use std::fmt;
use tracing::field::{Field, Visit};
use tracing_subscriber::fmt::format::FormatEvent;

// Thread-local buffer holding the current span's `_msg` value.
// Set by the HTTP middleware via `set_span_msg()` at span creation time.
// The formatter reads this value when the event message is empty (span close).
thread_local! {
    static SPAN_MSG: RefCell<String> = const { RefCell::new(String::new()) };
}

/// Set the `_msg` for the current span. Called by middleware at span creation.
pub fn set_span_msg(msg: String) {
    SPAN_MSG.with(|cell| cell.borrow_mut().clone_from(&msg));
}

/// Read and clear the current span's `_msg`. Called by the formatter.
fn take_span_msg() -> Option<String> {
    SPAN_MSG.with(|cell| {
        let mut s = cell.borrow_mut();
        if s.is_empty() {
            None
        } else {
            let out = s.clone();
            s.clear();
            Some(out)
        }
    })
}

/// Capture all event fields via the visitor pattern.
struct FieldCollector {
    message: Option<String>,
    fields: HashMap<String, serde_json::Value>,
}

impl FieldCollector {
    fn new() -> Self {
        Self {
            message: None,
            fields: HashMap::new(),
        }
    }
}

impl Visit for FieldCollector {
    fn record_str(&mut self, field: &Field, value: &str) {
        if field.name() == "message" {
            self.message = Some(value.to_string());
        } else {
            self.fields
                .insert(field.name().to_string(), serde_json::Value::String(value.to_string()));
        }
    }

    fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
        self.fields.insert(
            field.name().to_string(),
            serde_json::Value::String(format!("{:?}", value)),
        );
    }

    fn record_i64(&mut self, field: &Field, value: i64) {
        self.fields.insert(
            field.name().to_string(),
            serde_json::Value::Number(value.into()),
        );
    }

    fn record_u64(&mut self, field: &Field, value: u64) {
        use serde_json::Number;
        self.fields.insert(
            field.name().to_string(),
            serde_json::Value::Number(Number::from(value)),
        );
    }

    fn record_bool(&mut self, field: &Field, value: bool) {
        self.fields.insert(
            field.name().to_string(),
            serde_json::Value::Bool(value),
        );
    }
}

/// Custom `FormatEvent` that outputs JSON with `_msg` as the first field.
pub struct MsgJsonFormat;

impl<S, N> FormatEvent<S, N> for MsgJsonFormat
where
    S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
    N: for<'a> tracing_subscriber::fmt::format::FormatFields<'a> + 'static,
{
    fn format_event(
        &self,
        _ctx: &tracing_subscriber::fmt::FmtContext<'_, S, N>,
        mut writer: tracing_subscriber::fmt::format::Writer<'_>,
        event: &tracing::Event<'_>,
    ) -> std::fmt::Result {
        let mut collector = FieldCollector::new();
        event.record(&mut collector);

        // Derive _msg: event message → thread-local span msg → target fallback
        let msg = collector
            .message
            .clone()
            .or_else(take_span_msg)
            .unwrap_or_else(|| event.metadata().target().to_string());

        // Build ordered JSON: _msg first, then standard fields, then event fields
        let mut ordered = serde_json::Map::with_capacity(10 + collector.fields.len());
        ordered.insert("_msg".to_string(), serde_json::Value::String(msg));
        ordered.insert(
            "timestamp".to_string(),
            serde_json::Value::String(chrono::Utc::now().to_rfc3339()),
        );
        ordered.insert(
            "level".to_string(),
            serde_json::Value::String(format!("{}", event.metadata().level())),
        );
        ordered.insert(
            "target".to_string(),
            serde_json::Value::String(event.metadata().target().to_string()),
        );

        if let Some(module_path) = event.metadata().module_path() {
            ordered.insert(
                "module".to_string(),
                serde_json::Value::String(module_path.to_string()),
            );
        }
        if let Some(file) = event.metadata().file() {
            ordered.insert("file".to_string(), serde_json::Value::String(file.to_string()));
        }
        if let Some(line) = event.metadata().line() {
            ordered.insert("line".to_string(), serde_json::Value::Number(line.into()));
        }

        // Event fields (excluding `message` which is already in `_msg`)
        for (key, value) in &collector.fields {
            ordered.insert(key.clone(), value.clone());
        }

        // Use pretty JSON for human readability in terminals, compact for pipelines
        let json = if std::env::var("APP_LOG_PRETTY").as_deref() == Ok("true") {
            serde_json::to_string_pretty(&ordered)
        } else {
            serde_json::to_string(&ordered)
        };
        write!(writer, "{}", json.map_err(|_| fmt::Error)?)
    }
}

/// Log entry wrapper for VictoriaLogs-compatible output.
///
/// All fields serialize in a fixed order with `_msg` first.
#[derive(Serialize)]
pub struct VLogEntry {
    _msg: String,
    timestamp: String,
    level: String,
    target: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    module: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    file: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    line: Option<u32>,
    #[serde(flatten)]
    extra: serde_json::Map<String, serde_json::Value>,
}

impl VLogEntry {
    pub fn new(msg: String, level: &str, target: &str) -> Self {
        Self {
            _msg: msg,
            timestamp: chrono::Utc::now().to_rfc3339(),
            level: level.to_string(),
            target: target.to_string(),
            module: None,
            file: None,
            line: None,
            extra: serde_json::Map::new(),
        }
    }

    pub fn module(mut self, m: String) -> Self {
        self.module = Some(m);
        self
    }

    pub fn file(mut self, f: String) -> Self {
        self.file = Some(f);
        self
    }

    pub fn line(mut self, l: u32) -> Self {
        self.line = Some(l);
        self
    }

    pub fn field(mut self, key: String, value: serde_json::Value) -> Self {
        self.extra.insert(key, value);
        self
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn json_order_has_msg_first() {
        let entry = VLogEntry::new("GET /api/users".to_string(), "INFO", "my::module")
            .module("my::module".to_string())
            .field("status".to_string(), serde_json::json!(200))
            .field("duration_ms".to_string(), serde_json::json!(42));

        let json = serde_json::to_string(&entry).unwrap();
        assert!(json.starts_with(r#"{"_msg":"GET /api/users""#));
    }
}