grafos_observe/

dataplane.rs

//! Phase 219 data-plane span vocabulary.
//!
//! Producer-side data-plane operations (FBMU, FBBU, GPU session, runtime
//! capability cache) emit canonical [`ResourceSpan`]s using the helpers
//! in this module. Span names come from the [`PHASE_219_SPAN_NAMES`]
//! contract; lease IDs are kept in the span body, never on Prometheus
//! labels.
//!
//! The vocabulary lives here — not in `grafos-scheduler` — so that
//! data-plane producers (e.g. `grafos-runtime`, `grafos-std`) can emit
//! spans without depending on the scheduler crate. This aligns with the
//! v1.1 design (observability is substrate-level, not policy-level) and
//! with §3.3 invariant 2 (every edge has a trace context, so every edge
//! realization is a candidate emission site).
//!
//! [`PHASE_219_SPAN_NAMES`]: crate::PHASE_219_SPAN_NAMES

use alloc::string::ToString;

use grafos_core::ResourceKind;

use crate::event::OpType;
use crate::export::emit_span;
use crate::span::{ResourceSpan, SpanStatus};
use crate::trace::{TraceContext, TraceContextError};

fn resource_kind_label(kind: ResourceKind) -> &'static str {
    match kind {
        ResourceKind::Mem => "mem",
        ResourceKind::Block => "block",
        ResourceKind::Net => "net",
        ResourceKind::Cpu => "cpu",
        ResourceKind::Gpu => "gpu",
        ResourceKind::GpuMem => "gpumem",
        ResourceKind::Tasklet => "tasklet",
    }
}

fn optional_trace_context(
    traceparent: Option<&str>,
) -> Result<Option<TraceContext>, TraceContextError> {
    match traceparent {
        Some(value) => TraceContext::from_w3c_string(value).map(Some),
        None => Ok(None),
    }
}

/// Phase 219 data-plane span family. Each variant maps 1:1 to a span
/// name in [`PHASE_219_SPAN_NAMES`] so callers can refer to the family
/// by typed identity rather than stringly-typed names.
///
/// [`PHASE_219_SPAN_NAMES`]: crate::PHASE_219_SPAN_NAMES
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DataplaneSpanKind {
    /// `fabricbios.data.fbmu` — FBMU memory data-plane operation.
    Fbmu,
    /// `fabricbios.data.fbbu` — FBBU block data-plane operation.
    Fbbu,
    /// `fabricbios.data.gpu_session` — GPU session lifecycle/IO.
    GpuSession,
    /// `grafos.runtime.capability_cache` — capability cache hit/miss.
    CapabilityCache,
}

impl DataplaneSpanKind {
    /// Stable span name from the Phase 219 contract.
    pub fn span_name(self) -> &'static str {
        match self {
            DataplaneSpanKind::Fbmu => "fabricbios.data.fbmu",
            DataplaneSpanKind::Fbbu => "fabricbios.data.fbbu",
            DataplaneSpanKind::GpuSession => "fabricbios.data.gpu_session",
            DataplaneSpanKind::CapabilityCache => "grafos.runtime.capability_cache",
        }
    }
}

/// Map a producer-side [`OpType`] to the Phase 219 data-plane span
/// family it belongs to. `TaskletSubmit` is a control-plane submission
/// and has no data-plane span kind — returning `None` keeps the
/// vocabulary closed (callers must handle that variant explicitly).
///
/// The match is exhaustive: a future [`OpType`] variant produces a
/// compile error here, surfacing the missing mapping at build time
/// rather than at trace export time.
pub fn dataplane_span_kind_for_op_type(op_type: OpType) -> Option<DataplaneSpanKind> {
    match op_type {
        OpType::Read | OpType::Write => Some(DataplaneSpanKind::Fbmu),
        OpType::ReadBlock | OpType::WriteBlock => Some(DataplaneSpanKind::Fbbu),
        OpType::GpuSubmit => Some(DataplaneSpanKind::GpuSession),
        OpType::TaskletSubmit => None,
    }
}

