lmflow/
trace.rs

1//! 逐次调用的执行 trace:有界事件环 + Chrome Trace 导出。
2//!
3//! 默认关闭(`config.trace_capacity == 0`)。开启时,每次算子回调(Open/Process/Close)记
4//! 一条 span 进有界环 —— 环满丢最旧,故内存有界。可经 [`to_chrome_trace_json`] 导出成
5//! chrome://tracing / perfetto 可读的 "Trace Event Format"(complete `"X"` 事件:含 `ts` +
6//! `dur`),从而看到「哪个节点在什么时刻、在哪个线程上跑了多久」的时间线 —— 这是聚合统计
7//! (总时长/最大/百分位)回答不了的。
8//!
9//! 并发:多个 worker 线程并发记录,整个环用一把 `Mutex` 保护。trace 是调试/剖析模式,这点
10//! 争用可接受,换来的是显然正确、tsan 干净。**关闭时热路径零开销** —— `GraphShared::trace`
11//! 为 `None` 时,scheduler 连时钟都不读(见 `call_kernel`)。
12
13use std::collections::{HashMap, VecDeque};
14use std::fmt::Write as _;
15use std::sync::Mutex;
16use std::thread::ThreadId;
17
18use crate::timestamp::Timestamp;
19
20/// 算子回调阶段。
21#[derive(Clone, Copy, Debug, PartialEq, Eq)]
22pub enum TracePhase {
23    Open,
24    Process,
25    Close,
26}
27
28impl TracePhase {
29    pub fn as_str(self) -> &'static str {
30        match self {
31            TracePhase::Open => "open",
32            TracePhase::Process => "process",
33            TracePhase::Close => "close",
34        }
35    }
36}
37
38/// 一条执行 span。时间均相对图 epoch,单位微秒。
39#[derive(Clone, Copy, Debug)]
40pub struct TraceEvent {
41    /// 节点索引(导出时解析成节点名/算子名)。
42    pub node: u32,
43    pub phase: TracePhase,
44    /// 该图内稳定分配的小整数线程号(chrome 的泳道 tid)。
45    pub tid: u32,
46    /// 相对图 epoch 的起始微秒。
47    pub start_us: i64,
48    /// 本次回调时长(微秒)。
49    pub dur_us: i64,
50    /// 本次激活的对齐时间戳原值。Open 阶段是 `Unstarted`、Close 阶段是 `Done`,都不是流内
51    /// 值 —— 导出时按名字记(见 [`to_chrome_trace_json`]),不然会是贴着 `i64::MIN/MAX` 的数。
52    pub input_ts: i64,
53}
54
55struct Inner {
56    events: VecDeque<TraceEvent>,
57    capacity: usize,
58    tids: HashMap<ThreadId, u32>,
59    next_tid: u32,
60}
61
62/// 有界 trace 事件环(FIFO,满了丢最旧)。
63pub struct TraceRing {
64    inner: Mutex<Inner>,
65}
66
67impl TraceRing {
68    /// 建一个容量为 `capacity` 条的环。每条 [`TraceEvent`] 40 字节,故 `capacity` 直接决定
69    /// 内存上限(如 4096 条 ≈ 160 KB)。引擎只在 `trace_capacity > 0` 时构造,但本构造器是
70    /// 公开 API,故 `0` 按 `1` 处理 —— 关键是任何输入下都必须有界。
71    pub fn new(capacity: usize) -> Self {
72        Self {
73            inner: Mutex::new(Inner {
74                events: VecDeque::with_capacity(capacity.min(1024)),
75                capacity: capacity.max(1),
76                tids: HashMap::new(),
77                next_tid: 0,
78            }),
79        }
80    }
81
82    /// 记一条 span。满了丢最旧。线程号在本环内稳定分配(首次见到某线程即给一个递增小整数)。
83    pub fn record(&self, node: u32, phase: TracePhase, start_us: i64, dur_us: i64, input_ts: i64) {
84        let id = std::thread::current().id();
85        let mut g = self.inner.lock().expect("trace ring poisoned");
86        let tid = match g.tids.get(&id) {
87            Some(&t) => t,
88            None => {
89                let t = g.next_tid;
90                g.next_tid += 1;
91                g.tids.insert(id, t);
92                t
93            }
94        };
95        if g.events.len() >= g.capacity {
96            g.events.pop_front();
97        }
98        g.events.push_back(TraceEvent {
99            node,
100            phase,
101            tid,
102            start_us,
103            dur_us,
104            input_ts,
105        });
106    }
107
108    /// 当前环内容的快照(**不清空**:环自身有界,可多次导出)。
109    pub fn snapshot(&self) -> Vec<TraceEvent> {
110        let g = self.inner.lock().expect("trace ring poisoned");
111        g.events.iter().copied().collect()
112    }
113
114    /// 清空(供图 reset 重跑;事件、线程号映射一并归零,容量不变)。
115    pub fn clear(&self) {
116        let mut g = self.inner.lock().expect("trace ring poisoned");
117        g.events.clear();
118        g.tids.clear();
119        g.next_tid = 0;
120    }
121}
122
123/// 把 JSON 字符串值需要转义的字符转义(引号/反斜杠/控制字符)。节点名/算子名由宿主定,
124/// 可能含任意字符,故必须转义,否则产出的不是合法 JSON。
125fn push_json_escaped(out: &mut String, s: &str) {
126    out.push('"');
127    for c in s.chars() {
128        match c {
129            '"' => out.push_str("\\\""),
130            '\\' => out.push_str("\\\\"),
131            '\n' => out.push_str("\\n"),
132            '\r' => out.push_str("\\r"),
133            '\t' => out.push_str("\\t"),
134            c if (c as u32) < 0x20 => {
135                let _ = write!(out, "\\u{:04x}", c as u32);
136            }
137            c => out.push(c),
138        }
139    }
140    out.push('"');
141}
142
143/// 把 span 导出成 Chrome Trace Event Format(chrome://tracing / perfetto 可直接打开)。
144///
145/// `labels[node]` = `(节点名, 算子名)`。产出 complete `"X"` 事件,时间单位微秒,`pid = 1`,
146/// `tid` 为环内线程号,并为每条泳道补一条 `"M"` 元事件命名 —— 没有它查看器只显示裸 tid
147/// 数字,而「哪个线程」正是这个视图的看点。
148///
149/// 手写 JSON 而非用 `serde_json`(它本就是本 crate 的依赖):一次导出可能上万条 span,直接
150/// 写字符串省掉为每条 span 建中间 `Value` 树的开销。
151pub fn to_chrome_trace_json(events: &[TraceEvent], labels: &[(String, String)]) -> String {
152    let mut out = String::with_capacity(128 + events.len() * 128);
153    out.push_str("{\"traceEvents\":[");
154    let mut first = true;
155
156    // ---- 泳道命名:每条出现过的 tid 一条 thread_name 元事件 ----
157    let mut tids: Vec<u32> = events.iter().map(|e| e.tid).collect();
158    tids.sort_unstable();
159    tids.dedup();
160    for tid in tids {
161        if first {
162            first = false;
163            out.push_str(
164                "{\"name\":\"process_name\",\"ph\":\"M\",\"pid\":1,\
165                 \"args\":{\"name\":\"lm-flow graph\"}},",
166            );
167        } else {
168            out.push(',');
169        }
170        let _ = write!(
171            out,
172            "{{\"name\":\"thread_name\",\"ph\":\"M\",\"pid\":1,\"tid\":{tid},\
173             \"args\":{{\"name\":\"worker {tid}\"}}}}"
174        );
175    }
176
177    // ---- span 本体 ----
178    for e in events {
179        if first {
180            first = false;
181        } else {
182            out.push(',');
183        }
184        let (node_name, kernel_name) = labels
185            .get(e.node as usize)
186            .map(|(n, k)| (n.as_str(), k.as_str()))
187            .unwrap_or(("<unknown>", ""));
188        out.push('{');
189        out.push_str("\"name\":");
190        push_json_escaped(&mut out, node_name);
191        out.push_str(",\"cat\":");
192        push_json_escaped(&mut out, e.phase.as_str());
193        out.push_str(",\"ph\":\"X\",\"pid\":1,\"tid\":");
194        let _ = write!(out, "{}", e.tid);
195        out.push_str(",\"ts\":");
196        let _ = write!(out, "{}", e.start_us);
197        out.push_str(",\"dur\":");
198        let _ = write!(out, "{}", e.dur_us.max(0));
199        out.push_str(",\"args\":{\"kernel\":");
200        push_json_escaped(&mut out, kernel_name);
201        out.push_str(",\"phase\":");
202        push_json_escaped(&mut out, e.phase.as_str());
203        out.push_str(",\"input_ts\":");
204        let ts = Timestamp(e.input_ts);
205        if ts.is_range_value() {
206            let _ = write!(out, "{}", e.input_ts);
207        } else {
208            // 哨兵(Unset/Unstarted/PreStream/PostStream/Done…)贴着 i64::MIN/MAX,原样导出
209            // 会在查看器里显示成 -9223372036854775806 这种天文数字。用 Timestamp 既有的
210            // Display 记成名字,一眼看出是哪个哨兵 —— Open/Close 阶段拿到的就是它们。
211            push_json_escaped(&mut out, &ts.to_string());
212        }
213        out.push_str("}}");
214    }
215    out.push_str("],\"displayTimeUnit\":\"ms\"}");
216    out
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222
223    #[test]
224    fn ring_is_bounded_and_drops_oldest() {
225        let ring = TraceRing::new(2);
226        ring.record(0, TracePhase::Process, 0, 5, 100);
227        ring.record(1, TracePhase::Process, 10, 5, 200);
228        ring.record(2, TracePhase::Process, 20, 5, 300);
229        let snap = ring.snapshot();
230        assert_eq!(snap.len(), 2, "capacity 2 keeps only 2");
231        assert_eq!(snap[0].node, 1, "oldest (node 0) was evicted");
232        assert_eq!(snap[1].node, 2);
233    }
234
235    #[test]
236    fn clear_resets_events_and_tids() {
237        let ring = TraceRing::new(4);
238        ring.record(0, TracePhase::Open, 0, 1, 0);
239        ring.clear();
240        assert!(ring.snapshot().is_empty());
241        // 清空后线程号重新从 0 开始
242        ring.record(0, TracePhase::Process, 0, 1, 0);
243        assert_eq!(ring.snapshot()[0].tid, 0);
244    }
245
246    #[test]
247    fn chrome_json_escapes_and_wraps() {
248        let events = vec![TraceEvent {
249            node: 0,
250            phase: TracePhase::Process,
251            tid: 3,
252            start_us: 12,
253            dur_us: 34,
254            input_ts: 1000,
255        }];
256        let labels = vec![("no\"de".to_string(), "Kern\\el".to_string())];
257        let json = to_chrome_trace_json(&events, &labels);
258        assert!(json.starts_with("{\"traceEvents\":["));
259        assert!(json.contains("\"ts\":12"));
260        assert!(json.contains("\"dur\":34"));
261        assert!(json.contains("\"tid\":3"));
262        assert!(json.contains("no\\\"de"), "name quote escaped");
263        assert!(json.contains("Kern\\\\el"), "kernel backslash escaped");
264        assert!(json.contains("\"input_ts\":1000"));
265    }
266
267    #[test]
268    fn lanes_get_named_and_empty_stays_valid() {
269        // 每条泳道一条 thread_name 元事件,否则查看器只显示裸 tid 数字。
270        let ev = |tid| TraceEvent {
271            node: 0,
272            phase: TracePhase::Process,
273            tid,
274            start_us: 0,
275            dur_us: 1,
276            input_ts: 7,
277        };
278        let labels = vec![("n".to_string(), "K".to_string())];
279        let json = to_chrome_trace_json(&[ev(0), ev(2), ev(0)], &labels);
280        assert!(json.contains("\"process_name\""), "进程名: {json}");
281        assert_eq!(
282            json.matches("\"thread_name\"").count(),
283            2,
284            "两条泳道各一条,重复 tid 不重复发: {json}"
285        );
286        assert!(
287            json.contains("\"name\":\"worker 2\""),
288            "泳道 2 命名: {json}"
289        );
290
291        // 没有 span 时不得凭空产出元事件,仍是合法空 trace。
292        let empty = to_chrome_trace_json(&[], &labels);
293        assert!(
294            empty.contains("\"traceEvents\":[]"),
295            "空 trace 不带元事件: {empty}"
296        );
297    }
298
299    #[test]
300    fn sentinel_input_ts_exports_as_a_name() {
301        // Open/Close 阶段的 input_ts 是贴着 i64::MIN/MAX 的哨兵,原样导出是天文数字。
302        let ev = |ts: Timestamp| TraceEvent {
303            node: 0,
304            phase: TracePhase::Open,
305            tid: 0,
306            start_us: 0,
307            dur_us: 1,
308            input_ts: ts.0,
309        };
310        let labels = vec![("n".to_string(), "K".to_string())];
311
312        let json = to_chrome_trace_json(&[ev(Timestamp::unstarted())], &labels);
313        assert!(
314            json.contains("\"input_ts\":\"Unstarted\""),
315            "哨兵应记成名字: {json}"
316        );
317        assert!(
318            !json.contains(&format!("{}", i64::MIN + 1)),
319            "不该出现天文数字: {json}"
320        );
321
322        let json = to_chrome_trace_json(&[ev(Timestamp::done())], &labels);
323        assert!(json.contains("\"input_ts\":\"Done\""), "Done: {json}");
324
325        // 普通流内时间戳仍是数字(供查看器按数值筛选)。
326        let json = to_chrome_trace_json(&[ev(Timestamp(42))], &labels);
327        assert!(json.contains("\"input_ts\":42"), "普通值仍为数字: {json}");
328    }
329
330    #[test]
331    fn zero_capacity_ring_stays_bounded() {
332        // `TraceRing::new` 是公开 API,故 0 是可达输入(引擎自己只在 >0 时构造)。
333        // `capacity.max(1)` 就是为它兜底:必须仍然有界,不能无限增长。
334        let ring = TraceRing::new(0);
335        for i in 0..5 {
336            ring.record(i, TracePhase::Process, 0, 1, 0);
337        }
338        assert_eq!(ring.snapshot().len(), 1, "0 容量按 1 处理,仍然有界");
339    }
340}
lmflow/trace.rs

lmflow/
trace.rs
