steep/crates/steep-server/src/job_queue.rs

use std::{collections::VecDeque, time::Duration};

use futures_util::stream::{self, Stream};
use toasty::Db;
use tracing::error;

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.
#[derive(Clone)]
pub struct JobQueue {
    db: Db,
}

impl JobQueue {
    pub fn new(db: Db) -> Self {
        Self { db }
    }

    pub async fn push(&self, webhook: WebhookType) -> toasty::Result<Job> {
        let event_id = webhook.event_id();
        let repo = webhook.repository_full_name().to_string();
        let pr = webhook.pull_request_number();

        let mut db = self.db.clone();

        let job = toasty::create!(Job {
            event_id,
            repo,
            pr,
            payload: webhook,
            status: JobStatus::Pending,
        })
        .exec(&mut db)
        .await?;

        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<()> {
        let mut db = self.db.clone();

        Job::update_by_id(id)
            .status(JobStatus::Processing)
            .exec(&mut db)
            .await?;

        Ok(())
    }

    /// Retire un job de la file, une fois la review enregistrée.
    pub async fn remove(&self, id: u64) -> toasty::Result<()> {
        let mut db = self.db.clone();

        Job::delete_by_id(&mut db, id).await?;

        Ok(())
    }

    pub async fn list(&self, status: JobStatus) -> toasty::Result<Vec<Job>> {
        let mut db = self.db.clone();
        let jobs = Job::filter_by_status(status).exec(&mut db).await?;

        Ok(jobs)
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use futures_util::StreamExt;

    use super::*;
    use crate::gitea::{Comment, PullRequest, Repository, ReviewPayload};

    fn webhook(event_id: u64) -> WebhookType {
        WebhookType::Review(ReviewPayload {
            action: "created".to_string(),
            pull_request: PullRequest {
                number: 1,
                title: "Test".to_string(),
            },
            repository: Repository {
                full_name: "owner/repo".to_string(),
                clone_url: "https://example.com/owner/repo.git".to_string(),
            },
            comment: Comment {
                id: event_id,
                body: "@bot review".to_string(),
            },
        })
    }

    async fn next_job(queue: &JobQueue, event_id: u64) -> Job {
        let mut pending = std::pin::pin!(queue.pending());

        tokio::time::timeout(Duration::from_secs(10), async {
            loop {
                if let Some(job) = pending.next().await
                    && job.event_id == event_id
                {
                    break job;
                }
            }
        })
        .await
        .expect("timed out waiting for the job")
    }

    #[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());

        let event_id = 999_000_201;
        Job::delete_by_event_id(&mut db.clone(), event_id)
            .await
            .ok();

        queue.push(webhook(event_id)).await.unwrap();

        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)
                .first()
                .exec(&mut db.clone())
                .await
                .unwrap()
                .is_none()
        );
    }

    #[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 event_id = 999_000_202;
        Job::delete_by_event_id(&mut db.clone(), event_id)
            .await
            .ok();

        queue.push(webhook(event_id)).await.unwrap();
        assert!(queue.push(webhook(event_id)).await.is_err());

        Job::delete_by_event_id(&mut db.clone(), event_id)
            .await
            .unwrap();
    }
}