Skip to main content

rpfm_server/
jobs.rs

1//---------------------------------------------------------------------------//
2// Copyright (c) 2017-2026 Ismael Gutiérrez González. All rights reserved.
3//
4// This file is part of the Rusted PackFile Manager (RPFM) project,
5// which can be found here: https://github.com/Frodo45127/rpfm.
6//
7// This file is licensed under the MIT license, which can be found here:
8// https://github.com/Frodo45127/rpfm/blob/master/LICENSE.
9//---------------------------------------------------------------------------//
10
11//! State of the jobs of a session. See [`rpfm_ipc::api::jobs`].
12//!
13//! Jobs run on the session's background thread like any other request. The registry only keeps their
14//! state, so it can be checked, waited for and cancelled from outside the background thread, and
15//! broadcasts every change so connected clients get notified.
16
17use 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
32/// Amount of ended jobs whose state is kept. Older ones are forgotten.
33const MAX_ENDED_JOBS: usize = 100;
34
35/// Capacity of the channel broadcasting job changes. Slow listeners skip the changes they miss.
36const UPDATES_CAPACITY: usize = 64;
37
38/// State of the jobs of a session.
39#[derive(Debug)]
40pub struct JobRegistry {
41
42    /// ID of the next job.
43    next_id: AtomicU64,
44
45    /// State of each job, by ID.
46    jobs: Mutex<BTreeMap<u64, JobRecord>>,
47
48    /// Channel every job change is sent to.
49    updates: broadcast::Sender<JobStatus>,
50}
51
52/// State of a job, with when it was created and when it last changed.
53#[derive(Clone, Debug)]
54pub struct JobRecord {
55
56    /// State of the job.
57    pub status: JobStatus,
58
59    /// When the job was created.
60    pub created_at: SystemTime,
61
62    /// When the state of the job last changed.
63    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    /// Returns if any job is queued or running.
79    pub fn has_unfinished_jobs(&self) -> bool {
80        self.jobs.lock().unwrap().values().any(|record| !record.status.state.has_ended())
81    }
82
83    /// Registers a new queued job.
84    ///
85    /// # Returns
86    ///
87    /// The ID of the job.
88    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    /// Marks a queued job as running.
95    ///
96    /// # Returns
97    ///
98    /// `false` if the job was cancelled or is unknown, so it must not run.
99    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    /// Sets the step a running job is on.
104    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    /// Sets how far along a running job is, from 0 to 100.
109    ///
110    /// Only increases are kept, as parallel steps of a job may report their progress out of order.
111    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    /// Moves a running job back to the queue.
119    ///
120    /// # Returns
121    ///
122    /// `false` if the job isn't running.
123    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    /// Ends a job as cancelled, even if it already started.
128    pub fn finish_cancelled(&self, job: u64) {
129        self.transition(job, |state| (!state.has_ended()).then_some(JobState::Cancelled));
130    }
131
132    /// Ends a job with the outcome of its method.
133    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    /// Cancels a job that hasn't started yet.
143    ///
144    /// # Errors
145    ///
146    /// Fails if the job is unknown, or if it already started.
147    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    /// Returns the state of a job, or `None` if it's unknown.
157    pub fn status(&self, job: u64) -> Option<JobStatus> {
158        self.record(job).map(|record| record.status)
159    }
160
161    /// Returns the state of a job with when it was created and last changed, or `None` if it's unknown.
162    pub fn record(&self, job: u64) -> Option<JobRecord> {
163        self.jobs.lock().unwrap().get(&job).cloned()
164    }
165
166    /// Returns a receiver of every job change from now on.
167    pub fn subscribe(&self) -> broadcast::Receiver<JobStatus> {
168        self.updates.subscribe()
169    }
170
171    /// Waits for a job to end, up to `timeout`.
172    ///
173    /// # Returns
174    ///
175    /// The state of the job when it ended, or when the wait timed out. `None` if the job is unknown.
176    pub async fn wait(&self, job: u64, timeout: Duration) -> Option<JobStatus> {
177
178        // Subscribe before checking the state, so an end happening in between isn't missed.
179        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    /// Changes the state of a job with `change`, which returns the new state, or `None` to keep the current one.
195    ///
196    /// # Returns
197    ///
198    /// The new state of the job, or `None` if it's unknown or didn't change.
199    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    /// Stores the state of a job, forgetting the oldest ended jobs over [`MAX_ENDED_JOBS`], and notifies it.
213    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    /// Sends a job change to the listeners, if there are any.
229    fn notify(&self, status: JobStatus) {
230        let _ = self.updates.send(status);
231    }
232}
233
234/// Returns if a method is one of the job control methods, which are answered by [`handle_request`].
235pub fn is_job_control_method(method: &str) -> bool {
236    matches!(method, GetJobStatus::METHOD | WaitForJob::METHOD | CancelJob::METHOD)
237}
238
239/// Answers a job control request, without waiting for the session's other requests.
240///
241/// # Arguments
242///
243/// * `registry` - Jobs of the session.
244/// * `request` - A job control request. See [`is_job_control_method`].
245/// * `sender` - Where to send the response.
246pub 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
269/// Serializes a response.
270fn 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//-------------------------------------------------------------------------------//
275//                                   Tests
276//-------------------------------------------------------------------------------//
277
278#[cfg(test)]
279#[path = "jobs_test.rs"]
280mod jobs_test;