1use std::{
2 sync::{
3 Arc, RwLock,
4 atomic::{AtomicUsize, Ordering},
5 mpsc::{Receiver, Sender, TryRecvError},
6 },
7 task::Waker,
8 thread::JoinHandle,
9};
10
11use anyhow::Result;
12
13use crate::jobs::{
14 check_update::CheckUpdateResult, create_scratch::CreateScratchResult,
15 find_similar::FindSimilarResult, objdiff::ObjDiffResult, update::UpdateResult,
16};
17
18pub mod check_update;
19pub mod create_scratch;
20pub mod find_similar;
21pub mod objdiff;
22pub mod update;
23
24#[derive(Debug, Eq, PartialEq, Copy, Clone)]
25pub enum Job {
26 ObjDiff,
27 CheckUpdate,
28 Update,
29 CreateScratch,
30 FindSimilar,
31}
32pub static JOB_ID: AtomicUsize = AtomicUsize::new(0);
33
34#[derive(Default)]
35pub struct JobQueue {
36 pub jobs: Vec<JobState>,
37 pub results: Vec<JobResult>,
38}
39
40impl JobQueue {
41 #[inline]
43 pub fn push(&mut self, state: JobState) { self.jobs.push(state); }
44
45 #[inline]
47 pub fn push_once(&mut self, job: Job, func: impl FnOnce() -> JobState) {
48 if !self.is_running(job) {
49 self.push(func());
50 }
51 }
52
53 pub fn is_running(&self, kind: Job) -> bool {
55 self.jobs.iter().any(|j| j.kind == kind && j.handle.is_some())
56 }
57
58 pub fn any_running(&self) -> bool {
60 self.jobs.iter().any(|job| {
61 if let Some(handle) = &job.handle {
62 return !handle.is_finished();
63 }
64 false
65 })
66 }
67
68 pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut JobState> + '_ { self.jobs.iter_mut() }
70
71 pub fn iter_finished(
73 &mut self,
74 ) -> impl Iterator<Item = (&mut JobState, std::thread::Result<JobResult>)> + '_ {
75 self.jobs.iter_mut().filter_map(|job| {
76 if let Some(handle) = &job.handle {
77 if !handle.is_finished() {
78 return None;
79 }
80 let result = job.handle.take().unwrap().join();
81 return Some((job, result));
82 }
83 None
84 })
85 }
86
87 pub fn clear_finished(&mut self) {
89 self.jobs.retain(|job| {
90 !(job.handle.is_none() && job.context.status.read().unwrap().error.is_none())
91 });
92 }
93
94 pub fn clear_errored(&mut self) {
96 self.jobs.retain(|job| job.context.status.read().unwrap().error.is_none());
97 }
98
99 pub fn remove(&mut self, id: usize) { self.jobs.retain(|job| job.id != id); }
101
102 pub fn cancel_kind(&mut self, kind: Job) {
104 let ids: Vec<usize> = self
105 .jobs
106 .iter()
107 .filter(|j| j.kind == kind)
108 .map(|j| {
109 let _ = j.cancel.send(());
110 j.id
111 })
112 .collect();
113 for id in ids {
114 self.remove(id);
115 }
116 }
117
118 pub fn collect_results(&mut self) {
120 let mut results = vec![];
121 for (job, result) in self.iter_finished() {
122 match result {
123 Ok(result) => {
124 match result {
125 JobResult::None => {
126 }
128 _ => results.push(result),
129 }
130 }
131 Err(err) => {
132 let err = if let Some(msg) = err.downcast_ref::<&'static str>() {
133 anyhow::Error::msg(*msg)
134 } else if let Some(msg) = err.downcast_ref::<String>() {
135 anyhow::Error::msg(msg.clone())
136 } else {
137 anyhow::Error::msg("Thread panicked")
138 };
139 let result = job.context.status.write();
140 if let Ok(mut guard) = result {
141 guard.error = Some(err);
142 } else {
143 drop(result);
144 job.context.status = Arc::new(RwLock::new(JobStatus {
145 title: "Error".to_string(),
146 progress_percent: 0.0,
147 progress_items: None,
148 status: String::new(),
149 error: Some(err),
150 }));
151 }
152 }
153 }
154 }
155 self.results.append(&mut results);
156 self.clear_finished();
157 }
158}
159
160#[derive(Clone)]
161pub struct JobContext {
162 pub status: Arc<RwLock<JobStatus>>,
163 pub waker: Waker,
164}
165
166pub struct JobState {
167 pub id: usize,
168 pub kind: Job,
169 pub handle: Option<JoinHandle<JobResult>>,
170 pub context: JobContext,
171 pub cancel: Sender<()>,
172}
173
174#[derive(Default)]
175pub struct JobStatus {
176 pub title: String,
177 pub progress_percent: f32,
178 pub progress_items: Option<[u32; 2]>,
179 pub status: String,
180 pub error: Option<anyhow::Error>,
181}
182
183pub enum JobResult {
184 None,
185 ObjDiff(Option<Box<ObjDiffResult>>),
186 CheckUpdate(Option<Box<CheckUpdateResult>>),
187 Update(Box<UpdateResult>),
188 CreateScratch(Option<Box<CreateScratchResult>>),
189 FindSimilar(Option<Box<FindSimilarResult>>),
190}
191
192fn start_job(
193 waker: Waker,
194 title: &str,
195 kind: Job,
196 run: impl FnOnce(JobContext, Receiver<()>) -> Result<JobResult> + Send + 'static,
197) -> JobState {
198 let status = Arc::new(RwLock::new(JobStatus {
199 title: title.to_string(),
200 progress_percent: 0.0,
201 progress_items: None,
202 status: String::new(),
203 error: None,
204 }));
205 let context = JobContext { status: status.clone(), waker: waker.clone() };
206 let context_inner = JobContext { status: status.clone(), waker };
207 let (tx, rx) = std::sync::mpsc::channel();
208 let handle = std::thread::spawn(move || match run(context_inner, rx) {
209 Ok(state) => state,
210 Err(e) => {
211 if let Ok(mut w) = status.write() {
212 w.error = Some(e);
213 }
214 JobResult::None
215 }
216 });
217 let id = JOB_ID.fetch_add(1, Ordering::Relaxed);
218 JobState { id, kind, handle: Some(handle), context, cancel: tx }
219}
220
221fn update_status(
222 context: &JobContext,
223 str: String,
224 count: u32,
225 total: u32,
226 cancel: &Receiver<()>,
227) -> Result<()> {
228 let mut w =
229 context.status.write().map_err(|_| anyhow::Error::msg("Failed to lock job status"))?;
230 w.progress_items = Some([count, total]);
231 w.progress_percent = count as f32 / total as f32;
232 if should_cancel(cancel) {
233 w.status = "Cancelled".to_string();
234 return Err(anyhow::Error::msg("Cancelled"));
235 } else {
236 w.status = str;
237 }
238 drop(w);
239 context.waker.wake_by_ref();
240 Ok(())
241}
242
243fn should_cancel(rx: &Receiver<()>) -> bool {
244 match rx.try_recv() {
245 Ok(_) | Err(TryRecvError::Disconnected) => true,
246 Err(_) => false,
247 }
248}