saddle-observability 0.1.1

use std::time::Instant;

use saddle_core::{
    ApplicationId, CallContext, ErrorKind, ModuleId, OperationId, SaddleError, ServiceId, SpanId,
};
use serde_json::{Value, json};

use crate::{
    DomainEvent, EventLevel, InboundTrace, Observer, logger::LogRecord, trace::select_trace_id,
};

/// Framework boundaries that automatically produce start and finish records.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CallKind {
    ExternalRequest,
    Service,
    Database,
    Transaction,
}

impl CallKind {
    const fn as_str(self) -> &'static str {
        match self {
            Self::ExternalRequest => "external_request",
            Self::Service => "service",
            Self::Database => "database",
            Self::Transaction => "transaction",
        }
    }
}

/// Stable completion classification for framework call logs.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CallOutcome {
    Success,
    Failure,
    Abandoned,
}

impl CallOutcome {
    const fn as_str(self) -> &'static str {
        match self {
            Self::Success => "success",
            Self::Failure => "failure",
            Self::Abandoned => "abandoned",
        }
    }
}

/// An in-flight framework boundary.
///
/// Dropping it without calling [`ActiveCall::succeed`] or
/// [`ActiveCall::fail`] emits an `abandoned` finish record, which preserves a
/// balanced trace when a future is cancelled or unwinds.
pub struct ActiveCall {
    observer: Observer,
    context: CallContext,
    parent_span_id: Option<SpanId>,
    kind: CallKind,
    started_at: Instant,
    finished: bool,
}

impl Observer {
    /// Starts an external request, inheriting a valid trace id or creating one.
    pub fn start_external_call(
        &self,
        application: impl Into<ApplicationId>,
        module: impl Into<ModuleId>,
        service: impl Into<ServiceId>,
        operation: impl Into<OperationId>,
        inbound_trace_id: Option<&str>,
    ) -> (ActiveCall, InboundTrace) {
        let (trace_id, inbound) = select_trace_id(inbound_trace_id, || self.new_trace_id());
        let context = CallContext::new(
            application.into(),
            module.into(),
            service.into(),
            operation.into(),
            trace_id,
            self.new_span_id(),
        );
        (
            ActiveCall::start(
                self.clone(),
                context,
                None,
                CallKind::ExternalRequest,
                Some(inbound),
            ),
            inbound,
        )
    }

    /// Starts a Service, DB or transaction child call on the parent's trace.
    pub fn start_child_call(
        &self,
        parent: &CallContext,
        kind: CallKind,
        module: impl Into<ModuleId>,
        service: impl Into<ServiceId>,
        operation: impl Into<OperationId>,
    ) -> ActiveCall {
        let context = CallContext::new(
            parent.application().clone(),
            module.into(),
            service.into(),
            operation.into(),
            parent.trace_id(),
            self.new_span_id(),
        );
        ActiveCall::start(self.clone(), context, Some(parent.span_id()), kind, None)
    }

    /// Records an intentional domain event without accepting arbitrary objects.
    pub fn record_domain_event(&self, context: &CallContext, event: DomainEvent) {
        let mut record = LogRecord::new(event.level, "domain.event");
        add_context(&mut record, context, None);
        add_identity(&mut record, context);
        record
            .data
            .insert("name".to_owned(), Value::String(event.name));
        let fields = event
            .fields
            .into_iter()
            .map(|(name, value)| (name, value.into_json()))
            .collect();
        record
            .data
            .insert("fields".to_owned(), Value::Object(fields));
        self.emit(record);
    }
}

impl ActiveCall {
    fn start(
        observer: Observer,
        context: CallContext,
        parent_span_id: Option<SpanId>,
        kind: CallKind,
        inbound_trace: Option<InboundTrace>,
    ) -> Self {
        let mut record = LogRecord::new(EventLevel::Info, "framework.call.started");
        add_context(&mut record, &context, parent_span_id);
        add_call_identity(&mut record, &context, kind);
        if let Some(inbound_trace) = inbound_trace {
            record.data.insert(
                "trace_source".to_owned(),
                Value::String(inbound_trace.as_str().to_owned()),
            );
        }
        observer.emit(record);

        Self {
            observer,
            context,
            parent_span_id,
            kind,
            started_at: Instant::now(),
            finished: false,
        }
    }

