pagesift_gpu/
worker.rs

1//! GPU 工人：一个线程独占 CUDA 上下文，其他线程把工作以闭包形式交给它、等结果。
2
3use std::panic::{AssertUnwindSafe, catch_unwind};
4use std::sync::Arc;
5use std::sync::mpsc::{RecvTimeoutError, Sender, channel};
6use std::thread::JoinHandle;
7use std::time::Duration;
8
9use crate::exec::Ctx;
10use crate::{Gpu, GpuInfo};
11
12/// 工人线程上执行的一段工作。
13type Job = Box<dyn for<'c, 'g> FnOnce(&'c Ctx<'g>) + Send + 'static>;
14
15/// [`Worker::call_timeout`] 没等到结果的原因。
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub enum WaitError {
18    /// 到时间了工作还没做完（它留在队列里或继续跑，结果没人要了）。
19    TimedOut,
20    /// 工人线程没了，或工作 panic 了。
21    Failed(String),
22}
23
24/// 工人句柄，丢掉时线程退出并释放上下文。
25pub struct Worker {
26    tx: Option<Sender<Job>>,
27    thread: Option<JoinHandle<()>>,
28    info: GpuInfo,
29}
30
31impl Worker {
32    /// 打开第 `index` 块卡并启动工人线程（内核 JIT 在这里完成）。
33    /// 打不开时返回原因，调用方退回 CPU。
34    pub fn spawn(index: usize) -> Result<Worker, String> {
35        let (tx, rx) = channel::<Job>();
36        let (ready_tx, ready_rx) = channel();
37        let thread = std::thread::Builder::new()
38            .name("pagesift-gpu".into())
39            .spawn(move || {
40                let gpu = match Gpu::open_index(index).and_then(|g| g.info().map(|i| (g, i))) {
41                    Ok((g, info)) => {
42                        // 主线程等到这个消息才返回；它已经走了也没关系
43                        let _ = ready_tx.send(Ok(info));
44                        g
45                    }
46                    Err(e) => {
47                        let _ = ready_tx.send(Err(e));
48                        return;
49                    }
50                };
51                let ctx = Ctx::new(&gpu);
52                for job in rx {
53                    job(&ctx);
54                }
55            })
56            .map_err(|e| format!("failed to start the GPU worker thread: {e}"))?;
57        let info = ready_rx
58            .recv()
59            .map_err(|_| "the GPU worker thread exited during startup".to_owned())??;
60        Ok(Worker {
61            tx: Some(tx),
62            thread: Some(thread),
63            info,
64        })
65    }
66
67    /// 打开时查到的设备情况。
68    pub fn info(&self) -> &GpuInfo {
69        &self.info
70    }
71
72    /// 在工人线程上执行 `f` 并等结果。工作按提交顺序一个接一个执行。
73    pub fn call<R: Send>(
74        &self,
75        f: impl for<'c, 'g> FnOnce(&'c Ctx<'g>) -> R + Send,
76    ) -> Result<R, String> {
77        let (tx, rx) = channel();
78        let job = move |ctx: &Ctx<'_>| {
79            // 回复通道断开说明调用方已不在，这时没有人关心结果
80            let _ = tx.send(catch_unwind(AssertUnwindSafe(|| f(ctx))));
81        };
82        let boxed: Box<dyn for<'c, 'g> FnOnce(&'c Ctx<'g>) + Send + '_> = Box::new(job);
83        // SAFETY: 擦掉闭包借用的生命周期。下面的 recv 要等到闭包执行完（发回结果）或被丢弃
84        // （通道关闭）才返回，所以闭包借用的东西在它存续期间一直有效
85        let boxed: Job = unsafe { std::mem::transmute(boxed) };
86        let tx = self.tx.as_ref().expect("工人只在 Drop 里关闭");
87        tx.send(boxed)
88            .map_err(|_| "the GPU worker thread has exited".to_owned())?;
89        match rx.recv() {
90            Ok(Ok(r)) => Ok(r),
91            Ok(Err(_)) => Err("a job on the GPU worker panicked".into()),
92            Err(_) => Err("the GPU worker thread has exited".into()),
93        }
94    }
95}
96
97impl Worker {
98    /// 同 [`Worker::call`]，但最多等 `wait`。调用方到时间可以走开，所以闭包不能借东西（`'static`）。
99    pub fn call_timeout<R: Send + 'static>(
100        &self,
101        f: impl for<'c, 'g> FnOnce(&'c Ctx<'g>) -> R + Send + 'static,
102        wait: Duration,
103    ) -> Result<R, WaitError> {
104        let (tx, rx) = channel();
105        let job: Job = Box::new(move |ctx: &Ctx<'_>| {
106            // 回复通道断开说明调用方已不在，这时没有人关心结果
107            let _ = tx.send(catch_unwind(AssertUnwindSafe(|| f(ctx))));
108        });
109        let tx = self.tx.as_ref().expect("工人只在 Drop 里关闭");
110        tx.send(job)
111            .map_err(|_| WaitError::Failed("the GPU worker thread has exited".into()))?;
112        match rx.recv_timeout(wait) {
113            Ok(Ok(r)) => Ok(r),
114            Ok(Err(_)) => Err(WaitError::Failed("a job on the GPU worker panicked".into())),
115            Err(RecvTimeoutError::Timeout) => Err(WaitError::TimedOut),
116            Err(RecvTimeoutError::Disconnected) => {
117                Err(WaitError::Failed("the GPU worker thread has exited".into()))
118            }
119        }
120    }
121}
122
123impl Drop for Worker {
124    fn drop(&mut self) {
125        drop(self.tx.take());
126        if let Some(t) = self.thread.take() {
127            // 线程 panic 也不再往外传：Drop 里没有地方报
128            let _ = t.join();
129        }
130    }
131}
132
133/// 把工人包成可以共享的句柄。
134pub type SharedWorker = Arc<Worker>;
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139
140    /// 工作没做完时调用方能走开，之后的工作照常；没有显卡时跳过。
141    #[test]
142    fn call_timeout_walks_away_from_a_slow_job() {
143        let Ok(w) = Worker::spawn(0) else {
144            eprintln!("没有能用的 N 卡，跳过");
145            return;
146        };
147        let slow = w.call_timeout(
148            |_: &Ctx<'_>| std::thread::sleep(Duration::from_millis(200)),
149            Duration::from_millis(1),
150        );
151        assert_eq!(slow, Err(WaitError::TimedOut));
152        let next = w.call_timeout(|_: &Ctx<'_>| 7, Duration::from_secs(60));
153        assert_eq!(next, Ok(7));
154    }
155}
pagesift_gpu/worker.rs

pagesift_gpu/
worker.rs
