azums 0.1.3

High-performance job queue & streaming engine for Rust — from embedded to cloud
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

// crates/postgresflow/tests/storm_control.rs
mod common;

use chrono::Utc;

use azums::jobs::{AttemptsRepo, JobsRepo, PolicyDecisionsRepo};
use common::setup_db;
use serial_test::serial;
use uuid::Uuid;

async fn upsert_queue_policy(
    pool: &sqlx::PgPool,
    queue: &str,
    max_attempts_per_minute: i32,
    max_in_flight: i32,
    throttle_delay_ms: i32,
) {
    sqlx::query(
        r#"
        INSERT INTO queue_policies (queue, max_attempts_per_minute, max_in_flight, throttle_delay_ms)
        VALUES ($1, $2, $3, $4)
        ON CONFLICT (queue) DO UPDATE
        SET max_attempts_per_minute = EXCLUDED.max_attempts_per_minute,
            max_in_flight = EXCLUDED.max_in_flight,
            throttle_delay_ms = EXCLUDED.throttle_delay_ms
        "#,
    )
    .bind(queue)
    .bind(max_attempts_per_minute)
    .bind(max_in_flight)
    .bind(throttle_delay_ms)
    .execute(pool)
    .await
    .unwrap();
}

async fn insert_job_direct(pool: &sqlx::PgPool, queue: &str, job_type: &str) -> Uuid {
    sqlx::query_scalar(
        r#"
        INSERT INTO jobs (queue, job_type, payload_json, run_at, status, priority, max_attempts)
        VALUES ($1, $2, '{}'::jsonb, now(), 'queued', 0, 5)
        RETURNING id
        "#,
    )
    .bind(queue)
    .bind(job_type)
    .fetch_one(pool)
    .await
    .unwrap()
}

#[tokio::test]
#[serial]
async fn throttles_when_in_flight_exceeded() {
    let Some(pool) = setup_db().await else {
        return;
    };

    let jobs = JobsRepo::new(pool.clone());
    let policies = PolicyDecisionsRepo::new(pool.clone());

    upsert_queue_policy(&pool, "default", 10_000, 2, 500).await;

    let job1_id = insert_job_direct(&pool, "default", "job1").await;
    let job2_id = insert_job_direct(&pool, "default", "job2").await;

    let leased1 = jobs
        .lease_one_job("default", "worker-a", 10)
        .await
        .unwrap()
        .expect("should lease first job");
    assert_eq!(leased1.id, job1_id);

    upsert_queue_policy(&pool, "default", 10_000, 1, 500).await;

    let before = Utc::now();
    let leased2 = jobs.lease_one_job("default", "worker-b", 10).await.unwrap();

    assert!(leased2.is_none());

    let j2 = jobs.get_job(job2_id).await.unwrap().unwrap();
    assert!(j2.run_at > before, "expected run_at pushed forward");

    let rows = policies.list_for_job(job2_id).await.unwrap();
    assert!(!rows.is_empty(), "expected policy decision row");

    let last = rows.last().unwrap();
    assert_eq!(last.decision, "THROTTLED");
    assert_eq!(last.reason_code, "IN_FLIGHT_EXCEEDED");
}

#[tokio::test]
#[serial]
async fn throttles_when_retry_rate_exceeded() {
    let Some(pool) = setup_db().await else {
        return;
    };

    let jobs = JobsRepo::new(pool.clone());
    let attempts = AttemptsRepo::new(pool.clone());
    let policies = PolicyDecisionsRepo::new(pool.clone());

    upsert_queue_policy(&pool, "default", 1, 999_999, 500).await;

    let seed_job_id = insert_job_direct(&pool, "default", "seed").await;

    let leased_seed = jobs
        .lease_one_job("default", "worker-a", 10)
        .await
        .unwrap()
        .expect("should lease seed job");
    assert_eq!(leased_seed.id, seed_job_id);

    let _a1 = attempts
        .start_attempt(seed_job_id, "worker-a")
        .await
        .unwrap();

    let target_job_id = insert_job_direct(&pool, "default", "target").await;

    let before = Utc::now();
    let leased = jobs.lease_one_job("default", "worker-b", 10).await.unwrap();

    assert!(leased.is_none(), "expected throttle -> no lease");

    let j = jobs.get_job(target_job_id).await.unwrap().unwrap();
    assert!(j.run_at > before, "expected run_at pushed forward");

    let rows = policies.list_for_job(target_job_id).await.unwrap();
    assert!(!rows.is_empty(), "expected policy decision row");

    let last = rows.last().unwrap();
    assert_eq!(last.decision, "THROTTLED");
    assert_eq!(last.reason_code, "RETRY_RATE_EXCEEDED");
}