OXYGENHUB
# ab2d8b40f9f4 oxygenhub / crates / hub / tests / roles.rs
11.4 KiB · Rust · last changed Flow's job times in /metrics (refs oxygenflow#35, closes #212) 15 days ago source blame history raw
1
//! Roles (docs/roles.md, #272): a web-only hub starts none of the worker's
2
//! background tasks, a worker-only hub serves no pages, and on PostgreSQL
3
//! two hubs on one database run each singleton task once between them.
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
/// The worker's singleton tasks that run for as long as the process does.
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
/// Wait up to `limit` for `ok`, checking every 100 ms.
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
    // Only the transparency log's checkpointer, which goes with the log's
78
    // appends in the web role.
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
    // It serves pages.
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
    // Flow's job times (oxygenflow#35) and the work queue (#276).
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
    // It reads the transparency log but never writes it.
123
    assert!(app.tlog.append(&[]).is_err());
124
}
125
126
/// oxygenhub#275: a named lock is held by one holder at a time across the
127
/// database's sessions (on SQLite, the process); other names are free.
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
/// oxygenhub#276: the work queue. A keyed item is queued once; a claim
151
/// leases it (another claim skips it); a failure backs off and the last
152
/// attempt leaves it dead; done removes it. An item no worker knows is dead
153
/// at once.
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
    // A dead item does not hold its key: the same work queues again.
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
    // Repository 7 does not exist: nothing to refresh, so it is done.
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
/// Two hubs on one PostgreSQL database (one schema): exactly one leads
200
/// each singleton; when its connection drops it stops them, and when it
201
/// stops the other takes every one over within seconds.
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
    // A's connection is ended, as when the network drops: its locks free,
224
    // it stops what it led, and between them each runs once again.
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
    // The hub leading the scheduler stops; the other takes everything over.
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
}