137 lines
4.4 KiB
Rust
137 lines
4.4 KiB
Rust
use std::{collections::HashMap, sync::Arc, time::Duration};
|
|
|
|
use async_trait::async_trait;
|
|
use tokio::time::sleep;
|
|
use tracing::{error, info, warn};
|
|
|
|
use crate::{
|
|
jobs::{mark_job_failed, mark_job_succeeded, reserve_job, retry_job_after, JobQueueError},
|
|
models::Job,
|
|
state::AppState,
|
|
};
|
|
|
|
pub mod analyze;
|
|
pub mod index;
|
|
pub mod ocr;
|
|
pub mod thumbnails;
|
|
|
|
#[derive(Debug)]
|
|
pub enum JobExecution {
|
|
Success,
|
|
Retry { delay: Duration, error: String },
|
|
Failed { error: String },
|
|
}
|
|
|
|
#[async_trait]
|
|
pub trait JobHandler: Send + Sync {
|
|
fn job_type(&self) -> &'static str;
|
|
async fn handle(&self, state: Arc<AppState>, job: Job) -> JobExecution;
|
|
}
|
|
|
|
pub struct Worker {
|
|
state: Arc<AppState>,
|
|
handlers: HashMap<&'static str, Arc<dyn JobHandler>>,
|
|
poll_interval: Duration,
|
|
}
|
|
|
|
impl Worker {
|
|
pub fn new(
|
|
state: Arc<AppState>,
|
|
handlers: Vec<Arc<dyn JobHandler>>,
|
|
poll_interval: Duration,
|
|
) -> Self {
|
|
let map = handlers
|
|
.into_iter()
|
|
.map(|handler| (handler.job_type(), handler))
|
|
.collect();
|
|
Self {
|
|
state,
|
|
handlers: map,
|
|
poll_interval,
|
|
}
|
|
}
|
|
|
|
pub async fn run(&self) {
|
|
info!("worker started");
|
|
loop {
|
|
match self.tick().await {
|
|
Ok(true) => {}
|
|
Ok(false) => sleep(self.poll_interval).await,
|
|
Err(err) => {
|
|
error!(error = %err, "worker tick failed");
|
|
sleep(self.poll_interval).await;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn tick(&self) -> Result<bool, JobQueueError> {
|
|
let job_types: Vec<&str> = self.handlers.keys().copied().collect();
|
|
if job_types.is_empty() {
|
|
return Ok(false);
|
|
}
|
|
|
|
let mut conn = match self.state.db() {
|
|
Ok(conn) => conn,
|
|
Err(err) => {
|
|
error!(?err, "failed to obtain database connection in worker");
|
|
return Ok(false);
|
|
}
|
|
};
|
|
|
|
let job_opt = reserve_job(&mut conn, &job_types)?;
|
|
drop(conn);
|
|
|
|
if let Some(job) = job_opt {
|
|
if let Some(handler) = self.handlers.get(job.job_type.as_str()) {
|
|
let result = handler.handle(self.state.clone(), job.clone()).await;
|
|
match result {
|
|
JobExecution::Success => {
|
|
if let Ok(mut conn) = self.state.db() {
|
|
mark_job_succeeded(&mut conn, job.id)?;
|
|
info!(job_id = %job.id, job_type = %job.job_type, "job completed successfully");
|
|
} else {
|
|
error!("failed to mark job succeeded due to pool error");
|
|
}
|
|
}
|
|
JobExecution::Retry { delay, error } => {
|
|
warn!(job_id = %job.id, job_type = %job.job_type, %error, "job will retry");
|
|
if let Ok(mut conn) = self.state.db() {
|
|
retry_job_after(&mut conn, job.id, delay, &error)?;
|
|
} else {
|
|
error!("failed to requeue job for retry due to pool error");
|
|
}
|
|
}
|
|
JobExecution::Failed { error } => {
|
|
error!(job_id = %job.id, job_type = %job.job_type, %error, "job failed");
|
|
if let Ok(mut conn) = self.state.db() {
|
|
mark_job_failed(&mut conn, job.id, &error)?;
|
|
} else {
|
|
error!("failed to mark job failed due to pool error");
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
error!(job_type = %job.job_type, "no handler registered for job type");
|
|
if let Ok(mut conn) = self.state.db() {
|
|
mark_job_failed(&mut conn, job.id, "no handler registered")?;
|
|
} else {
|
|
error!("failed to mark job failed for missing handler due to pool error");
|
|
}
|
|
}
|
|
Ok(true)
|
|
} else {
|
|
Ok(false)
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn default_handlers() -> Vec<Arc<dyn JobHandler>> {
|
|
vec![
|
|
Arc::new(analyze::AnalyzeDocumentJob::new()),
|
|
Arc::new(thumbnails::GenerateThumbnailsJob::new()),
|
|
Arc::new(ocr::GenerateOcrTextJob::new()),
|
|
Arc::new(index::IndexDocumentTextJob::new()),
|
|
]
|
|
}
|