/// Stable operation label for a producer-side [`OpType`]. Mirrors
/// [`OpType`]'s [`fmt::Display`] impl but is `&'static str` so callers
/// can set it as a span attribute without allocation. The match is
/// exhaustive for the same reason as [`dataplane_span_kind_for_op_type`].
///
/// [`fmt::Display`]: core::fmt::Display
pub fn dataplane_operation_label(op_type: OpType) -> &'static str {
    match op_type {
        OpType::Read => "read",
        OpType::Write => "write",
        OpType::ReadBlock => "read_block",
        OpType::WriteBlock => "write_block",
        OpType::GpuSubmit => "gpu_submit",
        OpType::TaskletSubmit => "tasklet_submit",
    }
}

/// Build a data-plane span using the Phase 219 canonical span names.
/// The helper keeps the lease id in the span body (never a Prometheus
/// label) and routes resource kind / operation / outcome through stable
/// attribute keys so consumers see the same shape across families.
pub fn build_dataplane_span(
    trace_context: TraceContext,
    kind: DataplaneSpanKind,
    operation: &str,
    outcome: &str,
    resource_kind: ResourceKind,
    lease_id: Option<u128>,
    start_time_unix_us: u64,
    end_time_unix_us: u64,
) -> ResourceSpan {
    let mut span = ResourceSpan::new(kind.span_name(), trace_context);
    span.start_time_unix_us = start_time_unix_us;
    span.end_time_unix_us = end_time_unix_us;
    span.status = if outcome == "failed" || outcome == "rejected" {
        SpanStatus::Error(outcome.to_string())
    } else {
        SpanStatus::Ok
    };
    span.set_attribute("operation", operation);
    span.set_attribute("outcome", outcome);
    span.set_attribute("resource_kind", resource_kind_label(resource_kind));
    if let Some(lease_id) = lease_id {
        span.add_lease_id(lease_id);
    }
    span
}

/// Emit a data-plane span through grafos-observe's shared span exporter.
pub fn emit_dataplane_span(
    trace_context: TraceContext,
    kind: DataplaneSpanKind,
    operation: &str,
    outcome: &str,
    resource_kind: ResourceKind,
    lease_id: Option<u128>,
    start_time_unix_us: u64,
    end_time_unix_us: u64,
) {
    emit_span(build_dataplane_span(
        trace_context,
        kind,
        operation,
        outcome,
        resource_kind,
        lease_id,
        start_time_unix_us,
        end_time_unix_us,
    ));
}

/// Build a data-plane span from an optional W3C traceparent.
///
/// Absent trace context returns `Ok(None)`; invalid trace context
/// returns a typed error. The helper intentionally does not synthesize
/// root contexts — data-plane ops without an inbound trace are surfaced
/// as "no span" so callers can decide between linking and root-creation
/// explicitly (v1.1 §3.3 invariant 2 / no silent best-effort loss).
pub fn build_dataplane_span_from_traceparent(
    traceparent: Option<&str>,
    kind: DataplaneSpanKind,
    operation: &str,
    outcome: &str,
    resource_kind: ResourceKind,
    lease_id: Option<u128>,
    start_time_unix_us: u64,
    end_time_unix_us: u64,
) -> Result<Option<ResourceSpan>, TraceContextError> {
    Ok(optional_trace_context(traceparent)?.map(|trace_context| {
        build_dataplane_span(
            trace_context,
            kind,
            operation,
            outcome,
            resource_kind,
            lease_id,
            start_time_unix_us,
            end_time_unix_us,
        )
    }))
}

