Krolyxon
Commit a44bdf7

Branches

Tags

Refs →

remagesrcworker.rs

281 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
112pub struct Pool {
113 shared: Arc<Shared>,
114 threads: Vec<thread::JoinHandle<()>>,
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 let threads = (0..n)
133 .map(|i| {
134 let shared = shared.clone();
135 thread::Builder::new()
136 .name(format!("remage-worker-{i}"))
137 .spawn(move || worker(&shared))
138 .expect("spawn worker thread")
139 })
140 .collect();
141 Self { shared, threads }
142 }
143
144 /// Queue `work` under `key` unless that key is already running.
145 pub fn submit(&self, key: JobKey, prio: Priority, work: impl FnOnce() + Send + 'static) {
146 let mut q = self.shared.lock();
147 if q.running.contains(&key) {
148 return;
149 }
150 q.seq += 1;
151 let (seq, epoch) = (q.seq, q.epoch);
152 q.live.insert(key.clone(), seq);
153 q.heap.push(Entry {
154 prio,
155 seq,
156 epoch,
157 key,
158 work: Box::new(work),
159 });
160 drop(q);
161 self.shared.cv.notify_one();
162 }
163
164 /// Drop `key` if it is still queued (a running job finishes regardless).
165 pub fn cancel(&self, key: &JobKey) {
166 self.shared.lock().live.remove(key);
167 }
168
169 /// Invalidate every queued current/preload job.
170 pub fn bump_epoch(&self) {
171 self.shared.lock().epoch += 1;
172 }
173}
174
175impl Drop for Pool {
176 fn drop(&mut self) {
177 self.shared.lock().shutdown = true;
178 self.shared.cv.notify_all();
179 for t in self.threads.drain(..) {
180 let _ = t.join();
181 }
182 }
183}
184
185fn worker(shared: &Shared) {
186 loop {
187 let entry = {
188 let mut q = shared.lock();
189 loop {
190 if q.shutdown {
191 return;
192 }
193 if let Some(e) = q.pop() {
194 break e;
195 }
196 q = shared.cv.wait(q).unwrap_or_else(|e| e.into_inner());
197 }
198 };
199 let key = entry.key.clone();
200 if std::panic::catch_unwind(std::panic::AssertUnwindSafe(entry.work)).is_err() {
201 log::error!("worker job {key:?} panicked");
202 }
203 shared.lock().running.remove(&key);
204 }
205}
206
207#[cfg(test)]
208mod tests {
209 use super::*;
210 use crossbeam_channel::unbounded;
211 use std::time::Duration;
212
213 fn key(s: &str) -> JobKey {
214 JobKey::Decode(PathBuf::from(s))
215 }
216
217 #[test]
218 fn runs_by_priority_and_drops_stale_jobs() {
219 let pool = Pool::with_threads(1);
220 let (gate_tx, gate_rx) = unbounded::<()>();
221 let (tx, rx) = unbounded();
222 // Block the single worker so the rest queue up.
223 pool.submit(key("gate"), Priority::Current, move || {
224 let _ = gate_rx.recv();
225 });
226 std::thread::sleep(Duration::from_millis(50));
227 for (k, p) in [("thumb", Priority::Thumb), ("stale", Priority::Preload)] {
228 let tx = tx.clone();
229 pool.submit(key(k), p, move || tx.send(k).unwrap());
230 }
231 pool.bump_epoch();
232 for (k, p) in [("pre", Priority::Preload), ("cur", Priority::Current)] {
233 let tx = tx.clone();
234 pool.submit(key(k), p, move || tx.send(k).unwrap());
235 }
236 gate_tx.send(()).unwrap();
237 let got: Vec<_> = (0..3)
238 .map(|_| rx.recv_timeout(Duration::from_secs(2)).unwrap())
239 .collect();
240 assert_eq!(got, ["cur", "pre", "thumb"]);
241 assert!(rx.recv_timeout(Duration::from_millis(100)).is_err());
242 }
243
244 #[test]
245 fn resubmitting_replaces_the_queued_job() {
246 let pool = Pool::with_threads(1);
247 let (gate_tx, gate_rx) = unbounded::<()>();
248 let (tx, rx) = unbounded();
249 pool.submit(key("gate"), Priority::Current, move || {
250 let _ = gate_rx.recv();
251 });
252 std::thread::sleep(Duration::from_millis(50));
253 for tag in ["old", "new"] {
254 let tx = tx.clone();
255 pool.submit(key("a"), Priority::Preload, move || tx.send(tag).unwrap());
256 }
257 gate_tx.send(()).unwrap();
258 assert_eq!(rx.recv_timeout(Duration::from_secs(2)).unwrap(), "new");
259 assert!(rx.recv_timeout(Duration::from_millis(100)).is_err());
260 }
261
262 #[test]
263 fn cancelled_jobs_never_run() {
264 let pool = Pool::with_threads(1);
265 let (gate_tx, gate_rx) = unbounded::<()>();
266 let (tx, rx) = unbounded();
267 pool.submit(key("gate"), Priority::Current, move || {
268 let _ = gate_rx.recv();
269 });
270 std::thread::sleep(Duration::from_millis(50));
271 for k in ["a", "b"] {
272 let tx = tx.clone();
273 let job = JobKey::Thumb(PathBuf::from(k));
274 pool.submit(job, Priority::Thumb, move || tx.send(k).unwrap());
275 }
276 pool.cancel(&JobKey::Thumb(PathBuf::from("a")));
277 gate_tx.send(()).unwrap();
278 assert_eq!(rx.recv_timeout(Duration::from_secs(2)).unwrap(), "b");
279 assert!(rx.recv_timeout(Duration::from_millis(100)).is_err());
280 }
281}