2023-03-04 19:46:09 +00:00
|
|
|
use crate::errors::AsyncQueueError;
|
2023-03-10 22:41:34 +00:00
|
|
|
use crate::runnable::BackgroundTask;
|
|
|
|
use crate::task::{NewTask, Task, TaskHash, TaskId, TaskState};
|
2023-03-04 18:07:17 +00:00
|
|
|
use diesel::result::Error::QueryBuilderError;
|
|
|
|
use diesel_async::scoped_futures::ScopedFutureExt;
|
|
|
|
use diesel_async::AsyncConnection;
|
2023-03-05 00:19:35 +00:00
|
|
|
use diesel_async::{pg::AsyncPgConnection, pooled_connection::bb8::Pool};
|
2023-03-10 22:41:34 +00:00
|
|
|
use std::sync::Arc;
|
|
|
|
use std::time::Duration;
|
2023-03-04 18:07:17 +00:00
|
|
|
|
2023-03-10 22:41:34 +00:00
|
|
|
#[derive(Clone)]
|
|
|
|
pub struct Queue {
|
|
|
|
task_store: Arc<PgTaskStore>,
|
|
|
|
}
|
2023-03-04 18:07:17 +00:00
|
|
|
|
2023-03-10 22:41:34 +00:00
|
|
|
impl Queue {
|
|
|
|
pub(crate) fn new(task_store: Arc<PgTaskStore>) -> Self {
|
|
|
|
Queue { task_store }
|
|
|
|
}
|
2023-03-04 18:07:17 +00:00
|
|
|
|
2023-03-10 22:41:34 +00:00
|
|
|
pub async fn enqueue<BT>(&self, background_task: BT) -> Result<(), AsyncQueueError>
|
|
|
|
where
|
|
|
|
BT: BackgroundTask,
|
|
|
|
{
|
|
|
|
self.task_store
|
|
|
|
.create_task(NewTask::new(background_task, Duration::from_secs(10))?)
|
|
|
|
.await?;
|
|
|
|
Ok(())
|
|
|
|
}
|
2023-03-04 18:07:17 +00:00
|
|
|
}
|
|
|
|
|
2023-03-07 16:52:26 +00:00
|
|
|
/// An async queue that is used to manipulate tasks, it uses PostgreSQL as storage.
|
2023-03-07 15:41:20 +00:00
|
|
|
#[derive(Debug, Clone)]
|
2023-03-10 22:41:34 +00:00
|
|
|
pub struct PgTaskStore {
|
2023-03-04 18:07:17 +00:00
|
|
|
pool: Pool<AsyncPgConnection>,
|
|
|
|
}
|
|
|
|
|
2023-03-10 22:41:34 +00:00
|
|
|
impl PgTaskStore {
|
2023-03-07 15:41:20 +00:00
|
|
|
pub fn new(pool: Pool<AsyncPgConnection>) -> Self {
|
2023-03-10 22:41:34 +00:00
|
|
|
PgTaskStore { pool }
|
2023-03-04 18:07:17 +00:00
|
|
|
}
|
|
|
|
|
2023-03-10 22:41:34 +00:00
|
|
|
pub(crate) async fn pull_next_task(
|
|
|
|
&self,
|
|
|
|
queue_name: &str,
|
2023-03-04 18:07:17 +00:00
|
|
|
) -> Result<Option<Task>, AsyncQueueError> {
|
|
|
|
let mut connection = self
|
|
|
|
.pool
|
|
|
|
.get()
|
|
|
|
.await
|
|
|
|
.map_err(|e| QueryBuilderError(e.into()))?;
|
|
|
|
connection
|
|
|
|
.transaction::<Option<Task>, AsyncQueueError, _>(|conn| {
|
|
|
|
async move {
|
2023-03-10 22:41:34 +00:00
|
|
|
let Some(pending_task) = Task::fetch_next_pending(conn, queue_name).await else {
|
2023-03-04 18:07:17 +00:00
|
|
|
return Ok(None);
|
|
|
|
};
|
|
|
|
|
2023-03-07 16:52:26 +00:00
|
|
|
Task::set_running(conn, pending_task).await.map(Some)
|
2023-03-04 18:07:17 +00:00
|
|
|
}
|
2023-03-04 19:46:09 +00:00
|
|
|
.scope_boxed()
|
2023-03-04 18:07:17 +00:00
|
|
|
})
|
|
|
|
.await
|
|
|
|
}
|
|
|
|
|
2023-03-10 22:41:34 +00:00
|
|
|
pub(crate) async fn create_task(&self, new_task: NewTask) -> Result<Task, AsyncQueueError> {
|
2023-03-04 18:07:17 +00:00
|
|
|
let mut connection = self
|
|
|
|
.pool
|
|
|
|
.get()
|
|
|
|
.await
|
|
|
|
.map_err(|e| QueryBuilderError(e.into()))?;
|
2023-03-10 22:41:34 +00:00
|
|
|
Task::insert(&mut connection, new_task).await
|
2023-03-04 18:07:17 +00:00
|
|
|
}
|
|
|
|
|
2023-03-10 22:41:34 +00:00
|
|
|
pub(crate) async fn find_task_by_id(&self, id: TaskId) -> Result<Task, AsyncQueueError> {
|
2023-03-04 18:07:17 +00:00
|
|
|
let mut connection = self
|
|
|
|
.pool
|
|
|
|
.get()
|
|
|
|
.await
|
|
|
|
.map_err(|e| QueryBuilderError(e.into()))?;
|
2023-03-07 15:41:20 +00:00
|
|
|
Task::find_by_id(&mut connection, id).await
|
2023-03-04 18:07:17 +00:00
|
|
|
}
|
|
|
|
|
2023-03-10 22:41:34 +00:00
|
|
|
pub(crate) async fn set_task_state(
|
|
|
|
&self,
|
2023-03-07 15:41:20 +00:00
|
|
|
id: TaskId,
|
2023-03-10 22:41:34 +00:00
|
|
|
state: TaskState,
|
|
|
|
) -> Result<(), AsyncQueueError> {
|
2023-03-04 18:07:17 +00:00
|
|
|
let mut connection = self
|
|
|
|
.pool
|
|
|
|
.get()
|
|
|
|
.await
|
|
|
|
.map_err(|e| QueryBuilderError(e.into()))?;
|
2023-03-10 22:41:34 +00:00
|
|
|
match state {
|
|
|
|
TaskState::Done => Task::set_done(&mut connection, id).await?,
|
|
|
|
TaskState::Failed(error_msg) => {
|
|
|
|
Task::fail_with_message(&mut connection, id, &error_msg).await?
|
|
|
|
}
|
|
|
|
_ => return Ok(()),
|
|
|
|
};
|
|
|
|
Ok(())
|
2023-03-07 15:41:20 +00:00
|
|
|
}
|
|
|
|
|
2023-03-10 22:41:34 +00:00
|
|
|
pub(crate) async fn keep_task_alive(&self, id: TaskId) -> Result<(), AsyncQueueError> {
|
2023-03-07 15:41:20 +00:00
|
|
|
let mut connection = self
|
|
|
|
.pool
|
|
|
|
.get()
|
|
|
|
.await
|
|
|
|
.map_err(|e| QueryBuilderError(e.into()))?;
|
|
|
|
connection
|
|
|
|
.transaction::<(), AsyncQueueError, _>(|conn| {
|
|
|
|
async move {
|
|
|
|
let task = Task::find_by_id(conn, id).await?;
|
|
|
|
Task::set_running(conn, task).await?;
|
|
|
|
Ok(())
|
2023-03-07 16:52:26 +00:00
|
|
|
}
|
|
|
|
.scope_boxed()
|
|
|
|
})
|
|
|
|
.await
|
2023-03-04 18:07:17 +00:00
|
|
|
}
|
|
|
|
|
2023-03-10 22:41:34 +00:00
|
|
|
pub(crate) async fn remove_task(&self, id: TaskId) -> Result<u64, AsyncQueueError> {
|
2023-03-04 18:07:17 +00:00
|
|
|
let mut connection = self
|
|
|
|
.pool
|
|
|
|
.get()
|
|
|
|
.await
|
|
|
|
.map_err(|e| QueryBuilderError(e.into()))?;
|
2023-03-05 00:19:35 +00:00
|
|
|
let result = Task::remove(&mut connection, id).await?;
|
2023-03-04 18:07:17 +00:00
|
|
|
Ok(result)
|
|
|
|
}
|
|
|
|
|
2023-03-10 22:41:34 +00:00
|
|
|
pub(crate) async fn remove_all_tasks(&self) -> Result<u64, AsyncQueueError> {
|
2023-03-07 15:41:20 +00:00
|
|
|
let mut connection = self
|
|
|
|
.pool
|
|
|
|
.get()
|
|
|
|
.await
|
|
|
|
.map_err(|e| QueryBuilderError(e.into()))?;
|
|
|
|
Task::remove_all(&mut connection).await
|
2023-03-04 18:07:17 +00:00
|
|
|
}
|
|
|
|
|
2023-03-10 22:41:34 +00:00
|
|
|
pub(crate) async fn schedule_task_retry(
|
|
|
|
&self,
|
2023-03-07 15:41:20 +00:00
|
|
|
id: TaskId,
|
2023-03-04 18:07:17 +00:00
|
|
|
backoff_seconds: u32,
|
|
|
|
error: &str,
|
|
|
|
) -> Result<Task, AsyncQueueError> {
|
|
|
|
let mut connection = self
|
|
|
|
.pool
|
|
|
|
.get()
|
|
|
|
.await
|
|
|
|
.map_err(|e| QueryBuilderError(e.into()))?;
|
2023-03-07 15:41:20 +00:00
|
|
|
let task = Task::schedule_retry(&mut connection, id, backoff_seconds, error).await?;
|
2023-03-04 18:07:17 +00:00
|
|
|
Ok(task)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
#[cfg(test)]
|
|
|
|
mod async_queue_tests {
|
|
|
|
use super::*;
|
2023-03-10 22:41:34 +00:00
|
|
|
use crate::CurrentTask;
|
2023-03-04 18:07:17 +00:00
|
|
|
use async_trait::async_trait;
|
|
|
|
use diesel_async::pooled_connection::{bb8::Pool, AsyncDieselConnectionManager};
|
|
|
|
use diesel_async::AsyncPgConnection;
|
|
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
|
|
|
|
#[derive(Serialize, Deserialize)]
|
|
|
|
struct AsyncTask {
|
|
|
|
pub number: u16,
|
|
|
|
}
|
|
|
|
|
|
|
|
#[async_trait]
|
2023-03-10 22:41:34 +00:00
|
|
|
impl BackgroundTask for AsyncTask {
|
|
|
|
const TASK_NAME: &'static str = "AsyncUniqTask";
|
|
|
|
type AppData = ();
|
|
|
|
|
|
|
|
async fn run(&self, _task: CurrentTask, _: Self::AppData) -> Result<(), anyhow::Error> {
|
2023-03-04 18:07:17 +00:00
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
#[derive(Serialize, Deserialize)]
|
|
|
|
struct AsyncUniqTask {
|
|
|
|
pub number: u16,
|
|
|
|
}
|
|
|
|
|
|
|
|
#[async_trait]
|
2023-03-10 22:41:34 +00:00
|
|
|
impl BackgroundTask for AsyncUniqTask {
|
|
|
|
const TASK_NAME: &'static str = "AsyncUniqTask";
|
|
|
|
type AppData = ();
|
|
|
|
|
|
|
|
async fn run(&self, _task: CurrentTask, _: Self::AppData) -> Result<(), anyhow::Error> {
|
2023-03-04 18:07:17 +00:00
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
|
2023-03-07 15:41:20 +00:00
|
|
|
fn uniq(&self) -> Option<TaskHash> {
|
|
|
|
TaskHash::default_for_task(self).ok()
|
2023-03-04 18:07:17 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
#[derive(Serialize, Deserialize)]
|
|
|
|
struct AsyncTaskSchedule {
|
|
|
|
pub number: u16,
|
|
|
|
pub datetime: String,
|
|
|
|
}
|
|
|
|
|
|
|
|
#[async_trait]
|
2023-03-10 22:41:34 +00:00
|
|
|
impl BackgroundTask for AsyncTaskSchedule {
|
|
|
|
const TASK_NAME: &'static str = "AsyncUniqTask";
|
|
|
|
type AppData = ();
|
2023-03-04 18:07:17 +00:00
|
|
|
|
2023-03-10 22:41:34 +00:00
|
|
|
async fn run(&self, _task: CurrentTask, _: Self::AppData) -> Result<(), anyhow::Error> {
|
|
|
|
Ok(())
|
2023-03-04 18:07:17 +00:00
|
|
|
}
|
|
|
|
|
2023-03-10 22:41:34 +00:00
|
|
|
// fn cron(&self) -> Option<Scheduled> {
|
|
|
|
// let datetime = self.datetime.parse::<DateTime<Utc>>().ok()?;
|
|
|
|
// Some(Scheduled::ScheduleOnce(datetime))
|
|
|
|
// }
|
2023-03-04 18:07:17 +00:00
|
|
|
}
|
|
|
|
|
2023-03-07 15:41:20 +00:00
|
|
|
// #[tokio::test]
|
2023-03-10 22:41:34 +00:00
|
|
|
// async fn insert_task_creates_new_task() {
|
2023-03-07 15:41:20 +00:00
|
|
|
// let pool = pool().await;
|
2023-03-10 22:41:34 +00:00
|
|
|
// let mut queue = PgTaskStore::new(pool);
|
2023-03-07 15:41:20 +00:00
|
|
|
//
|
2023-03-10 22:41:34 +00:00
|
|
|
// let task = queue.create_task(AsyncTask { number: 1 }).await.unwrap();
|
2023-03-07 15:41:20 +00:00
|
|
|
//
|
2023-03-10 22:41:34 +00:00
|
|
|
// let metadata = task.payload.as_object().unwrap();
|
|
|
|
// let number = metadata["number"].as_u64();
|
|
|
|
// let type_task = metadata["type"].as_str();
|
|
|
|
//
|
|
|
|
// assert_eq!(Some(1), number);
|
|
|
|
// assert_eq!(Some("AsyncTask"), type_task);
|
2023-03-07 15:41:20 +00:00
|
|
|
//
|
2023-03-10 22:41:34 +00:00
|
|
|
// queue.remove_all_tasks().await.unwrap();
|
|
|
|
// }
|
|
|
|
//
|
|
|
|
// #[tokio::test]
|
|
|
|
// async fn update_task_state_test() {
|
|
|
|
// let pool = pool().await;
|
|
|
|
// let mut test = PgTaskStore::new(pool);
|
|
|
|
//
|
|
|
|
// let task = test.create_task(&AsyncTask { number: 1 }).await.unwrap();
|
2023-03-07 15:41:20 +00:00
|
|
|
//
|
|
|
|
// let metadata = task.payload.as_object().unwrap();
|
|
|
|
// let number = metadata["number"].as_u64();
|
|
|
|
// let type_task = metadata["type"].as_str();
|
2023-03-10 22:41:34 +00:00
|
|
|
// let id = task.id;
|
2023-03-07 15:41:20 +00:00
|
|
|
//
|
|
|
|
// assert_eq!(Some(1), number);
|
2023-03-10 22:41:34 +00:00
|
|
|
// assert_eq!(Some("AsyncTask"), type_task);
|
|
|
|
//
|
|
|
|
// let finished_task = test.set_task_state(task.id, TaskState::Done).await.unwrap();
|
|
|
|
//
|
|
|
|
// assert_eq!(id, finished_task.id);
|
|
|
|
// assert_eq!(TaskState::Done, finished_task.state());
|
2023-03-07 15:41:20 +00:00
|
|
|
//
|
|
|
|
// test.remove_all_tasks().await.unwrap();
|
|
|
|
// }
|
|
|
|
//
|
|
|
|
// #[tokio::test]
|
2023-03-10 22:41:34 +00:00
|
|
|
// async fn failed_task_query_test() {
|
2023-03-07 15:41:20 +00:00
|
|
|
// let pool = pool().await;
|
2023-03-10 22:41:34 +00:00
|
|
|
// let mut test = PgTaskStore::new(pool);
|
2023-03-07 15:41:20 +00:00
|
|
|
//
|
2023-03-10 22:41:34 +00:00
|
|
|
// let task = test.create_task(&AsyncTask { number: 1 }).await.unwrap();
|
2023-03-07 15:41:20 +00:00
|
|
|
//
|
2023-03-10 22:41:34 +00:00
|
|
|
// let metadata = task.payload.as_object().unwrap();
|
|
|
|
// let number = metadata["number"].as_u64();
|
|
|
|
// let type_task = metadata["type"].as_str();
|
|
|
|
// let id = task.id;
|
2023-03-07 15:41:20 +00:00
|
|
|
//
|
2023-03-10 22:41:34 +00:00
|
|
|
// assert_eq!(Some(1), number);
|
|
|
|
// assert_eq!(Some("AsyncTask"), type_task);
|
2023-03-07 15:41:20 +00:00
|
|
|
//
|
2023-03-10 22:41:34 +00:00
|
|
|
// let failed_task = test.set_task_state(task.id, TaskState::Failed("Some error".to_string())).await.unwrap();
|
2023-03-07 15:41:20 +00:00
|
|
|
//
|
2023-03-10 22:41:34 +00:00
|
|
|
// assert_eq!(id, failed_task.id);
|
|
|
|
// assert_eq!(Some("Some error"), failed_task.error_message.as_deref());
|
|
|
|
// assert_eq!(TaskState::Failed, failed_task.state());
|
2023-03-07 15:41:20 +00:00
|
|
|
//
|
|
|
|
// test.remove_all_tasks().await.unwrap();
|
|
|
|
// }
|
2023-03-10 22:41:34 +00:00
|
|
|
//
|
|
|
|
// #[tokio::test]
|
|
|
|
// async fn remove_all_tasks_test() {
|
|
|
|
// let pool = pool().await;
|
|
|
|
// let mut test = PgTaskStore::new(pool);
|
|
|
|
//
|
|
|
|
// let task = test.create_task(&AsyncTask { number: 1 }).await.unwrap();
|
|
|
|
//
|
|
|
|
// let metadata = task.payload.as_object().unwrap();
|
|
|
|
// let number = metadata["number"].as_u64();
|
|
|
|
// let type_task = metadata["type"].as_str();
|
|
|
|
//
|
|
|
|
// assert_eq!(Some(1), number);
|
|
|
|
// assert_eq!(Some("AsyncTask"), type_task);
|
|
|
|
//
|
|
|
|
// let task = test.create_task(&AsyncTask { number: 2 }).await.unwrap();
|
|
|
|
//
|
|
|
|
// let metadata = task.payload.as_object().unwrap();
|
|
|
|
// let number = metadata["number"].as_u64();
|
|
|
|
// let type_task = metadata["type"].as_str();
|
|
|
|
//
|
|
|
|
// assert_eq!(Some(2), number);
|
|
|
|
// assert_eq!(Some("AsyncTask"), type_task);
|
|
|
|
//
|
|
|
|
// let result = test.remove_all_tasks().await.unwrap();
|
|
|
|
// assert_eq!(2, result);
|
|
|
|
// }
|
|
|
|
//
|
|
|
|
// #[tokio::test]
|
|
|
|
// async fn pull_next_task_test() {
|
|
|
|
// let pool = pool().await;
|
|
|
|
// let mut queue = PgTaskStore::new(pool);
|
|
|
|
//
|
|
|
|
// let task = queue.create_task(&AsyncTask { number: 1 }).await.unwrap();
|
|
|
|
//
|
|
|
|
// let metadata = task.payload.as_object().unwrap();
|
|
|
|
// let number = metadata["number"].as_u64();
|
|
|
|
// let type_task = metadata["type"].as_str();
|
|
|
|
//
|
|
|
|
// assert_eq!(Some(1), number);
|
|
|
|
// assert_eq!(Some("AsyncTask"), type_task);
|
|
|
|
//
|
|
|
|
// let task = queue.create_task(&AsyncTask { number: 2 }).await.unwrap();
|
|
|
|
//
|
|
|
|
// let metadata = task.payload.as_object().unwrap();
|
|
|
|
// let number = metadata["number"].as_u64();
|
|
|
|
// let type_task = metadata["type"].as_str();
|
|
|
|
//
|
|
|
|
// assert_eq!(Some(2), number);
|
|
|
|
// assert_eq!(Some("AsyncTask"), type_task);
|
|
|
|
//
|
|
|
|
// let task = queue.pull_next_task(None).await.unwrap().unwrap();
|
|
|
|
//
|
|
|
|
// let metadata = task.payload.as_object().unwrap();
|
|
|
|
// let number = metadata["number"].as_u64();
|
|
|
|
// let type_task = metadata["type"].as_str();
|
|
|
|
//
|
|
|
|
// assert_eq!(Some(1), number);
|
|
|
|
// assert_eq!(Some("AsyncTask"), type_task);
|
|
|
|
//
|
|
|
|
// let task = queue.pull_next_task(None).await.unwrap().unwrap();
|
|
|
|
// let metadata = task.payload.as_object().unwrap();
|
|
|
|
// let number = metadata["number"].as_u64();
|
|
|
|
// let type_task = metadata["type"].as_str();
|
|
|
|
//
|
|
|
|
// assert_eq!(Some(2), number);
|
|
|
|
// assert_eq!(Some("AsyncTask"), type_task);
|
|
|
|
//
|
|
|
|
// queue.remove_all_tasks().await.unwrap();
|
|
|
|
// }
|
2023-03-04 18:07:17 +00:00
|
|
|
|
|
|
|
async fn pool() -> Pool<AsyncPgConnection> {
|
|
|
|
let manager = AsyncDieselConnectionManager::<AsyncPgConnection>::new(
|
2023-03-09 15:59:45 +00:00
|
|
|
option_env!("DATABASE_URL").expect("DATABASE_URL must be set"),
|
2023-03-04 18:07:17 +00:00
|
|
|
);
|
|
|
|
Pool::builder()
|
|
|
|
.max_size(1)
|
|
|
|
.min_idle(Some(1))
|
|
|
|
.build(manager)
|
|
|
|
.await
|
|
|
|
.unwrap()
|
|
|
|
}
|
|
|
|
}
|