diff --git a/.env b/.env index 6d0480c..b725b6b 100644 --- a/.env +++ b/.env @@ -1 +1 @@ -DATABASE_URL=postgres://postgres:postgres@localhost/fang +DATABASE_URL=postgres://postgres:password@localhost/fang diff --git a/Cargo.toml b/Cargo.toml index 05f92b3..ab10a16 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,24 +1,21 @@ [package] name = "fang" version = "0.10.2" -authors = ["Ayrat Badykov " , "Pepe Márquez "] +authors = [ + "Ayrat Badykov ", + "Pepe Márquez ", + "Rafael Caricio " +] description = "Background job processing library for Rust" -repository = "https://github.com/ayrat555/fang" +repository = "https://github.com/rafaelcaricio/fang" edition = "2021" license = "MIT" readme = "README.md" -rust-version = "1.62" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html +rust-version = "1.67" [lib] doctest = false -[features] -default = ["blocking", "asynk"] -blocking = ["diesel", "diesel-derive-enum", "dotenvy"] -asynk = ["bb8-postgres", "postgres-types", "tokio", "async-trait", "async-recursion"] - [dependencies] cron = "0.12" chrono = "0.4" @@ -36,36 +33,21 @@ uuid = { version = "1.1", features = ["v4"] } [dependencies.diesel] version = "2.0" features = ["postgres", "serde_json", "chrono", "uuid", "r2d2"] -optional = true [dependencies.diesel-derive-enum] version = "2.0.1" features = ["postgres"] -optional = true -[dependencies.dotenvy] -version = "0.15" -optional = true - -[dependencies.bb8-postgres] -version = "0.8" -features = ["with-serde_json-1" , "with-uuid-1" , "with-chrono-0_4"] -optional = true - -[dependencies.postgres-types] -version = "0.X.X" -features = ["derive"] -optional = true +[dependencies.diesel-async] +version = "0.2" +features = ["postgres", "bb8"] [dependencies.tokio] version = "1.25" features = ["rt", "time", "macros"] -optional = true [dependencies.async-trait] version = "0.1" -optional = true [dependencies.async-recursion] version = "1" -optional = true diff --git a/README.md b/README.md index a12f7ed..bf3baa6 100644 --- a/README.md +++ b/README.md @@ -45,7 +45,7 @@ fang = { version = "0.10" , features = ["asynk"], default-features = false } fang = { version = "0.10" } ``` -*Supports rustc 1.62+* +*Supports rustc 1.67+* 2. Create the `fang_tasks` table in the Postgres database. The migration can be found in [the migrations directory](https://github.com/ayrat555/fang/blob/master/migrations/2022-08-20-151615_create_fang_tasks/up.sql). diff --git a/diesel.toml b/diesel.toml index 73a18e3..f57985a 100644 --- a/diesel.toml +++ b/diesel.toml @@ -1,2 +1,2 @@ [print_schema] -file = "src/blocking/schema.rs" \ No newline at end of file +file = "src/schema.rs" diff --git a/fang_examples/asynk/simple_async_worker/Cargo.toml b/fang_examples/asynk/simple_async_worker/Cargo.toml index 8248aed..60c7dbf 100644 --- a/fang_examples/asynk/simple_async_worker/Cargo.toml +++ b/fang_examples/asynk/simple_async_worker/Cargo.toml @@ -3,10 +3,10 @@ name = "simple_async_worker" version = "0.1.0" edition = "2021" -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - [dependencies] -fang = { path = "../../../" , features = ["asynk"]} +fang = { path = "../../../" } env_logger = "0.9.0" log = "0.4.0" tokio = { version = "1", features = ["full"] } +diesel-async = { version = "0.2", features = ["postgres", "bb8"] } +diesel = { version = "2.0", features = ["postgres"] } diff --git a/fang_examples/asynk/simple_async_worker/src/lib.rs b/fang_examples/asynk/simple_async_worker/src/lib.rs index cfd269b..6e6ddd7 100644 --- a/fang_examples/asynk/simple_async_worker/src/lib.rs +++ b/fang_examples/asynk/simple_async_worker/src/lib.rs @@ -1,8 +1,8 @@ use fang::async_trait; -use fang::asynk::async_queue::AsyncQueueable; +use fang::queue::AsyncQueueable; use fang::serde::{Deserialize, Serialize}; use fang::typetag; -use fang::AsyncRunnable; +use fang::runnable::AsyncRunnable; use fang::FangError; use std::time::Duration; @@ -34,11 +34,11 @@ impl MyFailingTask { #[typetag::serde] impl AsyncRunnable for MyTask { async fn run(&self, queue: &mut dyn AsyncQueueable) -> Result<(), FangError> { - let new_task = MyTask::new(self.number + 1); - queue - .insert_task(&new_task as &dyn AsyncRunnable) - .await - .unwrap(); + // let new_task = MyTask::new(self.number + 1); + // queue + // .insert_task(&new_task as &dyn AsyncRunnable) + // .await + // .unwrap(); log::info!("the current number is {}", self.number); tokio::time::sleep(Duration::from_secs(3)).await; @@ -51,21 +51,24 @@ impl AsyncRunnable for MyTask { #[typetag::serde] impl AsyncRunnable for MyFailingTask { async fn run(&self, queue: &mut dyn AsyncQueueable) -> Result<(), FangError> { - let new_task = MyFailingTask::new(self.number + 1); - queue - .insert_task(&new_task as &dyn AsyncRunnable) - .await - .unwrap(); + // let new_task = MyFailingTask::new(self.number + 1); + // queue + // .insert_task(&new_task as &dyn AsyncRunnable) + // .await + // .unwrap(); log::info!("the current number is {}", self.number); tokio::time::sleep(Duration::from_secs(3)).await; - let b = true; - - if b { - panic!("Hello!"); - } else { - Ok(()) - } + log::info!("done.."); + // + // let b = true; + // + // if b { + // panic!("Hello!"); + // } else { + // Ok(()) + // } + Ok(()) } } diff --git a/fang_examples/asynk/simple_async_worker/src/main.rs b/fang_examples/asynk/simple_async_worker/src/main.rs index cbe9446..4f90f53 100644 --- a/fang_examples/asynk/simple_async_worker/src/main.rs +++ b/fang_examples/asynk/simple_async_worker/src/main.rs @@ -1,34 +1,44 @@ -use fang::asynk::async_queue::AsyncQueue; -use fang::asynk::async_queue::AsyncQueueable; -use fang::asynk::async_worker_pool::AsyncWorkerPool; -use fang::AsyncRunnable; -use fang::NoTls; +use fang::queue::AsyncQueue; +use fang::queue::AsyncQueueable; +use fang::worker_pool::AsyncWorkerPool; +use fang::runnable::AsyncRunnable; use simple_async_worker::MyFailingTask; use simple_async_worker::MyTask; use std::time::Duration; +use diesel_async::pg::AsyncPgConnection; +use diesel_async::pooled_connection::{bb8::Pool, AsyncDieselConnectionManager}; +use diesel::PgConnection; #[tokio::main] async fn main() { env_logger::init(); + let connection_url = "postgres://postgres:password@localhost/fang"; + log::info!("Starting..."); let max_pool_size: u32 = 3; + let manager = AsyncDieselConnectionManager::::new(connection_url); + let pool = Pool::builder() + .max_size(max_pool_size) + .min_idle(Some(1)) + .build(manager) + .await + .unwrap(); + let mut queue = AsyncQueue::builder() - .uri("postgres://postgres:postgres@localhost/fang") - .max_pool_size(max_pool_size) + .pool(pool) .build(); - queue.connect(NoTls).await.unwrap(); log::info!("Queue connected..."); - let mut pool: AsyncWorkerPool> = AsyncWorkerPool::builder() + let mut workers_pool: AsyncWorkerPool = AsyncWorkerPool::builder() .number_of_workers(10_u32) .queue(queue.clone()) .build(); log::info!("Pool created ..."); - pool.start().await; + workers_pool.start().await; log::info!("Workers started ..."); let task1 = MyTask::new(0); diff --git a/fang_examples/blocking/simple_cron_worker/.gitignore b/fang_examples/blocking/simple_cron_worker/.gitignore deleted file mode 100644 index 96ef6c0..0000000 --- a/fang_examples/blocking/simple_cron_worker/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -/target -Cargo.lock diff --git a/fang_examples/blocking/simple_cron_worker/Cargo.toml b/fang_examples/blocking/simple_cron_worker/Cargo.toml deleted file mode 100644 index b29f1c7..0000000 --- a/fang_examples/blocking/simple_cron_worker/Cargo.toml +++ /dev/null @@ -1,13 +0,0 @@ -[package] -name = "simple_cron_worker" -version = "0.1.0" -edition = "2021" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - -[dependencies] -fang = { path = "../../../" , features = ["blocking"]} -dotenv = "0.15.0" -env_logger = "0.9.0" -log = "0.4.0" -diesel = { version = "2", features = ["postgres", "r2d2"] } diff --git a/fang_examples/blocking/simple_cron_worker/README.md b/fang_examples/blocking/simple_cron_worker/README.md deleted file mode 100644 index f1c3f1b..0000000 --- a/fang_examples/blocking/simple_cron_worker/README.md +++ /dev/null @@ -1,3 +0,0 @@ -## Simple example - -The job described in this example enqueues a new job during its execution saving thread name of the current worker to its metadata. diff --git a/fang_examples/blocking/simple_cron_worker/src/lib.rs b/fang_examples/blocking/simple_cron_worker/src/lib.rs deleted file mode 100644 index ab44341..0000000 --- a/fang_examples/blocking/simple_cron_worker/src/lib.rs +++ /dev/null @@ -1,35 +0,0 @@ -use fang::runnable::Runnable; -use fang::serde::{Deserialize, Serialize}; -use fang::typetag; -use fang::FangError; -use fang::Queueable; -use fang::Scheduled; - -#[derive(Serialize, Deserialize)] -#[serde(crate = "fang::serde")] -pub struct MyCronTask {} - -#[typetag::serde] -impl Runnable for MyCronTask { - fn run(&self, _queue: &dyn Queueable) -> Result<(), FangError> { - log::info!("CRON !!!!!!!!!!!!!!!!!"); - - Ok(()) - } - - fn task_type(&self) -> String { - "cron_test".to_string() - } - - fn cron(&self) -> Option { - // sec min hour day of month month day of week year - // be careful works only with UTC hour. - // https://www.timeanddate.com/worldclock/timezone/utc - let expression = "0/20 * * * Aug-Sep * 2022/1"; - Some(Scheduled::CronPattern(expression.to_string())) - } - - fn uniq(&self) -> bool { - true - } -} diff --git a/fang_examples/blocking/simple_cron_worker/src/main.rs b/fang_examples/blocking/simple_cron_worker/src/main.rs deleted file mode 100644 index 76380f0..0000000 --- a/fang_examples/blocking/simple_cron_worker/src/main.rs +++ /dev/null @@ -1,43 +0,0 @@ -use diesel::r2d2; -use dotenv::dotenv; -use fang::PgConnection; -use fang::Queue; -use fang::Queueable; -use fang::RetentionMode; -use fang::WorkerPool; -use simple_cron_worker::MyCronTask; -use std::env; -use std::thread::sleep; -use std::time::Duration; - -pub fn connection_pool(pool_size: u32) -> r2d2::Pool> { - let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set"); - - let manager = r2d2::ConnectionManager::::new(database_url); - - r2d2::Pool::builder() - .max_size(pool_size) - .build(manager) - .unwrap() -} - -fn main() { - dotenv().ok(); - - env_logger::init(); - - let queue = Queue::builder().connection_pool(connection_pool(2)).build(); - - let mut worker_pool = WorkerPool::::builder() - .queue(queue) - .retention_mode(RetentionMode::KeepAll) - .number_of_workers(2_u32) - .task_type("cron_test".to_string()) - .build(); - - worker_pool.queue.schedule_task(&MyCronTask {}).unwrap(); - - worker_pool.start().unwrap(); - - sleep(Duration::from_secs(100)) -} diff --git a/fang_examples/blocking/simple_worker/.gitignore b/fang_examples/blocking/simple_worker/.gitignore deleted file mode 100644 index 96ef6c0..0000000 --- a/fang_examples/blocking/simple_worker/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -/target -Cargo.lock diff --git a/fang_examples/blocking/simple_worker/Cargo.toml b/fang_examples/blocking/simple_worker/Cargo.toml deleted file mode 100644 index 2bd62b0..0000000 --- a/fang_examples/blocking/simple_worker/Cargo.toml +++ /dev/null @@ -1,13 +0,0 @@ -[package] -name = "simple_worker" -version = "0.1.0" -edition = "2021" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - -[dependencies] -fang = { path = "../../../" , features = ["blocking"]} -dotenv = "0.15.0" -env_logger = "0.9.0" -log = "0.4.0" -diesel = { version = "2", features = ["postgres", "r2d2"] } diff --git a/fang_examples/blocking/simple_worker/README.md b/fang_examples/blocking/simple_worker/README.md deleted file mode 100644 index f1c3f1b..0000000 --- a/fang_examples/blocking/simple_worker/README.md +++ /dev/null @@ -1,3 +0,0 @@ -## Simple example - -The job described in this example enqueues a new job during its execution saving thread name of the current worker to its metadata. diff --git a/fang_examples/blocking/simple_worker/src/lib.rs b/fang_examples/blocking/simple_worker/src/lib.rs deleted file mode 100644 index 4129591..0000000 --- a/fang_examples/blocking/simple_worker/src/lib.rs +++ /dev/null @@ -1,96 +0,0 @@ -use fang::runnable::Runnable; -use fang::serde::{Deserialize, Serialize}; -use fang::typetag; -use fang::FangError; -use fang::Queueable; -use std::thread; -use std::time::Duration; - -#[derive(Serialize, Deserialize)] -#[serde(crate = "fang::serde")] -pub struct MyTask { - pub number: u16, - pub current_thread_name: String, -} - -impl MyTask { - pub fn new(number: u16) -> Self { - let handle = thread::current(); - let current_thread_name = handle.name().unwrap().to_string(); - - Self { - number, - current_thread_name, - } - } -} - -#[typetag::serde] -impl Runnable for MyTask { - fn run(&self, queue: &dyn Queueable) -> Result<(), FangError> { - let new_task = MyTask::new(self.number + 1); - - log::info!( - "The number is {}, thread name {}", - self.number, - self.current_thread_name - ); - - queue.insert_task(&new_task).unwrap(); - - thread::sleep(Duration::from_secs(2)); - - Ok(()) - } - - fn task_type(&self) -> String { - "worker_pool_test".to_string() - } -} - -#[derive(Serialize, Deserialize)] -#[serde(crate = "fang::serde")] -pub struct MyFailingTask { - pub number: u16, - pub current_thread_name: String, -} - -impl MyFailingTask { - pub fn new(number: u16) -> Self { - let handle = thread::current(); - let current_thread_name = handle.name().unwrap().to_string(); - Self { - number, - current_thread_name, - } - } -} - -#[typetag::serde] -impl Runnable for MyFailingTask { - fn run(&self, queue: &dyn Queueable) -> Result<(), FangError> { - let new_task = MyFailingTask::new(self.number + 1); - - queue.insert_task(&new_task).unwrap(); - - log::info!( - "Failing task number {}, Thread name:{}", - self.number, - self.current_thread_name - ); - - thread::sleep(Duration::from_secs(3)); - - let b = true; - - if b { - panic!("Hello!"); - } else { - Ok(()) - } - } - - fn task_type(&self) -> String { - "worker_pool_test".to_string() - } -} diff --git a/fang_examples/blocking/simple_worker/src/main.rs b/fang_examples/blocking/simple_worker/src/main.rs deleted file mode 100644 index 6df20f2..0000000 --- a/fang_examples/blocking/simple_worker/src/main.rs +++ /dev/null @@ -1,50 +0,0 @@ -use diesel::r2d2; -use dotenvy::dotenv; -use fang::PgConnection; -use fang::Queue; -use fang::Queueable; -use fang::RetentionMode; -use fang::WorkerPool; -use simple_worker::MyFailingTask; -use simple_worker::MyTask; -use std::env; -use std::thread::sleep; -use std::time::Duration; - -pub fn connection_pool(pool_size: u32) -> r2d2::Pool> { - let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set"); - - let manager = r2d2::ConnectionManager::::new(database_url); - - r2d2::Pool::builder() - .max_size(pool_size) - .build(manager) - .unwrap() -} - -fn main() { - dotenv().ok(); - - env_logger::init(); - - let queue = Queue::builder().connection_pool(connection_pool(3)).build(); - - let mut worker_pool = WorkerPool::::builder() - .queue(queue) - .retention_mode(RetentionMode::KeepAll) - .number_of_workers(3_u32) - .task_type("worker_pool_test".to_string()) - .build(); - - worker_pool.queue.insert_task(&MyTask::new(1)).unwrap(); - worker_pool.queue.insert_task(&MyTask::new(1000)).unwrap(); - - worker_pool - .queue - .insert_task(&MyFailingTask::new(5000)) - .unwrap(); - - worker_pool.start().unwrap(); - - sleep(Duration::from_secs(100)) -} diff --git a/src/asynk.rs b/src/asynk.rs deleted file mode 100644 index a75dd03..0000000 --- a/src/asynk.rs +++ /dev/null @@ -1,9 +0,0 @@ -pub mod async_queue; -pub mod async_runnable; -pub mod async_worker; -pub mod async_worker_pool; - -pub use async_queue::*; -pub use async_runnable::AsyncRunnable; -pub use async_worker::*; -pub use async_worker_pool::*; diff --git a/src/asynk/async_queue.rs b/src/asynk/async_queue.rs deleted file mode 100644 index 877bff1..0000000 --- a/src/asynk/async_queue.rs +++ /dev/null @@ -1,1250 +0,0 @@ -use crate::asynk::async_runnable::AsyncRunnable; -use crate::CronError; -use crate::Scheduled::*; -use async_trait::async_trait; -use bb8_postgres::bb8::Pool; -use bb8_postgres::bb8::RunError; -use bb8_postgres::tokio_postgres::row::Row; -use bb8_postgres::tokio_postgres::tls::{MakeTlsConnect, TlsConnect}; -use bb8_postgres::tokio_postgres::Socket; -use bb8_postgres::tokio_postgres::Transaction; -use bb8_postgres::PostgresConnectionManager; -use chrono::DateTime; -use chrono::Duration; -use chrono::Utc; -use cron::Schedule; -use postgres_types::{FromSql, ToSql}; -use sha2::{Digest, Sha256}; -use std::str::FromStr; -use thiserror::Error; -use typed_builder::TypedBuilder; -use uuid::Uuid; - -#[cfg(test)] -use bb8_postgres::tokio_postgres::tls::NoTls; - -const INSERT_TASK_QUERY: &str = include_str!("queries/insert_task.sql"); -const INSERT_TASK_UNIQ_QUERY: &str = include_str!("queries/insert_task_uniq.sql"); -const UPDATE_TASK_STATE_QUERY: &str = include_str!("queries/update_task_state.sql"); -const FAIL_TASK_QUERY: &str = include_str!("queries/fail_task.sql"); -const REMOVE_ALL_TASK_QUERY: &str = include_str!("queries/remove_all_tasks.sql"); -const REMOVE_ALL_SCHEDULED_TASK_QUERY: &str = - include_str!("queries/remove_all_scheduled_tasks.sql"); -const REMOVE_TASK_QUERY: &str = include_str!("queries/remove_task.sql"); -const REMOVE_TASK_BY_METADATA_QUERY: &str = include_str!("queries/remove_task_by_metadata.sql"); -const REMOVE_TASKS_TYPE_QUERY: &str = include_str!("queries/remove_tasks_type.sql"); -const FETCH_TASK_TYPE_QUERY: &str = include_str!("queries/fetch_task_type.sql"); -const FIND_TASK_BY_UNIQ_HASH_QUERY: &str = include_str!("queries/find_task_by_uniq_hash.sql"); -const FIND_TASK_BY_ID_QUERY: &str = include_str!("queries/find_task_by_id.sql"); -const RETRY_TASK_QUERY: &str = include_str!("queries/retry_task.sql"); - -pub const DEFAULT_TASK_TYPE: &str = "common"; - -#[derive(Debug, Eq, PartialEq, Clone, ToSql, FromSql)] -#[postgres(name = "fang_task_state")] -pub enum FangTaskState { - #[postgres(name = "new")] - New, - #[postgres(name = "in_progress")] - InProgress, - #[postgres(name = "failed")] - Failed, - #[postgres(name = "finished")] - Finished, - #[postgres(name = "retried")] - Retried, -} - -impl Default for FangTaskState { - fn default() -> Self { - FangTaskState::New - } -} - -#[derive(TypedBuilder, Debug, Eq, PartialEq, Clone)] -pub struct Task { - #[builder(setter(into))] - pub id: Uuid, - #[builder(setter(into))] - pub metadata: serde_json::Value, - #[builder(setter(into))] - pub error_message: Option, - #[builder(default, setter(into))] - pub state: FangTaskState, - #[builder(setter(into))] - pub task_type: String, - #[builder(setter(into))] - pub uniq_hash: Option, - #[builder(setter(into))] - pub retries: i32, - #[builder(setter(into))] - pub scheduled_at: DateTime, - #[builder(setter(into))] - pub created_at: DateTime, - #[builder(setter(into))] - pub updated_at: DateTime, -} - -#[derive(Debug, Error)] -pub enum AsyncQueueError { - #[error(transparent)] - PoolError(#[from] RunError), - #[error(transparent)] - PgError(#[from] bb8_postgres::tokio_postgres::Error), - #[error(transparent)] - SerdeError(#[from] serde_json::Error), - #[error(transparent)] - CronError(#[from] CronError), - #[error("returned invalid result (expected {expected:?}, found {found:?})")] - ResultError { expected: u64, found: u64 }, - #[error( - "AsyncQueue is not connected :( , call connect() method first and then perform operations" - )] - NotConnectedError, - #[error("Can not convert `std::time::Duration` to `chrono::Duration`")] - TimeError, - #[error("Can not perform this operation if task is not uniq, please check its definition in impl AsyncRunnable")] - TaskNotUniqError, -} - -impl From for AsyncQueueError { - fn from(error: cron::error::Error) -> Self { - AsyncQueueError::CronError(CronError::LibraryError(error)) - } -} - -/// This trait defines operations for an asynchronous queue. -/// The trait can be implemented for different storage backends. -/// For now, the trait is only implemented for PostgreSQL. More backends are planned to be implemented in the future. - -#[async_trait] -pub trait AsyncQueueable: Send { - /// This method should retrieve one task of the `task_type` type. If `task_type` is `None` it will try to - /// fetch a task of the type `common`. After fetching it should update the state of the task to - /// `FangTaskState::InProgress`. - /// - async fn fetch_and_touch_task( - &mut self, - task_type: Option, - ) -> Result, AsyncQueueError>; - - /// Enqueue a task to the queue, The task will be executed as soon as possible by the worker of the same type - /// created by an AsyncWorkerPool. - async fn insert_task(&mut self, task: &dyn AsyncRunnable) -> Result; - - /// The method will remove all tasks from the queue - async fn remove_all_tasks(&mut self) -> Result; - - /// Remove all tasks that are scheduled in the future. - async fn remove_all_scheduled_tasks(&mut self) -> Result; - - /// Remove a task by its id. - async fn remove_task(&mut self, id: Uuid) -> Result; - - /// Remove a task by its metadata (struct fields values) - async fn remove_task_by_metadata( - &mut self, - task: &dyn AsyncRunnable, - ) -> Result; - - /// Removes all tasks that have the specified `task_type`. - async fn remove_tasks_type(&mut self, task_type: &str) -> Result; - - /// Retrieve a task from storage by its `id`. - async fn find_task_by_id(&mut self, id: Uuid) -> Result; - - /// Update the state field of the specified task - /// See the `FangTaskState` enum for possible states. - async fn update_task_state( - &mut self, - task: Task, - state: FangTaskState, - ) -> Result; - - /// Update the state of a task to `FangTaskState::Failed` and set an error_message. - async fn fail_task(&mut self, task: Task, error_message: &str) - -> Result; - - /// Schedule a task. - async fn schedule_task(&mut self, task: &dyn AsyncRunnable) -> Result; - - async fn schedule_retry( - &mut self, - task: &Task, - backoff_seconds: u32, - error: &str, - ) -> Result; -} - -/// An async queue that can be used to enqueue tasks. -/// It uses a PostgreSQL storage. It must be connected to perform any operation. -/// To connect an `AsyncQueue` to PostgreSQL database call the `connect` method. -/// A Queue can be created with the TypedBuilder. -/// -/// ```rust -/// let mut queue = AsyncQueue::builder() -/// .uri("postgres://postgres:postgres@localhost/fang") -/// .max_pool_size(max_pool_size) -/// .build(); -/// ``` -/// - -#[derive(TypedBuilder, Debug, Clone)] -pub struct AsyncQueue -where - Tls: MakeTlsConnect + Clone + Send + Sync + 'static, - >::Stream: Send + Sync, - >::TlsConnect: Send, - <>::TlsConnect as TlsConnect>::Future: Send, -{ - #[builder(default=None, setter(skip))] - pool: Option>>, - #[builder(setter(into))] - uri: String, - #[builder(setter(into))] - max_pool_size: u32, - #[builder(default = false, setter(skip))] - connected: bool, -} - -#[cfg(test)] -#[derive(TypedBuilder)] -pub struct AsyncQueueTest<'a> { - #[builder(setter(into))] - pub transaction: Transaction<'a>, -} - -#[cfg(test)] -#[async_trait] -impl AsyncQueueable for AsyncQueueTest<'_> { - async fn find_task_by_id(&mut self, id: Uuid) -> Result { - let transaction = &mut self.transaction; - - AsyncQueue::::find_task_by_id_query(transaction, id).await - } - - async fn fetch_and_touch_task( - &mut self, - task_type: Option, - ) -> Result, AsyncQueueError> { - let transaction = &mut self.transaction; - - AsyncQueue::::fetch_and_touch_task_query(transaction, task_type).await - } - - async fn insert_task(&mut self, task: &dyn AsyncRunnable) -> Result { - let transaction = &mut self.transaction; - - let metadata = serde_json::to_value(task)?; - - let task: Task = if !task.uniq() { - AsyncQueue::::insert_task_query( - transaction, - metadata, - &task.task_type(), - Utc::now(), - ) - .await? - } else { - AsyncQueue::::insert_task_if_not_exist_query( - transaction, - metadata, - &task.task_type(), - Utc::now(), - ) - .await? - }; - Ok(task) - } - - async fn schedule_task(&mut self, task: &dyn AsyncRunnable) -> Result { - let transaction = &mut self.transaction; - - let metadata = serde_json::to_value(task)?; - - let scheduled_at = match task.cron() { - Some(scheduled) => match scheduled { - CronPattern(cron_pattern) => { - let schedule = Schedule::from_str(&cron_pattern)?; - let mut iterator = schedule.upcoming(Utc); - - iterator - .next() - .ok_or(AsyncQueueError::CronError(CronError::NoTimestampsError))? - } - ScheduleOnce(datetime) => datetime, - }, - None => { - return Err(AsyncQueueError::CronError( - CronError::TaskNotSchedulableError, - )); - } - }; - - let task: Task = if !task.uniq() { - AsyncQueue::::insert_task_query( - transaction, - metadata, - &task.task_type(), - scheduled_at, - ) - .await? - } else { - AsyncQueue::::insert_task_if_not_exist_query( - transaction, - metadata, - &task.task_type(), - scheduled_at, - ) - .await? - }; - - Ok(task) - } - async fn remove_all_tasks(&mut self) -> Result { - let transaction = &mut self.transaction; - - AsyncQueue::::remove_all_tasks_query(transaction).await - } - - async fn remove_all_scheduled_tasks(&mut self) -> Result { - let transaction = &mut self.transaction; - - AsyncQueue::::remove_all_scheduled_tasks_query(transaction).await - } - - async fn remove_task(&mut self, id: Uuid) -> Result { - let transaction = &mut self.transaction; - - AsyncQueue::::remove_task_query(transaction, id).await - } - - async fn remove_task_by_metadata( - &mut self, - task: &dyn AsyncRunnable, - ) -> Result { - if task.uniq() { - let transaction = &mut self.transaction; - - AsyncQueue::::remove_task_by_metadata_query(transaction, task).await - } else { - Err(AsyncQueueError::TaskNotUniqError) - } - } - - async fn remove_tasks_type(&mut self, task_type: &str) -> Result { - let transaction = &mut self.transaction; - - AsyncQueue::::remove_tasks_type_query(transaction, task_type).await - } - - async fn update_task_state( - &mut self, - task: Task, - state: FangTaskState, - ) -> Result { - let transaction = &mut self.transaction; - - AsyncQueue::::update_task_state_query(transaction, task, state).await - } - - async fn fail_task( - &mut self, - task: Task, - error_message: &str, - ) -> Result { - let transaction = &mut self.transaction; - - AsyncQueue::::fail_task_query(transaction, task, error_message).await - } - - async fn schedule_retry( - &mut self, - task: &Task, - backoff_seconds: u32, - error: &str, - ) -> Result { - let transaction = &mut self.transaction; - - AsyncQueue::::schedule_retry_query(transaction, task, backoff_seconds, error).await - } -} - -impl AsyncQueue -where - Tls: MakeTlsConnect + Clone + Send + Sync + 'static, - >::Stream: Send + Sync, - >::TlsConnect: Send, - <>::TlsConnect as TlsConnect>::Future: Send, -{ - /// Check if the connection with db is established - pub fn check_if_connection(&self) -> Result<(), AsyncQueueError> { - if self.connected { - Ok(()) - } else { - Err(AsyncQueueError::NotConnectedError) - } - } - - /// Connect to the db if not connected - pub async fn connect(&mut self, tls: Tls) -> Result<(), AsyncQueueError> { - let manager = PostgresConnectionManager::new_from_stringlike(self.uri.clone(), tls)?; - - let pool = Pool::builder() - .max_size(self.max_pool_size) - .build(manager) - .await?; - - self.pool = Some(pool); - self.connected = true; - Ok(()) - } - - async fn remove_all_tasks_query( - transaction: &mut Transaction<'_>, - ) -> Result { - Self::execute_query(transaction, REMOVE_ALL_TASK_QUERY, &[], None).await - } - - async fn remove_all_scheduled_tasks_query( - transaction: &mut Transaction<'_>, - ) -> Result { - Self::execute_query( - transaction, - REMOVE_ALL_SCHEDULED_TASK_QUERY, - &[&Utc::now()], - None, - ) - .await - } - - async fn remove_task_query( - transaction: &mut Transaction<'_>, - id: Uuid, - ) -> Result { - Self::execute_query(transaction, REMOVE_TASK_QUERY, &[&id], Some(1)).await - } - - async fn remove_task_by_metadata_query( - transaction: &mut Transaction<'_>, - task: &dyn AsyncRunnable, - ) -> Result { - let metadata = serde_json::to_value(task)?; - - let uniq_hash = Self::calculate_hash(metadata.to_string()); - - Self::execute_query( - transaction, - REMOVE_TASK_BY_METADATA_QUERY, - &[&uniq_hash], - None, - ) - .await - } - - async fn remove_tasks_type_query( - transaction: &mut Transaction<'_>, - task_type: &str, - ) -> Result { - Self::execute_query(transaction, REMOVE_TASKS_TYPE_QUERY, &[&task_type], None).await - } - - async fn find_task_by_id_query( - transaction: &mut Transaction<'_>, - id: Uuid, - ) -> Result { - let row: Row = transaction.query_one(FIND_TASK_BY_ID_QUERY, &[&id]).await?; - - let task = Self::row_to_task(row); - Ok(task) - } - - async fn fail_task_query( - transaction: &mut Transaction<'_>, - task: Task, - error_message: &str, - ) -> Result { - let updated_at = Utc::now(); - - let row: Row = transaction - .query_one( - FAIL_TASK_QUERY, - &[ - &FangTaskState::Failed, - &error_message, - &updated_at, - &task.id, - ], - ) - .await?; - let failed_task = Self::row_to_task(row); - Ok(failed_task) - } - - async fn schedule_retry_query( - transaction: &mut Transaction<'_>, - task: &Task, - backoff_seconds: u32, - error: &str, - ) -> Result { - let now = Utc::now(); - let scheduled_at = now + Duration::seconds(backoff_seconds as i64); - let retries = task.retries + 1; - - let row: Row = transaction - .query_one( - RETRY_TASK_QUERY, - &[&error, &retries, &scheduled_at, &now, &task.id], - ) - .await?; - let failed_task = Self::row_to_task(row); - Ok(failed_task) - } - - async fn fetch_and_touch_task_query( - transaction: &mut Transaction<'_>, - task_type: Option, - ) -> Result, AsyncQueueError> { - let task_type = match task_type { - Some(passed_task_type) => passed_task_type, - None => DEFAULT_TASK_TYPE.to_string(), - }; - - let task = match Self::get_task_type_query(transaction, &task_type).await { - Ok(some_task) => Some(some_task), - Err(_) => None, - }; - let result_task = if let Some(some_task) = task { - Some( - Self::update_task_state_query(transaction, some_task, FangTaskState::InProgress) - .await?, - ) - } else { - None - }; - Ok(result_task) - } - - async fn get_task_type_query( - transaction: &mut Transaction<'_>, - task_type: &str, - ) -> Result { - let row: Row = transaction - .query_one(FETCH_TASK_TYPE_QUERY, &[&task_type, &Utc::now()]) - .await?; - - let task = Self::row_to_task(row); - - Ok(task) - } - - async fn update_task_state_query( - transaction: &mut Transaction<'_>, - task: Task, - state: FangTaskState, - ) -> Result { - let updated_at = Utc::now(); - - let row: Row = transaction - .query_one(UPDATE_TASK_STATE_QUERY, &[&state, &updated_at, &task.id]) - .await?; - let task = Self::row_to_task(row); - Ok(task) - } - - async fn insert_task_query( - transaction: &mut Transaction<'_>, - metadata: serde_json::Value, - task_type: &str, - scheduled_at: DateTime, - ) -> Result { - let row: Row = transaction - .query_one(INSERT_TASK_QUERY, &[&metadata, &task_type, &scheduled_at]) - .await?; - let task = Self::row_to_task(row); - Ok(task) - } - - async fn insert_task_uniq_query( - transaction: &mut Transaction<'_>, - metadata: serde_json::Value, - task_type: &str, - scheduled_at: DateTime, - ) -> Result { - let uniq_hash = Self::calculate_hash(metadata.to_string()); - - let row: Row = transaction - .query_one( - INSERT_TASK_UNIQ_QUERY, - &[&metadata, &task_type, &uniq_hash, &scheduled_at], - ) - .await?; - - let task = Self::row_to_task(row); - Ok(task) - } - - async fn execute_query( - transaction: &mut Transaction<'_>, - query: &str, - params: &[&(dyn ToSql + Sync)], - expected_result_count: Option, - ) -> Result { - let result = transaction.execute(query, params).await?; - - if let Some(expected_result) = expected_result_count { - if result != expected_result { - return Err(AsyncQueueError::ResultError { - expected: expected_result, - found: result, - }); - } - } - Ok(result) - } - - async fn insert_task_if_not_exist_query( - transaction: &mut Transaction<'_>, - metadata: serde_json::Value, - task_type: &str, - scheduled_at: DateTime, - ) -> Result { - match Self::find_task_by_uniq_hash_query(transaction, &metadata).await { - Some(task) => Ok(task), - None => { - Self::insert_task_uniq_query(transaction, metadata, task_type, scheduled_at).await - } - } - } - - fn calculate_hash(json: String) -> String { - let mut hasher = Sha256::new(); - hasher.update(json.as_bytes()); - let result = hasher.finalize(); - hex::encode(result) - } - - async fn find_task_by_uniq_hash_query( - transaction: &mut Transaction<'_>, - metadata: &serde_json::Value, - ) -> Option { - let uniq_hash = Self::calculate_hash(metadata.to_string()); - - let result = transaction - .query_one(FIND_TASK_BY_UNIQ_HASH_QUERY, &[&uniq_hash]) - .await; - - match result { - Ok(row) => Some(Self::row_to_task(row)), - Err(_) => None, - } - } - - fn row_to_task(row: Row) -> Task { - let id: Uuid = row.get("id"); - let metadata: serde_json::Value = row.get("metadata"); - - let error_message: Option = row.try_get("error_message").ok(); - - let uniq_hash: Option = row.try_get("uniq_hash").ok(); - let state: FangTaskState = row.get("state"); - let task_type: String = row.get("task_type"); - let retries: i32 = row.get("retries"); - let created_at: DateTime = row.get("created_at"); - let updated_at: DateTime = row.get("updated_at"); - let scheduled_at: DateTime = row.get("scheduled_at"); - - Task::builder() - .id(id) - .metadata(metadata) - .error_message(error_message) - .state(state) - .uniq_hash(uniq_hash) - .task_type(task_type) - .retries(retries) - .created_at(created_at) - .updated_at(updated_at) - .scheduled_at(scheduled_at) - .build() - } -} - -#[async_trait] -impl AsyncQueueable for AsyncQueue -where - Tls: MakeTlsConnect + Clone + Send + Sync + 'static, - >::Stream: Send + Sync, - >::TlsConnect: Send, - <>::TlsConnect as TlsConnect>::Future: Send, -{ - async fn find_task_by_id(&mut self, id: Uuid) -> Result { - let mut connection = self.pool.as_ref().unwrap().get().await?; - let mut transaction = connection.transaction().await?; - - let task = Self::find_task_by_id_query(&mut transaction, id).await?; - - transaction.commit().await?; - - Ok(task) - } - - async fn fetch_and_touch_task( - &mut self, - task_type: Option, - ) -> Result, AsyncQueueError> { - self.check_if_connection()?; - let mut connection = self.pool.as_ref().unwrap().get().await?; - let mut transaction = connection.transaction().await?; - - let task = Self::fetch_and_touch_task_query(&mut transaction, task_type).await?; - - transaction.commit().await?; - - Ok(task) - } - - async fn insert_task(&mut self, task: &dyn AsyncRunnable) -> Result { - self.check_if_connection()?; - let mut connection = self.pool.as_ref().unwrap().get().await?; - let mut transaction = connection.transaction().await?; - - let metadata = serde_json::to_value(task)?; - - let task: Task = if !task.uniq() { - Self::insert_task_query(&mut transaction, metadata, &task.task_type(), Utc::now()) - .await? - } else { - Self::insert_task_if_not_exist_query( - &mut transaction, - metadata, - &task.task_type(), - Utc::now(), - ) - .await? - }; - - transaction.commit().await?; - - Ok(task) - } - - async fn schedule_task(&mut self, task: &dyn AsyncRunnable) -> Result { - self.check_if_connection()?; - let mut connection = self.pool.as_ref().unwrap().get().await?; - let mut transaction = connection.transaction().await?; - let metadata = serde_json::to_value(task)?; - - let scheduled_at = match task.cron() { - Some(scheduled) => match scheduled { - CronPattern(cron_pattern) => { - let schedule = Schedule::from_str(&cron_pattern)?; - let mut iterator = schedule.upcoming(Utc); - iterator - .next() - .ok_or(AsyncQueueError::CronError(CronError::NoTimestampsError))? - } - ScheduleOnce(datetime) => datetime, - }, - None => { - return Err(AsyncQueueError::CronError( - CronError::TaskNotSchedulableError, - )); - } - }; - - let task: Task = if !task.uniq() { - Self::insert_task_query(&mut transaction, metadata, &task.task_type(), scheduled_at) - .await? - } else { - Self::insert_task_if_not_exist_query( - &mut transaction, - metadata, - &task.task_type(), - scheduled_at, - ) - .await? - }; - transaction.commit().await?; - Ok(task) - } - - async fn remove_all_tasks(&mut self) -> Result { - self.check_if_connection()?; - let mut connection = self.pool.as_ref().unwrap().get().await?; - let mut transaction = connection.transaction().await?; - - let result = Self::remove_all_tasks_query(&mut transaction).await?; - - transaction.commit().await?; - - Ok(result) - } - - async fn remove_all_scheduled_tasks(&mut self) -> Result { - self.check_if_connection()?; - let mut connection = self.pool.as_ref().unwrap().get().await?; - let mut transaction = connection.transaction().await?; - - let result = Self::remove_all_scheduled_tasks_query(&mut transaction).await?; - - transaction.commit().await?; - - Ok(result) - } - - async fn remove_task(&mut self, id: Uuid) -> Result { - self.check_if_connection()?; - let mut connection = self.pool.as_ref().unwrap().get().await?; - let mut transaction = connection.transaction().await?; - - let result = Self::remove_task_query(&mut transaction, id).await?; - - transaction.commit().await?; - - Ok(result) - } - - async fn remove_task_by_metadata( - &mut self, - task: &dyn AsyncRunnable, - ) -> Result { - if task.uniq() { - self.check_if_connection()?; - let mut connection = self.pool.as_ref().unwrap().get().await?; - let mut transaction = connection.transaction().await?; - - let result = Self::remove_task_by_metadata_query(&mut transaction, task).await?; - - transaction.commit().await?; - - Ok(result) - } else { - Err(AsyncQueueError::TaskNotUniqError) - } - } - - async fn remove_tasks_type(&mut self, task_type: &str) -> Result { - self.check_if_connection()?; - let mut connection = self.pool.as_ref().unwrap().get().await?; - let mut transaction = connection.transaction().await?; - - let result = Self::remove_tasks_type_query(&mut transaction, task_type).await?; - - transaction.commit().await?; - - Ok(result) - } - - async fn update_task_state( - &mut self, - task: Task, - state: FangTaskState, - ) -> Result { - self.check_if_connection()?; - let mut connection = self.pool.as_ref().unwrap().get().await?; - let mut transaction = connection.transaction().await?; - - let task = Self::update_task_state_query(&mut transaction, task, state).await?; - transaction.commit().await?; - - Ok(task) - } - - async fn fail_task( - &mut self, - task: Task, - error_message: &str, - ) -> Result { - self.check_if_connection()?; - let mut connection = self.pool.as_ref().unwrap().get().await?; - let mut transaction = connection.transaction().await?; - - let task = Self::fail_task_query(&mut transaction, task, error_message).await?; - transaction.commit().await?; - - Ok(task) - } - - async fn schedule_retry( - &mut self, - task: &Task, - backoff_seconds: u32, - error: &str, - ) -> Result { - self.check_if_connection()?; - let mut connection = self.pool.as_ref().unwrap().get().await?; - let mut transaction = connection.transaction().await?; - - let task = - Self::schedule_retry_query(&mut transaction, task, backoff_seconds, error).await?; - transaction.commit().await?; - - Ok(task) - } -} - -#[cfg(test)] -mod async_queue_tests { - use super::AsyncQueueTest; - use super::AsyncQueueable; - use super::FangTaskState; - use super::Task; - use crate::asynk::AsyncRunnable; - use crate::FangError; - use crate::Scheduled; - use async_trait::async_trait; - use bb8_postgres::bb8::Pool; - use bb8_postgres::tokio_postgres::NoTls; - use bb8_postgres::PostgresConnectionManager; - use chrono::DateTime; - use chrono::Duration; - use chrono::SubsecRound; - use chrono::Utc; - use serde::{Deserialize, Serialize}; - - #[derive(Serialize, Deserialize)] - struct AsyncTask { - pub number: u16, - } - - #[typetag::serde] - #[async_trait] - impl AsyncRunnable for AsyncTask { - async fn run(&self, _queueable: &mut dyn AsyncQueueable) -> Result<(), FangError> { - Ok(()) - } - } - - #[derive(Serialize, Deserialize)] - struct AsyncUniqTask { - pub number: u16, - } - - #[typetag::serde] - #[async_trait] - impl AsyncRunnable for AsyncUniqTask { - async fn run(&self, _queueable: &mut dyn AsyncQueueable) -> Result<(), FangError> { - Ok(()) - } - - fn uniq(&self) -> bool { - true - } - } - - #[derive(Serialize, Deserialize)] - struct AsyncTaskSchedule { - pub number: u16, - pub datetime: String, - } - - #[typetag::serde] - #[async_trait] - impl AsyncRunnable for AsyncTaskSchedule { - async fn run(&self, _queueable: &mut dyn AsyncQueueable) -> Result<(), FangError> { - Ok(()) - } - - fn cron(&self) -> Option { - let datetime = self.datetime.parse::>().ok()?; - Some(Scheduled::ScheduleOnce(datetime)) - } - } - - #[tokio::test] - async fn insert_task_creates_new_task() { - let pool = pool().await; - let mut connection = pool.get().await.unwrap(); - let transaction = connection.transaction().await.unwrap(); - - let mut test = AsyncQueueTest::builder().transaction(transaction).build(); - - let task = insert_task(&mut test, &AsyncTask { number: 1 }).await; - - let metadata = task.metadata.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); - test.transaction.rollback().await.unwrap(); - } - - #[tokio::test] - async fn update_task_state_test() { - let pool = pool().await; - let mut connection = pool.get().await.unwrap(); - let transaction = connection.transaction().await.unwrap(); - - let mut test = AsyncQueueTest::builder().transaction(transaction).build(); - - let task = insert_task(&mut test, &AsyncTask { number: 1 }).await; - - let metadata = task.metadata.as_object().unwrap(); - let number = metadata["number"].as_u64(); - let type_task = metadata["type"].as_str(); - let id = task.id; - - assert_eq!(Some(1), number); - assert_eq!(Some("AsyncTask"), type_task); - - let finished_task = test - .update_task_state(task, FangTaskState::Finished) - .await - .unwrap(); - - assert_eq!(id, finished_task.id); - assert_eq!(FangTaskState::Finished, finished_task.state); - - test.transaction.rollback().await.unwrap(); - } - - #[tokio::test] - async fn failed_task_query_test() { - let pool = pool().await; - let mut connection = pool.get().await.unwrap(); - let transaction = connection.transaction().await.unwrap(); - - let mut test = AsyncQueueTest::builder().transaction(transaction).build(); - - let task = insert_task(&mut test, &AsyncTask { number: 1 }).await; - - let metadata = task.metadata.as_object().unwrap(); - let number = metadata["number"].as_u64(); - let type_task = metadata["type"].as_str(); - let id = task.id; - - assert_eq!(Some(1), number); - assert_eq!(Some("AsyncTask"), type_task); - - let failed_task = test.fail_task(task, "Some error").await.unwrap(); - - assert_eq!(id, failed_task.id); - assert_eq!(Some("Some error"), failed_task.error_message.as_deref()); - assert_eq!(FangTaskState::Failed, failed_task.state); - - test.transaction.rollback().await.unwrap(); - } - - #[tokio::test] - async fn remove_all_tasks_test() { - let pool = pool().await; - let mut connection = pool.get().await.unwrap(); - let transaction = connection.transaction().await.unwrap(); - - let mut test = AsyncQueueTest::builder().transaction(transaction).build(); - - let task = insert_task(&mut test, &AsyncTask { number: 1 }).await; - - let metadata = task.metadata.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 = insert_task(&mut test, &AsyncTask { number: 2 }).await; - - let metadata = task.metadata.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); - - test.transaction.rollback().await.unwrap(); - } - - #[tokio::test] - async fn schedule_task_test() { - let pool = pool().await; - let mut connection = pool.get().await.unwrap(); - let transaction = connection.transaction().await.unwrap(); - - let mut test = AsyncQueueTest::builder().transaction(transaction).build(); - - let datetime = (Utc::now() + Duration::seconds(7)).round_subsecs(0); - - let task = &AsyncTaskSchedule { - number: 1, - datetime: datetime.to_string(), - }; - - let task = test.schedule_task(task).await.unwrap(); - - let metadata = task.metadata.as_object().unwrap(); - let number = metadata["number"].as_u64(); - let type_task = metadata["type"].as_str(); - - assert_eq!(Some(1), number); - assert_eq!(Some("AsyncTaskSchedule"), type_task); - assert_eq!(task.scheduled_at, datetime); - } - - #[tokio::test] - async fn remove_all_scheduled_tasks_test() { - let pool = pool().await; - let mut connection = pool.get().await.unwrap(); - let transaction = connection.transaction().await.unwrap(); - - let mut test = AsyncQueueTest::builder().transaction(transaction).build(); - - let datetime = (Utc::now() + Duration::seconds(7)).round_subsecs(0); - - let task1 = &AsyncTaskSchedule { - number: 1, - datetime: datetime.to_string(), - }; - - let task2 = &AsyncTaskSchedule { - number: 2, - datetime: datetime.to_string(), - }; - - test.schedule_task(task1).await.unwrap(); - test.schedule_task(task2).await.unwrap(); - - let number = test.remove_all_scheduled_tasks().await.unwrap(); - - assert_eq!(2, number); - } - - #[tokio::test] - async fn fetch_and_touch_test() { - let pool = pool().await; - let mut connection = pool.get().await.unwrap(); - let transaction = connection.transaction().await.unwrap(); - - let mut test = AsyncQueueTest::builder().transaction(transaction).build(); - - let task = insert_task(&mut test, &AsyncTask { number: 1 }).await; - - let metadata = task.metadata.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 = insert_task(&mut test, &AsyncTask { number: 2 }).await; - - let metadata = task.metadata.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 = test.fetch_and_touch_task(None).await.unwrap().unwrap(); - - let metadata = task.metadata.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.fetch_and_touch_task(None).await.unwrap().unwrap(); - let metadata = task.metadata.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); - - test.transaction.rollback().await.unwrap(); - } - - #[tokio::test] - async fn remove_tasks_type_test() { - let pool = pool().await; - let mut connection = pool.get().await.unwrap(); - let transaction = connection.transaction().await.unwrap(); - - let mut test = AsyncQueueTest::builder().transaction(transaction).build(); - - let task = insert_task(&mut test, &AsyncTask { number: 1 }).await; - - let metadata = task.metadata.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 = insert_task(&mut test, &AsyncTask { number: 2 }).await; - - let metadata = task.metadata.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_tasks_type("mytype").await.unwrap(); - assert_eq!(0, result); - - let result = test.remove_tasks_type("common").await.unwrap(); - assert_eq!(2, result); - - test.transaction.rollback().await.unwrap(); - } - - #[tokio::test] - async fn remove_tasks_by_metadata() { - let pool = pool().await; - let mut connection = pool.get().await.unwrap(); - let transaction = connection.transaction().await.unwrap(); - - let mut test = AsyncQueueTest::builder().transaction(transaction).build(); - - let task = insert_task(&mut test, &AsyncUniqTask { number: 1 }).await; - - let metadata = task.metadata.as_object().unwrap(); - let number = metadata["number"].as_u64(); - let type_task = metadata["type"].as_str(); - - assert_eq!(Some(1), number); - assert_eq!(Some("AsyncUniqTask"), type_task); - - let task = insert_task(&mut test, &AsyncUniqTask { number: 2 }).await; - - let metadata = task.metadata.as_object().unwrap(); - let number = metadata["number"].as_u64(); - let type_task = metadata["type"].as_str(); - - assert_eq!(Some(2), number); - assert_eq!(Some("AsyncUniqTask"), type_task); - - let result = test - .remove_task_by_metadata(&AsyncUniqTask { number: 0 }) - .await - .unwrap(); - assert_eq!(0, result); - - let result = test - .remove_task_by_metadata(&AsyncUniqTask { number: 1 }) - .await - .unwrap(); - assert_eq!(1, result); - - test.transaction.rollback().await.unwrap(); - } - - async fn insert_task(test: &mut AsyncQueueTest<'_>, task: &dyn AsyncRunnable) -> Task { - test.insert_task(task).await.unwrap() - } - - async fn pool() -> Pool> { - let pg_mgr = PostgresConnectionManager::new_from_stringlike( - "postgres://postgres:postgres@localhost/fang", - NoTls, - ) - .unwrap(); - - Pool::builder().build(pg_mgr).await.unwrap() - } -} diff --git a/src/asynk/async_worker.rs b/src/asynk/async_worker.rs deleted file mode 100644 index b66af7d..0000000 --- a/src/asynk/async_worker.rs +++ /dev/null @@ -1,580 +0,0 @@ -use crate::asynk::async_queue::AsyncQueueable; -use crate::asynk::async_queue::FangTaskState; -use crate::asynk::async_queue::Task; -use crate::asynk::async_queue::DEFAULT_TASK_TYPE; -use crate::asynk::async_runnable::AsyncRunnable; -use crate::FangError; -use crate::Scheduled::*; -use crate::{RetentionMode, SleepParams}; -use log::error; -use typed_builder::TypedBuilder; - -/// it executes tasks only of task_type type, it sleeps when there are no tasks in the queue -#[derive(TypedBuilder)] -pub struct AsyncWorker -where - AQueue: AsyncQueueable + Clone + Sync + 'static, -{ - #[builder(setter(into))] - pub queue: AQueue, - #[builder(default=DEFAULT_TASK_TYPE.to_string(), setter(into))] - pub task_type: String, - #[builder(default, setter(into))] - pub sleep_params: SleepParams, - #[builder(default, setter(into))] - pub retention_mode: RetentionMode, -} - -impl AsyncWorker -where - AQueue: AsyncQueueable + Clone + Sync + 'static, -{ - async fn run(&mut self, task: Task, runnable: Box) -> Result<(), FangError> { - let result = runnable.run(&mut self.queue).await; - - match result { - Ok(_) => self.finalize_task(task, &result).await?, - - Err(ref error) => { - if task.retries < runnable.max_retries() { - let backoff_seconds = runnable.backoff(task.retries as u32); - - self.queue - .schedule_retry(&task, backoff_seconds, &error.description) - .await?; - } else { - self.finalize_task(task, &result).await?; - } - } - } - - Ok(()) - } - - async fn finalize_task( - &mut self, - task: Task, - result: &Result<(), FangError>, - ) -> Result<(), FangError> { - match self.retention_mode { - RetentionMode::KeepAll => match result { - Ok(_) => { - self.queue - .update_task_state(task, FangTaskState::Finished) - .await?; - } - Err(error) => { - self.queue.fail_task(task, &error.description).await?; - } - }, - RetentionMode::RemoveAll => { - self.queue.remove_task(task.id).await?; - } - RetentionMode::RemoveFinished => match result { - Ok(_) => { - self.queue.remove_task(task.id).await?; - } - Err(error) => { - self.queue.fail_task(task, &error.description).await?; - } - }, - }; - - Ok(()) - } - - async fn sleep(&mut self) { - self.sleep_params.maybe_increase_sleep_period(); - - tokio::time::sleep(self.sleep_params.sleep_period).await; - } - - pub(crate) async fn run_tasks(&mut self) -> Result<(), FangError> { - loop { - //fetch task - match self - .queue - .fetch_and_touch_task(Some(self.task_type.clone())) - .await - { - Ok(Some(task)) => { - let actual_task: Box = - serde_json::from_value(task.metadata.clone()).unwrap(); - - // check if task is scheduled or not - if let Some(CronPattern(_)) = actual_task.cron() { - // program task - self.queue.schedule_task(&*actual_task).await?; - } - self.sleep_params.maybe_reset_sleep_period(); - // run scheduled task - self.run(task, actual_task).await?; - } - Ok(None) => { - self.sleep().await; - } - - Err(error) => { - error!("Failed to fetch a task {:?}", error); - - self.sleep().await; - } - }; - } - } -} - -#[cfg(test)] -#[derive(TypedBuilder)] -pub struct AsyncWorkerTest<'a> { - #[builder(setter(into))] - pub queue: &'a mut dyn AsyncQueueable, - #[builder(default=DEFAULT_TASK_TYPE.to_string(), setter(into))] - pub task_type: String, - #[builder(default, setter(into))] - pub sleep_params: SleepParams, - #[builder(default, setter(into))] - pub retention_mode: RetentionMode, -} - -#[cfg(test)] -impl<'a> AsyncWorkerTest<'a> { - pub async fn run( - &mut self, - task: Task, - runnable: Box, - ) -> Result<(), FangError> { - let result = runnable.run(self.queue).await; - - match result { - Ok(_) => self.finalize_task(task, &result).await?, - - Err(ref error) => { - if task.retries < runnable.max_retries() { - let backoff_seconds = runnable.backoff(task.retries as u32); - - self.queue - .schedule_retry(&task, backoff_seconds, &error.description) - .await?; - } else { - self.finalize_task(task, &result).await?; - } - } - } - - Ok(()) - } - - async fn finalize_task( - &mut self, - task: Task, - result: &Result<(), FangError>, - ) -> Result<(), FangError> { - match self.retention_mode { - RetentionMode::KeepAll => match result { - Ok(_) => { - self.queue - .update_task_state(task, FangTaskState::Finished) - .await?; - } - Err(error) => { - self.queue.fail_task(task, &error.description).await?; - } - }, - RetentionMode::RemoveAll => match result { - Ok(_) => { - self.queue.remove_task(task.id).await?; - } - Err(_error) => { - self.queue.remove_task(task.id).await?; - } - }, - RetentionMode::RemoveFinished => match result { - Ok(_) => { - self.queue.remove_task(task.id).await?; - } - Err(error) => { - self.queue.fail_task(task, &error.description).await?; - } - }, - }; - - Ok(()) - } - - pub async fn sleep(&mut self) { - self.sleep_params.maybe_increase_sleep_period(); - - tokio::time::sleep(self.sleep_params.sleep_period).await; - } - - pub async fn run_tasks_until_none(&mut self) -> Result<(), FangError> { - loop { - match self - .queue - .fetch_and_touch_task(Some(self.task_type.clone())) - .await - { - Ok(Some(task)) => { - let actual_task: Box = - serde_json::from_value(task.metadata.clone()).unwrap(); - - // check if task is scheduled or not - if let Some(CronPattern(_)) = actual_task.cron() { - // program task - self.queue.schedule_task(&*actual_task).await?; - } - self.sleep_params.maybe_reset_sleep_period(); - // run scheduled task - self.run(task, actual_task).await?; - } - Ok(None) => { - return Ok(()); - } - Err(error) => { - error!("Failed to fetch a task {:?}", error); - - self.sleep().await; - } - }; - } - } -} - -#[cfg(test)] -mod async_worker_tests { - use super::AsyncWorkerTest; - use crate::asynk::async_queue::AsyncQueueTest; - use crate::asynk::async_queue::AsyncQueueable; - use crate::asynk::async_queue::FangTaskState; - use crate::asynk::async_worker::Task; - use crate::asynk::AsyncRunnable; - use crate::FangError; - use crate::RetentionMode; - use crate::Scheduled; - use async_trait::async_trait; - use bb8_postgres::bb8::Pool; - use bb8_postgres::tokio_postgres::NoTls; - use bb8_postgres::PostgresConnectionManager; - use chrono::Duration; - use chrono::Utc; - use serde::{Deserialize, Serialize}; - - #[derive(Serialize, Deserialize)] - struct WorkerAsyncTask { - pub number: u16, - } - - #[typetag::serde] - #[async_trait] - impl AsyncRunnable for WorkerAsyncTask { - async fn run(&self, _queueable: &mut dyn AsyncQueueable) -> Result<(), FangError> { - Ok(()) - } - } - - #[derive(Serialize, Deserialize)] - struct WorkerAsyncTaskSchedule { - pub number: u16, - } - - #[typetag::serde] - #[async_trait] - impl AsyncRunnable for WorkerAsyncTaskSchedule { - async fn run(&self, _queueable: &mut dyn AsyncQueueable) -> Result<(), FangError> { - Ok(()) - } - fn cron(&self) -> Option { - Some(Scheduled::ScheduleOnce(Utc::now() + Duration::seconds(1))) - } - } - - #[derive(Serialize, Deserialize)] - struct AsyncFailedTask { - pub number: u16, - } - - #[typetag::serde] - #[async_trait] - impl AsyncRunnable for AsyncFailedTask { - async fn run(&self, _queueable: &mut dyn AsyncQueueable) -> Result<(), FangError> { - let message = format!("number {} is wrong :(", self.number); - - Err(FangError { - description: message, - }) - } - - fn max_retries(&self) -> i32 { - 0 - } - } - - #[derive(Serialize, Deserialize, Clone)] - struct AsyncRetryTask {} - - #[typetag::serde] - #[async_trait] - impl AsyncRunnable for AsyncRetryTask { - async fn run(&self, _queueable: &mut dyn AsyncQueueable) -> Result<(), FangError> { - let message = "Failed".to_string(); - - Err(FangError { - description: message, - }) - } - - fn max_retries(&self) -> i32 { - 2 - } - } - - #[derive(Serialize, Deserialize)] - struct AsyncTaskType1 {} - - #[typetag::serde] - #[async_trait] - impl AsyncRunnable for AsyncTaskType1 { - async fn run(&self, _queueable: &mut dyn AsyncQueueable) -> Result<(), FangError> { - Ok(()) - } - - fn task_type(&self) -> String { - "type1".to_string() - } - } - - #[derive(Serialize, Deserialize)] - struct AsyncTaskType2 {} - - #[typetag::serde] - #[async_trait] - impl AsyncRunnable for AsyncTaskType2 { - async fn run(&self, _queueable: &mut dyn AsyncQueueable) -> Result<(), FangError> { - Ok(()) - } - - fn task_type(&self) -> String { - "type2".to_string() - } - } - - #[tokio::test] - async fn execute_and_finishes_task() { - let pool = pool().await; - let mut connection = pool.get().await.unwrap(); - let transaction = connection.transaction().await.unwrap(); - - let mut test = AsyncQueueTest::builder().transaction(transaction).build(); - let actual_task = WorkerAsyncTask { number: 1 }; - - let task = insert_task(&mut test, &actual_task).await; - let id = task.id; - - let mut worker = AsyncWorkerTest::builder() - .queue(&mut test as &mut dyn AsyncQueueable) - .retention_mode(RetentionMode::KeepAll) - .build(); - - worker.run(task, Box::new(actual_task)).await.unwrap(); - let task_finished = test.find_task_by_id(id).await.unwrap(); - assert_eq!(id, task_finished.id); - assert_eq!(FangTaskState::Finished, task_finished.state); - test.transaction.rollback().await.unwrap(); - } - - #[tokio::test] - async fn schedule_task_test() { - let pool = pool().await; - let mut connection = pool.get().await.unwrap(); - let transaction = connection.transaction().await.unwrap(); - - let mut test = AsyncQueueTest::builder().transaction(transaction).build(); - - let actual_task = WorkerAsyncTaskSchedule { number: 1 }; - - let task = test.schedule_task(&actual_task).await.unwrap(); - - let id = task.id; - - let mut worker = AsyncWorkerTest::builder() - .queue(&mut test as &mut dyn AsyncQueueable) - .retention_mode(RetentionMode::KeepAll) - .build(); - - worker.run_tasks_until_none().await.unwrap(); - - let task = worker.queue.find_task_by_id(id).await.unwrap(); - - assert_eq!(id, task.id); - assert_eq!(FangTaskState::New, task.state); - - tokio::time::sleep(core::time::Duration::from_secs(3)).await; - - worker.run_tasks_until_none().await.unwrap(); - - let task = test.find_task_by_id(id).await.unwrap(); - assert_eq!(id, task.id); - assert_eq!(FangTaskState::Finished, task.state); - } - - #[tokio::test] - async fn retries_task_test() { - let pool = pool().await; - let mut connection = pool.get().await.unwrap(); - let transaction = connection.transaction().await.unwrap(); - - let mut test = AsyncQueueTest::builder().transaction(transaction).build(); - - let actual_task = AsyncRetryTask {}; - - let task = test.insert_task(&actual_task).await.unwrap(); - - let id = task.id; - - let mut worker = AsyncWorkerTest::builder() - .queue(&mut test as &mut dyn AsyncQueueable) - .retention_mode(RetentionMode::KeepAll) - .build(); - - worker.run_tasks_until_none().await.unwrap(); - - let task = worker.queue.find_task_by_id(id).await.unwrap(); - - assert_eq!(id, task.id); - assert_eq!(FangTaskState::Retried, task.state); - assert_eq!(1, task.retries); - - tokio::time::sleep(core::time::Duration::from_secs(5)).await; - worker.run_tasks_until_none().await.unwrap(); - - let task = worker.queue.find_task_by_id(id).await.unwrap(); - - assert_eq!(id, task.id); - assert_eq!(FangTaskState::Retried, task.state); - assert_eq!(2, task.retries); - - tokio::time::sleep(core::time::Duration::from_secs(10)).await; - worker.run_tasks_until_none().await.unwrap(); - - let task = test.find_task_by_id(id).await.unwrap(); - assert_eq!(id, task.id); - assert_eq!(FangTaskState::Failed, task.state); - assert_eq!("Failed".to_string(), task.error_message.unwrap()); - } - - #[tokio::test] - async fn saves_error_for_failed_task() { - let pool = pool().await; - let mut connection = pool.get().await.unwrap(); - let transaction = connection.transaction().await.unwrap(); - - let mut test = AsyncQueueTest::builder().transaction(transaction).build(); - let failed_task = AsyncFailedTask { number: 1 }; - - let task = insert_task(&mut test, &failed_task).await; - let id = task.id; - - let mut worker = AsyncWorkerTest::builder() - .queue(&mut test as &mut dyn AsyncQueueable) - .retention_mode(RetentionMode::KeepAll) - .build(); - - worker.run(task, Box::new(failed_task)).await.unwrap(); - let task_finished = test.find_task_by_id(id).await.unwrap(); - - assert_eq!(id, task_finished.id); - assert_eq!(FangTaskState::Failed, task_finished.state); - assert_eq!( - "number 1 is wrong :(".to_string(), - task_finished.error_message.unwrap() - ); - test.transaction.rollback().await.unwrap(); - } - - #[tokio::test] - async fn executes_task_only_of_specific_type() { - let pool = pool().await; - let mut connection = pool.get().await.unwrap(); - let transaction = connection.transaction().await.unwrap(); - - let mut test = AsyncQueueTest::builder().transaction(transaction).build(); - - let task1 = insert_task(&mut test, &AsyncTaskType1 {}).await; - let task12 = insert_task(&mut test, &AsyncTaskType1 {}).await; - let task2 = insert_task(&mut test, &AsyncTaskType2 {}).await; - - let id1 = task1.id; - let id12 = task12.id; - let id2 = task2.id; - - let mut worker = AsyncWorkerTest::builder() - .queue(&mut test as &mut dyn AsyncQueueable) - .task_type("type1".to_string()) - .retention_mode(RetentionMode::KeepAll) - .build(); - - worker.run_tasks_until_none().await.unwrap(); - let task1 = test.find_task_by_id(id1).await.unwrap(); - let task12 = test.find_task_by_id(id12).await.unwrap(); - let task2 = test.find_task_by_id(id2).await.unwrap(); - - assert_eq!(id1, task1.id); - assert_eq!(id12, task12.id); - assert_eq!(id2, task2.id); - assert_eq!(FangTaskState::Finished, task1.state); - assert_eq!(FangTaskState::Finished, task12.state); - assert_eq!(FangTaskState::New, task2.state); - test.transaction.rollback().await.unwrap(); - } - - #[tokio::test] - async fn remove_when_finished() { - let pool = pool().await; - let mut connection = pool.get().await.unwrap(); - let transaction = connection.transaction().await.unwrap(); - - let mut test = AsyncQueueTest::builder().transaction(transaction).build(); - - let task1 = insert_task(&mut test, &AsyncTaskType1 {}).await; - let task12 = insert_task(&mut test, &AsyncTaskType1 {}).await; - let task2 = insert_task(&mut test, &AsyncTaskType2 {}).await; - - let _id1 = task1.id; - let _id12 = task12.id; - let id2 = task2.id; - - let mut worker = AsyncWorkerTest::builder() - .queue(&mut test as &mut dyn AsyncQueueable) - .task_type("type1".to_string()) - .build(); - - worker.run_tasks_until_none().await.unwrap(); - let task = test - .fetch_and_touch_task(Some("type1".to_string())) - .await - .unwrap(); - assert_eq!(None, task); - - let task2 = test - .fetch_and_touch_task(Some("type2".to_string())) - .await - .unwrap() - .unwrap(); - assert_eq!(id2, task2.id); - - test.transaction.rollback().await.unwrap(); - } - async fn insert_task(test: &mut AsyncQueueTest<'_>, task: &dyn AsyncRunnable) -> Task { - test.insert_task(task).await.unwrap() - } - async fn pool() -> Pool> { - let pg_mgr = PostgresConnectionManager::new_from_stringlike( - "postgres://postgres:postgres@localhost/fang", - NoTls, - ) - .unwrap(); - - Pool::builder().build(pg_mgr).await.unwrap() - } -} diff --git a/src/asynk/queries/fail_task.sql b/src/asynk/queries/fail_task.sql deleted file mode 100644 index 1719286..0000000 --- a/src/asynk/queries/fail_task.sql +++ /dev/null @@ -1 +0,0 @@ -UPDATE "fang_tasks" SET "state" = $1 , "error_message" = $2 , "updated_at" = $3 WHERE id = $4 RETURNING * diff --git a/src/asynk/queries/fetch_task_type.sql b/src/asynk/queries/fetch_task_type.sql deleted file mode 100644 index e055820..0000000 --- a/src/asynk/queries/fetch_task_type.sql +++ /dev/null @@ -1 +0,0 @@ -SELECT * FROM fang_tasks WHERE task_type = $1 AND state in ('new', 'retried') AND $2 >= scheduled_at ORDER BY created_at ASC, scheduled_at ASC LIMIT 1 FOR UPDATE SKIP LOCKED diff --git a/src/asynk/queries/find_task_by_id.sql b/src/asynk/queries/find_task_by_id.sql deleted file mode 100644 index 608166f..0000000 --- a/src/asynk/queries/find_task_by_id.sql +++ /dev/null @@ -1 +0,0 @@ -SELECT * FROM fang_tasks WHERE id = $1 diff --git a/src/asynk/queries/find_task_by_uniq_hash.sql b/src/asynk/queries/find_task_by_uniq_hash.sql deleted file mode 100644 index cb53f45..0000000 --- a/src/asynk/queries/find_task_by_uniq_hash.sql +++ /dev/null @@ -1 +0,0 @@ -SELECT * FROM fang_tasks WHERE uniq_hash = $1 AND state in ('new', 'retried') LIMIT 1 diff --git a/src/asynk/queries/insert_task.sql b/src/asynk/queries/insert_task.sql deleted file mode 100644 index 514d921..0000000 --- a/src/asynk/queries/insert_task.sql +++ /dev/null @@ -1 +0,0 @@ -INSERT INTO "fang_tasks" ("metadata", "task_type", "scheduled_at") VALUES ($1, $2, $3) RETURNING * diff --git a/src/asynk/queries/insert_task_uniq.sql b/src/asynk/queries/insert_task_uniq.sql deleted file mode 100644 index 0817383..0000000 --- a/src/asynk/queries/insert_task_uniq.sql +++ /dev/null @@ -1 +0,0 @@ -INSERT INTO "fang_tasks" ("metadata", "task_type" , "uniq_hash", "scheduled_at") VALUES ($1, $2 , $3, $4) RETURNING * diff --git a/src/asynk/queries/remove_all_scheduled_tasks.sql b/src/asynk/queries/remove_all_scheduled_tasks.sql deleted file mode 100644 index 61a5b6b..0000000 --- a/src/asynk/queries/remove_all_scheduled_tasks.sql +++ /dev/null @@ -1 +0,0 @@ -DELETE FROM "fang_tasks" WHERE scheduled_at > $1 diff --git a/src/asynk/queries/remove_all_tasks.sql b/src/asynk/queries/remove_all_tasks.sql deleted file mode 100644 index eaecbba..0000000 --- a/src/asynk/queries/remove_all_tasks.sql +++ /dev/null @@ -1 +0,0 @@ -DELETE FROM "fang_tasks" diff --git a/src/asynk/queries/remove_task.sql b/src/asynk/queries/remove_task.sql deleted file mode 100644 index b6da69f..0000000 --- a/src/asynk/queries/remove_task.sql +++ /dev/null @@ -1 +0,0 @@ -DELETE FROM "fang_tasks" WHERE id = $1 diff --git a/src/asynk/queries/remove_task_by_metadata.sql b/src/asynk/queries/remove_task_by_metadata.sql deleted file mode 100644 index 94324e2..0000000 --- a/src/asynk/queries/remove_task_by_metadata.sql +++ /dev/null @@ -1 +0,0 @@ -DELETE FROM "fang_tasks" WHERE uniq_hash = $1 diff --git a/src/asynk/queries/remove_tasks_type.sql b/src/asynk/queries/remove_tasks_type.sql deleted file mode 100644 index e4de9c0..0000000 --- a/src/asynk/queries/remove_tasks_type.sql +++ /dev/null @@ -1 +0,0 @@ -DELETE FROM "fang_tasks" WHERE task_type = $1 diff --git a/src/asynk/queries/retry_task.sql b/src/asynk/queries/retry_task.sql deleted file mode 100644 index f26267c..0000000 --- a/src/asynk/queries/retry_task.sql +++ /dev/null @@ -1 +0,0 @@ -UPDATE "fang_tasks" SET "state" = 'retried' , "error_message" = $1, "retries" = $2, scheduled_at = $3, "updated_at" = $4 WHERE id = $5 RETURNING * diff --git a/src/asynk/queries/update_task_state.sql b/src/asynk/queries/update_task_state.sql deleted file mode 100644 index e2e2d94..0000000 --- a/src/asynk/queries/update_task_state.sql +++ /dev/null @@ -1 +0,0 @@ -UPDATE "fang_tasks" SET "state" = $1 , "updated_at" = $2 WHERE id = $3 RETURNING * diff --git a/src/blocking.rs b/src/blocking.rs deleted file mode 100644 index 69f3b07..0000000 --- a/src/blocking.rs +++ /dev/null @@ -1,14 +0,0 @@ -mod error; -pub mod fang_task_state; -pub mod queue; -pub mod runnable; -pub mod schema; -pub mod worker; -pub mod worker_pool; - -pub use fang_task_state::FangTaskState; -pub use queue::*; -pub use runnable::Runnable; -pub use schema::*; -pub use worker::*; -pub use worker_pool::*; diff --git a/src/blocking/error.rs b/src/blocking/error.rs deleted file mode 100644 index 333b868..0000000 --- a/src/blocking/error.rs +++ /dev/null @@ -1,31 +0,0 @@ -use crate::blocking::queue::QueueError; -use crate::FangError; -use diesel::r2d2::PoolError; -use diesel::result::Error as DieselError; -use std::io::Error as IoError; - -impl From for FangError { - fn from(error: IoError) -> Self { - let description = format!("{error:?}"); - FangError { description } - } -} - -impl From for FangError { - fn from(error: QueueError) -> Self { - let description = format!("{error:?}"); - FangError { description } - } -} - -impl From for FangError { - fn from(error: DieselError) -> Self { - Self::from(QueueError::DieselError(error)) - } -} - -impl From for FangError { - fn from(error: PoolError) -> Self { - Self::from(QueueError::PoolError(error)) - } -} diff --git a/src/blocking/queue.rs b/src/blocking/queue.rs deleted file mode 100644 index 9095729..0000000 --- a/src/blocking/queue.rs +++ /dev/null @@ -1,922 +0,0 @@ -use crate::fang_task_state::FangTaskState; -use crate::runnable::Runnable; -use crate::schema::fang_tasks; -use crate::CronError; -use crate::Scheduled::*; -use chrono::DateTime; -use chrono::Duration; -use chrono::Utc; -use cron::Schedule; -use diesel::pg::PgConnection; -use diesel::prelude::*; -use diesel::r2d2; -use diesel::r2d2::ConnectionManager; -use diesel::r2d2::PoolError; -use diesel::r2d2::PooledConnection; -use diesel::result::Error as DieselError; -use sha2::Digest; -use sha2::Sha256; -use std::str::FromStr; -use thiserror::Error; -use typed_builder::TypedBuilder; -use uuid::Uuid; - -#[cfg(test)] -use dotenvy::dotenv; -#[cfg(test)] -use std::env; - -pub type PoolConnection = PooledConnection>; - -#[derive(Queryable, Identifiable, Debug, Eq, PartialEq, Clone, TypedBuilder)] -#[diesel(table_name = fang_tasks)] -pub struct Task { - #[builder(setter(into))] - pub id: Uuid, - #[builder(setter(into))] - pub metadata: serde_json::Value, - #[builder(setter(into))] - pub error_message: Option, - #[builder(setter(into))] - pub state: FangTaskState, - #[builder(setter(into))] - pub task_type: String, - #[builder(setter(into))] - pub uniq_hash: Option, - #[builder(setter(into))] - pub retries: i32, - #[builder(setter(into))] - pub scheduled_at: DateTime, - #[builder(setter(into))] - pub created_at: DateTime, - #[builder(setter(into))] - pub updated_at: DateTime, -} - -#[derive(Insertable, Debug, Eq, PartialEq, Clone, TypedBuilder)] -#[diesel(table_name = fang_tasks)] -pub struct NewTask { - #[builder(setter(into))] - metadata: serde_json::Value, - #[builder(setter(into))] - task_type: String, - #[builder(setter(into))] - uniq_hash: Option, - #[builder(setter(into))] - scheduled_at: DateTime, -} - -#[derive(Debug, Error)] -pub enum QueueError { - #[error(transparent)] - DieselError(#[from] DieselError), - #[error(transparent)] - PoolError(#[from] PoolError), - #[error(transparent)] - CronError(#[from] CronError), - #[error("Can not perform this operation if task is not uniq, please check its definition in impl Runnable")] - TaskNotUniqError, -} - -impl From for QueueError { - fn from(error: cron::error::Error) -> Self { - QueueError::CronError(CronError::LibraryError(error)) - } -} - -/// This trait defines operations for a synchronous queue. -/// The trait can be implemented for different storage backends. -/// For now, the trait is only implemented for PostgreSQL. More backends are planned to be implemented in the future. -pub trait Queueable { - /// This method should retrieve one task of the `task_type` type. If `task_type` is `None` it will try to - /// fetch a task of the type `common`. After fetching it should update the state of the task to - /// `FangTaskState::InProgress`. - fn fetch_and_touch_task(&self, task_type: String) -> Result, QueueError>; - - /// Enqueue a task to the queue, The task will be executed as soon as possible by the worker of the same type - /// created by an `WorkerPool`. - fn insert_task(&self, params: &dyn Runnable) -> Result; - - /// The method will remove all tasks from the queue - fn remove_all_tasks(&self) -> Result; - - /// Remove all tasks that are scheduled in the future. - fn remove_all_scheduled_tasks(&self) -> Result; - - /// Removes all tasks that have the specified `task_type`. - fn remove_tasks_of_type(&self, task_type: &str) -> Result; - - /// Remove a task by its id. - fn remove_task(&self, id: Uuid) -> Result; - - /// To use this function task has to be uniq. uniq() has to return true. - /// If task is not uniq this function will not do anything. - /// Remove a task by its metadata (struct fields values) - fn remove_task_by_metadata(&self, task: &dyn Runnable) -> Result; - - fn find_task_by_id(&self, id: Uuid) -> Option; - - /// Update the state field of the specified task - /// See the `FangTaskState` enum for possible states. - fn update_task_state(&self, task: &Task, state: FangTaskState) -> Result; - - /// Update the state of a task to `FangTaskState::Failed` and set an error_message. - fn fail_task(&self, task: &Task, error: &str) -> Result; - - /// Schedule a task. - fn schedule_task(&self, task: &dyn Runnable) -> Result; - - fn schedule_retry( - &self, - task: &Task, - backoff_in_seconds: u32, - error: &str, - ) -> Result; -} - -/// An async queue that can be used to enqueue tasks. -/// It uses a PostgreSQL storage. It must be connected to perform any operation. -/// To connect a `Queue` to the PostgreSQL database call the `get_connection` method. -/// A Queue can be created with the TypedBuilder. -/// -/// ```rust -/// // Set DATABASE_URL enviroment variable if you would like to try this function. -/// pub fn connection_pool(pool_size: u32) -> r2d2::Pool> { -/// let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set"); -/// -/// let manager = r2d2::ConnectionManager::::new(database_url); -/// -/// r2d2::Pool::builder() -/// .max_size(pool_size) -/// .build(manager) -/// .unwrap() -/// } -/// -/// let queue = Queue::builder().connection_pool(connection_pool(3)).build(); -/// ``` -/// -#[derive(Clone, TypedBuilder)] -pub struct Queue { - #[builder(setter(into))] - pub connection_pool: r2d2::Pool>, -} - -impl Queueable for Queue { - fn fetch_and_touch_task(&self, task_type: String) -> Result, QueueError> { - let mut connection = self.get_connection()?; - - Self::fetch_and_touch_query(&mut connection, task_type) - } - - fn insert_task(&self, params: &dyn Runnable) -> Result { - let mut connection = self.get_connection()?; - - Self::insert_query(&mut connection, params, Utc::now()) - } - fn schedule_task(&self, params: &dyn Runnable) -> Result { - let mut connection = self.get_connection()?; - - Self::schedule_task_query(&mut connection, params) - } - - fn remove_all_scheduled_tasks(&self) -> Result { - let mut connection = self.get_connection()?; - - Self::remove_all_scheduled_tasks_query(&mut connection) - } - - fn remove_all_tasks(&self) -> Result { - let mut connection = self.get_connection()?; - - Self::remove_all_tasks_query(&mut connection) - } - - fn remove_tasks_of_type(&self, task_type: &str) -> Result { - let mut connection = self.get_connection()?; - - Self::remove_tasks_of_type_query(&mut connection, task_type) - } - - fn remove_task(&self, id: Uuid) -> Result { - let mut connection = self.get_connection()?; - - Self::remove_task_query(&mut connection, id) - } - - /// To use this function task has to be uniq. uniq() has to return true. - /// If task is not uniq this function will not do anything. - fn remove_task_by_metadata(&self, task: &dyn Runnable) -> Result { - if task.uniq() { - let mut connection = self.get_connection()?; - - Self::remove_task_by_metadata_query(&mut connection, task) - } else { - Err(QueueError::TaskNotUniqError) - } - } - - fn update_task_state(&self, task: &Task, state: FangTaskState) -> Result { - let mut connection = self.get_connection()?; - - Self::update_task_state_query(&mut connection, task, state) - } - - fn fail_task(&self, task: &Task, error: &str) -> Result { - let mut connection = self.get_connection()?; - - Self::fail_task_query(&mut connection, task, error) - } - - fn find_task_by_id(&self, id: Uuid) -> Option { - let mut connection = self.get_connection().unwrap(); - - Self::find_task_by_id_query(&mut connection, id) - } - - fn schedule_retry( - &self, - task: &Task, - backoff_seconds: u32, - error: &str, - ) -> Result { - let mut connection = self.get_connection()?; - - Self::schedule_retry_query(&mut connection, task, backoff_seconds, error) - } -} - -impl Queue { - /// Connect to the db if not connected - pub fn get_connection(&self) -> Result { - let result = self.connection_pool.get(); - - if let Err(err) = result { - log::error!("Failed to get a db connection {:?}", err); - return Err(QueueError::PoolError(err)); - } - - Ok(result.unwrap()) - } - - pub fn schedule_task_query( - connection: &mut PgConnection, - params: &dyn Runnable, - ) -> Result { - let scheduled_at = match params.cron() { - Some(scheduled) => match scheduled { - CronPattern(cron_pattern) => { - let schedule = Schedule::from_str(&cron_pattern)?; - let mut iterator = schedule.upcoming(Utc); - - iterator - .next() - .ok_or(QueueError::CronError(CronError::NoTimestampsError))? - } - ScheduleOnce(datetime) => datetime, - }, - None => { - return Err(QueueError::CronError(CronError::TaskNotSchedulableError)); - } - }; - - Self::insert_query(connection, params, scheduled_at) - } - - fn calculate_hash(json: String) -> String { - let mut hasher = Sha256::new(); - hasher.update(json.as_bytes()); - let result = hasher.finalize(); - hex::encode(result) - } - - pub fn insert_query( - connection: &mut PgConnection, - params: &dyn Runnable, - scheduled_at: DateTime, - ) -> Result { - if !params.uniq() { - let new_task = NewTask::builder() - .scheduled_at(scheduled_at) - .uniq_hash(None) - .task_type(params.task_type()) - .metadata(serde_json::to_value(params).unwrap()) - .build(); - - Ok(diesel::insert_into(fang_tasks::table) - .values(new_task) - .get_result::(connection)?) - } else { - let metadata = serde_json::to_value(params).unwrap(); - - let uniq_hash = Self::calculate_hash(metadata.to_string()); - - match Self::find_task_by_uniq_hash_query(connection, &uniq_hash) { - Some(task) => Ok(task), - None => { - let new_task = NewTask::builder() - .scheduled_at(scheduled_at) - .uniq_hash(Some(uniq_hash)) - .task_type(params.task_type()) - .metadata(serde_json::to_value(params).unwrap()) - .build(); - - Ok(diesel::insert_into(fang_tasks::table) - .values(new_task) - .get_result::(connection)?) - } - } - } - } - - pub fn fetch_task_query(connection: &mut PgConnection, task_type: String) -> Option { - Self::fetch_task_of_type_query(connection, &task_type) - } - - pub fn fetch_and_touch_query( - connection: &mut PgConnection, - task_type: String, - ) -> Result, QueueError> { - connection.transaction::, QueueError, _>(|conn| { - let found_task = Self::fetch_task_query(conn, task_type); - - if found_task.is_none() { - return Ok(None); - } - - match Self::update_task_state_query( - conn, - &found_task.unwrap(), - FangTaskState::InProgress, - ) { - Ok(updated_task) => Ok(Some(updated_task)), - Err(err) => Err(err), - } - }) - } - - pub fn find_task_by_id_query(connection: &mut PgConnection, id: Uuid) -> Option { - fang_tasks::table - .filter(fang_tasks::id.eq(id)) - .first::(connection) - .ok() - } - - pub fn remove_all_tasks_query(connection: &mut PgConnection) -> Result { - Ok(diesel::delete(fang_tasks::table).execute(connection)?) - } - - pub fn remove_all_scheduled_tasks_query( - connection: &mut PgConnection, - ) -> Result { - let query = fang_tasks::table.filter(fang_tasks::scheduled_at.gt(Utc::now())); - - Ok(diesel::delete(query).execute(connection)?) - } - - pub fn remove_tasks_of_type_query( - connection: &mut PgConnection, - task_type: &str, - ) -> Result { - let query = fang_tasks::table.filter(fang_tasks::task_type.eq(task_type)); - - Ok(diesel::delete(query).execute(connection)?) - } - - pub fn remove_task_by_metadata_query( - connection: &mut PgConnection, - task: &dyn Runnable, - ) -> Result { - let metadata = serde_json::to_value(task).unwrap(); - - let uniq_hash = Self::calculate_hash(metadata.to_string()); - - let query = fang_tasks::table.filter(fang_tasks::uniq_hash.eq(uniq_hash)); - - Ok(diesel::delete(query).execute(connection)?) - } - - pub fn remove_task_query(connection: &mut PgConnection, id: Uuid) -> Result { - let query = fang_tasks::table.filter(fang_tasks::id.eq(id)); - - Ok(diesel::delete(query).execute(connection)?) - } - - pub fn update_task_state_query( - connection: &mut PgConnection, - task: &Task, - state: FangTaskState, - ) -> Result { - Ok(diesel::update(task) - .set(( - fang_tasks::state.eq(state), - fang_tasks::updated_at.eq(Self::current_time()), - )) - .get_result::(connection)?) - } - - pub fn fail_task_query( - connection: &mut PgConnection, - task: &Task, - error: &str, - ) -> Result { - Ok(diesel::update(task) - .set(( - fang_tasks::state.eq(FangTaskState::Failed), - fang_tasks::error_message.eq(error), - fang_tasks::updated_at.eq(Self::current_time()), - )) - .get_result::(connection)?) - } - - fn current_time() -> DateTime { - Utc::now() - } - - #[cfg(test)] - pub fn connection_pool(pool_size: u32) -> r2d2::Pool> { - dotenv().ok(); - - let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set"); - - let manager = r2d2::ConnectionManager::::new(database_url); - - r2d2::Pool::builder() - .max_size(pool_size) - .build(manager) - .unwrap() - } - - fn fetch_task_of_type_query(connection: &mut PgConnection, task_type: &str) -> Option { - fang_tasks::table - .order(fang_tasks::created_at.asc()) - .order(fang_tasks::scheduled_at.asc()) - .limit(1) - .filter(fang_tasks::scheduled_at.le(Utc::now())) - .filter(fang_tasks::state.eq_any(vec![FangTaskState::New, FangTaskState::Retried])) - .filter(fang_tasks::task_type.eq(task_type)) - .for_update() - .skip_locked() - .get_result::(connection) - .ok() - } - - fn find_task_by_uniq_hash_query( - connection: &mut PgConnection, - uniq_hash: &str, - ) -> Option { - fang_tasks::table - .filter(fang_tasks::uniq_hash.eq(uniq_hash)) - .filter(fang_tasks::state.eq_any(vec![FangTaskState::New, FangTaskState::Retried])) - .first::(connection) - .ok() - } - - pub fn schedule_retry_query( - connection: &mut PgConnection, - task: &Task, - backoff_seconds: u32, - error: &str, - ) -> Result { - let now = Self::current_time(); - let scheduled_at = now + Duration::seconds(backoff_seconds as i64); - - let task = diesel::update(task) - .set(( - fang_tasks::state.eq(FangTaskState::Retried), - fang_tasks::error_message.eq(error), - fang_tasks::retries.eq(task.retries + 1), - fang_tasks::scheduled_at.eq(scheduled_at), - fang_tasks::updated_at.eq(now), - )) - .get_result::(connection)?; - - Ok(task) - } -} - -#[cfg(test)] -mod queue_tests { - use super::Queue; - use super::Queueable; - use crate::chrono::SubsecRound; - use crate::fang_task_state::FangTaskState; - use crate::runnable::Runnable; - use crate::runnable::COMMON_TYPE; - use crate::typetag; - use crate::FangError; - use crate::Scheduled; - use chrono::DateTime; - use chrono::Duration; - use chrono::Utc; - use diesel::connection::Connection; - use diesel::result::Error; - use serde::{Deserialize, Serialize}; - - #[derive(Serialize, Deserialize)] - struct PepeTask { - pub number: u16, - } - - #[typetag::serde] - impl Runnable for PepeTask { - fn run(&self, _queue: &dyn Queueable) -> Result<(), FangError> { - println!("the number is {}", self.number); - - Ok(()) - } - fn uniq(&self) -> bool { - true - } - } - - #[derive(Serialize, Deserialize)] - struct AyratTask { - pub number: u16, - } - - #[typetag::serde] - impl Runnable for AyratTask { - fn run(&self, _queue: &dyn Queueable) -> Result<(), FangError> { - println!("the number is {}", self.number); - - Ok(()) - } - fn uniq(&self) -> bool { - true - } - - fn task_type(&self) -> String { - "weirdo".to_string() - } - } - - #[derive(Serialize, Deserialize)] - struct ScheduledPepeTask { - pub number: u16, - pub datetime: String, - } - - #[typetag::serde] - impl Runnable for ScheduledPepeTask { - fn run(&self, _queue: &dyn Queueable) -> Result<(), FangError> { - println!("the number is {}", self.number); - - Ok(()) - } - fn uniq(&self) -> bool { - true - } - - fn task_type(&self) -> String { - "scheduled".to_string() - } - - fn cron(&self) -> Option { - let datetime = self.datetime.parse::>().ok()?; - Some(Scheduled::ScheduleOnce(datetime)) - } - } - - #[test] - fn insert_task_test() { - let task = PepeTask { number: 10 }; - - let pool = Queue::connection_pool(5); - - let queue = Queue::builder().connection_pool(pool).build(); - - let mut queue_pooled_connection = queue.connection_pool.get().unwrap(); - - queue_pooled_connection.test_transaction::<(), Error, _>(|conn| { - let task = Queue::insert_query(conn, &task, Utc::now()).unwrap(); - - let metadata = task.metadata.as_object().unwrap(); - let number = metadata["number"].as_u64(); - let type_task = metadata["type"].as_str(); - - assert_eq!(task.error_message, None); - assert_eq!(FangTaskState::New, task.state); - assert_eq!(Some(10), number); - assert_eq!(Some("PepeTask"), type_task); - - Ok(()) - }); - } - - #[test] - fn fetch_task_fetches_the_oldest_task() { - let task1 = PepeTask { number: 10 }; - let task2 = PepeTask { number: 11 }; - - let pool = Queue::connection_pool(5); - - let queue = Queue::builder().connection_pool(pool).build(); - - let mut queue_pooled_connection = queue.connection_pool.get().unwrap(); - - queue_pooled_connection.test_transaction::<(), Error, _>(|conn| { - let task1 = Queue::insert_query(conn, &task1, Utc::now()).unwrap(); - let _task2 = Queue::insert_query(conn, &task2, Utc::now()).unwrap(); - - let found_task = Queue::fetch_task_query(conn, COMMON_TYPE.to_string()).unwrap(); - assert_eq!(found_task.id, task1.id); - Ok(()) - }); - } - - #[test] - fn update_task_state_test() { - let task = PepeTask { number: 10 }; - - let pool = Queue::connection_pool(5); - - let queue = Queue::builder().connection_pool(pool).build(); - - let mut queue_pooled_connection = queue.connection_pool.get().unwrap(); - - queue_pooled_connection.test_transaction::<(), Error, _>(|conn| { - let task = Queue::insert_query(conn, &task, Utc::now()).unwrap(); - - let found_task = - Queue::update_task_state_query(conn, &task, FangTaskState::Finished).unwrap(); - - let metadata = found_task.metadata.as_object().unwrap(); - let number = metadata["number"].as_u64(); - let type_task = metadata["type"].as_str(); - - assert_eq!(found_task.id, task.id); - assert_eq!(found_task.state, FangTaskState::Finished); - assert_eq!(Some(10), number); - assert_eq!(Some("PepeTask"), type_task); - - Ok(()) - }); - } - - #[test] - fn fail_task_updates_state_field_and_sets_error_message() { - let task = PepeTask { number: 10 }; - - let pool = Queue::connection_pool(5); - - let queue = Queue::builder().connection_pool(pool).build(); - - let mut queue_pooled_connection = queue.connection_pool.get().unwrap(); - - queue_pooled_connection.test_transaction::<(), Error, _>(|conn| { - let task = Queue::insert_query(conn, &task, Utc::now()).unwrap(); - - let error = "Failed".to_string(); - - let found_task = Queue::fail_task_query(conn, &task, &error).unwrap(); - - let metadata = found_task.metadata.as_object().unwrap(); - let number = metadata["number"].as_u64(); - let type_task = metadata["type"].as_str(); - - assert_eq!(found_task.id, task.id); - assert_eq!(found_task.state, FangTaskState::Failed); - assert_eq!(Some(10), number); - assert_eq!(Some("PepeTask"), type_task); - assert_eq!(found_task.error_message.unwrap(), error); - - Ok(()) - }); - } - - #[test] - fn fetch_and_touch_updates_state() { - let task = PepeTask { number: 10 }; - - let pool = Queue::connection_pool(5); - - let queue = Queue::builder().connection_pool(pool).build(); - - let mut queue_pooled_connection = queue.connection_pool.get().unwrap(); - - queue_pooled_connection.test_transaction::<(), Error, _>(|conn| { - let task = Queue::insert_query(conn, &task, Utc::now()).unwrap(); - - let found_task = Queue::fetch_and_touch_query(conn, COMMON_TYPE.to_string()) - .unwrap() - .unwrap(); - - let metadata = found_task.metadata.as_object().unwrap(); - let number = metadata["number"].as_u64(); - let type_task = metadata["type"].as_str(); - - assert_eq!(found_task.id, task.id); - assert_eq!(found_task.state, FangTaskState::InProgress); - assert_eq!(Some(10), number); - assert_eq!(Some("PepeTask"), type_task); - - Ok(()) - }); - } - - #[test] - fn fetch_and_touch_returns_none() { - let pool = Queue::connection_pool(5); - - let queue = Queue::builder().connection_pool(pool).build(); - - let mut queue_pooled_connection = queue.connection_pool.get().unwrap(); - - queue_pooled_connection.test_transaction::<(), Error, _>(|conn| { - let found_task = Queue::fetch_and_touch_query(conn, COMMON_TYPE.to_string()).unwrap(); - - assert_eq!(None, found_task); - - Ok(()) - }); - } - - #[test] - fn insert_task_uniq_test() { - let task = PepeTask { number: 10 }; - - let pool = Queue::connection_pool(5); - - let queue = Queue::builder().connection_pool(pool).build(); - - let mut queue_pooled_connection = queue.connection_pool.get().unwrap(); - - queue_pooled_connection.test_transaction::<(), Error, _>(|conn| { - let task1 = Queue::insert_query(conn, &task, Utc::now()).unwrap(); - let task2 = Queue::insert_query(conn, &task, Utc::now()).unwrap(); - - assert_eq!(task2.id, task1.id); - Ok(()) - }); - } - - #[test] - fn schedule_task_test() { - let pool = Queue::connection_pool(5); - let queue = Queue::builder().connection_pool(pool).build(); - - let mut queue_pooled_connection = queue.connection_pool.get().unwrap(); - - queue_pooled_connection.test_transaction::<(), Error, _>(|conn| { - let datetime = (Utc::now() + Duration::seconds(7)).round_subsecs(0); - - let task = &ScheduledPepeTask { - number: 10, - datetime: datetime.to_string(), - }; - - let task = Queue::schedule_task_query(conn, task).unwrap(); - - let metadata = task.metadata.as_object().unwrap(); - let number = metadata["number"].as_u64(); - let type_task = metadata["type"].as_str(); - - assert_eq!(Some(10), number); - assert_eq!(Some("ScheduledPepeTask"), type_task); - assert_eq!(task.scheduled_at, datetime); - - Ok(()) - }); - } - - #[test] - fn remove_all_scheduled_tasks_test() { - let pool = Queue::connection_pool(5); - let queue = Queue::builder().connection_pool(pool).build(); - - let mut queue_pooled_connection = queue.connection_pool.get().unwrap(); - - queue_pooled_connection.test_transaction::<(), Error, _>(|conn| { - let datetime = (Utc::now() + Duration::seconds(7)).round_subsecs(0); - - let task1 = &ScheduledPepeTask { - number: 10, - datetime: datetime.to_string(), - }; - - let task2 = &ScheduledPepeTask { - number: 11, - datetime: datetime.to_string(), - }; - - Queue::schedule_task_query(conn, task1).unwrap(); - Queue::schedule_task_query(conn, task2).unwrap(); - - let number = Queue::remove_all_scheduled_tasks_query(conn).unwrap(); - - assert_eq!(2, number); - - Ok(()) - }); - } - - #[test] - fn remove_all_tasks_test() { - let task1 = PepeTask { number: 10 }; - let task2 = PepeTask { number: 11 }; - - let pool = Queue::connection_pool(5); - - let queue = Queue::builder().connection_pool(pool).build(); - - let mut queue_pooled_connection = queue.connection_pool.get().unwrap(); - - queue_pooled_connection.test_transaction::<(), Error, _>(|conn| { - let task1 = Queue::insert_query(conn, &task1, Utc::now()).unwrap(); - let task2 = Queue::insert_query(conn, &task2, Utc::now()).unwrap(); - - let result = Queue::remove_all_tasks_query(conn).unwrap(); - - assert_eq!(2, result); - assert_eq!(None, Queue::find_task_by_id_query(conn, task1.id)); - assert_eq!(None, Queue::find_task_by_id_query(conn, task2.id)); - - Ok(()) - }); - } - - #[test] - fn remove_task() { - let task1 = PepeTask { number: 10 }; - let task2 = PepeTask { number: 11 }; - - let pool = Queue::connection_pool(5); - - let queue = Queue::builder().connection_pool(pool).build(); - - let mut queue_pooled_connection = queue.connection_pool.get().unwrap(); - - queue_pooled_connection.test_transaction::<(), Error, _>(|conn| { - let task1 = Queue::insert_query(conn, &task1, Utc::now()).unwrap(); - let task2 = Queue::insert_query(conn, &task2, Utc::now()).unwrap(); - - assert!(Queue::find_task_by_id_query(conn, task1.id).is_some()); - assert!(Queue::find_task_by_id_query(conn, task2.id).is_some()); - - Queue::remove_task_query(conn, task1.id).unwrap(); - - assert!(Queue::find_task_by_id_query(conn, task1.id).is_none()); - assert!(Queue::find_task_by_id_query(conn, task2.id).is_some()); - - Ok(()) - }); - } - - #[test] - fn remove_task_of_type() { - let task1 = PepeTask { number: 10 }; - let task2 = AyratTask { number: 10 }; - - let pool = Queue::connection_pool(5); - - let queue = Queue::builder().connection_pool(pool).build(); - - let mut queue_pooled_connection = queue.connection_pool.get().unwrap(); - - queue_pooled_connection.test_transaction::<(), Error, _>(|conn| { - let task1 = Queue::insert_query(conn, &task1, Utc::now()).unwrap(); - let task2 = Queue::insert_query(conn, &task2, Utc::now()).unwrap(); - - assert!(Queue::find_task_by_id_query(conn, task1.id).is_some()); - assert!(Queue::find_task_by_id_query(conn, task2.id).is_some()); - - Queue::remove_tasks_of_type_query(conn, "weirdo").unwrap(); - - assert!(Queue::find_task_by_id_query(conn, task1.id).is_some()); - assert!(Queue::find_task_by_id_query(conn, task2.id).is_none()); - - Ok(()) - }); - } - - #[test] - fn remove_task_by_metadata() { - let m_task1 = PepeTask { number: 10 }; - let m_task2 = PepeTask { number: 11 }; - let m_task3 = AyratTask { number: 10 }; - - let pool = Queue::connection_pool(5); - - let queue = Queue::builder().connection_pool(pool).build(); - - let mut queue_pooled_connection = queue.connection_pool.get().unwrap(); - - queue_pooled_connection.test_transaction::<(), Error, _>(|conn| { - let task1 = Queue::insert_query(conn, &m_task1, Utc::now()).unwrap(); - let task2 = Queue::insert_query(conn, &m_task2, Utc::now()).unwrap(); - let task3 = Queue::insert_query(conn, &m_task3, Utc::now()).unwrap(); - - assert!(Queue::find_task_by_id_query(conn, task1.id).is_some()); - assert!(Queue::find_task_by_id_query(conn, task2.id).is_some()); - assert!(Queue::find_task_by_id_query(conn, task3.id).is_some()); - - Queue::remove_task_by_metadata_query(conn, &m_task1).unwrap(); - - assert!(Queue::find_task_by_id_query(conn, task1.id).is_none()); - assert!(Queue::find_task_by_id_query(conn, task2.id).is_some()); - assert!(Queue::find_task_by_id_query(conn, task3.id).is_some()); - - Ok(()) - }); - } -} diff --git a/src/blocking/runnable.rs b/src/blocking/runnable.rs deleted file mode 100644 index 01a6f69..0000000 --- a/src/blocking/runnable.rs +++ /dev/null @@ -1,53 +0,0 @@ -use crate::queue::Queueable; -use crate::FangError; -use crate::Scheduled; - -pub const COMMON_TYPE: &str = "common"; -pub const RETRIES_NUMBER: i32 = 20; - -/// Implement this trait to run your custom tasks. -#[typetag::serde(tag = "type")] -pub trait Runnable { - /// Execute the task. This method should define its logic - fn run(&self, _queueable: &dyn Queueable) -> Result<(), FangError>; - - /// Define the type of the task. - /// The `common` task type is used by default - fn task_type(&self) -> String { - COMMON_TYPE.to_string() - } - - /// If set to true, no new tasks with the same metadata will be inserted - /// By default it is set to false. - fn uniq(&self) -> bool { - false - } - - /// This method defines if a task is periodic or it should be executed once in the future. - /// - /// Be careful it works only with the UTC timezone. - /** - ```rust - fn cron(&self) -> Option { - let expression = "0/20 * * * Aug-Sep * 2022/1"; - Some(Scheduled::CronPattern(expression.to_string())) - } - ``` - */ - /// In order to schedule a task once, use the `Scheduled::ScheduleOnce` enum variant. - fn cron(&self) -> Option { - None - } - - /// Define the maximum number of retries the task will be retried. - /// By default the number of retries is 20. - fn max_retries(&self) -> i32 { - RETRIES_NUMBER - } - - /// Define the backoff mode - /// By default, it is exponential, 2^(attempt) - fn backoff(&self, attempt: u32) -> u32 { - u32::pow(2, attempt) - } -} diff --git a/src/blocking/schema.rs b/src/blocking/schema.rs deleted file mode 100644 index 4f9dff2..0000000 --- a/src/blocking/schema.rs +++ /dev/null @@ -1,25 +0,0 @@ -// @generated automatically by Diesel CLI. - -pub mod sql_types { - #[derive(diesel::sql_types::SqlType)] - #[diesel(postgres_type(name = "fang_task_state"))] - pub struct FangTaskState; -} - -diesel::table! { - use diesel::sql_types::*; - use super::sql_types::FangTaskState; - - fang_tasks (id) { - id -> Uuid, - metadata -> Jsonb, - error_message -> Nullable, - state -> FangTaskState, - task_type -> Varchar, - uniq_hash -> Nullable, - retries -> Int4, - scheduled_at -> Timestamptz, - created_at -> Timestamptz, - updated_at -> Timestamptz, - } -} diff --git a/src/blocking/worker.rs b/src/blocking/worker.rs deleted file mode 100644 index a114f7e..0000000 --- a/src/blocking/worker.rs +++ /dev/null @@ -1,411 +0,0 @@ -#![allow(clippy::borrowed_box)] -#![allow(clippy::unnecessary_unwrap)] - -use crate::fang_task_state::FangTaskState; -use crate::queue::Queueable; -use crate::queue::Task; -use crate::runnable::Runnable; -use crate::runnable::COMMON_TYPE; -use crate::FangError; -use crate::Scheduled::*; -use crate::{RetentionMode, SleepParams}; -use log::error; -use std::thread; -use typed_builder::TypedBuilder; - -/// A executioner of tasks, it executes tasks only of one given task_type, it sleeps when they are -/// not tasks to be executed. -#[derive(TypedBuilder)] -pub struct Worker -where - BQueue: Queueable + Clone + Sync + Send + 'static, -{ - #[builder(setter(into))] - pub queue: BQueue, - #[builder(default=COMMON_TYPE.to_string(), setter(into))] - pub task_type: String, - #[builder(default, setter(into))] - pub sleep_params: SleepParams, - #[builder(default, setter(into))] - pub retention_mode: RetentionMode, -} - -impl Worker -where - BQueue: Queueable + Clone + Sync + Send + 'static, -{ - pub fn run(&self, task: Task) { - let runnable: Box = serde_json::from_value(task.metadata.clone()).unwrap(); - let result = runnable.run(&self.queue); - - match result { - Ok(_) => self.finalize_task(task, &result), - Err(ref error) => { - if task.retries < runnable.max_retries() { - let backoff_seconds = runnable.backoff(task.retries as u32); - - self.queue - .schedule_retry(&task, backoff_seconds, &error.description) - .expect("Failed to retry"); - } else { - self.finalize_task(task, &result); - } - } - } - } - - pub(crate) fn run_tasks(&mut self) -> Result<(), FangError> { - loop { - match self.queue.fetch_and_touch_task(self.task_type.clone()) { - Ok(Some(task)) => { - let actual_task: Box = - serde_json::from_value(task.metadata.clone()).unwrap(); - - // check if task is scheduled or not - if let Some(CronPattern(_)) = actual_task.cron() { - // program task - self.queue.schedule_task(&*actual_task)?; - } - - self.maybe_reset_sleep_period(); - self.run(task); - } - Ok(None) => { - self.sleep(); - } - - Err(error) => { - error!("Failed to fetch a task {:?}", error); - - self.sleep(); - } - }; - } - } - - #[cfg(test)] - pub fn run_tasks_until_none(&mut self) -> Result<(), FangError> { - loop { - match self.queue.fetch_and_touch_task(self.task_type.clone()) { - Ok(Some(task)) => { - let actual_task: Box = - serde_json::from_value(task.metadata.clone()).unwrap(); - - // check if task is scheduled or not - if let Some(CronPattern(_)) = actual_task.cron() { - // program task - self.queue.schedule_task(&*actual_task)?; - } - - self.maybe_reset_sleep_period(); - self.run(task); - } - Ok(None) => { - return Ok(()); - } - Err(error) => { - error!("Failed to fetch a task {:?}", error); - - self.sleep(); - } - }; - } - } - - pub fn maybe_reset_sleep_period(&mut self) { - self.sleep_params.maybe_reset_sleep_period(); - } - - fn sleep(&mut self) { - self.sleep_params.maybe_increase_sleep_period(); - - thread::sleep(self.sleep_params.sleep_period); - } - - fn finalize_task(&self, task: Task, result: &Result<(), FangError>) { - match self.retention_mode { - RetentionMode::KeepAll => { - match result { - Ok(_) => self - .queue - .update_task_state(&task, FangTaskState::Finished) - .unwrap(), - Err(error) => self.queue.fail_task(&task, &error.description).unwrap(), - }; - } - - RetentionMode::RemoveAll => { - self.queue.remove_task(task.id).unwrap(); - } - - RetentionMode::RemoveFinished => match result { - Ok(_) => { - self.queue.remove_task(task.id).unwrap(); - } - Err(error) => { - self.queue.fail_task(&task, &error.description).unwrap(); - } - }, - } - } -} - -#[cfg(test)] -mod worker_tests { - use super::RetentionMode; - use super::Runnable; - use super::Worker; - use crate::fang_task_state::FangTaskState; - use crate::queue::Queue; - use crate::queue::Queueable; - use crate::typetag; - use crate::FangError; - use chrono::Utc; - use serde::{Deserialize, Serialize}; - - #[derive(Serialize, Deserialize)] - struct WorkerTaskTest { - pub number: u16, - } - - #[typetag::serde] - impl Runnable for WorkerTaskTest { - fn run(&self, _queue: &dyn Queueable) -> Result<(), FangError> { - println!("the number is {}", self.number); - - Ok(()) - } - - fn task_type(&self) -> String { - "worker_task".to_string() - } - } - - #[derive(Serialize, Deserialize)] - struct FailedTask { - pub number: u16, - } - - #[typetag::serde] - impl Runnable for FailedTask { - fn run(&self, _queue: &dyn Queueable) -> Result<(), FangError> { - let message = format!("the number is {}", self.number); - - Err(FangError { - description: message, - }) - } - - fn max_retries(&self) -> i32 { - 0 - } - - fn task_type(&self) -> String { - "F_task".to_string() - } - } - - #[derive(Serialize, Deserialize)] - struct RetryTask { - pub number: u16, - } - - #[typetag::serde] - impl Runnable for RetryTask { - fn run(&self, _queue: &dyn Queueable) -> Result<(), FangError> { - let message = format!("Saving Pepe. Attempt {}", self.number); - - Err(FangError { - description: message, - }) - } - - fn max_retries(&self) -> i32 { - 2 - } - - fn task_type(&self) -> String { - "Retry_task".to_string() - } - } - - #[derive(Serialize, Deserialize)] - struct TaskType1 {} - - #[typetag::serde] - impl Runnable for TaskType1 { - fn run(&self, _queue: &dyn Queueable) -> Result<(), FangError> { - Ok(()) - } - - fn task_type(&self) -> String { - "type1".to_string() - } - } - - #[derive(Serialize, Deserialize)] - struct TaskType2 {} - - #[typetag::serde] - impl Runnable for TaskType2 { - fn run(&self, _queue: &dyn Queueable) -> Result<(), FangError> { - Ok(()) - } - - fn task_type(&self) -> String { - "type2".to_string() - } - } - - // Worker tests has to commit because the worker operations commits - #[test] - #[ignore] - fn executes_and_finishes_task() { - let task = WorkerTaskTest { number: 10 }; - - let pool = Queue::connection_pool(5); - - let queue = Queue::builder().connection_pool(pool).build(); - - let worker = Worker::::builder() - .queue(queue) - .retention_mode(RetentionMode::KeepAll) - .task_type(task.task_type()) - .build(); - let mut pooled_connection = worker.queue.connection_pool.get().unwrap(); - - let task = Queue::insert_query(&mut pooled_connection, &task, Utc::now()).unwrap(); - - assert_eq!(FangTaskState::New, task.state); - - // this operation commits and thats why need to commit this test - worker.run(task.clone()); - - let found_task = Queue::find_task_by_id_query(&mut pooled_connection, task.id).unwrap(); - - assert_eq!(FangTaskState::Finished, found_task.state); - - Queue::remove_tasks_of_type_query(&mut pooled_connection, "worker_task").unwrap(); - } - - #[test] - #[ignore] - fn executes_task_only_of_specific_type() { - let task1 = TaskType1 {}; - let task2 = TaskType2 {}; - - let pool = Queue::connection_pool(5); - - let queue = Queue::builder().connection_pool(pool).build(); - - let mut worker = Worker::::builder() - .queue(queue) - .task_type(task1.task_type()) - .retention_mode(RetentionMode::KeepAll) - .build(); - - let mut pooled_connection = worker.queue.connection_pool.get().unwrap(); - - let task1 = Queue::insert_query(&mut pooled_connection, &task1, Utc::now()).unwrap(); - let task2 = Queue::insert_query(&mut pooled_connection, &task2, Utc::now()).unwrap(); - - assert_eq!(FangTaskState::New, task1.state); - assert_eq!(FangTaskState::New, task2.state); - - worker.run_tasks_until_none().unwrap(); - - std::thread::sleep(std::time::Duration::from_millis(1000)); - - let found_task1 = Queue::find_task_by_id_query(&mut pooled_connection, task1.id).unwrap(); - assert_eq!(FangTaskState::Finished, found_task1.state); - - let found_task2 = Queue::find_task_by_id_query(&mut pooled_connection, task2.id).unwrap(); - assert_eq!(FangTaskState::New, found_task2.state); - - Queue::remove_tasks_of_type_query(&mut pooled_connection, "type1").unwrap(); - Queue::remove_tasks_of_type_query(&mut pooled_connection, "type2").unwrap(); - } - - #[test] - #[ignore] - fn saves_error_for_failed_task() { - let task = FailedTask { number: 10 }; - - let pool = Queue::connection_pool(5); - - let queue = Queue::builder().connection_pool(pool).build(); - - let worker = Worker::::builder() - .queue(queue) - .retention_mode(RetentionMode::KeepAll) - .task_type(task.task_type()) - .build(); - - let mut pooled_connection = worker.queue.connection_pool.get().unwrap(); - - let task = Queue::insert_query(&mut pooled_connection, &task, Utc::now()).unwrap(); - - assert_eq!(FangTaskState::New, task.state); - - worker.run(task.clone()); - - let found_task = Queue::find_task_by_id_query(&mut pooled_connection, task.id).unwrap(); - - assert_eq!(FangTaskState::Failed, found_task.state); - assert_eq!( - "the number is 10".to_string(), - found_task.error_message.unwrap() - ); - - Queue::remove_tasks_of_type_query(&mut pooled_connection, "F_task").unwrap(); - } - - #[test] - #[ignore] - fn retries_task() { - let task = RetryTask { number: 10 }; - - let pool = Queue::connection_pool(5); - - let queue = Queue::builder().connection_pool(pool).build(); - - let mut worker = Worker::::builder() - .queue(queue) - .retention_mode(RetentionMode::KeepAll) - .task_type(task.task_type()) - .build(); - - let mut pooled_connection = worker.queue.connection_pool.get().unwrap(); - - let task = Queue::insert_query(&mut pooled_connection, &task, Utc::now()).unwrap(); - - assert_eq!(FangTaskState::New, task.state); - - worker.run(task.clone()); - - std::thread::sleep(std::time::Duration::from_millis(1000)); - - let found_task = Queue::find_task_by_id_query(&mut pooled_connection, task.id).unwrap(); - - assert_eq!(FangTaskState::Retried, found_task.state); - assert_eq!(1, found_task.retries); - - worker.run_tasks_until_none().unwrap(); - - std::thread::sleep(std::time::Duration::from_millis(14000)); - - worker.run_tasks_until_none().unwrap(); - - let found_task = Queue::find_task_by_id_query(&mut pooled_connection, task.id).unwrap(); - - assert_eq!(FangTaskState::Failed, found_task.state); - assert_eq!(2, found_task.retries); - - assert_eq!( - "Saving Pepe. Attempt 10".to_string(), - found_task.error_message.unwrap() - ); - - Queue::remove_tasks_of_type_query(&mut pooled_connection, "Retry_task").unwrap(); - } -} diff --git a/src/blocking/worker_pool.rs b/src/blocking/worker_pool.rs deleted file mode 100644 index 266e6a7..0000000 --- a/src/blocking/worker_pool.rs +++ /dev/null @@ -1,122 +0,0 @@ -use crate::queue::Queueable; -use crate::worker::Worker; -use crate::FangError; -use crate::RetentionMode; -use crate::SleepParams; -use log::error; -use log::info; -use std::thread; -use typed_builder::TypedBuilder; - -#[derive(Clone, TypedBuilder)] -pub struct WorkerPool -where - BQueue: Queueable + Clone + Sync + Send + 'static, -{ - /// the AsyncWorkerPool uses a queue to control the tasks that will be executed. - #[builder(setter(into))] - pub queue: BQueue, - /// sleep_params controls how much time a worker will sleep while waiting for tasks - /// execute. - #[builder(setter(into), default)] - pub sleep_params: SleepParams, - /// retention_mode controls if tasks should be persisted after execution - #[builder(setter(into), default)] - pub retention_mode: RetentionMode, - /// the number of workers of the AsyncWorkerPool. - #[builder(setter(into))] - pub number_of_workers: u32, - /// The type of tasks that will be executed by `AsyncWorkerPool`. - #[builder(setter(into), default)] - pub task_type: String, -} - -#[derive(Clone, TypedBuilder)] -pub struct WorkerThread -where - BQueue: Queueable + Clone + Sync + Send + 'static, -{ - pub name: String, - pub restarts: u64, - pub worker_pool: WorkerPool, -} - -#[derive(Clone)] -pub struct WorkerParams { - pub retention_mode: Option, - pub sleep_params: Option, - pub task_type: Option, -} - -impl WorkerPool -where - BQueue: Queueable + Clone + Sync + Send + 'static, -{ - /// Starts the configured number of workers - /// This is necessary in order to execute tasks. - pub fn start(&mut self) -> Result<(), FangError> { - for idx in 1..self.number_of_workers + 1 { - let name = format!("worker_{}{idx}", self.task_type); - - let worker_thread = WorkerThread::builder() - .name(name.clone()) - .restarts(0) - .worker_pool(self.clone()) - .build(); - - worker_thread.spawn()?; - } - Ok(()) - } -} - -impl WorkerThread -where - BQueue: Queueable + Clone + Sync + Send + 'static, -{ - fn spawn(self) -> Result<(), FangError> { - info!( - "starting a worker thread {}, number of restarts {}", - self.name, self.restarts - ); - - let builder = thread::Builder::new().name(self.name.clone()); - - builder - .spawn(move || { - let mut worker: Worker = Worker::builder() - .queue(self.worker_pool.queue.clone()) - .task_type(self.worker_pool.task_type.clone()) - .retention_mode(self.worker_pool.retention_mode.clone()) - .sleep_params(self.worker_pool.sleep_params.clone()) - .build(); - - // Run worker - if let Err(error) = worker.run_tasks() { - error!( - "Error executing tasks in worker '{}': {:?}", - self.name, error - ); - } - }) - .map_err(FangError::from)?; - - Ok(()) - } -} - -impl Drop for WorkerThread -where - BQueue: Queueable + Clone + Sync + Send + 'static, -{ - fn drop(&mut self) { - self.restarts += 1; - - error!( - "Worker {} stopped. Restarting. The number of restarts {}", - self.name, self.restarts, - ); - - self.clone().spawn().unwrap(); - } -} diff --git a/src/errors.rs b/src/errors.rs new file mode 100644 index 0000000..7687245 --- /dev/null +++ b/src/errors.rs @@ -0,0 +1,48 @@ +use thiserror::Error; + +/// An error that can happen during executing of tasks +#[derive(Debug)] +pub struct FangError { + /// A description of an error + pub description: String, +} + +/// List of error types that can occur while working with cron schedules. +#[derive(Debug, Error)] +pub enum CronError { + /// A problem occured during cron schedule parsing. + #[error(transparent)] + LibraryError(#[from] cron::error::Error), + /// [`Scheduled`] enum variant is not provided + #[error("You have to implement method `cron()` in your AsyncRunnable")] + TaskNotSchedulableError, + /// The next execution can not be determined using the current [`Scheduled::CronPattern`] + #[error("No timestamps match with this cron pattern")] + NoTimestampsError, +} + +#[derive(Debug, Error)] +pub enum AsyncQueueError { + #[error(transparent)] + PgError(#[from] diesel::result::Error), + #[error(transparent)] + SerdeError(#[from] serde_json::Error), + #[error(transparent)] + CronError(#[from] CronError), + #[error("returned invalid result (expected {expected:?}, found {found:?})")] + ResultError { expected: u64, found: u64 }, + #[error( + "AsyncQueue is not connected :( , call connect() method first and then perform operations" + )] + NotConnectedError, + #[error("Can not convert `std::time::Duration` to `chrono::Duration`")] + TimeError, + #[error("Can not perform this operation if task is not uniq, please check its definition in impl AsyncRunnable")] + TaskNotUniqError, +} + +impl From for AsyncQueueError { + fn from(error: cron::error::Error) -> Self { + AsyncQueueError::CronError(CronError::LibraryError(error)) + } +} diff --git a/src/blocking/fang_task_state.rs b/src/fang_task_state.rs similarity index 100% rename from src/blocking/fang_task_state.rs rename to src/fang_task_state.rs diff --git a/src/lib.rs b/src/lib.rs index afbe235..739ddbf 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,7 +1,7 @@ #![doc = include_str!("../README.md")] use std::time::Duration; -use thiserror::Error; +use chrono::{DateTime, Utc}; use typed_builder::TypedBuilder; /// Represents a schedule for scheduled tasks. @@ -19,20 +19,6 @@ pub enum Scheduled { ScheduleOnce(DateTime), } -/// List of error types that can occur while working with cron schedules. -#[derive(Debug, Error)] -pub enum CronError { - /// A problem occured during cron schedule parsing. - #[error(transparent)] - LibraryError(#[from] cron::error::Error), - /// [`Scheduled`] enum variant is not provided - #[error("You have to implement method `cron()` in your AsyncRunnable")] - TaskNotSchedulableError, - /// The next execution can not be determined using the current [`Scheduled::CronPattern`] - #[error("No timestamps match with this cron pattern")] - NoTimestampsError, -} - /// All possible options for retaining tasks in the db after their execution. /// /// The default mode is [`RetentionMode::RemoveAll`] @@ -94,54 +80,12 @@ impl Default for SleepParams { } } -/// An error that can happen during executing of tasks -#[derive(Debug)] -pub struct FangError { - /// A description of an error - pub description: String, -} - -#[doc(hidden)] -#[cfg(feature = "blocking")] -extern crate diesel; - -#[doc(hidden)] -#[cfg(feature = "blocking")] -pub use diesel::pg::PgConnection; - -#[doc(hidden)] -pub use typetag; - -#[doc(hidden)] -pub extern crate serde; - -#[doc(hidden)] -pub extern crate chrono; - -#[doc(hidden)] -pub use serde_derive::{Deserialize, Serialize}; - -#[doc(hidden)] -pub use chrono::DateTime; -#[doc(hidden)] -pub use chrono::Utc; - -#[cfg(feature = "blocking")] -pub mod blocking; - -#[cfg(feature = "blocking")] -pub use blocking::*; - -#[cfg(feature = "asynk")] -pub mod asynk; - -#[cfg(feature = "asynk")] -pub use asynk::*; - -#[cfg(feature = "asynk")] -#[doc(hidden)] -pub use bb8_postgres::tokio_postgres::tls::NoTls; - -#[cfg(feature = "asynk")] -#[doc(hidden)] -pub use async_trait::async_trait; +pub mod fang_task_state; +pub mod schema; +pub mod task; +pub mod queue; +mod queries; +pub mod errors; +pub mod runnable; +pub mod worker; +pub mod worker_pool; diff --git a/src/queries.rs b/src/queries.rs new file mode 100644 index 0000000..82248d3 --- /dev/null +++ b/src/queries.rs @@ -0,0 +1,205 @@ +use crate::runnable::AsyncRunnable; +use crate::fang_task_state::FangTaskState; +use crate::schema::fang_tasks; +use crate::errors::CronError; +use crate::Scheduled::*; +use crate::task::{DEFAULT_TASK_TYPE, Task}; +use async_trait::async_trait; +use chrono::DateTime; +use chrono::Duration; +use chrono::Utc; +use cron::Schedule; +use diesel::prelude::*; +use diesel::result::Error::QueryBuilderError; +use diesel::ExpressionMethods; +use diesel_async::scoped_futures::ScopedFutureExt; +use diesel_async::AsyncConnection; +use diesel_async::{pg::AsyncPgConnection, pooled_connection::bb8::Pool, pooled_connection::bb8::PooledConnection, RunQueryDsl}; +use diesel_async::pooled_connection::PoolableConnection; +use sha2::{Digest, Sha256}; +use std::str::FromStr; +use typed_builder::TypedBuilder; +use uuid::Uuid; +use crate::task::NewTask; +use crate::errors::AsyncQueueError; + + +impl Task { + pub async fn remove_all_scheduled_tasks( + connection: &mut AsyncPgConnection, + ) -> Result { + let query = fang_tasks::table.filter(fang_tasks::scheduled_at.gt(Utc::now())); + Ok(diesel::delete(query).execute(connection).await? as u64) + } + + pub async fn remove_task( + connection: &mut AsyncPgConnection, + id: Uuid, + ) -> Result { + let query = fang_tasks::table.filter(fang_tasks::id.eq(id)); + Ok(diesel::delete(query).execute(connection).await? as u64) + } + + pub async fn remove_task_by_metadata( + connection: &mut AsyncPgConnection, + task: &dyn AsyncRunnable, + ) -> Result { + let metadata = serde_json::to_value(task)?; + + let uniq_hash = Self::calculate_hash(metadata.to_string()); + + let query = fang_tasks::table.filter(fang_tasks::uniq_hash.eq(uniq_hash)); + + Ok(diesel::delete(query).execute(connection).await? as u64) + } + + pub async fn remove_tasks_type( + connection: &mut AsyncPgConnection, + task_type: &str, + ) -> Result { + let query = fang_tasks::table.filter(fang_tasks::task_type.eq(task_type)); + Ok(diesel::delete(query).execute(connection).await? as u64) + } + + pub async fn find_task_by_id( + connection: &mut AsyncPgConnection, + id: Uuid, + ) -> Result { + let task = fang_tasks::table + .filter(fang_tasks::id.eq(id)) + .first::(connection) + .await?; + Ok(task) + } + + pub async fn fail_task( + connection: &mut AsyncPgConnection, + task: Task, + error_message: &str, + ) -> Result { + Ok(diesel::update(&task) + .set(( + fang_tasks::state.eq(FangTaskState::Failed), + fang_tasks::error_message.eq(error_message), + fang_tasks::updated_at.eq(Utc::now()), + )) + .get_result::(connection) + .await?) + } + + pub async fn schedule_retry( + connection: &mut AsyncPgConnection, + task: &Task, + backoff_seconds: u32, + error: &str, + ) -> Result { + let now = Utc::now(); + let scheduled_at = now + Duration::seconds(backoff_seconds as i64); + + let task = diesel::update(task) + .set(( + fang_tasks::state.eq(FangTaskState::Retried), + fang_tasks::error_message.eq(error), + fang_tasks::retries.eq(task.retries + 1), + fang_tasks::scheduled_at.eq(scheduled_at), + fang_tasks::updated_at.eq(now), + )) + .get_result::(connection) + .await?; + + Ok(task) + } + + pub async fn fetch_task_of_type( + connection: &mut AsyncPgConnection, + task_type: Option, + ) -> Option { + fang_tasks::table + .order(fang_tasks::created_at.asc()) + .order(fang_tasks::scheduled_at.asc()) + .limit(1) + .filter(fang_tasks::scheduled_at.le(Utc::now())) + .filter(fang_tasks::state.eq_any(vec![FangTaskState::New, FangTaskState::Retried])) + .filter(fang_tasks::task_type.eq(task_type.unwrap_or_else(|| DEFAULT_TASK_TYPE.to_string()))) + .for_update() + .skip_locked() + .get_result::(connection) + .await + .ok() + } + + pub async fn update_task_state( + connection: &mut AsyncPgConnection, + task: Task, + state: FangTaskState, + ) -> Result { + let updated_at = Utc::now(); + Ok(diesel::update(&task) + .set(( + fang_tasks::state.eq(state), + fang_tasks::updated_at.eq(updated_at), + )) + .get_result::(connection) + .await?) + } + + pub async fn insert_task( + connection: &mut AsyncPgConnection, + params: &dyn AsyncRunnable, + scheduled_at: DateTime, + ) -> Result { + if !params.uniq() { + let new_task = NewTask::builder() + .scheduled_at(scheduled_at) + .uniq_hash(None) + .task_type(params.task_type()) + .metadata(serde_json::to_value(params).unwrap()) + .build(); + + Ok(diesel::insert_into(fang_tasks::table) + .values(new_task) + .get_result::(connection) + .await?) + } else { + let metadata = serde_json::to_value(params).unwrap(); + + let uniq_hash = Self::calculate_hash(metadata.to_string()); + + match Self::find_task_by_uniq_hash(connection, &uniq_hash).await { + Some(task) => Ok(task), + None => { + let new_task = NewTask::builder() + .scheduled_at(scheduled_at) + .uniq_hash(Some(uniq_hash)) + .task_type(params.task_type()) + .metadata(serde_json::to_value(params).unwrap()) + .build(); + + Ok(diesel::insert_into(fang_tasks::table) + .values(new_task) + .get_result::(connection) + .await?) + } + } + } + } + + fn calculate_hash(json: String) -> String { + let mut hasher = Sha256::new(); + hasher.update(json.as_bytes()); + let result = hasher.finalize(); + hex::encode(result) + } + + pub async fn find_task_by_uniq_hash( + connection: &mut AsyncPgConnection, + uniq_hash: &str, + ) -> Option { + fang_tasks::table + .filter(fang_tasks::uniq_hash.eq(uniq_hash)) + .filter(fang_tasks::state.eq_any(vec![FangTaskState::New, FangTaskState::Retried])) + .first::(connection) + .await + .ok() + } +} diff --git a/src/queue.rs b/src/queue.rs new file mode 100644 index 0000000..f9bb13c --- /dev/null +++ b/src/queue.rs @@ -0,0 +1,635 @@ +use crate::runnable::AsyncRunnable; +use crate::fang_task_state::FangTaskState; +use crate::schema::fang_tasks; +use crate::errors::CronError; +use crate::Scheduled::*; +use crate::task::{DEFAULT_TASK_TYPE, Task}; +use async_trait::async_trait; +use chrono::DateTime; +use chrono::Duration; +use chrono::Utc; +use crate::task::NewTask; +use cron::Schedule; +use diesel::result::Error::QueryBuilderError; +use diesel::ExpressionMethods; +use diesel_async::scoped_futures::ScopedFutureExt; +use diesel_async::AsyncConnection; +use diesel_async::{pg::AsyncPgConnection, pooled_connection::bb8::Pool, pooled_connection::AsyncDieselConnectionManager, RunQueryDsl}; +use sha2::{Sha256}; +use std::str::FromStr; +use diesel_async::pooled_connection::PoolableConnection; +use thiserror::Error; +use crate::errors::AsyncQueueError; +use typed_builder::TypedBuilder; +use uuid::Uuid; + + +/// This trait defines operations for an asynchronous queue. +/// The trait can be implemented for different storage backends. +/// For now, the trait is only implemented for PostgreSQL. More backends are planned to be implemented in the future. + +#[async_trait] +pub trait AsyncQueueable: Send { + /// This method should retrieve one task of the `task_type` type. If `task_type` is `None` it will try to + /// fetch a task of the type `common`. After fetching it should update the state of the task to + /// `FangTaskState::InProgress`. + /// + async fn fetch_and_touch_task( + &mut self, + task_type: Option, + ) -> Result, AsyncQueueError>; + + /// Enqueue a task to the queue, The task will be executed as soon as possible by the worker of the same type + /// created by an AsyncWorkerPool. + async fn insert_task(&mut self, task: &dyn AsyncRunnable) -> Result; + + /// The method will remove all tasks from the queue + async fn remove_all_tasks(&mut self) -> Result; + + /// Remove all tasks that are scheduled in the future. + async fn remove_all_scheduled_tasks(&mut self) -> Result; + + /// Remove a task by its id. + async fn remove_task(&mut self, id: Uuid) -> Result; + + /// Remove a task by its metadata (struct fields values) + async fn remove_task_by_metadata( + &mut self, + task: &dyn AsyncRunnable, + ) -> Result; + + /// Removes all tasks that have the specified `task_type`. + async fn remove_tasks_type(&mut self, task_type: &str) -> Result; + + /// Retrieve a task from storage by its `id`. + async fn find_task_by_id(&mut self, id: Uuid) -> Result; + + /// Update the state field of the specified task + /// See the `FangTaskState` enum for possible states. + async fn update_task_state( + &mut self, + task: Task, + state: FangTaskState, + ) -> Result; + + /// Update the state of a task to `FangTaskState::Failed` and set an error_message. + async fn fail_task(&mut self, task: Task, error_message: &str) + -> Result; + + /// Schedule a task. + async fn schedule_task(&mut self, task: &dyn AsyncRunnable) -> Result; + + async fn schedule_retry( + &mut self, + task: &Task, + backoff_seconds: u32, + error: &str, + ) -> Result; +} + +/// An async queue that can be used to enqueue tasks. +/// It uses a PostgreSQL storage. It must be connected to perform any operation. +/// To connect an `AsyncQueue` to PostgreSQL database call the `connect` method. +/// A Queue can be created with the TypedBuilder. +/// +/// ```rust +/// let mut queue = AsyncQueue::builder() +/// .uri("postgres://postgres:postgres@localhost/fang") +/// .max_pool_size(max_pool_size) +/// .build(); +/// ``` +/// +#[derive(TypedBuilder, Debug, Clone)] +pub struct AsyncQueue { + pool: Pool, +} + +#[async_trait] +impl AsyncQueueable for AsyncQueue { + async fn find_task_by_id(&mut self, id: Uuid) -> Result { + let mut connection = self + .pool + .get() + .await + .map_err(|e| QueryBuilderError(e.into()))?; + Task::find_task_by_id(&mut connection, id).await + } + + async fn fetch_and_touch_task( + &mut self, + task_type: Option, + ) -> Result, AsyncQueueError> { + let mut connection = self + .pool + .get() + .await + .map_err(|e| QueryBuilderError(e.into()))?; + connection + .transaction::, AsyncQueueError, _>(|conn| { + async move { + let Some(found_task) = Task::fetch_task_of_type(conn, task_type).await else { + return Ok(None); + }; + + match Task::update_task_state( + conn, + found_task, + FangTaskState::InProgress, + ) + .await + { + Ok(updated_task) => Ok(Some(updated_task)), + Err(err) => Err(err), + } + } + .scope_boxed() + }) + .await + } + + async fn insert_task(&mut self, task: &dyn AsyncRunnable) -> Result { + let mut connection = self + .pool + .get() + .await + .map_err(|e| QueryBuilderError(e.into()))?; + Ok(Task::insert_task(&mut connection, task, Utc::now()).await?) + } + + async fn schedule_task(&mut self, task: &dyn AsyncRunnable) -> Result { + let mut connection = self + .pool + .get() + .await + .map_err(|e| QueryBuilderError(e.into()))?; + let scheduled_at = match task.cron() { + Some(scheduled) => match scheduled { + CronPattern(cron_pattern) => { + let schedule = Schedule::from_str(&cron_pattern)?; + let mut iterator = schedule.upcoming(Utc); + iterator + .next() + .ok_or(AsyncQueueError::CronError(CronError::NoTimestampsError))? + } + ScheduleOnce(datetime) => datetime, + }, + None => { + return Err(AsyncQueueError::CronError( + CronError::TaskNotSchedulableError, + )); + } + }; + + Ok(Task::insert_task(&mut connection, task, scheduled_at).await?) + } + + async fn remove_all_tasks(&mut self) -> Result { + let mut connection = self + .pool + .get() + .await + .map_err(|e| QueryBuilderError(e.into()))?; + + Ok(diesel::delete(fang_tasks::table) + .execute(&mut connection) + .await? as u64) + } + + async fn remove_all_scheduled_tasks(&mut self) -> Result { + let mut connection = self + .pool + .get() + .await + .map_err(|e| QueryBuilderError(e.into()))?; + let result = Task::remove_all_scheduled_tasks(&mut connection).await?; + Ok(result) + } + + async fn remove_task(&mut self, id: Uuid) -> Result { + let mut connection = self + .pool + .get() + .await + .map_err(|e| QueryBuilderError(e.into()))?; + let result = Task::remove_task(&mut connection, id).await?; + Ok(result) + } + + async fn remove_task_by_metadata( + &mut self, + task: &dyn AsyncRunnable, + ) -> Result { + if task.uniq() { + let mut connection = self + .pool + .get() + .await + .map_err(|e| QueryBuilderError(e.into()))?; + let result = Task::remove_task_by_metadata(&mut connection, task).await?; + Ok(result) + } else { + Err(AsyncQueueError::TaskNotUniqError) + } + } + + async fn remove_tasks_type(&mut self, task_type: &str) -> Result { + let mut connection = self + .pool + .get() + .await + .map_err(|e| QueryBuilderError(e.into()))?; + let result = Task::remove_tasks_type(&mut connection, task_type).await?; + Ok(result) + } + + async fn update_task_state( + &mut self, + task: Task, + state: FangTaskState, + ) -> Result { + let mut connection = self + .pool + .get() + .await + .map_err(|e| QueryBuilderError(e.into()))?; + let task = Task::update_task_state(&mut connection, task, state).await?; + Ok(task) + } + + async fn fail_task( + &mut self, + task: Task, + error_message: &str, + ) -> Result { + let mut connection = self + .pool + .get() + .await + .map_err(|e| QueryBuilderError(e.into()))?; + let task = Task::fail_task(&mut connection, task, error_message).await?; + Ok(task) + } + + async fn schedule_retry( + &mut self, + task: &Task, + backoff_seconds: u32, + error: &str, + ) -> Result { + let mut connection = self + .pool + .get() + .await + .map_err(|e| QueryBuilderError(e.into()))?; + let task = + Task::schedule_retry(&mut connection, task, backoff_seconds, error).await?; + Ok(task) + } +} + +#[cfg(test)] +mod async_queue_tests { + use super::*; + use crate::schema::fang_tasks::task_type; + use crate::errors::FangError; + use crate::Scheduled; + use async_trait::async_trait; + use chrono::prelude::*; + use chrono::DateTime; + use chrono::Utc; + use diesel_async::pooled_connection::{bb8::Pool, AsyncDieselConnectionManager}; + use diesel_async::AsyncPgConnection; + use serde::{Deserialize, Serialize}; + + #[derive(Serialize, Deserialize)] + struct AsyncTask { + pub number: u16, + } + + #[typetag::serde] + #[async_trait] + impl AsyncRunnable for AsyncTask { + async fn run(&self, _queueable: &mut dyn AsyncQueueable) -> Result<(), FangError> { + Ok(()) + } + } + + #[derive(Serialize, Deserialize)] + struct AsyncUniqTask { + pub number: u16, + } + + #[typetag::serde] + #[async_trait] + impl AsyncRunnable for AsyncUniqTask { + async fn run(&self, _queueable: &mut dyn AsyncQueueable) -> Result<(), FangError> { + Ok(()) + } + + fn uniq(&self) -> bool { + true + } + } + + #[derive(Serialize, Deserialize)] + struct AsyncTaskSchedule { + pub number: u16, + pub datetime: String, + } + + #[typetag::serde] + #[async_trait] + impl AsyncRunnable for AsyncTaskSchedule { + async fn run(&self, _queueable: &mut dyn AsyncQueueable) -> Result<(), FangError> { + Ok(()) + } + + fn cron(&self) -> Option { + let datetime = self.datetime.parse::>().ok()?; + Some(Scheduled::ScheduleOnce(datetime)) + } + } + + #[tokio::test] + async fn insert_task_creates_new_task() { + let pool = pool().await; + let mut test = AsyncQueue::builder().pool(pool).build(); + + let task = insert_task(&mut test, &AsyncTask { number: 1 }).await; + + let metadata = task.metadata.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); + + test.remove_all_tasks().await.unwrap(); + } + + #[tokio::test] + async fn update_task_state_test() { + let pool = pool().await; + let mut test = AsyncQueue::builder().pool(pool).build(); + + let task = insert_task(&mut test, &AsyncTask { number: 1 }).await; + + let metadata = task.metadata.as_object().unwrap(); + let number = metadata["number"].as_u64(); + let type_task = metadata["type"].as_str(); + let id = task.id; + + assert_eq!(Some(1), number); + assert_eq!(Some("AsyncTask"), type_task); + + let finished_task = test + .update_task_state(task, FangTaskState::Finished) + .await + .unwrap(); + + assert_eq!(id, finished_task.id); + assert_eq!(FangTaskState::Finished, finished_task.state); + + test.remove_all_tasks().await.unwrap(); + } + + #[tokio::test] + async fn failed_task_query_test() { + let pool = pool().await; + let mut test = AsyncQueue::builder().pool(pool).build(); + + let task = insert_task(&mut test, &AsyncTask { number: 1 }).await; + + let metadata = task.metadata.as_object().unwrap(); + let number = metadata["number"].as_u64(); + let type_task = metadata["type"].as_str(); + let id = task.id; + + assert_eq!(Some(1), number); + assert_eq!(Some("AsyncTask"), type_task); + + let failed_task = test.fail_task(task, "Some error").await.unwrap(); + + assert_eq!(id, failed_task.id); + assert_eq!(Some("Some error"), failed_task.error_message.as_deref()); + assert_eq!(FangTaskState::Failed, failed_task.state); + + test.remove_all_tasks().await.unwrap(); + } + + #[tokio::test] + async fn remove_all_tasks_test() { + let pool = pool().await; + let mut test = AsyncQueue::builder().pool(pool.into()).build(); + + let task = insert_task(&mut test, &AsyncTask { number: 1 }).await; + + let metadata = task.metadata.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 = insert_task(&mut test, &AsyncTask { number: 2 }).await; + + let metadata = task.metadata.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 schedule_task_test() { + let pool = pool().await; + let mut test = AsyncQueue::builder().pool(pool).build(); + + let datetime = (Utc::now() + Duration::seconds(7)).round_subsecs(0); + + let task = &AsyncTaskSchedule { + number: 1, + datetime: datetime.to_string(), + }; + + let task = test.schedule_task(task).await.unwrap(); + + let metadata = task.metadata.as_object().unwrap(); + let number = metadata["number"].as_u64(); + let type_task = metadata["type"].as_str(); + + assert_eq!(Some(1), number); + assert_eq!(Some("AsyncTaskSchedule"), type_task); + assert_eq!(task.scheduled_at, datetime); + + test.remove_all_tasks().await.unwrap(); + } + + #[tokio::test] + async fn remove_all_scheduled_tasks_test() { + let pool = pool().await; + let mut test = AsyncQueue::builder().pool(pool).build(); + + let datetime = (Utc::now() + Duration::seconds(7)).round_subsecs(0); + + let task1 = &AsyncTaskSchedule { + number: 1, + datetime: datetime.to_string(), + }; + + let task2 = &AsyncTaskSchedule { + number: 2, + datetime: datetime.to_string(), + }; + + test.schedule_task(task1).await.unwrap(); + test.schedule_task(task2).await.unwrap(); + + let number = test.remove_all_scheduled_tasks().await.unwrap(); + + assert_eq!(2, number); + + test.remove_all_tasks().await.unwrap(); + } + + #[tokio::test] + async fn fetch_and_touch_test() { + let pool = pool().await; + let mut test = AsyncQueue::builder().pool(pool).build(); + + let task = insert_task(&mut test, &AsyncTask { number: 1 }).await; + + let metadata = task.metadata.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 = insert_task(&mut test, &AsyncTask { number: 2 }).await; + + let metadata = task.metadata.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 = test + .fetch_and_touch_task(None) + .await + .unwrap() + .unwrap(); + + let metadata = task.metadata.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 + .fetch_and_touch_task(None) + .await + .unwrap() + .unwrap(); + let metadata = task.metadata.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); + + test.remove_all_tasks().await.unwrap(); + } + + #[tokio::test] + async fn remove_tasks_type_test() { + let pool = pool().await; + let mut test = AsyncQueue::builder().pool(pool).build(); + + let task = insert_task(&mut test, &AsyncTask { number: 1 }).await; + + let metadata = task.metadata.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 = insert_task(&mut test, &AsyncTask { number: 2 }).await; + + let metadata = task.metadata.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_tasks_type("mytype").await.unwrap(); + assert_eq!(0, result); + + let result = test.remove_tasks_type("common").await.unwrap(); + assert_eq!(2, result); + + test.remove_all_tasks().await.unwrap(); + } + + #[tokio::test] + async fn remove_tasks_by_metadata() { + let pool = pool().await; + let mut test = AsyncQueue::builder().pool(pool).build(); + + let task = insert_task(&mut test, &AsyncUniqTask { number: 1 }).await; + + let metadata = task.metadata.as_object().unwrap(); + let number = metadata["number"].as_u64(); + let type_task = metadata["type"].as_str(); + + assert_eq!(Some(1), number); + assert_eq!(Some("AsyncUniqTask"), type_task); + + let task = insert_task(&mut test, &AsyncUniqTask { number: 2 }).await; + + let metadata = task.metadata.as_object().unwrap(); + let number = metadata["number"].as_u64(); + let type_task = metadata["type"].as_str(); + + assert_eq!(Some(2), number); + assert_eq!(Some("AsyncUniqTask"), type_task); + + let result = test + .remove_task_by_metadata(&AsyncUniqTask { number: 0 }) + .await + .unwrap(); + assert_eq!(0, result); + + let result = test + .remove_task_by_metadata(&AsyncUniqTask { number: 1 }) + .await + .unwrap(); + assert_eq!(1, result); + + test.remove_all_tasks().await.unwrap(); + } + + async fn insert_task(test: &mut AsyncQueue, task: &dyn AsyncRunnable) -> Task { + test.insert_task(task).await.unwrap() + } + + async fn pool() -> Pool { + let manager = AsyncDieselConnectionManager::::new( + "postgres://postgres:password@localhost/fang", + ); + Pool::builder() + .max_size(1) + .min_idle(Some(1)) + .build(manager) + .await + .unwrap() + } +} diff --git a/src/asynk/async_runnable.rs b/src/runnable.rs similarity index 79% rename from src/asynk/async_runnable.rs rename to src/runnable.rs index bde2bed..435acfc 100644 --- a/src/asynk/async_runnable.rs +++ b/src/runnable.rs @@ -1,10 +1,8 @@ -use crate::async_queue::AsyncQueueError; -use crate::asynk::async_queue::AsyncQueueable; -use crate::FangError; +use crate::errors::AsyncQueueError; +use crate::queue::AsyncQueueable; +use crate::errors::FangError; use crate::Scheduled; use async_trait::async_trait; -use bb8_postgres::bb8::RunError; -use bb8_postgres::tokio_postgres::Error as TokioPostgresError; use serde_json::Error as SerdeError; const COMMON_TYPE: &str = "common"; @@ -19,18 +17,6 @@ impl From for FangError { } } -impl From for FangError { - fn from(error: TokioPostgresError) -> Self { - Self::from(AsyncQueueError::PgError(error)) - } -} - -impl From> for FangError { - fn from(error: RunError) -> Self { - Self::from(AsyncQueueError::PoolError(error)) - } -} - impl From for FangError { fn from(error: SerdeError) -> Self { Self::from(AsyncQueueError::SerdeError(error)) diff --git a/src/task.rs b/src/task.rs new file mode 100644 index 0000000..68af926 --- /dev/null +++ b/src/task.rs @@ -0,0 +1,51 @@ +use crate::fang_task_state::FangTaskState; +use crate::schema::fang_tasks; +use chrono::DateTime; +use chrono::Duration; +use chrono::Utc; +use cron::Schedule; +use diesel::prelude::*; +use sha2::{Digest, Sha256}; +use thiserror::Error; +use typed_builder::TypedBuilder; +use uuid::Uuid; + +pub const DEFAULT_TASK_TYPE: &str = "common"; + +#[derive(Queryable, Identifiable, Debug, Eq, PartialEq, Clone, TypedBuilder)] +#[diesel(table_name = fang_tasks)] +pub struct Task { + #[builder(setter(into))] + pub id: Uuid, + #[builder(setter(into))] + pub metadata: serde_json::Value, + #[builder(setter(into))] + pub error_message: Option, + #[builder(setter(into))] + pub state: FangTaskState, + #[builder(setter(into))] + pub task_type: String, + #[builder(setter(into))] + pub uniq_hash: Option, + #[builder(setter(into))] + pub retries: i32, + #[builder(setter(into))] + pub scheduled_at: DateTime, + #[builder(setter(into))] + pub created_at: DateTime, + #[builder(setter(into))] + pub updated_at: DateTime, +} + +#[derive(Insertable, Debug, Eq, PartialEq, Clone, TypedBuilder)] +#[diesel(table_name = fang_tasks)] +pub struct NewTask { + #[builder(setter(into))] + metadata: serde_json::Value, + #[builder(setter(into))] + task_type: String, + #[builder(setter(into))] + uniq_hash: Option, + #[builder(setter(into))] + scheduled_at: DateTime, +} diff --git a/src/worker.rs b/src/worker.rs new file mode 100644 index 0000000..85bc1c3 --- /dev/null +++ b/src/worker.rs @@ -0,0 +1,453 @@ +use crate::queue::AsyncQueueable; +use crate::task::Task; +use crate::task::DEFAULT_TASK_TYPE; +use crate::runnable::AsyncRunnable; +use crate::fang_task_state::FangTaskState; +use crate::errors::FangError; +use crate::Scheduled::*; +use crate::{RetentionMode, SleepParams}; +use log::error; +use typed_builder::TypedBuilder; + +/// it executes tasks only of task_type type, it sleeps when there are no tasks in the queue +#[derive(TypedBuilder)] +pub struct AsyncWorker +where + AQueue: AsyncQueueable + Clone + Sync + 'static, +{ + #[builder(setter(into))] + pub queue: AQueue, + #[builder(default=DEFAULT_TASK_TYPE.to_string(), setter(into))] + pub task_type: String, + #[builder(default, setter(into))] + pub sleep_params: SleepParams, + #[builder(default, setter(into))] + pub retention_mode: RetentionMode, +} + +impl AsyncWorker +where + AQueue: AsyncQueueable + Clone + Sync + 'static, +{ + async fn run(&mut self, task: Task, runnable: Box) -> Result<(), FangError> { + let result = runnable.run(&mut self.queue).await; + + match result { + Ok(_) => self.finalize_task(task, &result).await?, + + Err(ref error) => { + if task.retries < runnable.max_retries() { + let backoff_seconds = runnable.backoff(task.retries as u32); + + self.queue + .schedule_retry(&task, backoff_seconds, &error.description) + .await?; + } else { + self.finalize_task(task, &result).await?; + } + } + } + + Ok(()) + } + + async fn finalize_task( + &mut self, + task: Task, + result: &Result<(), FangError>, + ) -> Result<(), FangError> { + match self.retention_mode { + RetentionMode::KeepAll => match result { + Ok(_) => { + self.queue + .update_task_state(task, FangTaskState::Finished) + .await?; + } + Err(error) => { + self.queue.fail_task(task, &error.description).await?; + } + }, + RetentionMode::RemoveAll => { + self.queue.remove_task(task.id).await?; + } + RetentionMode::RemoveFinished => match result { + Ok(_) => { + self.queue.remove_task(task.id).await?; + } + Err(error) => { + self.queue.fail_task(task, &error.description).await?; + } + }, + }; + + Ok(()) + } + + async fn sleep(&mut self) { + self.sleep_params.maybe_increase_sleep_period(); + + tokio::time::sleep(self.sleep_params.sleep_period).await; + } + + pub(crate) async fn run_tasks(&mut self) -> Result<(), FangError> { + loop { + //fetch task + match self.queue.fetch_and_touch_task(Some(self.task_type.clone())).await { + Ok(Some(task)) => { + let actual_task: Box = + serde_json::from_value(task.metadata.clone()).unwrap(); + + // check if task is scheduled or not + if let Some(CronPattern(_)) = actual_task.cron() { + // program task + self.queue.schedule_task(&*actual_task).await?; + } + self.sleep_params.maybe_reset_sleep_period(); + // run scheduled task + self.run(task, actual_task).await?; + } + Ok(None) => { + self.sleep().await; + } + + Err(error) => { + error!("Failed to fetch a task {:?}", error); + + self.sleep().await; + } + }; + } + } +} + +#[cfg(test)] +mod async_worker_tests { + use super::*; + use crate::queue::AsyncQueueable; + use crate::worker::Task; + use crate::errors::FangError; + use crate::RetentionMode; + use crate::Scheduled; + use async_trait::async_trait; + use chrono::Duration; + use chrono::Utc; + use serde::{Deserialize, Serialize}; + + #[derive(Serialize, Deserialize)] + struct WorkerAsyncTask { + pub number: u16, + } + + #[typetag::serde] + #[async_trait] + impl AsyncRunnable for WorkerAsyncTask { + async fn run(&self, _queueable: &mut dyn AsyncQueueable) -> Result<(), FangError> { + Ok(()) + } + } + + #[derive(Serialize, Deserialize)] + struct WorkerAsyncTaskSchedule { + pub number: u16, + } + + #[typetag::serde] + #[async_trait] + impl AsyncRunnable for WorkerAsyncTaskSchedule { + async fn run(&self, _queueable: &mut dyn AsyncQueueable) -> Result<(), FangError> { + Ok(()) + } + fn cron(&self) -> Option { + Some(Scheduled::ScheduleOnce(Utc::now() + Duration::seconds(1))) + } + } + + #[derive(Serialize, Deserialize)] + struct AsyncFailedTask { + pub number: u16, + } + + #[typetag::serde] + #[async_trait] + impl AsyncRunnable for AsyncFailedTask { + async fn run(&self, _queueable: &mut dyn AsyncQueueable) -> Result<(), FangError> { + let message = format!("number {} is wrong :(", self.number); + + Err(FangError { + description: message, + }) + } + + fn max_retries(&self) -> i32 { + 0 + } + } + + #[derive(Serialize, Deserialize, Clone)] + struct AsyncRetryTask {} + + #[typetag::serde] + #[async_trait] + impl AsyncRunnable for AsyncRetryTask { + async fn run(&self, _queueable: &mut dyn AsyncQueueable) -> Result<(), FangError> { + let message = "Failed".to_string(); + + Err(FangError { + description: message, + }) + } + + fn max_retries(&self) -> i32 { + 2 + } + } + + #[derive(Serialize, Deserialize)] + struct AsyncTaskType1 {} + + #[typetag::serde] + #[async_trait] + impl AsyncRunnable for AsyncTaskType1 { + async fn run(&self, _queueable: &mut dyn AsyncQueueable) -> Result<(), FangError> { + Ok(()) + } + + fn task_type(&self) -> String { + "type1".to_string() + } + } + + #[derive(Serialize, Deserialize)] + struct AsyncTaskType2 {} + + #[typetag::serde] + #[async_trait] + impl AsyncRunnable for AsyncTaskType2 { + async fn run(&self, _queueable: &mut dyn AsyncQueueable) -> Result<(), FangError> { + Ok(()) + } + + fn task_type(&self) -> String { + "type2".to_string() + } + } + + // #[tokio::test] + // async fn execute_and_finishes_task() { + // let pool = pool().await; + // let mut connection = pool.get().await.unwrap(); + // let transaction = connection.transaction().await.unwrap(); + // + // let mut test = AsyncQueueTest::builder().transaction(transaction).build(); + // let actual_task = WorkerAsyncTask { number: 1 }; + // + // let task = insert_task(&mut test, &actual_task).await; + // let id = task.id; + // + // let mut worker = AsyncWorkerTest::builder() + // .queue(&mut test as &mut dyn AsyncQueueable) + // .retention_mode(RetentionMode::KeepAll) + // .build(); + // + // worker.run(task, Box::new(actual_task)).await.unwrap(); + // let task_finished = test.find_task_by_id(id).await.unwrap(); + // assert_eq!(id, task_finished.id); + // assert_eq!(FangTaskState::Finished, task_finished.state); + // test.transaction.rollback().await.unwrap(); + // } + // + // #[tokio::test] + // async fn schedule_task_test() { + // let pool = pool().await; + // let mut connection = pool.get().await.unwrap(); + // let transaction = connection.transaction().await.unwrap(); + // + // let mut test = AsyncQueueTest::builder().transaction(transaction).build(); + // + // let actual_task = WorkerAsyncTaskSchedule { number: 1 }; + // + // let task = test.schedule_task(&actual_task).await.unwrap(); + // + // let id = task.id; + // + // let mut worker = AsyncWorkerTest::builder() + // .queue(&mut test as &mut dyn AsyncQueueable) + // .retention_mode(RetentionMode::KeepAll) + // .build(); + // + // worker.run_tasks_until_none().await.unwrap(); + // + // let task = worker.queue.find_task_by_id(id).await.unwrap(); + // + // assert_eq!(id, task.id); + // assert_eq!(FangTaskState::New, task.state); + // + // tokio::time::sleep(core::time::Duration::from_secs(3)).await; + // + // worker.run_tasks_until_none().await.unwrap(); + // + // let task = test.find_task_by_id(id).await.unwrap(); + // assert_eq!(id, task.id); + // assert_eq!(FangTaskState::Finished, task.state); + // } + // + // #[tokio::test] + // async fn retries_task_test() { + // let pool = pool().await; + // let mut connection = pool.get().await.unwrap(); + // let transaction = connection.transaction().await.unwrap(); + // + // let mut test = AsyncQueueTest::builder().transaction(transaction).build(); + // + // let actual_task = AsyncRetryTask {}; + // + // let task = test.insert_task(&actual_task).await.unwrap(); + // + // let id = task.id; + // + // let mut worker = AsyncWorkerTest::builder() + // .queue(&mut test as &mut dyn AsyncQueueable) + // .retention_mode(RetentionMode::KeepAll) + // .build(); + // + // worker.run_tasks_until_none().await.unwrap(); + // + // let task = worker.queue.find_task_by_id(id).await.unwrap(); + // + // assert_eq!(id, task.id); + // assert_eq!(FangTaskState::Retried, task.state); + // assert_eq!(1, task.retries); + // + // tokio::time::sleep(core::time::Duration::from_secs(5)).await; + // worker.run_tasks_until_none().await.unwrap(); + // + // let task = worker.queue.find_task_by_id(id).await.unwrap(); + // + // assert_eq!(id, task.id); + // assert_eq!(FangTaskState::Retried, task.state); + // assert_eq!(2, task.retries); + // + // tokio::time::sleep(core::time::Duration::from_secs(10)).await; + // worker.run_tasks_until_none().await.unwrap(); + // + // let task = test.find_task_by_id(id).await.unwrap(); + // assert_eq!(id, task.id); + // assert_eq!(FangTaskState::Failed, task.state); + // assert_eq!("Failed".to_string(), task.error_message.unwrap()); + // } + // + // #[tokio::test] + // async fn saves_error_for_failed_task() { + // let pool = pool().await; + // let mut connection = pool.get().await.unwrap(); + // let transaction = connection.transaction().await.unwrap(); + // + // let mut test = AsyncQueueTest::builder().transaction(transaction).build(); + // let failed_task = AsyncFailedTask { number: 1 }; + // + // let task = insert_task(&mut test, &failed_task).await; + // let id = task.id; + // + // let mut worker = AsyncWorkerTest::builder() + // .queue(&mut test as &mut dyn AsyncQueueable) + // .retention_mode(RetentionMode::KeepAll) + // .build(); + // + // worker.run(task, Box::new(failed_task)).await.unwrap(); + // let task_finished = test.find_task_by_id(id).await.unwrap(); + // + // assert_eq!(id, task_finished.id); + // assert_eq!(FangTaskState::Failed, task_finished.state); + // assert_eq!( + // "number 1 is wrong :(".to_string(), + // task_finished.error_message.unwrap() + // ); + // test.transaction.rollback().await.unwrap(); + // } + // + // #[tokio::test] + // async fn executes_task_only_of_specific_type() { + // let pool = pool().await; + // let mut connection = pool.get().await.unwrap(); + // let transaction = connection.transaction().await.unwrap(); + // + // let mut test = AsyncQueueTest::builder().transaction(transaction).build(); + // + // let task1 = insert_task(&mut test, &AsyncTaskType1 {}).await; + // let task12 = insert_task(&mut test, &AsyncTaskType1 {}).await; + // let task2 = insert_task(&mut test, &AsyncTaskType2 {}).await; + // + // let id1 = task1.id; + // let id12 = task12.id; + // let id2 = task2.id; + // + // let mut worker = AsyncWorkerTest::builder() + // .queue(&mut test as &mut dyn AsyncQueueable) + // .task_type("type1".to_string()) + // .retention_mode(RetentionMode::KeepAll) + // .build(); + // + // worker.run_tasks_until_none().await.unwrap(); + // let task1 = test.find_task_by_id(id1).await.unwrap(); + // let task12 = test.find_task_by_id(id12).await.unwrap(); + // let task2 = test.find_task_by_id(id2).await.unwrap(); + // + // assert_eq!(id1, task1.id); + // assert_eq!(id12, task12.id); + // assert_eq!(id2, task2.id); + // assert_eq!(FangTaskState::Finished, task1.state); + // assert_eq!(FangTaskState::Finished, task12.state); + // assert_eq!(FangTaskState::New, task2.state); + // test.transaction.rollback().await.unwrap(); + // } + // + // #[tokio::test] + // async fn remove_when_finished() { + // let pool = pool().await; + // let mut connection = pool.get().await.unwrap(); + // let transaction = connection.transaction().await.unwrap(); + // + // let mut test = AsyncQueueTest::builder().transaction(transaction).build(); + // + // let task1 = insert_task(&mut test, &AsyncTaskType1 {}).await; + // let task12 = insert_task(&mut test, &AsyncTaskType1 {}).await; + // let task2 = insert_task(&mut test, &AsyncTaskType2 {}).await; + // + // let _id1 = task1.id; + // let _id12 = task12.id; + // let id2 = task2.id; + // + // let mut worker = AsyncWorkerTest::builder() + // .queue(&mut test as &mut dyn AsyncQueueable) + // .task_type("type1".to_string()) + // .build(); + // + // worker.run_tasks_until_none().await.unwrap(); + // let task = test + // .fetch_and_touch_task(Some("type1".to_string())) + // .await + // .unwrap(); + // assert_eq!(None, task); + // + // let task2 = test + // .fetch_and_touch_task(Some("type2".to_string())) + // .await + // .unwrap() + // .unwrap(); + // assert_eq!(id2, task2.id); + // + // test.transaction.rollback().await.unwrap(); + // } + // async fn insert_task(test: &mut AsyncQueueTest<'_>, task: &dyn AsyncRunnable) -> Task { + // test.insert_task(task).await.unwrap() + // } + // async fn pool() -> Pool> { + // let pg_mgr = PostgresConnectionManager::new_from_stringlike( + // "postgres://postgres:postgres@localhost/fang", + // NoTls, + // ) + // .unwrap(); + // + // Pool::builder().build(pg_mgr).await.unwrap() + // } +} diff --git a/src/asynk/async_worker_pool.rs b/src/worker_pool.rs similarity index 94% rename from src/asynk/async_worker_pool.rs rename to src/worker_pool.rs index 784440e..f8a5bb3 100644 --- a/src/asynk/async_worker_pool.rs +++ b/src/worker_pool.rs @@ -1,7 +1,7 @@ -use crate::asynk::async_queue::AsyncQueueable; -use crate::asynk::async_queue::DEFAULT_TASK_TYPE; -use crate::asynk::async_worker::AsyncWorker; -use crate::FangError; +use crate::queue::AsyncQueueable; +use crate::task::DEFAULT_TASK_TYPE; +use crate::worker::AsyncWorker; +use crate::errors::FangError; use crate::{RetentionMode, SleepParams}; use async_recursion::async_recursion; use log::error;