grafos_observe/

span.rs

//! Resource-attributed spans for grafOS distributed tracing.
//!
//! A [`ResourceSpan`] extends standard OTel span data with grafOS-specific
//! resource attribution: lease IDs, operation counts, byte throughput, and
//! lease cost. This metadata makes grafOS traces strictly richer than
//! standard application traces.

use alloc::string::String;
use alloc::vec::Vec;

use crate::event::{OpType, ResourceType};
use crate::trace::{SpanId, TraceContext, TraceId};

/// Status of a completed span.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SpanStatus {
    /// The operation completed without error.
    Ok,
    /// The operation failed with the given error message.
    Error(String),
}

/// Key for operation count tracking: (resource type, operation type).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct OpKey {
    pub resource_type: ResourceType,
    pub op_type: OpType,
}

/// A resource-attributed span carrying grafOS-specific metadata.
///
/// Populated automatically by lease API instrumentation hooks and the
/// `#[grafos::instrument]` proc macro. Exported to OTel collectors via
/// the `otlp` exporter (requires `otlp` feature).
#[derive(Debug, Clone)]
pub struct ResourceSpan {
    /// Human-readable name for this span (usually the function name).
    pub name: String,

    /// Trace context (trace_id + span_id + flags).
    pub trace_context: TraceContext,

    /// Parent span ID (if this span has a parent in the same service).
    pub parent_span_id: Option<SpanId>,

    /// Start time as Unix microseconds.
    pub start_time_unix_us: u64,

    /// End time as Unix microseconds.
    pub end_time_unix_us: u64,

    /// Completion status.
    pub status: SpanStatus,

    // -- Resource attribution --
    /// Lease IDs touched during this span.
    pub lease_ids: Vec<u128>,

    /// Operation counts by (resource_type, op_type).
    pub ops: Vec<(OpKey, u64)>,

    /// Total bytes read across all lease operations in this span.
    pub bytes_read: u64,

    /// Total bytes written across all lease operations in this span.
    pub bytes_written: u64,

    /// Sum of (leased_bytes * seconds_held) for all leases in this span.
    pub lease_cost_byte_secs: u64,

    /// Time spent waiting for lease acquisition, in microseconds.
    pub lease_acquire_wait_us: u64,

    /// User-defined key-value attributes.
    pub attributes: Vec<(String, String)>,
}

impl ResourceSpan {
    /// Create a new span with the given name and trace context.
    pub fn new(name: &str, trace_context: TraceContext) -> Self {
        ResourceSpan {
            name: String::from(name),
            trace_context,
            parent_span_id: None,
            start_time_unix_us: 0,
            end_time_unix_us: 0,
            status: SpanStatus::Ok,
            lease_ids: Vec::new(),
            ops: Vec::new(),
            bytes_read: 0,
            bytes_written: 0,
            lease_cost_byte_secs: 0,
            lease_acquire_wait_us: 0,
            attributes: Vec::new(),
        }
    }

    /// Duration in microseconds.
    pub fn duration_us(&self) -> u64 {
        self.end_time_unix_us
            .saturating_sub(self.start_time_unix_us)
    }

    /// Record an operation.
    pub fn record_op(&mut self, resource_type: ResourceType, op_type: OpType, count: u64) {
        let key = OpKey {
            resource_type,
            op_type,
        };
        if let Some(entry) = self.ops.iter_mut().find(|(k, _)| *k == key) {
            entry.1 += count;
        } else {
            self.ops.push((key, count));
        }
    }

    /// Get the operation count for a specific (resource_type, op_type) pair.
    pub fn op_count(&self, resource_type: ResourceType, op_type: OpType) -> u64 {
        let key = OpKey {
            resource_type,
            op_type,
        };
        self.ops
            .iter()
            .find(|(k, _)| *k == key)
            .map(|(_, c)| *c)
            .unwrap_or(0)
    }

    /// Add a lease ID to the span's lease list.
    pub fn add_lease_id(&mut self, lease_id: u128) {
        if !self.lease_ids.contains(&lease_id) {
            self.lease_ids.push(lease_id);
        }
    }

    /// Add a user-defined attribute.
    pub fn set_attribute(&mut self, key: &str, value: &str) {
        if let Some(entry) = self.attributes.iter_mut().find(|(k, _)| k == key) {
            entry.1 = String::from(value);
        } else {
            self.attributes
                .push((String::from(key), String::from(value)));
        }
    }

    /// Convenience: trace ID from the context.
    pub fn trace_id(&self) -> TraceId {
        self.trace_context.trace_id
    }

    /// Convenience: span ID from the context.
    pub fn span_id(&self) -> SpanId {
        self.trace_context.span_id
    }
}

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

    fn test_ctx() -> TraceContext {
        let mut bytes = [0u8; 24];
        for (i, b) in bytes.iter_mut().enumerate() {
            *b = (i as u8).wrapping_add(0x42);
        }
        TraceContext::new_root(&bytes)
    }

    #[test]
    fn new_span_has_correct_defaults() {
        let ctx = test_ctx();
        let span = ResourceSpan::new("test_op", ctx);
        assert_eq!(span.name, "test_op");
        assert_eq!(span.trace_id(), ctx.trace_id);
        assert_eq!(span.status, SpanStatus::Ok);
        assert!(span.lease_ids.is_empty());
        assert!(span.ops.is_empty());
        assert_eq!(span.bytes_read, 0);
        assert_eq!(span.bytes_written, 0);
    }

    #[test]
    fn duration_calculation() {
        let mut span = ResourceSpan::new("d", test_ctx());
        span.start_time_unix_us = 1000;
        span.end_time_unix_us = 2500;
        assert_eq!(span.duration_us(), 1500);
    }

    #[test]
    fn record_op_accumulates() {
        let mut span = ResourceSpan::new("ops", test_ctx());
        span.record_op(ResourceType::Mem, OpType::Write, 5);
        span.record_op(ResourceType::Mem, OpType::Write, 3);
        span.record_op(ResourceType::Block, OpType::ReadBlock, 1);

        assert_eq!(span.op_count(ResourceType::Mem, OpType::Write), 8);
        assert_eq!(span.op_count(ResourceType::Block, OpType::ReadBlock), 1);
        assert_eq!(span.op_count(ResourceType::Gpu, OpType::GpuSubmit), 0);
    }

    #[test]
    fn add_lease_id_deduplicates() {
        let mut span = ResourceSpan::new("leases", test_ctx());
        span.add_lease_id(1);
        span.add_lease_id(2);
        span.add_lease_id(1); // duplicate
        assert_eq!(span.lease_ids.len(), 2);
    }

    #[test]
    fn set_attribute_updates_existing() {
        let mut span = ResourceSpan::new("attrs", test_ctx());
        span.set_attribute("env", "dev");
        span.set_attribute("env", "prod");
        assert_eq!(span.attributes.len(), 1);
        assert_eq!(span.attributes[0].1, "prod");
    }

    #[test]
    fn span_status_error() {
        let mut span = ResourceSpan::new("fail", test_ctx());
        span.status = SpanStatus::Error(String::from("connection refused"));
        match &span.status {
            SpanStatus::Error(msg) => assert_eq!(msg, "connection refused"),
            _ => panic!("expected error status"),
        }
    }
}