1use serde::Serialize;
18use serde_json::Value;
19use tokio::sync::{broadcast, mpsc::UnboundedSender};
20use tokio::time::{timeout_at, Duration, Instant};
21
22use std::collections::BTreeMap;
23use std::sync::{Arc, Mutex};
24use std::sync::atomic::{AtomicU64, Ordering};
25use std::time::SystemTime;
26
27use rpfm_ipc::api::{ApiError, Done, Request, RpcOutcome, RpcRequest, RpcResponse};
28use rpfm_ipc::api::jobs::{CancelJob, DEFAULT_WAIT_SECS, GetJobStatus, JobState, JobStatus, WaitForJob};
29
30use crate::api::parse_params;
31
32const MAX_ENDED_JOBS: usize = 100;
34
35const UPDATES_CAPACITY: usize = 64;
37
38#[derive(Debug)]
40pub struct JobRegistry {
41
42 next_id: AtomicU64,
44
45 jobs: Mutex<BTreeMap<u64, JobRecord>>,
47
48 updates: broadcast::Sender<JobStatus>,
50}
51
52#[derive(Clone, Debug)]
54pub struct JobRecord {
55
56 pub status: JobStatus,
58
59 pub created_at: SystemTime,
61
62 pub updated_at: SystemTime,
64}
65
66impl Default for JobRegistry {
67 fn default() -> Self {
68 Self {
69 next_id: AtomicU64::new(1),
70 jobs: Mutex::new(BTreeMap::new()),
71 updates: broadcast::channel(UPDATES_CAPACITY).0,
72 }
73 }
74}
75
76impl JobRegistry {
77
78 pub fn has_unfinished_jobs(&self) -> bool {
80 self.jobs.lock().unwrap().values().any(|record| !record.status.state.has_ended())
81 }
82
83 pub fn create(&self, method: &str) -> u64 {
89 let job = self.next_id.fetch_add(1, Ordering::SeqCst);
90 self.set(JobStatus { job, method: method.to_owned(), state: JobState::Queued });
91 job
92 }
93
94 pub fn start(&self, job: u64) -> bool {
100 self.transition(job, |state| matches!(state, JobState::Queued).then_some(JobState::Running { stage: None, progress: None })).is_some()
101 }
102
103 pub fn set_stage(&self, job: u64, stage: &str) {
105 self.transition(job, |state| matches!(state, JobState::Running { .. }).then(|| JobState::Running { stage: Some(stage.to_owned()), progress: None }));
106 }
107
108 pub fn set_progress(&self, job: u64, progress: u8) {
112 self.transition(job, |state| match state {
113 JobState::Running { stage, progress: current } if current.is_none_or(|current| current < progress) => Some(JobState::Running { stage: stage.clone(), progress: Some(progress) }),
114 _ => None,
115 });
116 }
117
118 pub fn requeue(&self, job: u64) -> bool {
124 self.transition(job, |state| matches!(state, JobState::Running { .. }).then_some(JobState::Queued)).is_some()
125 }
126
127 pub fn finish_cancelled(&self, job: u64) {
129 self.transition(job, |state| (!state.has_ended()).then_some(JobState::Cancelled));
130 }
131
132 pub fn finish(&self, job: u64, outcome: RpcOutcome) {
134 let state = match outcome {
135 RpcOutcome::Result(result) => JobState::Finished { result },
136 RpcOutcome::Error(error) => JobState::Failed { error },
137 };
138
139 self.transition(job, |_| Some(state));
140 }
141
142 pub fn cancel(&self, job: u64) -> Result<JobStatus, ApiError> {
148 let already = |state: &JobState| if state.has_ended() { "ended" } else { "started" };
149 match self.status(job) {
150 None => Err(ApiError::JobNotFound(job)),
151 Some(status) => self.transition(job, |state| matches!(state, JobState::Queued).then_some(JobState::Cancelled))
152 .ok_or_else(|| ApiError::InvalidParams(format!("The job {job} already {}, so it can't be cancelled.", already(&status.state)))),
153 }
154 }
155
156 pub fn status(&self, job: u64) -> Option<JobStatus> {
158 self.record(job).map(|record| record.status)
159 }
160
161 pub fn record(&self, job: u64) -> Option<JobRecord> {
163 self.jobs.lock().unwrap().get(&job).cloned()
164 }
165
166 pub fn subscribe(&self) -> broadcast::Receiver<JobStatus> {
168 self.updates.subscribe()
169 }
170
171 pub async fn wait(&self, job: u64, timeout: Duration) -> Option<JobStatus> {
177
178 let mut updates = self.subscribe();
180 let deadline = Instant::now() + timeout;
181 loop {
182 let status = self.status(job)?;
183 if status.state.has_ended() {
184 return Some(status);
185 }
186
187 match timeout_at(deadline, updates.recv()).await {
188 Ok(Ok(_)) | Ok(Err(broadcast::error::RecvError::Lagged(_))) => continue,
189 Ok(Err(broadcast::error::RecvError::Closed)) | Err(_) => return self.status(job),
190 }
191 }
192 }
193
194 fn transition(&self, job: u64, change: impl FnOnce(&JobState) -> Option<JobState>) -> Option<JobStatus> {
200 let status = {
201 let mut jobs = self.jobs.lock().unwrap();
202 let record = jobs.get_mut(&job)?;
203 record.status.state = change(&record.status.state)?;
204 record.updated_at = SystemTime::now();
205 record.status.clone()
206 };
207
208 self.notify(status.clone());
209 Some(status)
210 }
211
212 fn set(&self, status: JobStatus) {
214 {
215 let mut jobs = self.jobs.lock().unwrap();
216 let now = SystemTime::now();
217 jobs.insert(status.job, JobRecord { status: status.clone(), created_at: now, updated_at: now });
218
219 let ended = jobs.values().filter(|record| record.status.state.has_ended()).map(|record| record.status.job).collect::<Vec<_>>();
220 for job in ended.iter().take(ended.len().saturating_sub(MAX_ENDED_JOBS)) {
221 jobs.remove(job);
222 }
223 }
224
225 self.notify(status);
226 }
227
228 fn notify(&self, status: JobStatus) {
230 let _ = self.updates.send(status);
231 }
232}
233
234pub fn is_job_control_method(method: &str) -> bool {
236 matches!(method, GetJobStatus::METHOD | WaitForJob::METHOD | CancelJob::METHOD)
237}
238
239pub fn handle_request(registry: Arc<JobRegistry>, request: RpcRequest, sender: UnboundedSender<RpcResponse>) {
247 tokio::spawn(async move {
248 let result = match request.method.as_str() {
249 GetJobStatus::METHOD => parse_params::<GetJobStatus>(request.params)
250 .and_then(|request| registry.status(request.job).ok_or(ApiError::JobNotFound(request.job)))
251 .and_then(to_json),
252 WaitForJob::METHOD => match parse_params::<WaitForJob>(request.params) {
253 Ok(request) => {
254 let timeout = Duration::from_secs(request.timeout_secs.unwrap_or(DEFAULT_WAIT_SECS));
255 registry.wait(request.job, timeout).await.ok_or(ApiError::JobNotFound(request.job)).and_then(to_json)
256 }
257 Err(error) => Err(error),
258 },
259 CancelJob::METHOD => parse_params::<CancelJob>(request.params)
260 .and_then(|request| registry.cancel(request.job))
261 .and_then(|_| to_json(Done {})),
262 method => Err(ApiError::MethodNotFound(method.to_owned())),
263 };
264
265 let _ = sender.send(RpcResponse::new(request.id, result));
266 });
267}
268
269fn to_json(value: impl Serialize) -> Result<Value, ApiError> {
271 serde_json::to_value(value).map_err(|error| ApiError::Internal(error.to_string()))
272}
273
274#[cfg(test)]
279#[path = "jobs_test.rs"]
280mod jobs_test;