mod runner; mod spawn; mod utils; pub use runner::*; pub use spawn::*; pub use utils::OwnedTask; #[cfg(test)] mod tests { use super::*; use std::env; use std::future::Future; use std::ops::Deref; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Once}; use std::time::Duration; use futures::channel::mpsc; use futures::StreamExt; use sqlx::{Pool, Postgres}; use tokio::sync::{Mutex, MutexGuard}; use tokio::task; struct TestGuard(MutexGuard<'static, ()>, T); impl Deref for TestGuard { type Target = T; fn deref(&self) -> &T { &self.1 } } async fn test_pool() -> TestGuard> { static INIT_LOGGER: Once = Once::new(); static TEST_MUTEX: Mutex<()> = Mutex::const_new(()); let guard = TEST_MUTEX.lock().await; let _ = dotenv::dotenv(); INIT_LOGGER.call_once(|| pretty_env_logger::init()); let pool = Pool::connect(&env::var("DATABASE_URL").unwrap()) .await .unwrap(); sqlx::query("TRUNCATE TABLE mq_payloads") .execute(&pool) .await .unwrap(); sqlx::query("DELETE FROM mq_msgs WHERE id != uuid_nil()") .execute(&pool) .await .unwrap(); TestGuard(guard, pool) } async fn test_task_runner( pool: &Pool, f: impl (Fn(CurrentTask) -> F) + Send + Sync + 'static, ) -> (OwnedTask, Arc) where F::Output: Send + 'static, { let counter = Arc::new(AtomicUsize::new(0)); let counter2 = counter.clone(); let runner = TaskRunnerOptions::new(pool, move |task| { counter2.fetch_add(1, Ordering::SeqCst); task::spawn(f(task)); }) .run() .await .unwrap(); (runner, counter) } async fn pause() { pause_ms(50).await; } async fn pause_ms(ms: u64) { tokio::time::sleep(Duration::from_millis(ms)).await; } #[tokio::test] async fn it_can_spawn_task() { let pool = &*test_pool().await; let (_runner, counter) = test_task_runner(&pool, |mut task| async move { task.complete().await }).await; assert_eq!(counter.load(Ordering::SeqCst), 0); TaskBuilder::new("foo").spawn(pool).await.unwrap(); pause().await; assert_eq!(counter.load(Ordering::SeqCst), 1); } #[tokio::test] async fn it_runs_tasks_in_order() { let pool = &*test_pool().await; let (tx, mut rx) = mpsc::unbounded(); let (_runner, counter) = test_task_runner(&pool, move |task| { let tx = tx.clone(); async move { tx.unbounded_send(task).unwrap(); } }) .await; assert_eq!(counter.load(Ordering::SeqCst), 0); TaskBuilder::new("foo") .set_ordered(true) .spawn(pool) .await .unwrap(); TaskBuilder::new("bar") .set_ordered(true) .spawn(pool) .await .unwrap(); pause().await; assert_eq!(counter.load(Ordering::SeqCst), 1); let mut task = rx.next().await.unwrap(); task.complete().await.unwrap(); pause().await; assert_eq!(counter.load(Ordering::SeqCst), 2); } #[tokio::test] async fn it_runs_tasks_in_parallel() { let pool = &*test_pool().await; let (tx, mut rx) = mpsc::unbounded(); let (_runner, counter) = test_task_runner(&pool, move |task| { let tx = tx.clone(); async move { tx.unbounded_send(task).unwrap(); } }) .await; assert_eq!(counter.load(Ordering::SeqCst), 0); TaskBuilder::new("foo").spawn(pool).await.unwrap(); TaskBuilder::new("bar").spawn(pool).await.unwrap(); pause().await; assert_eq!(counter.load(Ordering::SeqCst), 2); for _ in 0..2 { let mut task = rx.next().await.unwrap(); task.complete().await.unwrap(); } } #[tokio::test] async fn it_retries_failed_tasks() { let pool = &*test_pool().await; let (_runner, counter) = test_task_runner(&pool, move |_| async {}).await; let backoff = 100; assert_eq!(counter.load(Ordering::SeqCst), 0); TaskBuilder::new("foo") .set_retry_backoff(Duration::from_millis(backoff)) .set_retries(2) .spawn(pool) .await .unwrap(); // First attempt pause().await; assert_eq!(counter.load(Ordering::SeqCst), 1); // Second attempt pause_ms(backoff).await; pause().await; assert_eq!(counter.load(Ordering::SeqCst), 2); // Third attempt pause_ms(backoff * 2).await; pause().await; assert_eq!(counter.load(Ordering::SeqCst), 3); // No more attempts pause_ms(backoff * 5).await; assert_eq!(counter.load(Ordering::SeqCst), 3); } }