use crate::{
config::QueueConfig,
db,
error::QueueError,
executor::{ProcessedJob, QueueExecutor},
types::{QueueJob, QueueJobDetails, QueuePriority, QueueStats},
JobHandler, QueueEventEmitter,
};
use rusqlite::Connection;
use std::sync::{Arc, Mutex};
pub struct QueueManager {
db: Arc<Mutex<Connection>>,
executor: Arc<QueueExecutor>,
config: QueueConfig,
}
impl QueueManager {
pub fn new(config: QueueConfig) -> Result<Self, QueueError> {
let conn = db::open_database(config.db_path.as_deref())?;
let requeued = db::requeue_interrupted(&conn)?;
if requeued > 0 {
tracing::info!(count = requeued, "Requeued interrupted jobs");
}
let db = Arc::new(Mutex::new(conn));
let executor = Arc::new(QueueExecutor::new(config.clone(), Arc::clone(&db)));
Ok(Self {
db,
executor,
config,
})
}
pub fn add<H>(&self, job: QueueJob<H>) -> Result<String, QueueError>
where
H: JobHandler,
{
let conn = self.db.lock()?;
let data = serde_json::to_value(&job.data)?;
db::insert_job_full(
&conn,
&job.id,
job.priority.as_i32(),
&data,
job.trace_id.as_deref(),
job.attempt_id.as_ref().map(|a| a.as_str()),
job.trial_id.as_ref().map(|t| t.as_str()),
)?;
Ok(job.id)
}
pub fn cancel(&self, job_id: &str) -> Result<(), QueueError> {
let conn = self.db.lock()?;
db::cancel_job(&conn, job_id)?;
Ok(())
}
pub fn reorder(&self, job_id: &str, new_priority: QueuePriority) -> Result<(), QueueError> {
let conn = self.db.lock()?;
match db::reorder_pending(&conn, job_id, new_priority.as_i32()) {
Ok(true) => Ok(()),
Ok(false) => {
let status = db::get_job(&conn, job_id)
.ok()
.flatten()
.map(|j| j.2)
.unwrap_or_else(|| "unknown".to_string());
Err(QueueError::Other(format!(
"Can only reorder pending jobs (job {} is {})",
job_id, status
)))
}
Err(_) => Err(QueueError::NotFound(job_id.to_string())),
}
}
pub fn pause(&self) {
self.executor.pause();
}
pub fn resume(&self) {
self.executor.resume();
}
pub fn is_paused(&self) -> bool {
self.executor.is_paused()
}
pub fn list_jobs(&self) -> Result<Vec<(String, String)>, QueueError> {
let conn = self.db.lock()?;
let jobs = db::list_all_jobs(&conn)?;
Ok(jobs
.into_iter()
.map(|(id, status, _)| (id, status))
.collect())
}
pub fn list_jobs_with_data(&self) -> Result<Vec<(String, String, String)>, QueueError> {
let conn = self.db.lock()?;
Ok(db::list_all_jobs(&conn)?)
}
pub fn get_job_details(&self, job_id: &str) -> Result<Option<QueueJobDetails>, QueueError> {
let conn = self.db.lock()?;
Ok(db::get_job_details(&conn, job_id)?)
}
pub fn prune(&self, days: u32) -> Result<u32, QueueError> {
let conn = self.db.lock()?;
Ok(db::prune_old_jobs(&conn, days)?)
}
pub async fn process_one<H>(
&self,
event_emitter: &Arc<dyn QueueEventEmitter>,
) -> Result<Option<ProcessedJob>, QueueError>
where
H: JobHandler,
{
self.executor.process_one::<H>(event_emitter).await
}
pub fn count_by_status(&self) -> Result<QueueStats, QueueError> {
let conn = self.db.lock()?;
Ok(db::count_by_status(&conn)?)
}
pub fn shutdown(&self) {
self.executor.shutdown();
}
pub fn is_shutdown(&self) -> bool {
self.executor.is_shutdown()
}
pub fn worker_id(&self) -> &str {
&self.config.worker_id
}
pub fn spawn<H>(self, event_emitter: Arc<dyn QueueEventEmitter>) -> Arc<Self>
where
H: JobHandler + 'static,
{
let manager = Arc::new(self);
let executor = Arc::clone(&manager.executor);
executor.spawn::<H>(event_emitter);
manager
}
pub fn spawn_on<H>(
self,
event_emitter: Arc<dyn QueueEventEmitter>,
handle: &tokio::runtime::Handle,
) -> Arc<Self>
where
H: JobHandler + 'static,
{
let manager = Arc::new(self);
let executor = Arc::clone(&manager.executor);
executor.spawn_on::<H>(event_emitter, handle);
manager
}
}