Skip to main content

z39_solver/job/
mod.rs

1use std::path::PathBuf;
2use std::sync::Arc;
3use dashmap::DashMap;
4use serde::{Deserialize, Serialize};
5use tokio::sync::Mutex;
6use tokio::task::JoinHandle;
7
8use crate::solver::{self, SolveStatus};
9
10#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
11#[serde(rename_all = "snake_case")]
12pub enum JobStatus {
13    Pending,
14    Running,
15    Done,
16    Error,
17    Cancelled,
18}
19
20#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct Job {
22    pub id: String,
23    pub status: JobStatus,
24    pub label: String,
25    pub result: Option<String>,
26    pub duration_ms: Option<u64>,
27}
28
29pub struct JobManager {
30    jobs: Arc<DashMap<String, Job>>,
31    handles: Arc<Mutex<Vec<(String, JoinHandle<()>)>>>,
32    z3_bin: PathBuf,
33}
34
35impl JobManager {
36    pub fn new(z3_bin: PathBuf) -> Self {
37        Self {
38            jobs: Arc::new(DashMap::new()),
39            handles: Arc::new(Mutex::new(Vec::new())),
40            z3_bin,
41        }
42    }
43
44    #[allow(dead_code)]
45    pub fn z3_bin(&self) -> &PathBuf { &self.z3_bin }
46
47    pub async fn submit_with_bin(&self, label: String, smt_input: String, timeout_secs: u64, z3_bin: PathBuf) -> String {
48        let id = uuid::Uuid::new_v4().to_string()[..8].to_string();
49
50        let job = Job {
51            id: id.clone(),
52            status: JobStatus::Running,
53            label,
54            result: None,
55            duration_ms: None,
56        };
57        self.jobs.insert(id.clone(), job);
58
59        let jobs = self.jobs.clone();
60        let jid = id.clone();
61
62        let handle = tokio::spawn(async move {
63            let result = solver::solve(&z3_bin, &smt_input, timeout_secs).await;
64
65            if let Some(mut j) = jobs.get_mut(&jid) {
66                j.status = match &result.status {
67                    SolveStatus::Timeout | SolveStatus::Error(_) => JobStatus::Error,
68                    _ => JobStatus::Done,
69                };
70                j.result = Some(result.to_compact());
71                j.duration_ms = Some(result.duration_ms);
72            }
73        });
74
75        let mut handles = self.handles.lock().await;
76        handles.push((id.clone(), handle));
77        id
78    }
79
80    pub fn status(&self, job_id: &str) -> Option<Job> {
81        self.jobs.get(job_id).map(|j| j.clone())
82    }
83
84    pub fn result(&self, job_id: &str) -> Option<String> {
85        self.jobs.get(job_id).and_then(|j| j.result.clone())
86    }
87
88    pub async fn cancel(&self, job_id: &str) -> bool {
89        if let Some(mut j) = self.jobs.get_mut(job_id) {
90            if j.status == JobStatus::Running || j.status == JobStatus::Pending {
91                j.status = JobStatus::Cancelled;
92                let mut handles = self.handles.lock().await;
93                handles.retain(|(id, h)| {
94                    if id == job_id { h.abort(); false } else { true }
95                });
96                return true;
97            }
98        }
99        false
100    }
101
102    #[allow(dead_code)]
103    pub fn list(&self) -> Vec<String> {
104        self.jobs.iter().map(|j| format!("{} {:?} {}", j.id, j.status, j.label)).collect()
105    }
106}