    pub fn context(&self) -> &CallContext {
        &self.context
    }

    pub fn succeed(mut self) {
        self.finish(CallOutcome::Success, None);
    }

    pub fn fail(mut self, error: &SaddleError) {
        self.finish(CallOutcome::Failure, Some(error));
    }

    fn finish(&mut self, outcome: CallOutcome, error: Option<&SaddleError>) {
        if self.finished {
            return;
        }

        let level = match outcome {
            CallOutcome::Success => EventLevel::Info,
            CallOutcome::Failure | CallOutcome::Abandoned => EventLevel::Error,
        };
        let mut record = LogRecord::new(level, "framework.call.finished");
        add_context(&mut record, &self.context, self.parent_span_id);
        add_call_identity(&mut record, &self.context, self.kind);
        record.data.insert(
            "outcome".to_owned(),
            Value::String(outcome.as_str().to_owned()),
        );
        record.data.insert(
            "elapsed_us".to_owned(),
            json!(u64::try_from(self.started_at.elapsed().as_micros()).unwrap_or(u64::MAX)),
        );
        if let Some(error) = error {
            record.data.insert(
                "error_kind".to_owned(),
                Value::String(error_kind(error.kind()).to_owned()),
            );
            record.data.insert(
                "error_code".to_owned(),
                Value::String(error.code().to_owned()),
            );
        }
        self.observer.emit(record);
        self.finished = true;
    }
}

impl Drop for ActiveCall {
    fn drop(&mut self) {
        self.finish(CallOutcome::Abandoned, None);
    }
}

fn add_context(record: &mut LogRecord, context: &CallContext, parent_span_id: Option<SpanId>) {
    record.trace_id = Some(context.trace_id().to_string());
    record.span_id = Some(context.span_id().to_string());
    record.parent_span_id = parent_span_id.map(|span_id| span_id.to_string());
}

fn add_call_identity(record: &mut LogRecord, context: &CallContext, kind: CallKind) {
    record.data.insert(
        "call_kind".to_owned(),
        Value::String(kind.as_str().to_owned()),
    );
    add_identity(record, context);
}

fn add_identity(record: &mut LogRecord, context: &CallContext) {
    record.data.insert(
        "application".to_owned(),
        Value::String(context.application().to_string()),
    );
    record.data.insert(
        "module".to_owned(),
        Value::String(context.module().to_string()),
    );
    record.data.insert(
        "service".to_owned(),
        Value::String(context.service().to_string()),
    );
    record.data.insert(
        "operation".to_owned(),
        Value::String(context.operation().to_string()),
    );
}

const fn error_kind(kind: ErrorKind) -> &'static str {
    match kind {
        ErrorKind::InvalidArgument => "invalid_argument",
        ErrorKind::NotFound => "not_found",
        ErrorKind::Conflict => "conflict",
        ErrorKind::Business => "business",
        ErrorKind::Unavailable => "unavailable",
        ErrorKind::Infrastructure => "infrastructure",
        ErrorKind::Internal => "internal",
        _ => "unknown",
    }
}

#[cfg(test)]
mod tests {
    use std::{
        future::Future,
        io,
        sync::{Arc, Mutex},
        task::{Context, Poll, Wake, Waker},
        thread,
    };

    use saddle_core::{ErrorKind, SaddleError};

    use super::*;
    use crate::ObserverConfig;

    struct ThreadWaker(thread::Thread);

    impl Wake for ThreadWaker {
        fn wake(self: Arc<Self>) {
            self.0.unpark();
        }
    }

    fn block_on<T>(future: impl Future<Output = T>) -> T {
        let waker = Waker::from(Arc::new(ThreadWaker(thread::current())));
        let mut context = Context::from_waker(&waker);
        let mut future = std::pin::pin!(future);
        loop {
            match future.as_mut().poll(&mut context) {
                Poll::Ready(output) => return output,
                Poll::Pending => thread::park(),
            }
        }
    }

    #[derive(Clone, Default)]
    struct Capture(Arc<Mutex<Vec<u8>>>);

    impl io::Write for Capture {
        fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
            self.0.lock().unwrap().extend_from_slice(bytes);
            Ok(bytes.len())
        }

