rigtop_core/
worker.rs

1//! Background thread that owns the data source.
2//!
3//! Polling (and, for remote sources, waiting on the network) happens off the
4//! UI thread so that input stays responsive. The UI sends [`Command`]s and
5//! receives [`Event`]s over channels.
6
7use std::sync::mpsc::{self, Receiver, RecvTimeoutError, Sender};
8use std::thread::{self, JoinHandle};
9use std::time::{Duration, Instant};
10
11use crate::{DataSource, MIN_POLL_INTERVAL, ProcessSignal, Snapshot};
12
13/// Requests from the UI to the worker.
14#[derive(Debug)]
15pub enum Command {
16    /// Send a signal to a process on the monitored host.
17    Signal { pid: u32, signal: ProcessSignal },
18    /// Change the polling interval.
19    SetInterval(Duration),
20}
21
22/// Messages from the worker to the UI.
23#[derive(Debug)]
24pub enum Event {
25    /// A new measurement.
26    Snapshot(Box<Snapshot>),
27    /// Outcome of a [`Command::Signal`].
28    SignalResult {
29        pid: u32,
30        signal: ProcessSignal,
31        result: Result<(), String>,
32    },
33    /// Polling failed; the worker keeps trying.
34    Error(String),
35}
36
37/// Handle to the worker thread. Dropping it stops the thread.
38pub struct Worker {
39    commands: Option<Sender<Command>>,
40    events: Receiver<Event>,
41    handle: Option<JoinHandle<()>>,
42}
43
44impl Worker {
45    /// Starts the worker. `make_source` runs on the worker thread, so slow
46    /// initialization (NVML, warm-up sample, SSH connection) does not delay
47    /// the first frame.
48    pub fn spawn<S, F>(make_source: F, interval: Duration) -> Self
49    where
50        S: DataSource + 'static,
51        F: FnOnce() -> S + Send + 'static,
52    {
53        Self::spawn_with_notify(make_source, interval, || {})
54    }
55
56    /// Like [`Worker::spawn`], and calls `notify` after every event so that
57    /// an event-driven UI (e.g. egui's `request_repaint`) can wake up instead
58    /// of polling.
59    pub fn spawn_with_notify<S, F, N>(make_source: F, interval: Duration, notify: N) -> Self
60    where
61        S: DataSource + 'static,
62        F: FnOnce() -> S + Send + 'static,
63        N: Fn() + Send + 'static,
64    {
65        let (command_tx, command_rx) = mpsc::channel();
66        let (event_tx, event_rx) = mpsc::channel();
67        let handle = thread::Builder::new()
68            .name("rigtop-collector".to_owned())
69            .spawn(move || {
70                let events = Notifying {
71                    sender: event_tx,
72                    notify,
73                };
74                run(make_source(), interval, &command_rx, &events)
75            })
76            .expect("failed to spawn the collector thread");
77        Self {
78            commands: Some(command_tx),
79            events: event_rx,
80            handle: Some(handle),
81        }
82    }
83
84    /// Sends a command; silently ignored if the worker has stopped.
85    pub fn send(&self, command: Command) {
86        if let Some(tx) = &self.commands {
87            let _ = tx.send(command);
88        }
89    }
90
91    /// Returns the next pending event without blocking.
92    pub fn try_recv(&self) -> Option<Event> {
93        self.events.try_recv().ok()
94    }
95}
96
97impl Drop for Worker {
98    fn drop(&mut self) {
99        // Closing the command channel makes the worker loop exit.
100        self.commands.take();
101        if let Some(handle) = self.handle.take() {
102            let _ = handle.join();
103        }
104    }
105}
106
107/// Worker loop: poll on schedule, execute commands in between.
108/// Event sender that calls a notification hook after each send.
109struct Notifying<N: Fn()> {
110    sender: Sender<Event>,
111    notify: N,
112}
113
114impl<N: Fn()> Notifying<N> {
115    /// Sends an event; fails when the UI has gone away.
116    fn send(&self, event: Event) -> Result<(), mpsc::SendError<Event>> {
117        self.sender.send(event)?;
118        (self.notify)();
119        Ok(())
120    }
121}
122
123fn run<S: DataSource, N: Fn()>(
124    mut source: S,
125    mut interval: Duration,
126    commands: &Receiver<Command>,
127    events: &Notifying<N>,
128) {
129    // The source has just taken its warm-up sample; the first real sample
130    // only needs the minimum interval, so the UI fills quickly.
131    let mut last_poll = Instant::now();
132    let mut next_poll = last_poll + MIN_POLL_INTERVAL;
133    loop {
134        let timeout = next_poll.saturating_duration_since(Instant::now());
135        match commands.recv_timeout(timeout) {
136            Ok(Command::Signal { pid, signal }) => {
137                let result = source
138                    .signal_process(pid, signal)
139                    .map_err(|e| e.to_string());
140                if events
141                    .send(Event::SignalResult {
142                        pid,
143                        signal,
144                        result,
145                    })
146                    .is_err()
147                {
148                    return;
149                }
150            }
151            Ok(Command::SetInterval(new_interval)) => {
152                interval = new_interval.max(MIN_POLL_INTERVAL);
153                next_poll = last_poll + interval;
154            }
155            Err(RecvTimeoutError::Timeout) => {
156                let event = match source.poll() {
157                    Ok(snapshot) => Event::Snapshot(Box::new(snapshot)),
158                    Err(err) => Event::Error(err.to_string()),
159                };
160                last_poll = Instant::now();
161                // Keep a steady cadence, but never try to catch up on
162                // missed ticks (e.g. after the machine was suspended).
163                next_poll = (next_poll + interval).max(last_poll + MIN_POLL_INTERVAL);
164                if events.send(event).is_err() {
165                    return;
166                }
167            }
168            Err(RecvTimeoutError::Disconnected) => return,
169        }
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176    use crate::SourceError;
177
178    /// A source that counts polls and records signals.
179    struct FakeSource {
180        polls: u64,
181        signals: Sender<(u32, ProcessSignal)>,
182    }
183
184    impl DataSource for FakeSource {
185        fn poll(&mut self) -> Result<Snapshot, SourceError> {
186            self.polls += 1;
187            Ok(Snapshot {
188                timestamp_ms: self.polls,
189                ..Default::default()
190            })
191        }
192
193        fn signal_process(&mut self, pid: u32, signal: ProcessSignal) -> Result<(), SourceError> {
194            self.signals.send((pid, signal)).unwrap();
195            if pid == 1 {
196                Err(SourceError::Disconnected("denied".to_owned()))
197            } else {
198                Ok(())
199            }
200        }
201    }
202
203    fn recv(worker: &Worker) -> Event {
204        let deadline = Instant::now() + Duration::from_secs(5);
205        loop {
206            if let Some(event) = worker.try_recv() {
207                return event;
208            }
209            assert!(Instant::now() < deadline, "no event from worker");
210            thread::sleep(Duration::from_millis(5));
211        }
212    }
213
214    #[test]
215    fn notify_runs_after_each_event() {
216        use std::sync::Arc;
217        use std::sync::atomic::{AtomicUsize, Ordering};
218
219        let (signal_tx, _signal_rx) = mpsc::channel();
220        let count = Arc::new(AtomicUsize::new(0));
221        let counter = Arc::clone(&count);
222        let worker = Worker::spawn_with_notify(
223            move || FakeSource {
224                polls: 0,
225                signals: signal_tx,
226            },
227            MIN_POLL_INTERVAL,
228            move || {
229                counter.fetch_add(1, Ordering::SeqCst);
230            },
231        );
232        recv(&worker);
233        // The event is sent before the notification; wait for the latter.
234        let deadline = Instant::now() + Duration::from_secs(5);
235        while count.load(Ordering::SeqCst) == 0 {
236            assert!(Instant::now() < deadline, "notify was not called");
237            thread::sleep(Duration::from_millis(5));
238        }
239    }
240
241    #[test]
242    fn delivers_snapshots_and_signal_results() {
243        let (signal_tx, signal_rx) = mpsc::channel();
244        let worker = Worker::spawn(
245            move || FakeSource {
246                polls: 0,
247                signals: signal_tx,
248            },
249            MIN_POLL_INTERVAL,
250        );
251
252        match recv(&worker) {
253            Event::Snapshot(s) => assert_eq!(s.timestamp_ms, 1),
254            other => panic!("unexpected {other:?}"),
255        }
256
257        worker.send(Command::Signal {
258            pid: 1,
259            signal: ProcessSignal::Kill,
260        });
261        // Skip snapshots until the signal result arrives.
262        let result = loop {
263            if let Event::SignalResult { pid, result, .. } = recv(&worker) {
264                assert_eq!(pid, 1);
265                break result;
266            }
267        };
268        assert!(result.unwrap_err().contains("denied"));
269        assert_eq!(signal_rx.recv().unwrap(), (1, ProcessSignal::Kill));
270        // Dropping the worker joins the thread without hanging.
271        drop(worker);
272    }
273}
rigtop_core/worker.rs

rigtop_core/
worker.rs