/// Emit a data-plane span from an optional W3C traceparent.
pub fn emit_dataplane_span_from_traceparent(
    traceparent: Option<&str>,
    kind: DataplaneSpanKind,
    operation: &str,
    outcome: &str,
    resource_kind: ResourceKind,
    lease_id: Option<u128>,
    start_time_unix_us: u64,
    end_time_unix_us: u64,
) -> Result<(), TraceContextError> {
    if let Some(span) = build_dataplane_span_from_traceparent(
        traceparent,
        kind,
        operation,
        outcome,
        resource_kind,
        lease_id,
        start_time_unix_us,
        end_time_unix_us,
    )? {
        emit_span(span);
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::PHASE_219_SPAN_NAMES;
    use alloc::format;
    use alloc::vec;

    fn test_ctx() -> TraceContext {
        TraceContext::new_root(&[0x42; 24])
    }

    #[test]
    fn dataplane_span_kind_maps_to_phase_219_contract_names() {
        for kind in [
            DataplaneSpanKind::Fbmu,
            DataplaneSpanKind::Fbbu,
            DataplaneSpanKind::GpuSession,
            DataplaneSpanKind::CapabilityCache,
        ] {
            assert!(
                PHASE_219_SPAN_NAMES.contains(&kind.span_name()),
                "data-plane span name {} must be in the Phase 219 contract",
                kind.span_name(),
            );
        }
        assert_eq!(DataplaneSpanKind::Fbmu.span_name(), "fabricbios.data.fbmu");
        assert_eq!(DataplaneSpanKind::Fbbu.span_name(), "fabricbios.data.fbbu");
        assert_eq!(
            DataplaneSpanKind::GpuSession.span_name(),
            "fabricbios.data.gpu_session"
        );
        assert_eq!(
            DataplaneSpanKind::CapabilityCache.span_name(),
            "grafos.runtime.capability_cache"
        );
    }

    #[test]
    fn dataplane_span_records_lease_id_resource_kind_and_outcome() {
        let span = build_dataplane_span(
            test_ctx(),
            DataplaneSpanKind::Fbmu,
            "bind",
            "bound",
            ResourceKind::Mem,
            Some(0x1234),
            50,
            150,
        );

        assert_eq!(span.name, "fabricbios.data.fbmu");
        assert_eq!(span.duration_us(), 100);
        assert_eq!(span.lease_ids, vec![0x1234]);
        assert!(matches!(span.status, SpanStatus::Ok));
        assert_eq!(
            span.attributes
                .iter()
                .find(|(k, _)| k == "operation")
                .map(|(_, v)| v.as_str()),
            Some("bind")
        );
        assert_eq!(
            span.attributes
                .iter()
                .find(|(k, _)| k == "outcome")
                .map(|(_, v)| v.as_str()),
            Some("bound")
        );
        assert_eq!(
            span.attributes
                .iter()
                .find(|(k, _)| k == "resource_kind")
                .map(|(_, v)| v.as_str()),
            Some("mem")
        );
    }

    #[test]
    fn dataplane_span_marks_failed_and_rejected_outcomes_as_error() {
        let failed = build_dataplane_span(
            test_ctx(),
            DataplaneSpanKind::Fbbu,
            "read",
            "failed",
            ResourceKind::Block,
            Some(0x9),
            1,
            2,
        );
        assert!(matches!(failed.status, SpanStatus::Error(_)));

        let rejected = build_dataplane_span(
            test_ctx(),
            DataplaneSpanKind::CapabilityCache,
            "lookup",
            "rejected",
            ResourceKind::Net,
            None,
            1,
            2,
        );
        assert!(matches!(rejected.status, SpanStatus::Error(_)));
        assert!(
            rejected.lease_ids.is_empty(),
            "no lease id provided → span lease list stays empty"
        );
    }

    #[test]
    fn dataplane_span_uses_gpu_session_name_and_gpu_label() {
        let span = build_dataplane_span(
            test_ctx(),
            DataplaneSpanKind::GpuSession,
            "open",
            "ok",
            ResourceKind::Gpu,
            Some(0xab),
            10,
            30,
        );
        assert_eq!(span.name, "fabricbios.data.gpu_session");
        assert_eq!(
            span.attributes
                .iter()
                .find(|(k, _)| k == "resource_kind")
                .map(|(_, v)| v.as_str()),
            Some("gpu")
        );
    }

    #[test]
    fn dataplane_traceparent_helper_rejects_invalid_context() {
        let err = build_dataplane_span_from_traceparent(
            Some("not-a-traceparent"),
            DataplaneSpanKind::Fbmu,
            "bind",
            "bound",
            ResourceKind::Mem,
            Some(1),
            1,
            2,
        )
        .expect_err("invalid traceparent must not silently drop the span");

        assert!(matches!(err, TraceContextError::InvalidFormat));
    }

    #[test]
    fn dataplane_traceparent_helper_omits_absent_context() {
        let span = build_dataplane_span_from_traceparent(
            None,
            DataplaneSpanKind::CapabilityCache,
            "lookup",
            "hit",
            ResourceKind::Net,
            None,
            1,
            2,
        )
        .expect("absent traceparent is a valid no-span case");

        assert!(span.is_none());
    }

    #[test]
    fn dataplane_traceparent_helper_uses_supplied_context() {
        let ctx = test_ctx();
        let traceparent = ctx.to_w3c_string();
        let span = build_dataplane_span_from_traceparent(
            Some(&traceparent),
            DataplaneSpanKind::Fbbu,
            "write",
            "ok",
            ResourceKind::Block,
            Some(0x77),
            10,
            40,
        )
        .expect("valid traceparent")
        .expect("span should be built when traceparent is present");

        assert_eq!(span.name, "fabricbios.data.fbbu");
        assert_eq!(span.trace_context.trace_id, ctx.trace_id);
        assert_eq!(span.trace_context.span_id, ctx.span_id);
        assert_eq!(span.lease_ids, vec![0x77]);
        assert_eq!(
            span.attributes
                .iter()
                .find(|(k, _)| k == "resource_kind")
                .map(|(_, v)| v.as_str()),
            Some("block")
        );
    }

    #[test]
    fn dataplane_span_kind_for_op_type_maps_memory_ops_to_fbmu() {
        assert_eq!(
            dataplane_span_kind_for_op_type(OpType::Read),
            Some(DataplaneSpanKind::Fbmu)
        );
        assert_eq!(
            dataplane_span_kind_for_op_type(OpType::Write),
            Some(DataplaneSpanKind::Fbmu)
        );
    }

    #[test]
    fn dataplane_span_kind_for_op_type_maps_block_ops_to_fbbu() {
        assert_eq!(
            dataplane_span_kind_for_op_type(OpType::ReadBlock),
            Some(DataplaneSpanKind::Fbbu)
        );
        assert_eq!(
            dataplane_span_kind_for_op_type(OpType::WriteBlock),
            Some(DataplaneSpanKind::Fbbu)
        );
    }

    #[test]
    fn dataplane_span_kind_for_op_type_maps_gpu_submit_to_gpu_session() {
        assert_eq!(
            dataplane_span_kind_for_op_type(OpType::GpuSubmit),
            Some(DataplaneSpanKind::GpuSession)
        );
    }

    #[test]
    fn dataplane_span_kind_for_op_type_returns_none_for_tasklet_submit() {
        // TaskletSubmit is a control-plane submission, not a data-plane
        // span family in the Phase 219 contract. Callers must handle the
        // None case explicitly rather than defaulting to a data-plane
        // span name that doesn't belong.
        assert_eq!(dataplane_span_kind_for_op_type(OpType::TaskletSubmit), None);
    }

    #[test]
    fn dataplane_operation_label_matches_op_type_display_strings() {
        assert_eq!(dataplane_operation_label(OpType::Read), "read");
        assert_eq!(dataplane_operation_label(OpType::Write), "write");
        assert_eq!(dataplane_operation_label(OpType::ReadBlock), "read_block");
        assert_eq!(dataplane_operation_label(OpType::WriteBlock), "write_block");
        assert_eq!(dataplane_operation_label(OpType::GpuSubmit), "gpu_submit");
        assert_eq!(
            dataplane_operation_label(OpType::TaskletSubmit),
            "tasklet_submit"
        );
    }

    #[test]
    fn dataplane_operation_label_agrees_with_op_type_display_impl() {
        // Pin that our &'static str label never drifts from the
        // grafos-observe OpType Display impl. If a future variant
        // diverges, this test surfaces the drift at build time.
        for op_type in [
            OpType::Read,
            OpType::Write,
            OpType::ReadBlock,
            OpType::WriteBlock,
            OpType::GpuSubmit,
            OpType::TaskletSubmit,
        ] {
            assert_eq!(
                dataplane_operation_label(op_type),
                format!("{}", op_type),
                "&'static str label drifted from Display for {op_type:?}",
            );
        }
    }
}