forge-ops-tracker 0.8.0

// One trace's worth of spans (a request's, or a job's, own call tree), sharing a single trace id.
// Held in a `thread_local!`, the same choice the breadcrumb trail and CURRENT_USER already made and
// with the same caveat: it follows a thread, not an async task, so it suits this crate's
// synchronous, thread-per-request-shaped design. Nesting comes from a stack of open span ids: a
// span opened while another is open becomes its child, and anything else parents under the root.
//
// A trace is sent only when its root span took at least `Configuration.trace_capture_threshold`,
// decided here once the root finishes, so a fast request costs nothing on the wire.

use std::cell::RefCell;
use std::collections::HashMap;
use std::hash::{BuildHasher, Hasher};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use crate::configuration::Configuration;
use crate::event_builder::format_unix_timestamp;
use crate::pii_scrubber::{json_string, Value};

pub const MAX_SPANS: usize = 500;

// The kinds the ingestion API accepts; anything else would fail validation for the whole trace,
// so an unknown kind is sent as "other" instead.
const KINDS: [&str; 7] = [
    "controller",
    "service",
    "database",
    "redis",
    "http",
    "job",
    "other",
];

struct Span {
    span_id: String,
    parent_span_id: Option<String>,
    name: String,
    kind: String,
    started_at: SystemTime,
    duration_ms: f64,
    data: HashMap<String, Value>,
}

pub struct SpanBuffer {
    trace_id: String,
    root_span_id: String,
    spans: Vec<Span>,
    open: Vec<String>,
    environment: String,
    release: Option<String>,
}

thread_local! {
    static TRACE: RefCell<Option<SpanBuffer>> = const { RefCell::new(None) };
}

impl SpanBuffer {
    fn new(config: &Configuration) -> Self {
        SpanBuffer {
            trace_id: random_hex(2),
            root_span_id: random_hex(1),
            spans: Vec::new(),
            open: Vec::new(),
            environment: config.environment.clone(),
            release: config.release.clone(),
        }
    }

    fn current_parent(&self) -> String {
        self.open
            .last()
            .cloned()
            .unwrap_or_else(|| self.root_span_id.clone())
    }

    fn record(&mut self, span: Span) {
        if self.spans.len() >= MAX_SPANS - 1 {
            return; // leave room for the root
        }
        self.spans.push(span);
    }

    fn span_json(&self, span: &Span) -> String {
        let kind = if KINDS.contains(&span.kind.as_str()) {
            span.kind.as_str()
        } else {
            "other"
        };
        let data = Value::Object(span.data.clone()).to_json();
        format!(
            "{{\"span_id\":{},\"parent_span_id\":{},\"name\":{},\"kind\":{},\"started_at\":{},\"duration_ms\":{},\"environment\":{},\"release\":{},\"data\":{}}}",
            json_string(&span.span_id),
            span.parent_span_id
                .as_deref()
                .map(json_string)
                .unwrap_or_else(|| "null".to_string()),
            json_string(&span.name),
            json_string(kind),
            json_string(&timestamp(span.started_at)),
            (span.duration_ms * 100.0).round() / 100.0,
            json_string(&self.environment),
            self.release
                .as_deref()
                .map(json_string)
                .unwrap_or_else(|| "null".to_string()),
            data
        )
    }
}

/// Starts a fresh trace on the calling thread, discarding any earlier one. Returns false (and
/// starts nothing) when a trace is already open here, which lets a nested `trace` degrade to a
/// span; also starts nothing when tracing is off or reporting isn't enabled, but still returns
/// true, since the caller then owns a (no-op) root.
pub fn begin(config: &Configuration) -> bool {
    TRACE.with(|trace| {
        let mut trace = trace.borrow_mut();
        if trace.is_some() {
            return false;
        }
        if config.track_tracing && config.is_enabled() {
            *trace = Some(SpanBuffer::new(config));
        }
        true
    })
}

#[cfg(test)]
pub fn is_active() -> bool {
    TRACE.with(|trace| trace.borrow().is_some())
}

/// Ends the calling thread's trace, always clearing it, and returns the whole request body
/// (`{"trace_id":...,"spans":[...]}`) when the root took at least `threshold`.
pub fn end(
    threshold: Duration,
    root_name: &str,
    root_kind: &str,
    started_at: SystemTime,
    duration_ms: f64,
) -> Option<String> {
    let buffer = TRACE.with(|trace| trace.borrow_mut().take())?;
    if duration_ms < threshold.as_secs_f64() * 1000.0 {
        return None;
    }

    let root = Span {
        span_id: buffer.root_span_id.clone(),
        parent_span_id: None,
        name: root_name.to_string(),
        kind: root_kind.to_string(),
        started_at,
        duration_ms,
        data: HashMap::new(),
    };
    let spans: Vec<String> = std::iter::once(&root)
        .chain(buffer.spans.iter())
        .map(|span| buffer.span_json(span))
        .collect();
    Some(format!(
        "{{\"trace_id\":{},\"spans\":[{}]}}",
        json_string(&buffer.trace_id),
        spans.join(",")
    ))
}

