2

backend/src/workers/mod.rs

@@ -0,0 +1,131 @@
+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 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)?;
+                        } 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()),
+    ]
+}