274 lines
8.0 KiB
Rust
274 lines
8.0 KiB
Rust
use std::sync::Arc;
|
|
|
|
use toasty::{Db, stmt::Value};
|
|
use tokio::sync::Notify;
|
|
|
|
use crate::{
|
|
gitea::WebhookType,
|
|
schemas::{Job, JobStatus},
|
|
};
|
|
|
|
/// 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,
|
|
work: Arc::new(Notify::new()),
|
|
}
|
|
}
|
|
|
|
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?;
|
|
|
|
// `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)
|
|
}
|
|
|
|
/// 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();
|
|
|
|
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(())
|
|
}
|
|
|
|
/// 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();
|
|
|
|
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 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(),
|
|
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(),
|
|
},
|
|
})
|
|
}
|
|
|
|
/// 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(());
|
|
|
|
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
|
|
.unwrap();
|
|
|
|
(guard, JobQueue::new(db.clone()), db)
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[ignore = "requires DATABASE_URL"]
|
|
async fn a_pushed_job_is_claimed_then_removed() {
|
|
let (_guard, queue, db) = exclusive_queue().await;
|
|
|
|
queue.push(webhook(FIRST_EVENT)).await.unwrap();
|
|
assert!(queue.is_queued(FIRST_EVENT).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);
|
|
|
|
queue.remove(job.id).await.unwrap();
|
|
|
|
assert!(
|
|
Job::filter_by_event_id(FIRST_EVENT)
|
|
.first()
|
|
.exec(&mut db.clone())
|
|
.await
|
|
.unwrap()
|
|
.is_none()
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[ignore = "requires DATABASE_URL"]
|
|
async fn a_duplicate_event_id_is_rejected() {
|
|
let (_guard, queue, _db) = exclusive_queue().await;
|
|
|
|
queue.push(webhook(FIRST_EVENT + 1)).await.unwrap();
|
|
assert!(queue.push(webhook(FIRST_EVENT + 1)).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;
|
|
|
|
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();
|
|
}
|
|
}
|