/// Opens a span under whatever is open (or the root), returning its id, or None outside a trace.
pub fn open_span() -> Option<(String, String)> {
    TRACE.with(|trace| {
        trace.borrow_mut().as_mut().map(|buffer| {
            let parent = buffer.current_parent();
            let id = random_hex(1);
            buffer.open.push(id.clone());
            (id, parent)
        })
    })
}

pub fn close_span(
    id: String,
    parent: String,
    name: &str,
    kind: &str,
    started_at: SystemTime,
    duration_ms: f64,
    data: HashMap<String, Value>,
) {
    TRACE.with(|trace| {
        if let Some(buffer) = trace.borrow_mut().as_mut() {
            buffer.open.retain(|open| open != &id);
            buffer.record(Span {
                span_id: id,
                parent_span_id: Some(parent),
                name: name.to_string(),
                kind: kind.to_string(),
                started_at,
                duration_ms,
                data,
            });
        }
    });
}

/// Records an already-finished span under whatever is currently open; a no-op outside a trace.
pub fn record_leaf(
    name: &str,
    kind: &str,
    started_at: SystemTime,
    duration_ms: f64,
    data: HashMap<String, Value>,
) {
    TRACE.with(|trace| {
        if let Some(buffer) = trace.borrow_mut().as_mut() {
            let parent = buffer.current_parent();
            buffer.record(Span {
                span_id: random_hex(1),
                parent_span_id: Some(parent),
                name: name.to_string(),
                kind: kind.to_string(),
                started_at,
                duration_ms,
                data,
            });
        }
    });
}

/// ISO 8601 with milliseconds, UTC.
fn timestamp(time: SystemTime) -> String {
    let since = time.duration_since(UNIX_EPOCH).unwrap_or_default();
    let whole = format_unix_timestamp(since.as_secs());
    format!(
        "{}.{:03}Z",
        whole.trim_end_matches('Z'),
        since.subsec_millis()
    )
}

