clean project structure

crates/steep-server/src/job_queue.rs

@@ -1,26 +1,30 @@
-use std::{collections::VecDeque, time::Duration};
+use std::sync::Arc;
 
-use futures_util::stream::{self, Stream};
-use toasty::Db;
-use tracing::error;
+use toasty::{Db, stmt::Value};
+use tokio::sync::Notify;
 
 use crate::{
     gitea::WebhookType,
     schemas::{Job, JobStatus},
 };
 
-/// Délai entre deux interrogations de la table quand la file est vide.
-const POLL_INTERVAL: Duration = Duration::from_secs(1);
-
-/// File des reviews à faire.
+/// The reviews to do, stored in PostgreSQL.
+///
+/// Steep runs as a single instance, so every job is pushed by this process: the
+/// worker is woken up by [`JobQueue::push`] instead of polling the table. The table
+/// only makes the queue survive a restart.
 #[derive(Clone)]
 pub struct JobQueue {
     db: Db,
+    work: Arc<Notify>,
 }
 
 impl JobQueue {
     pub fn new(db: Db) -> Self {
-        Self { db }
+        Self {
+            db,
+            work: Arc::new(Notify::new()),
+        }
     }
 
     pub async fn push(&self, webhook: WebhookType) -> toasty::Result<Job> {
@@ -40,54 +44,76 @@ impl JobQueue {
         .exec(&mut db)
         .await?;
 
+        // `notify_one` keeps a permit when nobody is waiting yet, so a job pushed
+        // while the worker is busy is not missed.
+        self.work.notify_one();
+
         Ok(job)
     }
 
-    pub fn pending(&self) -> impl Stream<Item = Job> {
-        let db = self.db.clone();
-
-        stream::unfold(VecDeque::new(), move |mut buffer| {
-            let db = db.clone();
-
-            async move {
-                let mut db = db;
-
-                loop {
-                    if let Some(job) = buffer.pop_front() {
-                        return Some((job, buffer));
-                    }
-
-                    match Job::filter_by_status(JobStatus::Pending)
-                        .exec(&mut db)
-                        .await
-                    {
-                        Ok(jobs) if jobs.is_empty() => {
-                            tokio::time::sleep(POLL_INTERVAL).await;
-                        }
-                        Ok(jobs) => buffer.extend(jobs),
-                        Err(err) => {
-                            error!(%err, "failed to poll pending jobs");
-                            tokio::time::sleep(POLL_INTERVAL).await;
-                        }
-                    }
-                }
-            }
-        })
-    }
-
-    /// Marque un job comme en cours de traitement.
-    pub async fn mark_processing(&self, id: u64) -> toasty::Result<()> {
+    /// Whether a job is already queued for this event.
+    pub async fn is_queued(&self, event_id: u64) -> toasty::Result<bool> {
         let mut db = self.db.clone();
 
-        Job::update_by_id(id)
-            .status(JobStatus::Processing)
+        Ok(Job::filter_by_event_id(event_id)
+            .first()
+            .exec(&mut db)
+            .await?
+            .is_some())
+    }
+
+    /// Puts back in the queue the jobs a previous process left `processing`.
+    ///
+    /// To call once at startup, before the worker starts: it is only correct because
+    /// no other instance can be working on those jobs.
+    pub async fn requeue_interrupted(&self) -> toasty::Result<()> {
+        let mut db = self.db.clone();
+
+        Job::update_by_status(JobStatus::Processing)
+            .status(JobStatus::Pending)
             .exec(&mut db)
             .await?;
 
         Ok(())
     }
 
-    /// Retire un job de la file, une fois la review enregistrée.
+    /// Takes the oldest pending job and marks it `processing`, or returns `None` when
+    /// the queue is empty.
+    ///
+    /// The selection and the update are a single statement, so a job cannot be taken
+    /// twice.
+    pub async fn claim_next(&self) -> toasty::Result<Option<Job>> {
+        let mut db = self.db.clone();
+
+        let rows = toasty::sql::query(
+            "UPDATE jobs
+             SET status = 'processing'
+             WHERE id = (
+                 SELECT id FROM jobs
+                 WHERE status = 'pending'
+                 ORDER BY id
+                 LIMIT 1
+             )
+             RETURNING id",
+        )
+        .exec(&mut db)
+        .await?;
+
+        let Some(Value::Record(row)) = rows.into_iter().next() else {
+            return Ok(None);
+        };
+
+        let id = u64::try_from(row.as_slice()[0].clone())?;
+
+        Ok(Some(Job::get_by_id(&mut db, id).await?))
+    }
+
+    /// Waits until a job may have been pushed.
+    pub async fn wait_for_work(&self) {
+        self.work.notified().await;
+    }
+
+    /// Removes a job from the queue, once its review is recorded.
     pub async fn remove(&self, id: u64) -> toasty::Result<()> {
         let mut db = self.db.clone();
 
@@ -108,11 +134,13 @@ impl JobQueue {
 mod tests {
     use std::time::Duration;
 
-    use futures_util::StreamExt;
+    use tokio::sync::{Mutex, MutexGuard};
 
     use super::*;
     use crate::gitea::{Comment, PullRequest, Repository, ReviewPayload};
 
+    const FIRST_EVENT: u64 = 999_000_200;
+
     fn webhook(event_id: u64) -> WebhookType {
         WebhookType::Review(ReviewPayload {
             action: "created".to_string(),
@@ -131,41 +159,40 @@ mod tests {
         })
     }
 
-    async fn next_job(queue: &JobQueue, event_id: u64) -> Job {
-        let mut pending = std::pin::pin!(queue.pending());
+    /// The tests claim whichever job is next, so they must not overlap: one would
+    /// claim the jobs of another.
+    async fn exclusive_queue() -> (MutexGuard<'static, ()>, JobQueue, Db) {
+        static LOCK: Mutex<()> = Mutex::const_new(());
 
-        tokio::time::timeout(Duration::from_secs(10), async {
-            loop {
-                if let Some(job) = pending.next().await
-                    && job.event_id == event_id
-                {
-                    break job;
-                }
-            }
-        })
+        let guard = LOCK.lock().await;
+        let db = crate::db::test_connect().await;
+
+        toasty::sql::statement(
+            "DELETE FROM jobs WHERE event_id >= 999000200 AND event_id < 999000300",
+        )
+        .exec(&mut db.clone())
         .await
-        .expect("timed out waiting for the job")
+        .unwrap();
+
+        (guard, JobQueue::new(db.clone()), db)
     }
 
     #[tokio::test]
     #[ignore = "requires DATABASE_URL"]
-    async fn a_pushed_job_is_pending_then_removed() {
-        let db = crate::db::test_connect().await;
-        let queue = JobQueue::new(db.clone());
+    async fn a_pushed_job_is_claimed_then_removed() {
+        let (_guard, queue, db) = exclusive_queue().await;
 
-        let event_id = 999_000_201;
-        Job::delete_by_event_id(&mut db.clone(), event_id)
-            .await
-            .ok();
+        queue.push(webhook(FIRST_EVENT)).await.unwrap();
+        assert!(queue.is_queued(FIRST_EVENT).await.unwrap());
 
-        queue.push(webhook(event_id)).await.unwrap();
+        let job = queue.claim_next().await.unwrap().expect("a job to claim");
+        assert_eq!(job.event_id, FIRST_EVENT);
+        assert_eq!(job.status, JobStatus::Processing);
 
-        let job = next_job(&queue, event_id).await;
-        queue.mark_processing(job.id).await.unwrap();
         queue.remove(job.id).await.unwrap();
 
         assert!(
-            Job::filter_by_event_id(event_id)
+            Job::filter_by_event_id(FIRST_EVENT)
                 .first()
                 .exec(&mut db.clone())
                 .await
@@ -177,19 +204,70 @@ mod tests {
     #[tokio::test]
     #[ignore = "requires DATABASE_URL"]
     async fn a_duplicate_event_id_is_rejected() {
-        let db = crate::db::test_connect().await;
-        let queue = JobQueue::new(db.clone());
+        let (_guard, queue, _db) = exclusive_queue().await;
 
-        let event_id = 999_000_202;
-        Job::delete_by_event_id(&mut db.clone(), event_id)
-            .await
-            .ok();
+        queue.push(webhook(FIRST_EVENT + 1)).await.unwrap();
+        assert!(queue.push(webhook(FIRST_EVENT + 1)).await.is_err());
+    }
 
-        queue.push(webhook(event_id)).await.unwrap();
-        assert!(queue.push(webhook(event_id)).await.is_err());
+    #[tokio::test]
+    #[ignore = "requires DATABASE_URL"]
+    async fn a_claimed_job_is_not_claimed_again() {
+        let (_guard, queue, _db) = exclusive_queue().await;
 
-        Job::delete_by_event_id(&mut db.clone(), event_id)
+        queue.push(webhook(FIRST_EVENT + 2)).await.unwrap();
+
+        assert!(queue.claim_next().await.unwrap().is_some());
+        assert!(queue.claim_next().await.unwrap().is_none());
+    }
+
+    #[tokio::test]
+    #[ignore = "requires DATABASE_URL"]
+    async fn jobs_are_claimed_oldest_first() {
+        let (_guard, queue, _db) = exclusive_queue().await;
+
+        queue.push(webhook(FIRST_EVENT + 3)).await.unwrap();
+        queue.push(webhook(FIRST_EVENT + 4)).await.unwrap();
+
+        let first = queue.claim_next().await.unwrap().unwrap();
+        let second = queue.claim_next().await.unwrap().unwrap();
+
+        assert_eq!(first.event_id, FIRST_EVENT + 3);
+        assert_eq!(second.event_id, FIRST_EVENT + 4);
+    }
+
+    #[tokio::test]
+    #[ignore = "requires DATABASE_URL"]
+    async fn an_interrupted_job_is_requeued() {
+        let (_guard, queue, _db) = exclusive_queue().await;
+
+        queue.push(webhook(FIRST_EVENT + 5)).await.unwrap();
+        let job = queue.claim_next().await.unwrap().unwrap();
+        assert!(queue.claim_next().await.unwrap().is_none());
+
+        // The process died here: on restart, the job goes back to the queue.
+        queue.requeue_interrupted().await.unwrap();
+
+        let again = queue.claim_next().await.unwrap().expect("the job back");
+        assert_eq!(again.id, job.id);
+    }
+
+    #[tokio::test]
+    #[ignore = "requires DATABASE_URL"]
+    async fn waiting_for_work_returns_as_soon_as_a_job_is_pushed() {
+        let (_guard, queue, _db) = exclusive_queue().await;
+
+        let waiter = {
+            let queue = queue.clone();
+            tokio::spawn(async move { queue.wait_for_work().await })
+        };
+
+        tokio::time::sleep(Duration::from_millis(100)).await;
+        queue.push(webhook(FIRST_EVENT + 6)).await.unwrap();
+
+        tokio::time::timeout(Duration::from_secs(1), waiter)
             .await
+            .expect("the worker was not woken up")
             .unwrap();
     }
 }