| 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 | |
| 8 | use std::collections::{BinaryHeap, HashMap, HashSet}; |
| 9 | use std::path::PathBuf; |
| 10 | use std::sync::{Arc, Condvar, Mutex, MutexGuard}; |
| 11 | use std::thread; |
| 12 | |
| 13 | #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] |
| 14 | pub 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 | |
| 25 | impl Priority { |
| 26 | fn epoch_sensitive(self) -> bool { |
| 27 | matches!(self, Priority::Current | Priority::Preload) |
| 28 | } |
| 29 | } |
| 30 | |
| 31 | #[derive(Clone, Debug, PartialEq, Eq, Hash)] |
| 32 | pub enum JobKey { |
| 33 | Decode(PathBuf), |
| 34 | Thumb(PathBuf), |
| 35 | Refine(PathBuf), |
| 36 | } |
| 37 | |
| 38 | type Work = Box<dyn FnOnce() + Send>; |
| 39 | |
| 40 | struct Entry { |
| 41 | prio: Priority, |
| 42 | seq: u64, |
| 43 | epoch: u64, |
| 44 | key: JobKey, |
| 45 | work: Work, |
| 46 | } |
| 47 | |
| 48 | impl PartialEq for Entry { |
| 49 | fn eq(&self, other: &Self) -> bool { |
| 50 | self.prio == other.prio && self.seq == other.seq |
| 51 | } |
| 52 | } |
| 53 | |
| 54 | impl Eq for Entry {} |
| 55 | |
| 56 | impl PartialOrd for Entry { |
| 57 | fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> { |
| 58 | Some(self.cmp(other)) |
| 59 | } |
| 60 | } |
| 61 | |
| 62 | impl 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)] |
| 73 | struct 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 | |
| 83 | impl 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 | |
| 100 | struct Shared { |
| 101 | queue: Mutex<Queue>, |
| 102 | cv: Condvar, |
| 103 | } |
| 104 | |
| 105 | impl 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 | pub struct Pool { |
| 113 | shared: Arc<Shared>, |
| 114 | threads: Vec<thread::JoinHandle<()>>, |
| 115 | } |
| 116 | |
| 117 | impl 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 | |
| 175 | impl 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 | |
| 185 | fn 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)] |
| 208 | mod 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 | } |