Krolyxon
Commit 1d9227a

Branches

Tags

Refs →

remagesrcworker.rs

278 lines · 8.0 KB · Rust

Raw Blame History Download
1//! Priority job pool with epoch cancellation.
2//!
3//! Jobs are keyed: submitting a key that is already queued replaces the queued job (new
4//! priority, new epoch), and submitting one that is running is a no-op. Bumping the epoch makes
5//! every queued job of an epoch-sensitive priority stale, so mashing `n` never decodes the
6//! images that were skipped over.
7
8use std::collections::{BinaryHeap, HashMap, HashSet};
9use std::path::PathBuf;
10use std::sync::{Arc, Condvar, Mutex, MutexGuard};
11use std::thread;
12
13#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
14pub enum Priority {
15 /// The image on screen.
16 Current = 0,
17 /// Neighbours of the current image.
18 Preload = 1,
19 /// Thumbnails in view.
20 ThumbVisible = 2,
21 /// Other thumbnails.
22 Thumb = 3,
23}
24
25impl Priority {
26 fn epoch_sensitive(self) -> bool {
27 matches!(self, Priority::Current | Priority::Preload)
28 }
29}
30
31#[derive(Clone, Debug, PartialEq, Eq, Hash)]
32pub enum JobKey {
33 Decode(PathBuf),
34 Thumb(PathBuf),
35 Refine(PathBuf),
36}
37
38type Work = Box<dyn FnOnce() + Send>;
39
40struct Entry {
41 prio: Priority,
42 seq: u64,
43 epoch: u64,
44 key: JobKey,
45 work: Work,
46}
47
48impl PartialEq for Entry {
49 fn eq(&self, other: &Self) -> bool {
50 self.prio == other.prio && self.seq == other.seq
51 }
52}
53
54impl Eq for Entry {}
55
56impl PartialOrd for Entry {
57 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
58 Some(self.cmp(other))
59 }
60}
61
62impl Ord for Entry {
63 /// Max-heap order: lowest priority value first, then oldest submission.
64 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
65 other
66 .prio
67 .cmp(&self.prio)
68 .then_with(|| other.seq.cmp(&self.seq))
69 }
70}
71
72#[derive(Default)]
73struct Queue {
74 heap: BinaryHeap<Entry>,
75 /// Key → sequence number of its live (non-superseded) entry.
76 live: HashMap<JobKey, u64>,
77 running: HashSet<JobKey>,
78 seq: u64,
79 epoch: u64,
80 shutdown: bool,
81}
82
83impl Queue {
84 fn pop(&mut self) -> Option<Entry> {
85 while let Some(e) = self.heap.pop() {
86 if self.live.get(&e.key) != Some(&e.seq) {
87 continue; // superseded by a later submit
88 }
89 self.live.remove(&e.key);
90 if e.prio.epoch_sensitive() && e.epoch < self.epoch {
91 continue; // stale
92 }
93 self.running.insert(e.key.clone());
94 return Some(e);
95 }
96 None
97 }
98}
99
100struct Shared {
101 queue: Mutex<Queue>,
102 cv: Condvar,
103}
104
105impl Shared {
106 fn lock(&self) -> MutexGuard<'_, Queue> {
107 // A panicking job never holds the lock, so poisoning can't leave it inconsistent.
108 self.queue.lock().unwrap_or_else(|e| e.into_inner())
109 }
110}
111
112/// The workers run detached: nothing ever waits for them.
113pub struct Pool {
114 shared: Arc<Shared>,
115}
116
117impl Pool {
118 /// Start `nproc - 1` workers (at least one).
119 pub fn new() -> Self {
120 let n = thread::available_parallelism()
121 .map_or(2, |n| n.get())
122 .saturating_sub(1)
123 .max(1);
124 Self::with_threads(n)
125 }
126
127 pub fn with_threads(n: usize) -> Self {
128 let shared = Arc::new(Shared {
129 queue: Mutex::new(Queue::default()),
130 cv: Condvar::new(),
131 });
132 for i in 0..n {
133 let shared = shared.clone();
134 thread::Builder::new()
135 .name(format!("remage-worker-{i}"))
136 .spawn(move || worker(&shared))
137 .expect("spawn worker thread");
138 }
139 Self { shared }
140 }
141
142 /// Queue `work` under `key` unless that key is already running.
143 pub fn submit(&self, key: JobKey, prio: Priority, work: impl FnOnce() + Send + 'static) {
144 let mut q = self.shared.lock();
145 if q.running.contains(&key) {
146 return;
147 }
148 q.seq += 1;
149 let (seq, epoch) = (q.seq, q.epoch);
150 q.live.insert(key.clone(), seq);
151 q.heap.push(Entry {
152 prio,
153 seq,
154 epoch,
155 key,
156 work: Box::new(work),
157 });
158 drop(q);
159 self.shared.cv.notify_one();
160 }
161
162 /// Drop `key` if it is still queued (a running job finishes regardless).
163 pub fn cancel(&self, key: &JobKey) {
164 self.shared.lock().live.remove(key);
165 }
166
167 /// Invalidate every queued current/preload job.
168 pub fn bump_epoch(&self) {
169 self.shared.lock().epoch += 1;
170 }
171}
172
173impl Drop for Pool {
174 /// Idle workers stop; busy ones are left to it. This only happens on the way out, and
175 /// quitting must not wait for a decode nobody will see.
176 fn drop(&mut self) {
177 self.shared.lock().shutdown = true;
178 self.shared.cv.notify_all();
179 }
180}
181
182fn worker(shared: &Shared) {
183 loop {
184 let entry = {
185 let mut q = shared.lock();
186 loop {
187 if q.shutdown {
188 return;
189 }
190 if let Some(e) = q.pop() {
191 break e;
192 }
193 q = shared.cv.wait(q).unwrap_or_else(|e| e.into_inner());
194 }
195 };
196 let key = entry.key.clone();
197 if std::panic::catch_unwind(std::panic::AssertUnwindSafe(entry.work)).is_err() {
198 log::error!("worker job {key:?} panicked");
199 }
200 shared.lock().running.remove(&key);
201 }
202}
203
204#[cfg(test)]
205mod tests {
206 use super::*;
207 use crossbeam_channel::unbounded;
208 use std::time::Duration;
209
210 fn key(s: &str) -> JobKey {
211 JobKey::Decode(PathBuf::from(s))
212 }
213
214 #[test]
215 fn runs_by_priority_and_drops_stale_jobs() {
216 let pool = Pool::with_threads(1);
217 let (gate_tx, gate_rx) = unbounded::<()>();
218 let (tx, rx) = unbounded();
219 // Block the single worker so the rest queue up.
220 pool.submit(key("gate"), Priority::Current, move || {
221 let _ = gate_rx.recv();
222 });
223 std::thread::sleep(Duration::from_millis(50));
224 for (k, p) in [("thumb", Priority::Thumb), ("stale", Priority::Preload)] {
225 let tx = tx.clone();
226 pool.submit(key(k), p, move || tx.send(k).unwrap());
227 }
228 pool.bump_epoch();
229 for (k, p) in [("pre", Priority::Preload), ("cur", Priority::Current)] {
230 let tx = tx.clone();
231 pool.submit(key(k), p, move || tx.send(k).unwrap());
232 }
233 gate_tx.send(()).unwrap();
234 let got: Vec<_> = (0..3)
235 .map(|_| rx.recv_timeout(Duration::from_secs(2)).unwrap())
236 .collect();
237 assert_eq!(got, ["cur", "pre", "thumb"]);
238 assert!(rx.recv_timeout(Duration::from_millis(100)).is_err());
239 }
240
241 #[test]
242 fn resubmitting_replaces_the_queued_job() {
243 let pool = Pool::with_threads(1);
244 let (gate_tx, gate_rx) = unbounded::<()>();
245 let (tx, rx) = unbounded();
246 pool.submit(key("gate"), Priority::Current, move || {
247 let _ = gate_rx.recv();
248 });
249 std::thread::sleep(Duration::from_millis(50));
250 for tag in ["old", "new"] {
251 let tx = tx.clone();
252 pool.submit(key("a"), Priority::Preload, move || tx.send(tag).unwrap());
253 }
254 gate_tx.send(()).unwrap();
255 assert_eq!(rx.recv_timeout(Duration::from_secs(2)).unwrap(), "new");
256 assert!(rx.recv_timeout(Duration::from_millis(100)).is_err());
257 }
258
259 #[test]
260 fn cancelled_jobs_never_run() {
261 let pool = Pool::with_threads(1);
262 let (gate_tx, gate_rx) = unbounded::<()>();
263 let (tx, rx) = unbounded();
264 pool.submit(key("gate"), Priority::Current, move || {
265 let _ = gate_rx.recv();
266 });
267 std::thread::sleep(Duration::from_millis(50));
268 for k in ["a", "b"] {
269 let tx = tx.clone();
270 let job = JobKey::Thumb(PathBuf::from(k));
271 pool.submit(job, Priority::Thumb, move || tx.send(k).unwrap());
272 }
273 pool.cancel(&JobKey::Thumb(PathBuf::from("a")));
274 gate_tx.send(()).unwrap();
275 assert_eq!(rx.recv_timeout(Duration::from_secs(2)).unwrap(), "b");
276 assert!(rx.recv_timeout(Duration::from_millis(100)).is_err());
277 }
278}