lmflow/executor/
thread_pool.rs

1use std::collections::VecDeque;
2use std::sync::atomic::{AtomicBool, Ordering};
3use std::sync::{Arc, Condvar, Mutex, Weak};
4use std::thread::JoinHandle;
5use std::time::{Duration, Instant};
6
7use crate::config::StatsLevel;
8use crate::graph::{GraphInner, NodeId};
9
10use super::platform::{pin_current_thread_to, set_current_thread_rt_priority};
11use super::stats::{ExecutorStats, ExecutorStatsSnapshot};
12
13#[derive(Clone, Copy)]
14enum Task {
15    Run(NodeId),
16    WakeSource(NodeId, u64),
17}
18
19struct QueuedTask {
20    enqueued_at: Option<Instant>,
21    task: Task,
22}
23
24struct DelayedTask {
25    deadline: Instant,
26    task: Task,
27}
28
29#[derive(Default)]
30struct QueueState {
31    ready: VecDeque<QueuedTask>,
32    delayed: Vec<DelayedTask>,
33}
34
35struct Shared {
36    queue: Mutex<QueueState>,
37    cv: Condvar,
38    stop: AtomicBool,
39    stats: ExecutorStats,
40}
41
42impl Shared {
43    /// 取一个任务;返回 `None` 表示「已关停且队列排空」。
44    fn take(&self) -> Option<QueuedTask> {
45        let mut queue = self.queue.lock().unwrap_or_else(|error| error.into_inner());
46        loop {
47            let now = Instant::now();
48            while queue
49                .delayed
50                .first()
51                .is_some_and(|delayed| delayed.deadline <= now)
52            {
53                let delayed = queue.delayed.remove(0);
54                queue.ready.push_back(QueuedTask {
55                    enqueued_at: self.stats.full().then_some(now),
56                    task: delayed.task,
57                });
58                self.stats.enqueued();
59            }
60            if let Some(task) = queue.ready.pop_front() {
61                self.stats
62                    .started(task.enqueued_at.map(|enqueued_at| enqueued_at.elapsed()));
63                return Some(task);
64            }
65            if self.stop.load(Ordering::SeqCst) {
66                return None;
67            }
68            if let Some(delayed) = queue.delayed.first() {
69                let timeout = delayed.deadline.saturating_duration_since(now);
70                let (next, _) = self
71                    .cv
72                    .wait_timeout(queue, timeout)
73                    .unwrap_or_else(|error| error.into_inner());
74                queue = next;
75            } else {
76                queue = self
77                    .cv
78                    .wait(queue)
79                    .unwrap_or_else(|error| error.into_inner());
80            }
81        }
82    }
83}
84
85pub struct ThreadPool {
86    name: String,
87    num_threads: usize,
88    affinity: Vec<usize>,
89    priority: i32,
90    shared: Arc<Shared>,
91    threads: Mutex<Vec<JoinHandle<()>>>,
92}
93
94impl ThreadPool {
95    pub fn new(name: &str, num_threads: usize, affinity: Vec<usize>, priority: i32) -> Self {
96        Self::new_with_stats(name, num_threads, affinity, priority, StatsLevel::Full)
97    }
98
99    pub(crate) fn new_with_stats(
100        name: &str,
101        num_threads: usize,
102        affinity: Vec<usize>,
103        priority: i32,
104        stats_level: StatsLevel,
105    ) -> Self {
106        Self {
107            name: name.to_string(),
108            num_threads: num_threads.max(1),
109            affinity,
110            priority,
111            shared: Arc::new(Shared {
112                queue: Mutex::new(QueueState::default()),
113                cv: Condvar::new(),
114                stop: AtomicBool::new(false),
115                stats: ExecutorStats::new(stats_level),
116            }),
117            threads: Mutex::new(Vec::new()),
118        }
119    }
120
121    pub fn name(&self) -> &str {
122        &self.name
123    }
124
125    pub fn num_threads(&self) -> usize {
126        self.num_threads
127    }
128
129    pub fn affinity(&self) -> &[usize] {
130        &self.affinity
131    }
132
133    pub fn priority(&self) -> i32 {
134        self.priority
135    }
136
137    /// 拉起工作线程。必须在 `Arc<GraphInner>` 已存在之后调用。
138    pub fn start(&self, graph: Weak<GraphInner>) {
139        let mut handles = self
140            .threads
141            .lock()
142            .unwrap_or_else(|error| error.into_inner());
143        if !handles.is_empty() {
144            return;
145        }
146        for index in 0..self.num_threads {
147            let shared = self.shared.clone();
148            let graph = graph.clone();
149            let thread_name = format!("{}-{index}", self.name);
150            let cpu = if self.affinity.is_empty() {
151                None
152            } else {
153                Some(self.affinity[index % self.affinity.len()])
154            };
155            let priority = self.priority;
156            let handle = std::thread::Builder::new()
157                .name(thread_name)
158                .spawn(move || {
159                    if let Some(cpu) = cpu {
160                        pin_current_thread_to(cpu);
161                    }
162                    if priority > 0 {
163                        set_current_thread_rt_priority(priority);
164                    }
165                    worker(shared, graph);
166                })
167                .expect("failed to create worker thread");
168            handles.push(handle);
169        }
170    }
171
172    /// 投递一个就绪节点。关停后返回 `false`。
173    pub fn submit(&self, node: NodeId) -> bool {
174        let mut queue = self
175            .shared
176            .queue
177            .lock()
178            .unwrap_or_else(|error| error.into_inner());
179        if self.shared.stop.load(Ordering::SeqCst) {
180            return false;
181        }
182        queue.ready.push_back(QueuedTask {
183            enqueued_at: self.shared.stats.full().then(Instant::now),
184            task: Task::Run(node),
185        });
186        self.shared.stats.enqueued();
187        drop(queue);
188        self.shared.cv.notify_one();
189        true
190    }
191
192    pub fn submit_source_wake(&self, node: NodeId, generation: u64, delay: Duration) -> bool {
193        let mut queue = self
194            .shared
195            .queue
196            .lock()
197            .unwrap_or_else(|error| error.into_inner());
198        if self.shared.stop.load(Ordering::SeqCst) {
199            return false;
200        }
201        let task = Task::WakeSource(node, generation);
202        if delay.is_zero() {
203            queue.ready.push_back(QueuedTask {
204                enqueued_at: self.shared.stats.full().then(Instant::now),
205                task,
206            });
207            self.shared.stats.enqueued();
208        } else {
209            let now = Instant::now();
210            let mut bounded_delay = delay;
211            let deadline = loop {
212                if let Some(deadline) = now.checked_add(bounded_delay) {
213                    break deadline;
214                }
215                bounded_delay /= 2;
216            };
217            let index = queue
218                .delayed
219                .partition_point(|delayed| delayed.deadline <= deadline);
220            queue.delayed.insert(index, DelayedTask { deadline, task });
221        }
222        drop(queue);
223        self.shared.cv.notify_one();
224        true
225    }
226
227    pub fn pending(&self) -> usize {
228        self.shared.stats.queued()
229    }
230
231    pub fn clear_delayed(&self) {
232        let mut queue = self
233            .shared
234            .queue
235            .lock()
236            .unwrap_or_else(|error| error.into_inner());
237        queue.delayed.clear();
238        drop(queue);
239        self.shared.cv.notify_all();
240    }
241
242    pub(crate) fn stats(&self) -> ExecutorStatsSnapshot {
243        self.shared.stats.snapshot()
244    }
245
246    pub(crate) fn has_pending_work(&self) -> bool {
247        let queue = self
248            .shared
249            .queue
250            .lock()
251            .unwrap_or_else(|error| error.into_inner());
252        !queue.ready.is_empty() || !queue.delayed.is_empty() || self.shared.stats.has_running()
253    }
254
255    pub(crate) fn queued_nodes(&self) -> Vec<NodeId> {
256        self.shared
257            .queue
258            .lock()
259            .unwrap_or_else(|error| error.into_inner())
260            .ready
261            .iter()
262            .filter_map(|task| match task.task {
263                Task::Run(node) => Some(node),
264                Task::WakeSource(_, _) => None,
265            })
266            .collect()
267    }
268
269    pub fn reset_stats(&self) {
270        self.shared.stats.reset();
271    }
272
273    /// 关停并 join。幂等。
274    pub fn shutdown(&self) {
275        {
276            let _queue = self
277                .shared
278                .queue
279                .lock()
280                .unwrap_or_else(|error| error.into_inner());
281            self.shared.stop.store(true, Ordering::SeqCst);
282        }
283        self.shared.cv.notify_all();
284        let handles = {
285            let mut handles = self
286                .threads
287                .lock()
288                .unwrap_or_else(|error| error.into_inner());
289            std::mem::take(&mut *handles)
290        };
291        for handle in handles {
292            let _ = handle.join();
293        }
294    }
295}
296
297impl Drop for ThreadPool {
298    fn drop(&mut self) {
299        self.shutdown();
300    }
301}
302
303fn worker(shared: Arc<Shared>, graph: Weak<GraphInner>) {
304    while let Some(task) = shared.take() {
305        let Some(graph) = graph.upgrade() else {
306            break;
307        };
308        let started = shared.stats.full().then(Instant::now);
309        match task.task {
310            Task::Run(node) => graph.run_node_on_worker(node),
311            Task::WakeSource(node, generation) => graph.wake_source_on_worker(node, generation),
312        }
313        shared
314            .stats
315            .completed(started.map(|started| started.elapsed()));
316        graph.executor_task_completed();
317    }
318}
319
320#[cfg(test)]
321mod tests {
322    use super::*;
323
324    #[test]
325    fn zero_threads_becomes_one() {
326        let pool = ThreadPool::new("test", 0, vec![], 0);
327        assert_eq!(pool.num_threads(), 1);
328    }
329
330    #[test]
331    fn submit_before_start_is_queued() {
332        let pool = ThreadPool::new("test", 2, vec![], 0);
333        assert!(pool.submit(1));
334        assert_eq!(pool.pending(), 1);
335    }
336
337    #[test]
338    fn submit_after_shutdown_is_rejected() {
339        let pool = ThreadPool::new("test", 1, vec![], 0);
340        pool.shutdown();
341        assert!(!pool.submit(1));
342    }
343
344    #[test]
345    fn shutdown_is_idempotent_and_joins() {
346        let pool = ThreadPool::new("test", 3, vec![], 0);
347        pool.shutdown();
348        pool.shutdown();
349    }
350}
lmflow/executor/thread_pool.rs

lmflow/executor/
thread_pool.rs