        fn flush(&mut self) -> io::Result<()> {
            Ok(())
        }
    }

    fn records(capture: &Capture) -> Vec<Value> {
        String::from_utf8(capture.0.lock().unwrap().clone())
            .unwrap()
            .lines()
            .map(|line| serde_json::from_str(line).unwrap())
            .collect()
    }

    #[test]
    fn external_and_child_calls_share_trace_and_record_parentage() {
        let capture = Capture::default();
        let observer = Observer::with_writer(ObserverConfig::default(), capture.clone()).unwrap();
        let trace = "00112233445566778899aabbccddeeff";
        let (external, source) =
            observer.start_external_call("shop", "orders", "orders", "create", Some(trace));
        assert_eq!(source, InboundTrace::Inherited);
        assert_eq!(external.context().trace_id().to_string(), trace);

        let child = observer.start_child_call(
            external.context(),
            CallKind::Service,
            "users",
            "users",
            "validate",
        );
        assert_eq!(child.context().trace_id(), external.context().trace_id());
        assert_ne!(child.context().span_id(), external.context().span_id());
        child.succeed();
        external.succeed();
        block_on(observer.flush()).unwrap();

        let records = records(&capture);
        assert_eq!(records.len(), 4);
        assert_eq!(records[1]["parent_span_id"], records[0]["span_id"]);
        assert!(records.iter().all(|record| record["trace_id"] == trace));
    }

    #[test]
    fn invalid_inbound_trace_is_replaced_and_identified_in_start_record() {
        let capture = Capture::default();
        let observer = Observer::with_writer(ObserverConfig::default(), capture.clone()).unwrap();
        let (call, source) = observer.start_external_call(
            "shop",
            "orders",
            "orders",
            "create",
            Some("not-a-valid-trace"),
        );
        assert_eq!(source, InboundTrace::ReplacedInvalid);
        assert_ne!(call.context().trace_id().as_u128(), 0);
        call.succeed();
        block_on(observer.flush()).unwrap();

        let records = records(&capture);
        assert_eq!(records[0]["trace_source"], "replaced_invalid");
    }

    #[test]
    fn failure_logs_classification_but_not_sensitive_message() {
        let capture = Capture::default();
        let observer = Observer::with_writer(ObserverConfig::default(), capture.clone()).unwrap();
        let (call, _) = observer.start_external_call("shop", "orders", "orders", "create", None);
        let error = SaddleError::new(
            ErrorKind::Infrastructure,
            "db.query_failed",
            "password=must-not-be-logged",
        );
        call.fail(&error);
        block_on(observer.flush()).unwrap();

        let output = String::from_utf8(capture.0.lock().unwrap().clone()).unwrap();
        assert!(output.contains("db.query_failed"));
        assert!(output.contains("infrastructure"));
        assert!(!output.contains("must-not-be-logged"));
    }

    #[test]
    fn domain_event_contains_only_explicit_scalar_fields() {
        let capture = Capture::default();
        let observer = Observer::with_writer(ObserverConfig::default(), capture.clone()).unwrap();
        let (call, _) = observer.start_external_call("shop", "orders", "orders", "create", None);
        observer.record_domain_event(
            call.context(),
            DomainEvent::new("order.created")
                .unwrap()
                .field("order_id", "o-1")
                .unwrap()
                .field("item_count", 2)
                .unwrap(),
        );
        call.succeed();
        block_on(observer.flush()).unwrap();

        let records = records(&capture);
        assert_eq!(records[1]["event"], "domain.event");
        assert_eq!(records[1]["name"], "order.created");
        assert_eq!(records[1]["fields"]["order_id"], "o-1");
        assert_eq!(records[1]["fields"]["item_count"], 2);
    }

    #[test]
    fn dropped_call_is_recorded_as_abandoned() {
        let capture = Capture::default();
        let observer = Observer::with_writer(ObserverConfig::default(), capture.clone()).unwrap();
        let (call, _) = observer.start_external_call("shop", "orders", "orders", "create", None);
        drop(call);
        block_on(observer.flush()).unwrap();
        let records = records(&capture);
        assert_eq!(records[1]["outcome"], "abandoned");
    }
}