azums 0.2.0

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

use azums::jobs::JobsRepo;
use sqlx::PgPool;
use uuid::Uuid;

mod common;

async fn upsert_policy(
    pool: &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(pool: &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, 3)
        RETURNING id
        "#,
    )
    .bind(queue)
    .bind(job_type)
    .fetch_one(pool)
    .await
    .unwrap()
}

async fn mark_running(pool: &PgPool, job_id: Uuid) {
    sqlx::query(
        r#"
        UPDATE jobs
        SET status = 'running',
            locked_by = 'someone',
            locked_at = now(),
            lock_expires_at = now() + interval '60 seconds'
        WHERE id = $1
        "#,
    )
    .bind(job_id)
    .execute(pool)
    .await
    .unwrap();
}

#[tokio::test]
async fn writes_policy_decision_when_in_flight_exceeded() {
    let Some(pool) = common::setup_db().await else {
        return;
    };
    let jobs = JobsRepo::new(pool.clone());

    let queue = "q_policy";
    // allow no in-flight at all, so any running job triggers throttling
    upsert_policy(&pool, queue, 9999, 0, 500).await;

    // create one RUNNING job -> in_flight becomes 1
    let running_id = insert_job(&pool, queue, "running_job").await;
    mark_running(&pool, running_id).await;

    // create one QUEUED job that should be blocked
    let queued_id = insert_job(&pool, queue, "blocked_job").await;

    // try leasing - should return None because in-flight exceeded
    let leased = jobs.lease_one_job(queue, "worker-a", 10).await.unwrap();
    assert!(leased.is_none());

    // verify a policy decision row exists for the queued job
    let row: Option<(String, String)> = sqlx::query_as(
        r#"
        SELECT decision, reason_code
        FROM policy_decisions
        WHERE job_id = $1
        ORDER BY created_at DESC
        LIMIT 1
        "#,
    )
    .bind(queued_id)
    .fetch_optional(&pool)
    .await
    .unwrap();

    let (decision, reason_code) = row.expect("expected policy_decisions row, found none");
    assert_eq!(decision, "THROTTLED");
    assert_eq!(reason_code, "IN_FLIGHT_EXCEEDED");

    // also verify job got rescheduled into the future
    let job = jobs.get_job(queued_id).await.unwrap().unwrap();
    assert!(job.run_at > chrono::Utc::now());
}