| 1 |
|
| 2 |
|
| 3 |
|
| 4 |
|
| 5 |
|
| 6 |
|
| 7 |
|
| 8 |
|
| 9 |
|
| 10 |
|
| 11 |
|
| 12 |
use crate::harness::TestHarness; |
| 13 |
use makenotwork::db::UserId; |
| 14 |
use makenotwork::db::scan_jobs::{self, ScanTargetKind}; |
| 15 |
use makenotwork::storage::FileType; |
| 16 |
use uuid::Uuid; |
| 17 |
|
| 18 |
async fn seed_user(h: &TestHarness, name: &str) -> UserId { |
| 19 |
let hash = makenotwork::auth::hash_password("password123").expect("hash"); |
| 20 |
sqlx::query_scalar::<_, UserId>( |
| 21 |
"INSERT INTO users (username, email, password_hash, email_verified) |
| 22 |
VALUES ($1, $2, $3, true) RETURNING id", |
| 23 |
) |
| 24 |
.bind(name) |
| 25 |
.bind(format!("{name}@test.com")) |
| 26 |
.bind(&hash) |
| 27 |
.fetch_one(&h.db) |
| 28 |
.await |
| 29 |
.expect("seed user") |
| 30 |
} |
| 31 |
|
| 32 |
async fn status_of(h: &TestHarness, job: Uuid) -> String { |
| 33 |
sqlx::query_scalar::<_, String>("SELECT status FROM scan_jobs WHERE id = $1") |
| 34 |
.bind(job) |
| 35 |
.fetch_one(&h.db) |
| 36 |
.await |
| 37 |
.expect("status") |
| 38 |
} |
| 39 |
|
| 40 |
|
| 41 |
|
| 42 |
async fn enqueue_and_claim(h: &TestHarness, user: UserId, key: &str) -> Uuid { |
| 43 |
scan_jobs::enqueue( |
| 44 |
&h.db, |
| 45 |
ScanTargetKind::Item, |
| 46 |
Uuid::new_v4(), |
| 47 |
key, |
| 48 |
FileType::Download, |
| 49 |
user, |
| 50 |
1000, |
| 51 |
) |
| 52 |
.await |
| 53 |
.expect("enqueue"); |
| 54 |
let job = scan_jobs::claim_next(&h.db) |
| 55 |
.await |
| 56 |
.expect("claim") |
| 57 |
.expect("a queued job to claim"); |
| 58 |
assert_eq!(job.status, "running"); |
| 59 |
assert!(job.heartbeat_at.is_some(), "claim stamps heartbeat_at"); |
| 60 |
job.id |
| 61 |
} |
| 62 |
|
| 63 |
#[tokio::test] |
| 64 |
async fn reaper_spares_slow_but_alive_job() { |
| 65 |
let h = TestHarness::new().await; |
| 66 |
let user = seed_user(&h, "reaper_alive").await; |
| 67 |
let job = enqueue_and_claim(&h, user, "k/alive").await; |
| 68 |
|
| 69 |
|
| 70 |
|
| 71 |
|
| 72 |
sqlx::query( |
| 73 |
"UPDATE scan_jobs SET started_at = NOW() - interval '1 hour', heartbeat_at = NOW() WHERE id = $1", |
| 74 |
) |
| 75 |
.bind(job) |
| 76 |
.execute(&h.db) |
| 77 |
.await |
| 78 |
.expect("backdate started, fresh beat"); |
| 79 |
|
| 80 |
let reaped = scan_jobs::reap_stuck(&h.db, 300).await.expect("reap"); |
| 81 |
assert_eq!(reaped, 0, "a job with a fresh heartbeat must not be reaped"); |
| 82 |
assert_eq!(status_of(&h, job).await, "running"); |
| 83 |
} |
| 84 |
|
| 85 |
#[tokio::test] |
| 86 |
async fn reaper_reaps_crashed_worker() { |
| 87 |
let h = TestHarness::new().await; |
| 88 |
let user = seed_user(&h, "reaper_dead").await; |
| 89 |
let job = enqueue_and_claim(&h, user, "k/dead").await; |
| 90 |
|
| 91 |
|
| 92 |
sqlx::query("UPDATE scan_jobs SET heartbeat_at = NOW() - interval '1 hour' WHERE id = $1") |
| 93 |
.bind(job) |
| 94 |
.execute(&h.db) |
| 95 |
.await |
| 96 |
.expect("stale beat"); |
| 97 |
|
| 98 |
let reaped = scan_jobs::reap_stuck(&h.db, 300).await.expect("reap"); |
| 99 |
assert_eq!(reaped, 1, "a stale-heartbeat job must be reaped"); |
| 100 |
|
| 101 |
|
| 102 |
assert_eq!(status_of(&h, job).await, "queued"); |
| 103 |
} |
| 104 |
|
| 105 |
#[tokio::test] |
| 106 |
async fn reaper_falls_back_to_started_at_when_no_heartbeat() { |
| 107 |
|
| 108 |
|
| 109 |
let h = TestHarness::new().await; |
| 110 |
let user = seed_user(&h, "reaper_null").await; |
| 111 |
let job = enqueue_and_claim(&h, user, "k/null").await; |
| 112 |
|
| 113 |
sqlx::query( |
| 114 |
"UPDATE scan_jobs SET heartbeat_at = NULL, started_at = NOW() - interval '1 hour' WHERE id = $1", |
| 115 |
) |
| 116 |
.bind(job) |
| 117 |
.execute(&h.db) |
| 118 |
.await |
| 119 |
.expect("null beat, old start"); |
| 120 |
|
| 121 |
let reaped = scan_jobs::reap_stuck(&h.db, 300).await.expect("reap"); |
| 122 |
assert_eq!(reaped, 1, "null heartbeat falls back to started_at"); |
| 123 |
} |
| 124 |
|
| 125 |
#[tokio::test] |
| 126 |
async fn bump_heartbeat_refreshes_and_never_resurrects() { |
| 127 |
let h = TestHarness::new().await; |
| 128 |
let user = seed_user(&h, "reaper_bump").await; |
| 129 |
let job = enqueue_and_claim(&h, user, "k/bump").await; |
| 130 |
|
| 131 |
|
| 132 |
sqlx::query("UPDATE scan_jobs SET heartbeat_at = NOW() - interval '1 hour' WHERE id = $1") |
| 133 |
.bind(job) |
| 134 |
.execute(&h.db) |
| 135 |
.await |
| 136 |
.expect("stale"); |
| 137 |
|
| 138 |
scan_jobs::bump_heartbeat(&h.db, job).await.expect("bump"); |
| 139 |
let reaped = scan_jobs::reap_stuck(&h.db, 300).await.expect("reap"); |
| 140 |
assert_eq!(reaped, 0, "bump_heartbeat must refresh liveness"); |
| 141 |
assert_eq!(status_of(&h, job).await, "running"); |
| 142 |
|
| 143 |
|
| 144 |
|
| 145 |
sqlx::query("UPDATE scan_jobs SET status = 'done', heartbeat_at = NULL WHERE id = $1") |
| 146 |
.bind(job) |
| 147 |
.execute(&h.db) |
| 148 |
.await |
| 149 |
.expect("mark done"); |
| 150 |
scan_jobs::bump_heartbeat(&h.db, job) |
| 151 |
.await |
| 152 |
.expect("late bump"); |
| 153 |
let still_null: bool = |
| 154 |
sqlx::query_scalar("SELECT heartbeat_at IS NULL FROM scan_jobs WHERE id = $1") |
| 155 |
.bind(job) |
| 156 |
.fetch_one(&h.db) |
| 157 |
.await |
| 158 |
.expect("hb null check"); |
| 159 |
assert!( |
| 160 |
still_null, |
| 161 |
"bump must not touch a job that has left running" |
| 162 |
); |
| 163 |
} |
| 164 |
|