lmflow/
graph.rs

1//! The graph: topology, validation, execution and stream shutdown.
2//!
3//! Structurally this is an **index arena**: entities live in `Vec`s and refer to each other by
4//! `usize` id (see `docs/design.md` §6.1), which avoids a graph of self-referential pointers.
5//!
6//! Queues belong to a **consumer's input port**, not to an edge: one edge may feed several
7//! consumers and each of them must receive every packet, so a shared queue would make them steal
8//! packets from one another. An edge therefore only carries topology and closed-state, and packets
9//! are dispatched per consumer — cloning a reference, never the payload.
10
11use std::cell::UnsafeCell;
12use std::collections::{BTreeMap, BTreeSet, VecDeque};
13use std::ffi::c_void;
14use std::panic::{catch_unwind, AssertUnwindSafe};
15use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, AtomicUsize, Ordering};
16use std::sync::{Arc, Condvar, Mutex};
17use std::time::Instant;
18
19use crate::config::{GraphConfig, StatsLevel};
20use crate::context::Context;
21use crate::executor::Executor;
22use crate::kernel::{KernelInstance, PortTable};
23use crate::packet::Packet;
24use crate::runtime::{self, GraphShared};
25use crate::status::{Error, Result};
26use crate::timestamp::Timestamp;
27
28mod api;
29mod backpressure;
30mod build;
31mod dot;
32mod introspect;
33mod lifecycle;
34mod node;
35mod poller;
36mod scheduler;
37
38use backpressure::*;
39use node::*;
40pub use node::{DotView, State};
41use poller::PollerInner;
42pub use poller::{
43    OutputEvent, Poller, PollerBackpressureStatsSnapshot, PollerOptions, PollerOverflow,
44};
45
46pub type NodeId = usize;
47pub type EdgeId = usize;
48
49// ---------------------------------------------------------------- 图
50
51/// 「有进展」的通知状态:活动代数 + **当前阻塞在 condvar 上的宿主线程数**。
52///
53/// `waiters` 与 `gen` 同锁保护,这是能安全跳过 `notify_all` 的依据:notifier 在持锁时
54/// 读到 `waiters == 0`,则任何「正要等待」的线程此刻都还没拿到这把锁 —— 它随后会看到
55/// 递增后的 `gen`(≠ 它捕获的 `before`)从而根本不进入等待。故**不会丢唤醒**。
56///
57/// 为什么值得为此加一个计数:`notify_activity` 在 `dispatch` 里是**每包每条边**调一次,
58/// 而 `Condvar::notify_all` 即使没有任何等待者也会走一次 futex 系统调用
59/// (本机实测约 372 ns,与 `getpid()` 的 329 ns 同量级 —— 该机器系统调用被放大约 5 倍,
60/// 裸机约 60~80 ns)。没人在等的时候这纯属白付。
61#[derive(Debug, Default)]
62struct Activity {
63    gen: u64,
64    waiters: usize,
65}
66
67type WakeupCallback = Arc<dyn Fn() + Send + Sync>;
68
69pub struct GraphInner {
70    pub shared: Arc<GraphShared>,
71    nodes: Vec<Node>,
72    edges: Vec<Edge>,
73    graph_inputs: Vec<EdgeId>,
74    graph_outputs: Vec<EdgeId>,
75    input_by_name: BTreeMap<String, EdgeId>,
76    output_by_name: BTreeMap<String, EdgeId>,
77    edge_by_name: BTreeMap<String, EdgeId>,
78    state: Mutex<State>,
79    /// 所有执行器；委托执行器自持交还宿主线程执行的任务队列。
80    executors: Vec<Executor>,
81    /// 已投递到执行器、尚未跑完的任务数。用于「是否空闲」与终止判定。
82    in_flight: AtomicUsize,
83    /// 委托任务抽取游标。多个 DelegatingExecutor 之间轮询，避免前面的队列长期饿死后面的。
84    delegated_cursor: AtomicUsize,
85    /// 同一时刻只允许一个宿主线程执行委托任务，兑现 DelegatingExecutor 的零并发语义。
86    delegated_running: AtomicBool,
87    /// 有任何进展时唤醒阻塞中的宿主线程(取到输出、节点关闭、出错、任务入队)
88    activity: (Mutex<Activity>, Condvar),
89    /// 事件循环唤醒回调。只负责通知宿主“该进引擎推进了”，不得承载实际工作。
90    wakeup_callback: Mutex<Option<WakeupCallback>>,
91    /// 合并重复通知：宿主调用 `pump_step` 确认一次唤醒后重新武装。
92    wakeup_pending: AtomicBool,
93    /// 每次唤醒请求递增，避免活动恰好发生在 `pump_step(false)` 重新武装窗口时丢通知。
94    wakeup_generation: AtomicU64,
95    /// 暂停调度(调试/限速)。已在执行的算子不受影响。
96    paused: AtomicBool,
97    /// 因下游无损输入队列已满而保留 staging 的节点。下游每次出队后重试这些刷新。
98    blocked_flush_nodes: Mutex<BTreeSet<NodeId>>,
99    side_packets: Mutex<BTreeMap<String, Packet>>,
100    /// 各算子在 GetContract 里声明的必需 side packet:(名字, 声明它的节点)
101    required_side_packets: Vec<(String, String)>,
102    /// 计时基准。`Instant` 无法放进原子,故节点统计里存「相对本基准的微秒」。
103    epoch: Instant,
104    /// 本轮 `start` 的时刻(相对 `epoch` 的微秒 + 1);0 表示尚未开始。
105    run_started_us: AtomicI64,
106    /// Compact / Diagnostics 各自相邻两次导出之间的私有基线；不作为宿主查询 API 暴露。
107    dot_intervals: Mutex<dot::DotIntervalBaselines>,
108    /// 建图时确定，运行期间不变，故无需原子。
109    stats_level: StatsLevel,
110    e2e_stats: E2eStats,
111}
112
113impl GraphInner {
114    #[inline]
115    fn basic_stats(&self) -> bool {
116        self.stats_level != StatsLevel::Off
117    }
118
119    #[inline]
120    fn full_stats(&self) -> bool {
121        self.stats_level == StatsLevel::Full
122    }
123
124    #[inline]
125    pub(super) fn record_e2e(&self, pkt: &Packet) {
126        if !self.full_stats() || !pkt.mark_e2e_observed() {
127            return;
128        }
129        let ingress = pkt.ingress_us();
130        if ingress == 0 {
131            return;
132        }
133        self.e2e_stats
134            .record(self.epoch_us().saturating_sub(ingress).max(0) as u64);
135    }
136}
137
138/// A handle to a computation graph.
139///
140/// Build one with [`from_yaml`](Graph::from_yaml), attach output sinks with
141/// [`add_poller`](Graph::add_poller) (pull) or [`observe`](Graph::observe) (push), then
142/// [`start`](Graph::start). Feed packets in through [`input`](Graph::input); terminate by closing
143/// the inputs ([`close_all_inputs`](Graph::close_all_inputs)) and waiting for the pipeline to
144/// drain ([`wait_done`](Graph::wait_done)).
145///
146/// Every method takes `&self` — the handle is shared, not cloned. Dropping it shuts the graph down
147/// and joins the worker threads. After termination, [`reset`](Graph::reset) puts the graph back to
148/// a startable state while keeping already-opened kernel instances alive, so an expensive one-off
149/// such as loading a model is not repeated.
150pub struct Graph {
151    inner: Arc<GraphInner>,
152}
153
154impl Graph {
155    pub fn from_yaml(text: &str) -> Result<Self> {
156        let cfg = GraphConfig::from_yaml(text)?;
157        Self::from_config(cfg)
158    }
159
160    pub fn from_yaml_file(path: &str) -> Result<Self> {
161        let cfg = GraphConfig::from_yaml_file(path)?;
162        Self::from_config(cfg)
163    }
164
165    pub fn from_config(cfg: GraphConfig) -> Result<Self> {
166        let inner = GraphInner::build(cfg)?;
167        Ok(Self {
168            inner: Arc::new(inner),
169        })
170    }
171
172    pub fn state(&self) -> State {
173        *self.inner.state.lock().expect("state lock poisoned")
174    }
175
176    pub fn inner(&self) -> &Arc<GraphInner> {
177        &self.inner
178    }
179
180    /// 注入常量输入。必须在 `start` 之前。
181    pub fn set_side_packet(&self, name: &str, pkt: Packet) -> Result<()> {
182        if self.state() != State::Initialized {
183            return Err(Error::State(
184                "side packets must be injected before start".into(),
185            ));
186        }
187        crate::packet::validate_type_id(pkt.type_id()).map_err(|error| {
188            Error::InvalidArg(format!(
189                "side packet `{name}` carries an invalid type: {error}"
190            ))
191        })?;
192        self.inner
193            .side_packets
194            .lock()
195            .expect("side packet lock poisoned")
196            .insert(name.to_string(), pkt);
197        Ok(())
198    }
199
200    pub fn input(&self, port: &str) -> Result<Input> {
201        let edge =
202            *self.inner.input_by_name.get(port).ok_or_else(|| {
203                Error::NotFound(format!("graph input port `{port}` does not exist"))
204            })?;
205        Ok(Input {
206            graph: self.inner.clone(),
207            edge,
208        })
209    }
210
211    pub fn close_input(&self, port: &str) -> Result<()> {
212        let edge =
213            *self.inner.input_by_name.get(port).ok_or_else(|| {
214                Error::NotFound(format!("graph input port `{port}` does not exist"))
215            })?;
216        self.inner.close_edge(edge);
217        self.inner.set_state_draining_if_all_inputs_closed();
218        Ok(())
219    }
220
221    pub fn close_all_inputs(&self) {
222        for &e in &self.inner.graph_inputs {
223            self.inner.close_edge(e);
224        }
225        self.inner.set_state_draining_if_all_inputs_closed();
226    }
227
228    pub fn cancel(&self) {
229        self.inner.shared.cancel();
230        for node in &self.inner.nodes {
231            node.source_waiting.store(false, Ordering::SeqCst);
232            node.source_wait_reason.store(0, Ordering::Relaxed);
233            node.source_wake_deadline_us.store(0, Ordering::Relaxed);
234            node.source_wake_generation.fetch_add(1, Ordering::SeqCst);
235        }
236        for executor in &self.inner.executors {
237            executor.clear_delayed();
238        }
239        self.inner.resume_blocked_flushes();
240        self.inner.finish_all_backpressure_blocks();
241        self.inner.notify_activity();
242        self.inner.request_wakeup();
243    }
244
245    /// Wait for completion, checking the timeout between delegated callbacks.
246    /// A synchronous callback already executing on this thread cannot be preempted.
247    /// Timing out leaves pending work available for a later wait or cancellation.
248    pub fn wait_done_timeout(&self, timeout: std::time::Duration) -> Result<()> {
249        self.inner
250            .wait_done(Some(std::time::Instant::now() + timeout))
251    }
252
253    /// Wait for idle with a timeout checked between delegated callbacks.
254    /// A synchronous callback already executing on this thread cannot be preempted.
255    pub fn wait_until_idle_timeout(&self, timeout: std::time::Duration) -> Result<()> {
256        self.inner
257            .wait_until_idle(Some(std::time::Instant::now() + timeout))
258    }
259
260    pub fn pause(&self) {
261        self.inner.pause();
262    }
263
264    pub fn resume(&self) {
265        self.inner.resume();
266    }
267
268    pub fn is_paused(&self) -> bool {
269        self.inner.is_paused()
270    }
271
272    /// 已定义的线程池名字(按 YAML 顺序)。
273    pub fn executor_names(&self) -> Vec<&str> {
274        self.inner.executor_names()
275    }
276
277    /// 是否空闲:委托队列为空且所有执行器无在飞任务。
278    pub fn is_idle(&self) -> bool {
279        self.inner.is_idle_pub()
280    }
281
282    /// 在当前宿主线程上执行至多一个委托任务，或推进一次关流。
283    ///
284    /// 事件循环型宿主可反复调用它主动推进 `DelegatingExecutor`，而不必进入阻塞接口。
285    pub fn pump_step(&self) -> bool {
286        self.inner.pump_step_pub()
287    }
288
289    /// 在当前宿主线程最多推进 `max_steps` 个委托任务或关流步骤。
290    pub fn pump_steps(&self, max_steps: usize) -> usize {
291        (0..max_steps).take_while(|_| self.pump_step()).count()
292    }
293
294    /// Install an edge-triggered event-loop wakeup callback.
295    ///
296    /// The callback may run on any engine thread and must return quickly. It should only schedule
297    /// the host loop to call [`pump_step`](Self::pump_step) until it returns `false`.
298    ///
299    /// # There is only one slot
300    ///
301    /// A graph holds a *single* callback: installing again **replaces** the previous one rather
302    /// than adding to it, and the notification is graph-global — it says "something progressed",
303    /// not which port produced output or which queue freed a slot.
304    ///
305    /// So a host that wants several independent wakeup consumers (waiting for the run to finish,
306    /// for a port to emit, for an input queue to drain) must fan out itself: one owner installs
307    /// this callback and broadcasts to the waiters, each re-checking its own condition. Calling
308    /// this a second time from elsewhere silently displaces the first callback, and the symptom
309    /// is a graph that stops advancing — no error, no log.
310    pub fn set_wakeup_callback<F>(&self, callback: F)
311    where
312        F: Fn() + Send + Sync + 'static,
313    {
314        self.inner.set_wakeup_callback(Some(Arc::new(callback)));
315    }
316
317    /// Remove the event-loop wakeup callback.
318    pub fn clear_wakeup_callback(&self) {
319        self.inner.set_wakeup_callback(None);
320    }
321
322    /// 推模式订阅输出口。回调在派发该包的线程上执行(可能是线程池线程)。
323    ///
324    /// 必须在 `start` 之前注册,否则会漏掉已产出的包。
325    pub fn observe<F>(&self, port: &str, f: F) -> Result<()>
326    where
327        F: Fn(&Packet) + Send + Sync + 'static,
328    {
329        self.observe_inner(port, f, false)
330    }
331
332    /// 推模式订阅输出口，同时接收以空包编码的时间戳边界事件。
333    pub fn observe_with_timestamp_bounds<F>(&self, port: &str, f: F) -> Result<()>
334    where
335        F: Fn(&Packet) + Send + Sync + 'static,
336    {
337        self.observe_inner(port, f, true)
338    }
339
340    fn observe_inner<F>(&self, port: &str, f: F, observe_timestamp_bounds: bool) -> Result<()>
341    where
342        F: Fn(&Packet) + Send + Sync + 'static,
343    {
344        if self.state() != State::Initialized {
345            return Err(Error::State(
346                "observe must be called before start, otherwise already-produced packets are missed".into(),
347            ));
348        }
349        self.inner
350            .add_observer_fn(port, Arc::new(f), observe_timestamp_bounds)
351    }
352
353    /// 算子自报计数器的当前值(按图隔离)。
354    pub fn counter_value(&self, name: &str) -> i64 {
355        self.inner.shared.counter_value(name)
356    }
357
358    /// 已登记的计数器名。
359    pub fn counter_names(&self) -> Vec<String> {
360        self.inner.shared.counter_names()
361    }
362
363    /// 取出图级共享状态的句柄 —— 它比 `Graph` 句柄活得久,
364    /// 因此可以在图销毁**之后**仍然读取错误与计数器(测试与排障用)。
365    pub fn shared_for_inspection(&self) -> Arc<crate::runtime::GraphShared> {
366        self.inner.shared.clone()
367    }
368
369    pub fn node_count(&self) -> usize {
370        self.inner.nodes.len()
371    }
372    pub fn node_name(&self, i: usize) -> Option<&str> {
373        self.inner.nodes.get(i).map(|n| n.name.as_str())
374    }
375    pub fn input_port_names(&self) -> Vec<&str> {
376        self.inner
377            .graph_inputs
378            .iter()
379            .map(|&e| self.inner.edges[e].name.as_str())
380            .collect()
381    }
382    pub fn output_port_names(&self) -> Vec<&str> {
383        self.inner
384            .graph_outputs
385            .iter()
386            .map(|&e| self.inner.edges[e].name.as_str())
387            .collect()
388    }
389
390    /// 指定边的积压包数(该边所有消费者输入队列之和)。
391    pub fn queue_depth(&self, port: &str) -> Option<usize> {
392        let e = *self.inner.edge_by_name.get(port)?;
393        Some(self.inner.queue_depth(e))
394    }
395
396    pub fn dropped_count(&self, port: &str) -> Option<u64> {
397        let e = *self.inner.edge_by_name.get(port)?;
398        Some(self.inner.edges[e].dropped_count())
399    }
400
401    pub fn dump(&self) -> String {
402        self.inner.dump()
403    }
404
405    /// 导出 Graphviz DOT(拓扑 + 子图命名空间 cluster + 执行器/绑核图例)。见 `GraphInner::to_dot`。
406    pub fn to_dot(&self) -> String {
407        self.to_dot_with_view(DotView::Topology)
408    }
409
410    /// 导出指定详细程度的 Graphviz DOT。
411    pub fn to_dot_with_view(&self, view: DotView) -> String {
412        self.inner.to_dot(view)
413    }
414
415    /// 带节点状态与核心统计的紧凑视图,适合大型图持续刷新。
416    pub fn to_dot_compact(&self) -> String {
417        self.to_dot_with_view(DotView::Compact)
418    }
419
420    /// 同 [`to_dot`](Self::to_dot),但在每个节点标签上标出运行统计
421    /// (处理数 · 平均延迟 · 收/发包数 · 队列峰值 · 错误数),并把填充色换成
422    /// **按平均延迟的热力图**(绿=快 → 红=慢)—— 一眼看出瓶颈在哪个节点。
423    ///
424    /// 可在图运行期间随时调用(统计是原子读的快照),不必等跑完。
425    /// 注意:热力图占用了「按执行器上色」那一维,执行器仍以标签里的 `@name` 标出。
426    pub fn to_dot_with_stats(&self) -> String {
427        self.to_dot_with_view(DotView::Diagnostics)
428    }
429
430    /// 导出当前 trace 环内容为 Chrome Trace Event Format 的 JSON(chrome://tracing /
431    /// perfetto 可直接打开)。需在建图时设 `trace_capacity > 0`;未开启时返回合法的空
432    /// trace。见 [`GraphInner::to_chrome_trace`]。
433    pub fn to_chrome_trace(&self) -> String {
434        self.inner.to_chrome_trace()
435    }
436
437    pub fn node_stats(&self, i: usize) -> Option<NodeStatsSnapshot> {
438        self.inner.node_stats(i)
439    }
440
441    pub fn input_queue_stats(&self, node: usize, port: usize) -> Option<InputQueueStatsSnapshot> {
442        self.inner.input_queue_stats(node, port)
443    }
444}
445
446impl Drop for Graph {
447    /// 在宿主放开句柄时**就地关停线程池并 join**。
448    ///
449    /// 必须在这里做,不能只依赖 `GraphInner::drop`:工作线程为了执行节点会
450    /// `Weak::upgrade` 出一个临时强引用,若宿主的 Arc 已先释放,最后一个引用就落在
451    /// 工作线程手上 —— 于是 `GraphInner::drop` 在工作线程上运行,`shutdown` 变成
452    /// **join 自己**,得到 `EDEADLK`。先在宿主线程 join 完,worker 就不存在了。
453    fn drop(&mut self) {
454        // Input/Poller handles retain GraphInner. Stop accepting work and finish
455        // cancellation while executors can still drain their claimed tasks; merely
456        // joining workers leaves those handles attached to a permanently live graph.
457        if self.state() != State::Terminated {
458            self.cancel();
459            let _ = self.wait_done();
460        }
461        self.inner.shutdown_executors_pub();
462    }
463}
464
465impl std::fmt::Debug for Graph {
466    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
467        write!(
468            f,
469            "Graph{{state:{:?}, nodes:{}, inputs:{:?}, outputs:{:?}}}",
470            self.state(),
471            self.inner.nodes.len(),
472            self.input_port_names(),
473            self.output_port_names()
474        )
475    }
476}
477
478/// 图输入口句柄(热路径免按名字查表)。
479pub struct Input {
480    graph: Arc<GraphInner>,
481    edge: EdgeId,
482}
483
484impl Input {
485    /// Sends on the same input are serialized through dispatch. Reentrant sends
486    /// from that input's dispatch callback return a state error.
487    pub fn send(&self, pkt: Packet) -> Result<()> {
488        self.graph.send(self.edge, pkt, true)
489    }
490    /// Returns WouldBlock if another thread is dispatching on this input.
491    pub fn try_send(&self, pkt: Packet) -> Result<()> {
492        self.graph.send(self.edge, pkt, false)
493    }
494    pub fn backpressure_stats(&self) -> WatermarkBackpressureStatsSnapshot {
495        self.graph.watermark_backpressure_stats(self.edge)
496    }
497    /// Wait for the current send to finish before closing. From this input's
498    /// callback, closure is deferred until its current dispatch has finished.
499    pub fn close(&self) {
500        self.graph.close_edge(self.edge);
501        self.graph.set_state_draining_if_all_inputs_closed();
502    }
503}
504
505impl std::fmt::Debug for Input {
506    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
507        write!(f, "Input{{port:`{}`}}", self.graph.edges[self.edge].name)
508    }
509}
510
511/// `LMFlowNodeStats` 的 Rust 侧快照。
512#[derive(Debug, Clone)]
513pub struct NodeStatsSnapshot {
514    pub node_name: String,
515    pub kernel_name: String,
516    pub running: bool,
517    pub running_for_us: i64,
518    pub processed: u64,
519    pub errors: u64,
520    pub total_process_us: i64,
521    pub max_process_us: i64,
522    /// 从输入口取走的包数
523    pub packets_in: u64,
524    /// 产出并派发下游的包数
525    pub packets_out: u64,
526    /// 下游入队时观察到的队列深度峰值(高水位)
527    pub peak_queue_depth: usize,
528    pub queued: usize,
529}
530
531#[derive(Debug, Clone)]
532pub(super) struct ExecutorStatsSnapshot {
533    pub queued: usize,
534    pub running: usize,
535    pub peak_queued: usize,
536    pub completed: u64,
537    pub total_wait_us: u64,
538    pub total_execution_us: u64,
539    pub queued_for_us: u64,
540    pub saturated: bool,
541    pub queued_nodes: Vec<String>,
542}
543
544/// 节点单个输入口的队列与背压统计快照。
545#[derive(Debug, Clone)]
546pub struct InputQueueStatsSnapshot {
547    pub node_name: String,
548    pub port_name: String,
549    pub producer_name: Option<String>,
550    pub packet_capacity: Option<usize>,
551    pub queued_packets: usize,
552    pub queued_bytes: u64,
553    pub reserved_packets: usize,
554    pub peak_queued_packets: usize,
555    pub peak_queued_bytes: u64,
556    pub blocked: bool,
557    pub blocked_for_us: u64,
558    pub block_events: u64,
559    pub total_blocked_us: u64,
560}
561#[derive(Debug, Clone)]
562pub struct WatermarkBackpressureStatsSnapshot {
563    pub port_name: String,
564    pub packet_limit: usize,
565    pub total_queued_packets: usize,
566    pub blocked: bool,
567    pub active_waiters: usize,
568    pub blocked_for_us: u64,
569    pub block_events: u64,
570    pub total_blocked_us: u64,
571}
lmflow/graph.rs

lmflow/
graph.rs
