Krolyxon
Commit ab1dc3b

Branches

Tags

Refs →

remagesrcworker.rs

280 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}
36
37type Work = Box<dyn FnOnce() + Send>;
38
39struct Entry {
40 prio: Priority,
41 seq: u64,
42 epoch: u64,
43 key: JobKey,
44 work: Work,
45}
46
47impl PartialEq for Entry {
48 fn eq(&self, other: &Self) -> bool {
49 self.prio == other.prio && self.seq == other.seq
50 }
51}
52
53impl Eq for Entry {}
54
55impl PartialOrd for Entry {
56 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
57 Some(self.cmp(other))
58 }
59}
60
61impl Ord for Entry {
62 /// Max-heap order: lowest priority value first, then oldest submission.
63 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
64 other
65 .prio
66 .cmp(&self.prio)
67 .then_with(|| other.seq.cmp(&self.seq))
68 }
69}
70
71#[derive(Default)]
72struct Queue {
73 heap: BinaryHeap<Entry>,
74 /// Key → sequence number of its live (non-superseded) entry.
75 live: HashMap<JobKey, u64>,
76 running: HashSet<JobKey>,
77 seq: u64,
78 epoch: u64,
79 shutdown: bool,
80}
81
82impl Queue {
83 fn pop(&mut self) -> Option<Entry> {
84 while let Some(e) = self.heap.pop() {
85 if self.live.get(&e.key) != Some(&e.seq) {
86 continue; // superseded by a later submit
87 }
88 self.live.remove(&e.key);
89 if e.prio.epoch_sensitive() && e.epoch < self.epoch {
90 continue; // stale
91 }
92 self.running.insert(e.key.clone());
93 return Some(e);
94 }
95 None
96 }
97}
98
99struct Shared {
100 queue: Mutex<Queue>,
101 cv: Condvar,
102}
103
104impl Shared {
105 fn lock(&self) -> MutexGuard<'_, Queue> {
106 // A panicking job never holds the lock, so poisoning can't leave it inconsistent.
107 self.queue.lock().unwrap_or_else(|e| e.into_inner())
108 }
109}
110
111pub struct Pool {
112 shared: Arc<Shared>,
113 threads: Vec<thread::JoinHandle<()>>,
114}
115
116impl Pool {
117 /// Start `nproc - 1` workers (at least one).
118 pub fn new() -> Self {
119 let n = thread::available_parallelism()
120 .map_or(2, |n| n.get())
121 .saturating_sub(1)
122 .max(1);
123 Self::with_threads(n)
124 }
125
126 pub fn with_threads(n: usize) -> Self {
127 let shared = Arc::new(Shared {
128 queue: Mutex::new(Queue::default()),
129 cv: Condvar::new(),
130 });
131 let threads = (0..n)
132 .map(|i| {
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 .collect();
140 Self { shared, threads }
141 }
142
143 /// Queue `work` under `key` unless that key is already running.
144 pub fn submit(&self, key: JobKey, prio: Priority, work: impl FnOnce() + Send + 'static) {
145 let mut q = self.shared.lock();
146 if q.running.contains(&key) {
147 return;
148 }
149 q.seq += 1;
150 let (seq, epoch) = (q.seq, q.epoch);
151 q.live.insert(key.clone(), seq);
152 q.heap.push(Entry {
153 prio,
154 seq,
155 epoch,
156 key,
157 work: Box::new(work),
158 });
159 drop(q);
160 self.shared.cv.notify_one();
161 }
162
163 /// Drop `key` if it is still queued (a running job finishes regardless).
164 pub fn cancel(&self, key: &JobKey) {
165 self.shared.lock().live.remove(key);
166 }
167
168 /// Invalidate every queued current/preload job.
169 pub fn bump_epoch(&self) {
170 self.shared.lock().epoch += 1;
171 }
172}
173
174impl Drop for Pool {
175 fn drop(&mut self) {
176 self.shared.lock().shutdown = true;
177 self.shared.cv.notify_all();
178 for t in self.threads.drain(..) {
179 let _ = t.join();
180 }
181 }
182}
183
184fn worker(shared: &Shared) {
185 loop {
186 let entry = {
187 let mut q = shared.lock();
188 loop {
189 if q.shutdown {
190 return;
191 }
192 if let Some(e) = q.pop() {
193 break e;
194 }
195 q = shared.cv.wait(q).unwrap_or_else(|e| e.into_inner());
196 }
197 };
198 let key = entry.key.clone();
199 if std::panic::catch_unwind(std::panic::AssertUnwindSafe(entry.work)).is_err() {
200 log::error!("worker job {key:?} panicked");
201 }
202 shared.lock().running.remove(&key);
203 }
204}
205
206#[cfg(test)]
207mod tests {
208 use super::*;
209 use crossbeam_channel::unbounded;
210 use std::time::Duration;
211
212 fn key(s: &str) -> JobKey {
213 JobKey::Decode(PathBuf::from(s))
214 }
215
216 #[test]
217 fn runs_by_priority_and_drops_stale_jobs() {
218 let pool = Pool::with_threads(1);
219 let (gate_tx, gate_rx) = unbounded::<()>();
220 let (tx, rx) = unbounded();
221 // Block the single worker so the rest queue up.
222 pool.submit(key("gate"), Priority::Current, move || {
223 let _ = gate_rx.recv();
224 });
225 std::thread::sleep(Duration::from_millis(50));
226 for (k, p) in [("thumb", Priority::Thumb), ("stale", Priority::Preload)] {
227 let tx = tx.clone();
228 pool.submit(key(k), p, move || tx.send(k).unwrap());
229 }
230 pool.bump_epoch();
231 for (k, p) in [("pre", Priority::Preload), ("cur", Priority::Current)] {
232 let tx = tx.clone();
233 pool.submit(key(k), p, move || tx.send(k).unwrap());
234 }
235 gate_tx.send(()).unwrap();
236 let got: Vec<_> = (0..3)
237 .map(|_| rx.recv_timeout(Duration::from_secs(2)).unwrap())
238 .collect();
239 assert_eq!(got, ["cur", "pre", "thumb"]);
240 assert!(rx.recv_timeout(Duration::from_millis(100)).is_err());
241 }
242
243 #[test]
244 fn resubmitting_replaces_the_queued_job() {
245 let pool = Pool::with_threads(1);
246 let (gate_tx, gate_rx) = unbounded::<()>();
247 let (tx, rx) = unbounded();
248 pool.submit(key("gate"), Priority::Current, move || {
249 let _ = gate_rx.recv();
250 });
251 std::thread::sleep(Duration::from_millis(50));
252 for tag in ["old", "new"] {
253 let tx = tx.clone();
254 pool.submit(key("a"), Priority::Preload, move || tx.send(tag).unwrap());
255 }
256 gate_tx.send(()).unwrap();
257 assert_eq!(rx.recv_timeout(Duration::from_secs(2)).unwrap(), "new");
258 assert!(rx.recv_timeout(Duration::from_millis(100)).is_err());
259 }
260
261 #[test]
262 fn cancelled_jobs_never_run() {
263 let pool = Pool::with_threads(1);
264 let (gate_tx, gate_rx) = unbounded::<()>();
265 let (tx, rx) = unbounded();
266 pool.submit(key("gate"), Priority::Current, move || {
267 let _ = gate_rx.recv();
268 });
269 std::thread::sleep(Duration::from_millis(50));
270 for k in ["a", "b"] {
271 let tx = tx.clone();
272 let job = JobKey::Thumb(PathBuf::from(k));
273 pool.submit(job, Priority::Thumb, move || tx.send(k).unwrap());
274 }
275 pool.cancel(&JobKey::Thumb(PathBuf::from("a")));
276 gate_tx.send(()).unwrap();
277 assert_eq!(rx.recv_timeout(Duration::from_secs(2)).unwrap(), "b");
278 assert!(rx.recv_timeout(Duration::from_millis(100)).is_err());
279 }
280}