use awa::{BuildError, Client, JobArgs, JobContext, JobError, JobResult, QueueConfig};
use awa_model::{insert_with, migrations, InsertOpts};
use serde::{Deserialize, Serialize};
use sqlx::postgres::PgPoolOptions;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use std::time::Duration;
fn database_url() -> String {
std::env::var("DATABASE_URL")
.unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string())
}
async fn setup() -> sqlx::PgPool {
let pool = PgPoolOptions::new()
.max_connections(10)
.connect(&database_url())
.await
.expect("Failed to connect");
migrations::run(&pool).await.expect("Failed to migrate");
pool
}
async fn clean_queue(pool: &sqlx::PgPool, queue: &str) {
sqlx::query("DELETE FROM awa.jobs WHERE queue = $1")
.bind(queue)
.execute(pool)
.await
.expect("Failed to clean queue");
sqlx::query("DELETE FROM awa.queue_meta WHERE queue = $1")
.bind(queue)
.execute(pool)
.await
.expect("Failed to clean queue meta");
}
#[derive(Debug, Serialize, Deserialize, JobArgs)]
struct WeightedJob {
pub index: i64,
}
struct SlowWorker {
completed: Arc<AtomicU32>,
delay: Duration,
}
#[async_trait::async_trait]
impl awa::Worker for SlowWorker {
fn kind(&self) -> &'static str {
"weighted_job"
}
async fn perform(&self, _ctx: &JobContext) -> Result<JobResult, JobError> {
tokio::time::sleep(self.delay).await;
self.completed.fetch_add(1, Ordering::SeqCst);
Ok(JobResult::Completed)
}
}
#[tokio::test]
async fn test_hard_reserved_backward_compat() {
let pool = setup().await;
let queue = "wt_hard_compat";
clean_queue(&pool, queue).await;
for i in 0..10 {
insert_with(
&pool,
&WeightedJob { index: i },
InsertOpts {
queue: queue.into(),
..Default::default()
},
)
.await
.unwrap();
}
let completed = Arc::new(AtomicU32::new(0));
let client = Client::builder(pool.clone())
.queue(
queue,
QueueConfig {
max_workers: 5,
poll_interval: Duration::from_millis(50),
..Default::default()
},
)
.register_worker(SlowWorker {
completed: completed.clone(),
delay: Duration::from_millis(100),
})
.build()
.unwrap();
client.start().await.unwrap();
let start = std::time::Instant::now();
loop {
if completed.load(Ordering::SeqCst) >= 10 {
break;
}
if start.elapsed() > Duration::from_secs(5) {
break;
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
client.shutdown(Duration::from_secs(2)).await;
assert_eq!(completed.load(Ordering::SeqCst), 10);
}
#[tokio::test]
async fn test_global_cap_not_exceeded() {
let pool = setup().await;
let queue_a = "wt_cap_a";
let queue_b = "wt_cap_b";
clean_queue(&pool, queue_a).await;
clean_queue(&pool, queue_b).await;
for i in 0..20 {
insert_with(
&pool,
&WeightedJob { index: i },
InsertOpts {
queue: queue_a.into(),
..Default::default()
},
)
.await
.unwrap();
insert_with(
&pool,
&WeightedJob { index: i + 20 },
InsertOpts {
queue: queue_b.into(),
..Default::default()
},
)
.await
.unwrap();
}
let completed = Arc::new(AtomicU32::new(0));
let max_concurrent = Arc::new(AtomicU32::new(0));
let current_concurrent = Arc::new(AtomicU32::new(0));
struct ConcurrentTrackWorker {
completed: Arc<AtomicU32>,
max_concurrent: Arc<AtomicU32>,
current_concurrent: Arc<AtomicU32>,
}
#[async_trait::async_trait]
impl awa::Worker for ConcurrentTrackWorker {
fn kind(&self) -> &'static str {
"weighted_job"
}
async fn perform(&self, _ctx: &JobContext) -> Result<JobResult, JobError> {
let current = self.current_concurrent.fetch_add(1, Ordering::SeqCst) + 1;
self.max_concurrent.fetch_max(current, Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(200)).await;
self.current_concurrent.fetch_sub(1, Ordering::SeqCst);
self.completed.fetch_add(1, Ordering::SeqCst);
Ok(JobResult::Completed)
}
}
let global_max = 10u32;
let client = Client::builder(pool.clone())
.queue(
queue_a,
QueueConfig {
min_workers: 2,
weight: 1,
poll_interval: Duration::from_millis(50),
..Default::default()
},
)
.queue(
queue_b,
QueueConfig {
min_workers: 2,
weight: 1,
poll_interval: Duration::from_millis(50),
..Default::default()
},
)
.global_max_workers(global_max)
.register_worker(ConcurrentTrackWorker {
completed: completed.clone(),
max_concurrent: max_concurrent.clone(),
current_concurrent: current_concurrent.clone(),
})
.build()
.unwrap();
client.start().await.unwrap();
let start = std::time::Instant::now();
loop {
if completed.load(Ordering::SeqCst) >= 40 {
break;
}
if start.elapsed() > Duration::from_secs(15) {
break;
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
client.shutdown(Duration::from_secs(3)).await;
let max_seen = max_concurrent.load(Ordering::SeqCst);
assert_eq!(
completed.load(Ordering::SeqCst),
40,
"All 40 jobs should complete"
);
assert!(
max_seen <= global_max,
"Max concurrent ({max_seen}) should not exceed global_max_workers ({global_max})"
);
}
#[tokio::test]
async fn test_min_workers_exceed_global_rejected() {
let pool = setup().await;
let result = Client::builder(pool)
.queue(
"wt_val_a",
QueueConfig {
min_workers: 15,
..Default::default()
},
)
.queue(
"wt_val_b",
QueueConfig {
min_workers: 10,
..Default::default()
},
)
.global_max_workers(20)
.build();
assert!(matches!(
result,
Err(BuildError::MinWorkersExceedGlobal {
total_min: 25,
global_max: 20
})
));
}
#[tokio::test]
async fn test_idle_overflow_to_loaded_queue() {
let pool = setup().await;
let queue_a = "wt_overflow_a";
let queue_b = "wt_overflow_b";
clean_queue(&pool, queue_a).await;
clean_queue(&pool, queue_b).await;
for i in 0..30 {
insert_with(
&pool,
&WeightedJob { index: i },
InsertOpts {
queue: queue_a.into(),
..Default::default()
},
)
.await
.unwrap();
}
let completed = Arc::new(AtomicU32::new(0));
let client = Client::builder(pool.clone())
.queue(
queue_a,
QueueConfig {
min_workers: 5,
weight: 1,
poll_interval: Duration::from_millis(50),
..Default::default()
},
)
.queue(
queue_b,
QueueConfig {
min_workers: 5,
weight: 1,
poll_interval: Duration::from_millis(50),
..Default::default()
},
)
.global_max_workers(20)
.register_worker(SlowWorker {
completed: completed.clone(),
delay: Duration::from_millis(100),
})
.build()
.unwrap();
client.start().await.unwrap();
let start = std::time::Instant::now();
loop {
if completed.load(Ordering::SeqCst) >= 30 {
break;
}
if start.elapsed() > Duration::from_secs(10) {
break;
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
client.shutdown(Duration::from_secs(2)).await;
assert_eq!(completed.load(Ordering::SeqCst), 30);
}
#[tokio::test]
async fn test_floor_guarantee_under_load() {
let pool = setup().await;
let queue_a = "wt_floor_a";
let queue_b = "wt_floor_b";
clean_queue(&pool, queue_a).await;
clean_queue(&pool, queue_b).await;
for i in 0..40 {
insert_with(
&pool,
&WeightedJob { index: i },
InsertOpts {
queue: queue_a.into(),
..Default::default()
},
)
.await
.unwrap();
insert_with(
&pool,
&WeightedJob { index: i + 40 },
InsertOpts {
queue: queue_b.into(),
..Default::default()
},
)
.await
.unwrap();
}
let completed_a = Arc::new(AtomicU32::new(0));
let completed_b = Arc::new(AtomicU32::new(0));
let max_concurrent_a = Arc::new(AtomicU32::new(0));
let max_concurrent_b = Arc::new(AtomicU32::new(0));
let current_a = Arc::new(AtomicU32::new(0));
let current_b = Arc::new(AtomicU32::new(0));
struct FloorWorker {
completed_a: Arc<AtomicU32>,
completed_b: Arc<AtomicU32>,
max_concurrent_a: Arc<AtomicU32>,
max_concurrent_b: Arc<AtomicU32>,
current_a: Arc<AtomicU32>,
current_b: Arc<AtomicU32>,
}
#[async_trait::async_trait]
impl awa::Worker for FloorWorker {
fn kind(&self) -> &'static str {
"weighted_job"
}
async fn perform(&self, ctx: &JobContext) -> Result<JobResult, JobError> {
if ctx.job.queue == "wt_floor_a" {
let c = self.current_a.fetch_add(1, Ordering::SeqCst) + 1;
self.max_concurrent_a.fetch_max(c, Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(300)).await;
self.current_a.fetch_sub(1, Ordering::SeqCst);
self.completed_a.fetch_add(1, Ordering::SeqCst);
} else {
let c = self.current_b.fetch_add(1, Ordering::SeqCst) + 1;
self.max_concurrent_b.fetch_max(c, Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(300)).await;
self.current_b.fetch_sub(1, Ordering::SeqCst);
self.completed_b.fetch_add(1, Ordering::SeqCst);
}
Ok(JobResult::Completed)
}
}
let client = Client::builder(pool.clone())
.queue(
queue_a,
QueueConfig {
min_workers: 5,
weight: 1,
poll_interval: Duration::from_millis(50),
..Default::default()
},
)
.queue(
queue_b,
QueueConfig {
min_workers: 5,
weight: 1,
poll_interval: Duration::from_millis(50),
..Default::default()
},
)
.global_max_workers(20)
.register_worker(FloorWorker {
completed_a: completed_a.clone(),
completed_b: completed_b.clone(),
max_concurrent_a: max_concurrent_a.clone(),
max_concurrent_b: max_concurrent_b.clone(),
current_a: current_a.clone(),
current_b: current_b.clone(),
})
.build()
.unwrap();
client.start().await.unwrap();
let start = std::time::Instant::now();
loop {
let total = completed_a.load(Ordering::SeqCst) + completed_b.load(Ordering::SeqCst);
if total >= 60 {
break;
}
if start.elapsed() > Duration::from_secs(15) {
break;
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
client.shutdown(Duration::from_secs(3)).await;
let max_a = max_concurrent_a.load(Ordering::SeqCst);
let max_b = max_concurrent_b.load(Ordering::SeqCst);
assert!(
max_a >= 5,
"Queue A should reach at least 5 concurrent (min_workers), saw max of {max_a}"
);
assert!(
max_b >= 5,
"Queue B should reach at least 5 concurrent (min_workers), saw max of {max_b}"
);
}
#[tokio::test]
async fn test_weight_proportionality() {
let pool = setup().await;
let queue_a = "wt_prop_a";
let queue_b = "wt_prop_b";
clean_queue(&pool, queue_a).await;
clean_queue(&pool, queue_b).await;
for i in 0..120 {
insert_with(
&pool,
&WeightedJob { index: i },
InsertOpts {
queue: queue_a.into(),
..Default::default()
},
)
.await
.unwrap();
insert_with(
&pool,
&WeightedJob { index: i + 120 },
InsertOpts {
queue: queue_b.into(),
..Default::default()
},
)
.await
.unwrap();
}
let completed_a = Arc::new(AtomicU32::new(0));
let completed_b_counter = Arc::new(AtomicU32::new(0));
let max_concurrent_a = Arc::new(AtomicU32::new(0));
let max_concurrent_b = Arc::new(AtomicU32::new(0));
let current_a = Arc::new(AtomicU32::new(0));
let current_b = Arc::new(AtomicU32::new(0));
const SNAPSHOT_BASE: u32 = 40;
const SNAPSHOT_END: u32 = 160;
const SNAPSHOT_UNSET: u32 = u32::MAX;
let total_completed = Arc::new(AtomicU32::new(0));
let base_a = Arc::new(AtomicU32::new(SNAPSHOT_UNSET));
let base_b = Arc::new(AtomicU32::new(SNAPSHOT_UNSET));
let snapshot_a = Arc::new(AtomicU32::new(SNAPSHOT_UNSET));
let snapshot_b = Arc::new(AtomicU32::new(SNAPSHOT_UNSET));
struct ProportionWorker {
completed_a: Arc<AtomicU32>,
completed_b: Arc<AtomicU32>,
max_concurrent_a: Arc<AtomicU32>,
max_concurrent_b: Arc<AtomicU32>,
current_a: Arc<AtomicU32>,
current_b: Arc<AtomicU32>,
total_completed: Arc<AtomicU32>,
base_a: Arc<AtomicU32>,
base_b: Arc<AtomicU32>,
snapshot_a: Arc<AtomicU32>,
snapshot_b: Arc<AtomicU32>,
}
#[async_trait::async_trait]
impl awa::Worker for ProportionWorker {
fn kind(&self) -> &'static str {
"weighted_job"
}
async fn perform(&self, ctx: &JobContext) -> Result<JobResult, JobError> {
if ctx.job.queue == "wt_prop_a" {
let c = self.current_a.fetch_add(1, Ordering::SeqCst) + 1;
self.max_concurrent_a.fetch_max(c, Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(150)).await;
self.current_a.fetch_sub(1, Ordering::SeqCst);
self.completed_a.fetch_add(1, Ordering::SeqCst);
} else {
let c = self.current_b.fetch_add(1, Ordering::SeqCst) + 1;
self.max_concurrent_b.fetch_max(c, Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(150)).await;
self.current_b.fetch_sub(1, Ordering::SeqCst);
self.completed_b.fetch_add(1, Ordering::SeqCst);
}
match self.total_completed.fetch_add(1, Ordering::SeqCst) + 1 {
SNAPSHOT_BASE => {
self.base_a
.store(self.completed_a.load(Ordering::SeqCst), Ordering::SeqCst);
self.base_b
.store(self.completed_b.load(Ordering::SeqCst), Ordering::SeqCst);
}
SNAPSHOT_END => {
self.snapshot_a
.store(self.completed_a.load(Ordering::SeqCst), Ordering::SeqCst);
self.snapshot_b
.store(self.completed_b.load(Ordering::SeqCst), Ordering::SeqCst);
}
_ => {}
}
Ok(JobResult::Completed)
}
}
let client = Client::builder(pool.clone())
.queue(
queue_a,
QueueConfig {
min_workers: 2,
weight: 3,
poll_interval: Duration::from_millis(50),
..Default::default()
},
)
.queue(
queue_b,
QueueConfig {
min_workers: 2,
weight: 1,
poll_interval: Duration::from_millis(50),
..Default::default()
},
)
.global_max_workers(24)
.register_worker(ProportionWorker {
completed_a: completed_a.clone(),
completed_b: completed_b_counter.clone(),
max_concurrent_a: max_concurrent_a.clone(),
max_concurrent_b: max_concurrent_b.clone(),
current_a: current_a.clone(),
current_b: current_b.clone(),
total_completed: total_completed.clone(),
base_a: base_a.clone(),
base_b: base_b.clone(),
snapshot_a: snapshot_a.clone(),
snapshot_b: snapshot_b.clone(),
})
.build()
.unwrap();
client.start().await.unwrap();
let start = std::time::Instant::now();
loop {
if base_a.load(Ordering::SeqCst) != SNAPSHOT_UNSET
&& base_b.load(Ordering::SeqCst) != SNAPSHOT_UNSET
&& snapshot_a.load(Ordering::SeqCst) != SNAPSHOT_UNSET
&& snapshot_b.load(Ordering::SeqCst) != SNAPSHOT_UNSET
{
break;
}
if start.elapsed() > Duration::from_secs(20) {
break;
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
client.shutdown(Duration::from_secs(3)).await;
let snapshots_published = base_a.load(Ordering::SeqCst) != SNAPSHOT_UNSET
&& base_b.load(Ordering::SeqCst) != SNAPSHOT_UNSET
&& snapshot_a.load(Ordering::SeqCst) != SNAPSHOT_UNSET
&& snapshot_b.load(Ordering::SeqCst) != SNAPSHOT_UNSET;
let (ca, cb, window) = if snapshots_published {
(
snapshot_a.load(Ordering::SeqCst) - base_a.load(Ordering::SeqCst),
snapshot_b.load(Ordering::SeqCst) - base_b.load(Ordering::SeqCst),
"post-ramp-up window",
)
} else {
(
completed_a.load(Ordering::SeqCst),
completed_b_counter.load(Ordering::SeqCst),
"timeout fallback (full run)",
)
};
assert!(
ca > cb,
"Queue A (weight=3) should complete more jobs ({ca}) than B (weight=1, {cb}) in the {window}"
);
}
#[tokio::test]
async fn test_permit_before_claim_no_orphans() {
let pool = setup().await;
let queue = "wt_permit_claim";
clean_queue(&pool, queue).await;
for i in 0..20 {
insert_with(
&pool,
&WeightedJob { index: i },
InsertOpts {
queue: queue.into(),
..Default::default()
},
)
.await
.unwrap();
}
let completed = Arc::new(AtomicU32::new(0));
let client = Client::builder(pool.clone())
.queue(
queue,
QueueConfig {
min_workers: 3,
weight: 1,
poll_interval: Duration::from_millis(50),
..Default::default()
},
)
.global_max_workers(10)
.register_worker(SlowWorker {
completed: completed.clone(),
delay: Duration::from_millis(100),
})
.build()
.unwrap();
client.start().await.unwrap();
tokio::time::sleep(Duration::from_millis(500)).await;
client.shutdown(Duration::from_secs(5)).await;
let running: i64 =
sqlx::query_scalar("SELECT count(*) FROM awa.jobs WHERE queue = $1 AND state = 'running'")
.bind(queue)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(
running, 0,
"No jobs should be stuck in running state after shutdown"
);
}
#[tokio::test]
async fn test_health_check_weighted_mode() {
let pool = setup().await;
let queue = "wt_health";
clean_queue(&pool, queue).await;
let client = Client::builder(pool.clone())
.queue(
queue,
QueueConfig {
min_workers: 5,
weight: 3,
poll_interval: Duration::from_millis(200),
..Default::default()
},
)
.global_max_workers(20)
.register_worker(SlowWorker {
completed: Arc::new(AtomicU32::new(0)),
delay: Duration::from_millis(10),
})
.build()
.unwrap();
client.start().await.unwrap();
tokio::time::sleep(Duration::from_millis(300)).await;
let health = client.health_check().await;
let queue_health = health.queues.get(queue).expect("queue should be in health");
match &queue_health.capacity {
awa::QueueCapacity::Weighted {
min_workers,
weight,
..
} => {
assert_eq!(*min_workers, 5);
assert_eq!(*weight, 3);
}
other => panic!("Expected Weighted capacity, got: {other:?}"),
}
client.shutdown(Duration::from_secs(1)).await;
}
#[tokio::test]
async fn test_health_check_hard_reserved_mode() {
let pool = setup().await;
let queue = "wt_health_hard";
clean_queue(&pool, queue).await;
let client = Client::builder(pool.clone())
.queue(
queue,
QueueConfig {
max_workers: 42,
poll_interval: Duration::from_millis(200),
..Default::default()
},
)
.register_worker(SlowWorker {
completed: Arc::new(AtomicU32::new(0)),
delay: Duration::from_millis(10),
})
.build()
.unwrap();
client.start().await.unwrap();
tokio::time::sleep(Duration::from_millis(300)).await;
let health = client.health_check().await;
let queue_health = health.queues.get(queue).expect("queue should be in health");
match &queue_health.capacity {
awa::QueueCapacity::HardReserved { max_workers } => {
assert_eq!(*max_workers, 42);
}
other => panic!("Expected HardReserved capacity, got: {other:?}"),
}
client.shutdown(Duration::from_secs(1)).await;
}