/// `words` 64-bit words of hex (1 = a 16-char span id, 2 = a 32-char trace id). Rust's standard
/// library has no random number generator, so this mixes the per-process random `RandomState`
/// keys with the clock and a counter: an id only has to be unique within one project's traces, not
/// unpredictable, so this avoids a dependency for it.
fn random_hex(words: usize) -> String {
    static COUNTER: AtomicU64 = AtomicU64::new(0);
    (0..words)
        .map(|_| {
            let mut hasher = std::collections::hash_map::RandomState::new().build_hasher();
            hasher.write_u64(COUNTER.fetch_add(1, Ordering::Relaxed));
            hasher.write_u128(
                SystemTime::now()
                    .duration_since(UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_nanos(),
            );
            format!("{:016x}", hasher.finish())
        })
        .collect()
}

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

    fn config() -> Configuration {
        let mut config = Configuration::new();
        config.dsn = Some("https://key@tracker.example.com/api/v1/events".to_string());
        config.environment = "production".to_string();
        config.release = Some("abc123".to_string());
        config.trace_capture_threshold = Duration::from_millis(10);
        config
    }

    fn finish(config: &Configuration, name: &str, duration_ms: f64) -> Option<String> {
        end(
            config.trace_capture_threshold,
            name,
            "controller",
            SystemTime::now(),
            duration_ms,
        )
    }

    fn spans_of(body: &str) -> Vec<HashMap<String, String>> {
        // A tiny purpose-built reader, not a JSON parser: pulls the flat string fields this
        // test asserts on out of each span object.
        let start = body.find("\"spans\":[").unwrap() + "\"spans\":[".len();
        body[start..]
            .split("{\"span_id\"")
            .skip(1)
            .map(|chunk| {
                let chunk = format!("{{\"span_id\"{chunk}");
                let mut fields = HashMap::new();
                for key in ["span_id", "parent_span_id", "name", "kind", "started_at"] {
                    let needle = format!("\"{key}\":");
                    if let Some(at) = chunk.find(&needle) {
                        let rest = &chunk[at + needle.len()..];
                        let value = if let Some(stripped) = rest.strip_prefix('"') {
                            stripped.split('"').next().unwrap().to_string()
                        } else {
                            "null".to_string()
                        };
                        fields.insert(key.to_string(), value);
                    }
                }
                fields
            })
            .collect()
    }

    #[test]
    fn nests_spans_under_the_open_one_and_the_root_with_the_wire_shape() {
        let config = config();
        assert!(begin(&config));
        let (outer, parent) = open_span().unwrap();
        record_leaf(
            "SELECT users",
            "database",
            SystemTime::now(),
            3.0,
            HashMap::new(),
        );
        close_span(
            outer,
            parent,
            "charge",
            "service",
            SystemTime::now(),
            20.0,
            HashMap::new(),
        );
        record_leaf(
            "sibling",
            "database",
            SystemTime::now(),
            1.0,
            HashMap::new(),
        );

        let body = finish(&config, "GET /x", 1500.0).unwrap();
        assert!(!is_active(), "end clears the trace");
        let spans = spans_of(&body);
        let by_name: HashMap<_, _> = spans.iter().map(|s| (s["name"].clone(), s)).collect();

        assert!(body.contains("\"trace_id\":\""));
        assert_eq!(by_name["GET /x"]["parent_span_id"], "null");
        assert_eq!(by_name["GET /x"]["kind"], "controller");
        assert_eq!(
            by_name["SELECT users"]["parent_span_id"],
            by_name["charge"]["span_id"]
        );
        assert_eq!(
            by_name["charge"]["parent_span_id"],
            by_name["GET /x"]["span_id"]
        );
        assert_eq!(
            by_name["sibling"]["parent_span_id"],
            by_name["GET /x"]["span_id"]
        );
        assert_eq!(by_name["charge"]["span_id"].len(), 16);
        assert!(body.contains("\"environment\":\"production\""));
        assert!(body.contains("\"release\":\"abc123\""));
        assert!(body.contains("\"parent_span_id\":null"));
        let started = &by_name["charge"]["started_at"];
        assert_eq!(started.len(), 24, "started_at = {started}");
        assert!(started.ends_with('Z') && started.as_bytes()[19] == b'.');
    }

    #[test]
    fn trace_id_is_32_hex_and_ids_are_unique() {
        let config = config();
        assert!(begin(&config));
        let body = finish(&config, "a", 500.0).unwrap();
        let at = body.find("\"trace_id\":\"").unwrap() + "\"trace_id\":\"".len();
        let id = &body[at..at + 32];
        assert!(id.chars().all(|c| c.is_ascii_hexdigit()));
        assert_eq!(&body[at + 32..at + 33], "\"");

        let ids: std::collections::HashSet<String> = (0..1000).map(|_| random_hex(1)).collect();
        assert_eq!(ids.len(), 1000);
    }

    #[test]
    fn an_unknown_kind_is_sent_as_other_since_the_server_would_reject_the_whole_trace() {
        let config = config();
        begin(&config);
        record_leaf("q", "db", SystemTime::now(), 1.0, HashMap::new());
        record_leaf("r", "database", SystemTime::now(), 1.0, HashMap::new());
        let spans = spans_of(&finish(&config, "root", 500.0).unwrap());
        let by_name: HashMap<_, _> = spans.iter().map(|s| (s["name"].clone(), s)).collect();
        assert_eq!(by_name["q"]["kind"], "other");
        assert_eq!(by_name["r"]["kind"], "database");
    }

    #[test]
    fn nothing_is_sent_under_the_threshold_and_the_trace_is_cleared() {
        let config = config();
        begin(&config);
        assert_eq!(finish(&config, "GET /fast", 1.0), None);
        assert!(!is_active());
    }

    #[test]
    fn tracing_off_or_reporting_disabled_starts_no_trace() {
        let mut off = config();
        off.track_tracing = false;
        assert!(begin(&off), "the caller still owns a (no-op) root");
        assert!(!is_active());
        assert_eq!(finish(&off, "x", 500.0), None);

        let mut disabled = config();
        disabled.dsn = None;
        assert!(begin(&disabled));
        assert!(!is_active());
    }

    #[test]
    fn a_second_begin_on_the_same_thread_reports_a_trace_is_already_open() {
        let config = config();
        assert!(begin(&config));
        assert!(!begin(&config));
        finish(&config, "x", 500.0);
    }

    #[test]
    fn caps_a_trace_at_500_spans_including_the_root() {
        let config = config();
        begin(&config);
        for _ in 0..700 {
            record_leaf("q", "database", SystemTime::now(), 1.0, HashMap::new());
        }
        let body = finish(&config, "GET /x", 2000.0).unwrap();
        assert_eq!(body.matches("\"span_id\"").count(), 500);
    }

    #[test]
    fn the_trace_is_per_thread() {
        let config = config();
        begin(&config);
        std::thread::spawn(|| assert!(!is_active())).join().unwrap();
        finish(&config, "x", 500.0);
    }
}