| 1 |
|
| 2 |
|
| 3 |
|
| 4 |
|
| 5 |
use axum::body::Body; |
| 6 |
use axum::http::{header, Request, StatusCode}; |
| 7 |
use http_body_util::BodyExt; |
| 8 |
use oxygenhub::config::Config; |
| 9 |
use oxygenhub::Shared; |
| 10 |
use std::path::Path; |
| 11 |
use std::time::{Duration, Instant}; |
| 12 |
use tower::ServiceExt; |
| 13 |
|
| 14 |
|
| 15 |
const WORKER_SINGLETONS: [&str; 8] = [ |
| 16 |
"backups.scheduler", |
| 17 |
"diffusion.reconciler", |
| 18 |
"flow.placement", |
| 19 |
"flow.reaper", |
| 20 |
"flow.scheduler", |
| 21 |
"hooks.worker", |
| 22 |
"maintenance.nightly", |
| 23 |
"secrets.history_scan", |
| 24 |
]; |
| 25 |
|
| 26 |
fn hermetic() { |
| 27 |
static ONCE: std::sync::Once = std::sync::Once::new(); |
| 28 |
static DIR: std::sync::OnceLock<tempfile::TempDir> = std::sync::OnceLock::new(); |
| 29 |
ONCE.call_once(|| { |
| 30 |
let d = tempfile::tempdir().unwrap(); |
| 31 |
std::env::set_var("OXY_HOME", d.path()); |
| 32 |
DIR.set(d).ok(); |
| 33 |
}); |
| 34 |
} |
| 35 |
|
| 36 |
fn config(dir: &Path, roles: &[&str]) -> Config { |
| 37 |
let mut config = Config::load(dir).unwrap(); |
| 38 |
config.server.roles = roles.iter().map(|r| r.to_string()).collect(); |
| 39 |
config.registration.setup_code = "test-setup".into(); |
| 40 |
config |
| 41 |
} |
| 42 |
|
| 43 |
async fn get(app: &Shared, path: &str, token: Option<&str>) -> (StatusCode, String) { |
| 44 |
let mut req = Request::get(path); |
| 45 |
if let Some(t) = token { |
| 46 |
req = req.header(header::AUTHORIZATION, format!("Bearer {t}")); |
| 47 |
} |
| 48 |
let res = oxygenhub::router(app.clone()).oneshot(req.body(Body::empty()).unwrap()).await.unwrap(); |
| 49 |
let status = res.status(); |
| 50 |
let body = res.into_body().collect().await.unwrap().to_bytes(); |
| 51 |
(status, String::from_utf8_lossy(&body).into_owned()) |
| 52 |
} |
| 53 |
|
| 54 |
async fn health(app: &Shared) -> serde_json::Value { |
| 55 |
let (status, body) = get(app, "/health", None).await; |
| 56 |
assert_eq!(status, StatusCode::OK, "{body}"); |
| 57 |
serde_json::from_str(&body).unwrap() |
| 58 |
} |
| 59 |
|
| 60 |
|
| 61 |
async fn until(limit: Duration, mut ok: impl FnMut() -> bool) -> bool { |
| 62 |
let start = Instant::now(); |
| 63 |
while start.elapsed() < limit { |
| 64 |
if ok() { |
| 65 |
return true; |
| 66 |
} |
| 67 |
tokio::time::sleep(Duration::from_millis(100)).await; |
| 68 |
} |
| 69 |
ok() |
| 70 |
} |
| 71 |
|
| 72 |
#[tokio::test(flavor = "multi_thread")] |
| 73 |
async fn a_web_only_hub_starts_none_of_the_workers_tasks() { |
| 74 |
hermetic(); |
| 75 |
let dir = tempfile::tempdir().unwrap(); |
| 76 |
let app = oxygenhub::build(config(dir.path(), &["web"])).await.unwrap(); |
| 77 |
|
| 78 |
|
| 79 |
assert_eq!(app.leader.tasks(), ["tlog.checkpointer"]); |
| 80 |
assert!(until(Duration::from_secs(5), || app.leader.leading() == ["tlog.checkpointer"]).await, "{:?}", app.leader.leading()); |
| 81 |
let h = health(&app).await; |
| 82 |
assert_eq!(h["roles"], serde_json::json!(["web"])); |
| 83 |
assert_eq!(h["leading"], serde_json::json!(["tlog.checkpointer"])); |
| 84 |
|
| 85 |
let (status, body) = get(&app, "/explore", None).await; |
| 86 |
assert_eq!(status, StatusCode::OK, "{body}"); |
| 87 |
} |
| 88 |
|
| 89 |
#[tokio::test(flavor = "multi_thread")] |
| 90 |
async fn a_worker_only_hub_serves_its_health_and_metrics_and_no_pages() { |
| 91 |
hermetic(); |
| 92 |
let dir = tempfile::tempdir().unwrap(); |
| 93 |
let mut cfg = config(dir.path(), &["worker"]); |
| 94 |
cfg.server.metrics_token = "m".into(); |
| 95 |
let app = oxygenhub::build(cfg).await.unwrap(); |
| 96 |
let mut tasks = WORKER_SINGLETONS.to_vec(); |
| 97 |
tasks.extend(["diffusion.backfill", "images.learn_archs", "queue.worker"]); |
| 98 |
tasks.sort(); |
| 99 |
assert_eq!(app.leader.tasks(), tasks, "every worker task, and not the log's checkpointer"); |
| 100 |
assert!( |
| 101 |
until(Duration::from_secs(5), || WORKER_SINGLETONS.iter().all(|n| app.leader.leading().contains(n))).await, |
| 102 |
"{:?}", |
| 103 |
app.leader.leading() |
| 104 |
); |
| 105 |
let h = health(&app).await; |
| 106 |
assert_eq!(h["roles"], serde_json::json!(["worker"])); |
| 107 |
assert!(h["leading"].as_array().unwrap().iter().any(|n| n == "flow.scheduler"), "{h}"); |
| 108 |
for path in ["/", "/explore", "/login", "/api/v1/version", "/andy/demo", "/log/checkpoint"] { |
| 109 |
assert_eq!(get(&app, path, None).await.0, StatusCode::NOT_FOUND, "{path} is the web role's"); |
| 110 |
} |
| 111 |
let (status, body) = get(&app, "/metrics", Some("m")).await; |
| 112 |
assert_eq!(status, StatusCode::OK, "{body}"); |
| 113 |
|
| 114 |
for family in [ |
| 115 |
"# TYPE oxyhub_flow_job_queue_seconds histogram", |
| 116 |
"# TYPE oxyhub_flow_job_run_seconds histogram", |
| 117 |
"oxyhub_flow_jobs_reused_last_day 0", |
| 118 |
"oxyhub_queue_items{state=\"dead\"} 0", |
| 119 |
] { |
| 120 |
assert!(body.contains(family), "{family}: {body}"); |
| 121 |
} |
| 122 |
|
| 123 |
assert!(app.tlog.append(&[]).is_err()); |
| 124 |
} |
| 125 |
|
| 126 |
|
| 127 |
|
| 128 |
#[tokio::test(flavor = "multi_thread")] |
| 129 |
async fn an_exclusive_lock_has_one_holder_at_a_time() { |
| 130 |
hermetic(); |
| 131 |
let dir = tempfile::tempdir().unwrap(); |
| 132 |
let app = oxygenhub::build(config(dir.path(), &[])).await.unwrap(); |
| 133 |
let held = oxygenhub::leader::exclusive(&app.db, "autoland:7").await.unwrap(); |
| 134 |
let (tx, mut rx) = tokio::sync::oneshot::channel(); |
| 135 |
let db = app.db.clone(); |
| 136 |
tokio::spawn(async move { |
| 137 |
let g = oxygenhub::leader::exclusive(&db, "autoland:7").await.unwrap(); |
| 138 |
let _ = tx.send(()); |
| 139 |
drop(g); |
| 140 |
}); |
| 141 |
tokio::time::sleep(Duration::from_millis(300)).await; |
| 142 |
assert!(rx.try_recv().is_err(), "the second holder waits"); |
| 143 |
let other = tokio::time::timeout(Duration::from_secs(2), oxygenhub::leader::exclusive(&app.db, "autoland:8")).await; |
| 144 |
assert!(other.is_ok(), "another name is free"); |
| 145 |
drop(other); |
| 146 |
drop(held); |
| 147 |
assert!(tokio::time::timeout(Duration::from_secs(5), rx).await.is_ok(), "and takes it once it is let go"); |
| 148 |
} |
| 149 |
|
| 150 |
|
| 151 |
|
| 152 |
|
| 153 |
|
| 154 |
#[tokio::test(flavor = "multi_thread")] |
| 155 |
async fn the_work_queue_leases_retries_and_keeps_the_dead() { |
| 156 |
use oxygenhub::db::queue; |
| 157 |
hermetic(); |
| 158 |
let dir = tempfile::tempdir().unwrap(); |
| 159 |
let app = oxygenhub::build(config(dir.path(), &["web"])).await.unwrap(); |
| 160 |
let payload = serde_json::json!({ "repo_id": 7 }); |
| 161 |
assert!(queue::enqueue(&app.db, "mirror.refresh", &payload, "7", 0).await.unwrap()); |
| 162 |
assert!(!queue::enqueue(&app.db, "mirror.refresh", &payload, "7", 0).await.unwrap(), "once per key"); |
| 163 |
let got = queue::claim(&app.db, "a", 10, 60_000).await.unwrap(); |
| 164 |
assert_eq!((got.len(), got[0].attempts, got[0].claimed_by.as_str()), (1, 1, "a")); |
| 165 |
assert!(queue::claim(&app.db, "b", 10, 60_000).await.unwrap().is_empty(), "leased to a"); |
| 166 |
queue::failed(&app.db, &got[0], "the disk is full").await.unwrap(); |
| 167 |
assert!(queue::claim(&app.db, "b", 10, 60_000).await.unwrap().is_empty(), "backing off"); |
| 168 |
sqlx::query("UPDATE work_queue SET run_after = 0, attempts = $1").bind(queue::MAX_ATTEMPTS - 1).execute(&app.db).await.unwrap(); |
| 169 |
let last = queue::claim(&app.db, "b", 10, 60_000).await.unwrap(); |
| 170 |
assert_eq!(last[0].attempts, queue::MAX_ATTEMPTS); |
| 171 |
queue::failed(&app.db, &last[0], "still full").await.unwrap(); |
| 172 |
assert_eq!(queue::counts(&app.db).await.unwrap(), (0, 1), "dead after the last attempt"); |
| 173 |
|
| 174 |
assert!(queue::enqueue(&app.db, "mirror.refresh", &payload, "7", 0).await.unwrap()); |
| 175 |
let item = queue::claim(&app.db, "a", 10, 60_000).await.unwrap().remove(0); |
| 176 |
|
| 177 |
oxygenhub::queue::run(&app, item).await; |
| 178 |
assert_eq!(queue::counts(&app.db).await.unwrap(), (0, 1)); |
| 179 |
queue::enqueue(&app.db, "no.such.kind", &serde_json::json!({}), "", 0).await.unwrap(); |
| 180 |
let item = queue::claim(&app.db, "a", 10, 60_000).await.unwrap().remove(0); |
| 181 |
oxygenhub::queue::run(&app, item).await; |
| 182 |
assert_eq!(queue::counts(&app.db).await.unwrap(), (0, 2), "an unknown kind is dead at once"); |
| 183 |
} |
| 184 |
|
| 185 |
#[tokio::test(flavor = "multi_thread")] |
| 186 |
async fn a_hub_takes_every_role_by_default() { |
| 187 |
hermetic(); |
| 188 |
let dir = tempfile::tempdir().unwrap(); |
| 189 |
let app = oxygenhub::build(config(dir.path(), &[])).await.unwrap(); |
| 190 |
assert_eq!(app.leader.tasks().len(), 12, "{:?}", app.leader.tasks()); |
| 191 |
assert_eq!(health(&app).await["roles"], serde_json::json!(["web", "worker"])); |
| 192 |
assert_eq!(get(&app, "/explore", None).await.0, StatusCode::OK); |
| 193 |
let mut bad = config(dir.path(), &["web", "cron"]); |
| 194 |
bad.server.roles.push(String::new()); |
| 195 |
let e = oxygenhub::build(bad).await.err().expect("an unknown role is refused").to_string(); |
| 196 |
assert!(e.contains("unknown role \"cron\""), "{e}"); |
| 197 |
} |
| 198 |
|
| 199 |
|
| 200 |
|
| 201 |
|
| 202 |
#[cfg(feature = "postgres")] |
| 203 |
#[tokio::test(flavor = "multi_thread")] |
| 204 |
async fn two_hubs_on_one_database_run_each_singleton_once_and_fail_over() { |
| 205 |
hermetic(); |
| 206 |
let (da, dbb) = (tempfile::tempdir().unwrap(), tempfile::tempdir().unwrap()); |
| 207 |
let schema = format!("hub_roles_{}", da.path().file_name().unwrap().to_string_lossy()); |
| 208 |
let pg = |dir: &Path| { |
| 209 |
let mut c = config(dir, &[]); |
| 210 |
c.database.schema = schema.clone(); |
| 211 |
c |
| 212 |
}; |
| 213 |
let a = oxygenhub::build(pg(da.path())).await.unwrap(); |
| 214 |
let b = oxygenhub::build(pg(dbb.path())).await.unwrap(); |
| 215 |
let mut singletons = WORKER_SINGLETONS.to_vec(); |
| 216 |
singletons.push("tlog.checkpointer"); |
| 217 |
let once = |a: &Shared, b: &Shared| { |
| 218 |
let (la, lb) = (a.leader.leading(), b.leader.leading()); |
| 219 |
singletons.iter().all(|n| la.contains(n) != lb.contains(n)) |
| 220 |
}; |
| 221 |
assert!(until(Duration::from_secs(10), || once(&a, &b)).await, "a {:?}, b {:?}", a.leader.leading(), b.leader.leading()); |
| 222 |
|
| 223 |
|
| 224 |
|
| 225 |
let pid = a.leader.session_pid().await.expect("a has a session"); |
| 226 |
let ended: bool = sqlx::query_scalar("SELECT pg_terminate_backend($1)").bind(pid).fetch_one(&b.db).await.unwrap(); |
| 227 |
assert!(ended); |
| 228 |
let start = Instant::now(); |
| 229 |
loop { |
| 230 |
let reconnected = a.leader.session_pid().await.is_some_and(|p| p != pid); |
| 231 |
if reconnected && once(&a, &b) { |
| 232 |
break; |
| 233 |
} |
| 234 |
assert!(start.elapsed() < Duration::from_secs(15), "a {:?}, b {:?}", a.leader.leading(), b.leader.leading()); |
| 235 |
tokio::time::sleep(Duration::from_millis(100)).await; |
| 236 |
} |
| 237 |
|
| 238 |
|
| 239 |
let (gone, left) = if a.leader.leading().contains(&"flow.scheduler") { (&a, &b) } else { (&b, &a) }; |
| 240 |
gone.leader.resign().await; |
| 241 |
let start = Instant::now(); |
| 242 |
assert!( |
| 243 |
until(Duration::from_secs(10), || singletons.iter().all(|n| left.leader.leading().contains(n))).await, |
| 244 |
"left {:?}", |
| 245 |
left.leader.leading() |
| 246 |
); |
| 247 |
eprintln!("took over in {:.1} s", start.elapsed().as_secs_f64()); |
| 248 |
assert!(gone.leader.leading().is_empty(), "{:?}", gone.leader.leading()); |
| 249 |
assert_eq!(health(left).await["leading"].as_array().unwrap().len(), left.leader.leading().len()); |
| 250 |
} |