Krolyxon
Commit f2bd31a

Branches

Tags

Refs →

remagesrcworker.rs

256 lines · 7.1 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 #[allow(dead_code)] // thumbnails arrive in M3
21 ThumbVisible = 2,
22 /// Other thumbnails.
23 #[allow(dead_code)]
24 Thumb = 3,
25}
26
27impl Priority {
28 fn epoch_sensitive(self) -> bool {
29 matches!(self, Priority::Current | Priority::Preload)
30 }
31}
32
33#[derive(Clone, Debug, PartialEq, Eq, Hash)]
34pub enum JobKey {
35 Decode(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 /// Invalidate every queued current/preload job.
165 pub fn bump_epoch(&self) {
166 self.shared.lock().epoch += 1;
167 }
168}
169
170impl Drop for Pool {
171 fn drop(&mut self) {
172 self.shared.lock().shutdown = true;
173 self.shared.cv.notify_all();
174 for t in self.threads.drain(..) {
175 let _ = t.join();
176 }
177 }
178}
179
180fn worker(shared: &Shared) {
181 loop {
182 let entry = {
183 let mut q = shared.lock();
184 loop {
185 if q.shutdown {
186 return;
187 }
188 if let Some(e) = q.pop() {
189 break e;
190 }
191 q = shared.cv.wait(q).unwrap_or_else(|e| e.into_inner());
192 }
193 };
194 let key = entry.key.clone();
195 if std::panic::catch_unwind(std::panic::AssertUnwindSafe(entry.work)).is_err() {
196 log::error!("worker job {key:?} panicked");
197 }
198 shared.lock().running.remove(&key);
199 }
200}
201
202#[cfg(test)]
203mod tests {
204 use super::*;
205 use crossbeam_channel::unbounded;
206 use std::time::Duration;
207
208 fn key(s: &str) -> JobKey {
209 JobKey::Decode(PathBuf::from(s))
210 }
211
212 #[test]
213 fn runs_by_priority_and_drops_stale_jobs() {
214 let pool = Pool::with_threads(1);
215 let (gate_tx, gate_rx) = unbounded::<()>();
216 let (tx, rx) = unbounded();
217 // Block the single worker so the rest queue up.
218 pool.submit(key("gate"), Priority::Current, move || {
219 let _ = gate_rx.recv();
220 });
221 std::thread::sleep(Duration::from_millis(50));
222 for (k, p) in [("thumb", Priority::Thumb), ("stale", Priority::Preload)] {
223 let tx = tx.clone();
224 pool.submit(key(k), p, move || tx.send(k).unwrap());
225 }
226 pool.bump_epoch();
227 for (k, p) in [("pre", Priority::Preload), ("cur", Priority::Current)] {
228 let tx = tx.clone();
229 pool.submit(key(k), p, move || tx.send(k).unwrap());
230 }
231 gate_tx.send(()).unwrap();
232 let got: Vec<_> = (0..3)
233 .map(|_| rx.recv_timeout(Duration::from_secs(2)).unwrap())
234 .collect();
235 assert_eq!(got, ["cur", "pre", "thumb"]);
236 assert!(rx.recv_timeout(Duration::from_millis(100)).is_err());
237 }
238
239 #[test]
240 fn resubmitting_replaces_the_queued_job() {
241 let pool = Pool::with_threads(1);
242 let (gate_tx, gate_rx) = unbounded::<()>();
243 let (tx, rx) = unbounded();
244 pool.submit(key("gate"), Priority::Current, move || {
245 let _ = gate_rx.recv();
246 });
247 std::thread::sleep(Duration::from_millis(50));
248 for tag in ["old", "new"] {
249 let tx = tx.clone();
250 pool.submit(key("a"), Priority::Preload, move || tx.send(tag).unwrap());
251 }
252 gate_tx.send(()).unwrap();
253 assert_eq!(rx.recv_timeout(Duration::from_secs(2)).unwrap(), "new");
254 assert!(rx.recv_timeout(Duration::from_millis(100)).is_err());
255 }
256}