grafos_observe/
macro_support.rs

1//! Support module for the `#[grafos_instrument]` proc macro.
2//!
3//! This module is `#[doc(hidden)]` — it exists solely to provide the runtime
4//! side of the proc macro expansion. Do not depend on it directly.
5
6/// Log level for instrumented spans.
7#[derive(Debug, Clone, Copy)]
8pub enum Level {
9    #[allow(non_camel_case_types)]
10    debug,
11    #[allow(non_camel_case_types)]
12    info,
13    #[allow(non_camel_case_types)]
14    warn,
15}
16
17/// Record a completed instrumented span.
18///
19/// Called by the `#[grafos_instrument]` macro expansion at function exit.
20/// Formats the span data and emits it via the appropriate output channel.
21#[allow(clippy::too_many_arguments)]
22pub fn record_span(
23    function_name: &str,
24    level: Level,
25    duration_us: u64,
26    leases_held: usize,
27    total_bytes: u64,
28    nodes: usize,
29    fbmu_writes: u64,
30    fbmu_reads: u64,
31    fbbu_writes: u64,
32    fbbu_reads: u64,
33    gpu_submits: u64,
34    tasklet_submits: u64,
35) {
36    let msg = format_span(
37        function_name,
38        duration_us,
39        leases_held,
40        total_bytes,
41        nodes,
42        fbmu_writes,
43        fbmu_reads,
44        fbbu_writes,
45        fbbu_reads,
46        gpu_submits,
47        tasklet_submits,
48    );
49
50    #[cfg(feature = "tracing")]
51    {
52        match level {
53            Level::debug => tracing::debug!("{}", msg),
54            Level::info => tracing::info!("{}", msg),
55            Level::warn => tracing::warn!("{}", msg),
56        }
57    }
58
59    #[cfg(not(feature = "tracing"))]
60    {
61        let level_str = match level {
62            Level::debug => "DEBUG",
63            Level::info => "INFO",
64            Level::warn => "WARN",
65        };
66        #[cfg(feature = "std")]
67        eprintln!("[grafos::{level_str}] {msg}");
68        #[cfg(not(feature = "std"))]
69        let _ = (level_str, msg);
70    }
71
72    // Update global metrics
73    let m = crate::metrics::FabricMetrics::global();
74    m.op_latency.observe(duration_us);
75}
76
77#[allow(clippy::too_many_arguments)]
78fn format_span(
79    function_name: &str,
80    duration_us: u64,
81    leases_held: usize,
82    total_bytes: u64,
83    nodes: usize,
84    fbmu_writes: u64,
85    fbmu_reads: u64,
86    fbbu_writes: u64,
87    fbbu_reads: u64,
88    gpu_submits: u64,
89    tasklet_submits: u64,
90) -> String {
91    use alloc::format;
92    use alloc::string::String;
93
94    let duration_str = if duration_us >= 1_000_000 {
95        format!("{:.1}s", duration_us as f64 / 1_000_000.0)
96    } else if duration_us >= 1_000 {
97        format!("{:.1}ms", duration_us as f64 / 1_000.0)
98    } else {
99        format!("{duration_us}us")
100    };
101
102    let mut parts = String::with_capacity(256);
103    parts.push_str(&format!("{function_name}: {duration_str}"));
104
105    if leases_held > 0 || total_bytes > 0 {
106        let bytes_str = if total_bytes >= 1_073_741_824 {
107            format!("{:.1}GB", total_bytes as f64 / 1_073_741_824.0)
108        } else if total_bytes >= 1_048_576 {
109            format!("{:.1}MB", total_bytes as f64 / 1_048_576.0)
110        } else if total_bytes >= 1024 {
111            format!("{:.1}KB", total_bytes as f64 / 1024.0)
112        } else {
113            format!("{total_bytes}B")
114        };
115        parts.push_str(&format!(
116            ", {leases_held} leases ({bytes_str} across {nodes} nodes)"
117        ));
118    }
119
120    let mut ops = alloc::vec::Vec::new();
121    if fbmu_writes > 0 {
122        ops.push(format!("{fbmu_writes} FBMU writes"));
123    }
124    if fbmu_reads > 0 {
125        ops.push(format!("{fbmu_reads} FBMU reads"));
126    }
127    if fbbu_writes > 0 {
128        ops.push(format!("{fbbu_writes} FBBU writes"));
129    }
130    if fbbu_reads > 0 {
131        ops.push(format!("{fbbu_reads} FBBU reads"));
132    }
133    if gpu_submits > 0 {
134        ops.push(format!("{gpu_submits} GPU submits"));
135    }
136    if tasklet_submits > 0 {
137        ops.push(format!("{tasklet_submits} tasklet submits"));
138    }
139    if !ops.is_empty() {
140        parts.push_str(", ");
141        parts.push_str(&ops.join(", "));
142    }
143
144    parts
145}
146
147#[cfg(test)]
148mod tests {
149    use super::*;
150
151    #[test]
152    fn format_span_basic() {
153        let msg = format_span("my_func", 12_000, 0, 0, 0, 0, 0, 0, 0, 0, 0);
154        assert!(msg.contains("my_func"));
155        assert!(msg.contains("12.0ms"));
156    }
157
158    #[test]
159    fn format_span_with_leases() {
160        let msg = format_span(
161            "batch_op",
162            250,
163            3,
164            2_147_483_648, // 2GB
165            2,
166            47,
167            12,
168            0,
169            0,
170            0,
171            0,
172        );
173        assert!(msg.contains("batch_op"));
174        assert!(msg.contains("250us"));
175        assert!(msg.contains("3 leases"));
176        assert!(msg.contains("2.0GB"));
177        assert!(msg.contains("2 nodes"));
178        assert!(msg.contains("47 FBMU writes"));
179        assert!(msg.contains("12 FBMU reads"));
180    }
181
182    #[test]
183    fn format_span_seconds() {
184        let msg = format_span("slow", 2_500_000, 0, 0, 0, 0, 0, 0, 0, 0, 0);
185        assert!(msg.contains("2.5s"));
186    }
187
188    #[test]
189    fn format_span_all_op_types() {
190        let msg = format_span("full_op", 100, 1, 1024, 1, 1, 2, 3, 4, 5, 6);
191        assert!(msg.contains("1 FBMU writes"));
192        assert!(msg.contains("2 FBMU reads"));
193        assert!(msg.contains("3 FBBU writes"));
194        assert!(msg.contains("4 FBBU reads"));
195        assert!(msg.contains("5 GPU submits"));
196        assert!(msg.contains("6 tasklet submits"));
197    }
198
199    #[test]
200    fn record_span_does_not_panic() {
201        record_span("test_fn", Level::info, 100, 0, 0, 0, 0, 0, 0, 0, 0, 0);
202        record_span("test_fn", Level::debug, 100, 1, 4096, 1, 5, 3, 0, 0, 0, 0);
203        record_span("test_fn", Level::warn, 100, 0, 0, 0, 0, 0, 0, 0, 0, 0);
204    }
205}