OXYGENHUB
◆ v0.6.0 oxygenhub / crates / hub / tests / roles.rs
7.6 KiB · Rust · last changed Run the hub as web and worker processes, singletons under a leader lock 16 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; 7] =
16
    ["backups.scheduler", "diffusion.reconciler", "flow.placement", "flow.reaper", "flow.scheduler", "hooks.worker", "maintenance.nightly"];
17
18
fn hermetic() {
19
    static ONCE: std::sync::Once = std::sync::Once::new();
20
    static DIR: std::sync::OnceLock<tempfile::TempDir> = std::sync::OnceLock::new();
21
    ONCE.call_once(|| {
22
        let d = tempfile::tempdir().unwrap();
23
        std::env::set_var("OXY_HOME", d.path());
24
        DIR.set(d).ok();
25
    });
26
}
27
28
fn config(dir: &Path, roles: &[&str]) -> Config {
29
    let mut config = Config::load(dir).unwrap();
30
    config.server.roles = roles.iter().map(|r| r.to_string()).collect();
31
    config.registration.setup_code = "test-setup".into();
32
    config
33
}
34
35
async fn get(app: &Shared, path: &str, token: Option<&str>) -> (StatusCode, String) {
36
    let mut req = Request::get(path);
37
    if let Some(t) = token {
38
        req = req.header(header::AUTHORIZATION, format!("Bearer {t}"));
39
    }
40
    let res = oxygenhub::router(app.clone()).oneshot(req.body(Body::empty()).unwrap()).await.unwrap();
41
    let status = res.status();
42
    let body = res.into_body().collect().await.unwrap().to_bytes();
43
    (status, String::from_utf8_lossy(&body).into_owned())
44
}
45
46
async fn health(app: &Shared) -> serde_json::Value {
47
    let (status, body) = get(app, "/health", None).await;
48
    assert_eq!(status, StatusCode::OK, "{body}");
49
    serde_json::from_str(&body).unwrap()
50
}
51
52
/// Wait up to `limit` for `ok`, checking every 100 ms.
53
async fn until(limit: Duration, mut ok: impl FnMut() -> bool) -> bool {
54
    let start = Instant::now();
55
    while start.elapsed() < limit {
56
        if ok() {
57
            return true;
58
        }
59
        tokio::time::sleep(Duration::from_millis(100)).await;
60
    }
61
    ok()
62
}
63
64
#[tokio::test(flavor = "multi_thread")]
65
async fn a_web_only_hub_starts_none_of_the_workers_tasks() {
66
    hermetic();
67
    let dir = tempfile::tempdir().unwrap();
68
    let app = oxygenhub::build(config(dir.path(), &["web"])).await.unwrap();
69
    // Only the transparency log's checkpointer, which goes with the log's
70
    // appends in the web role.
71
    assert_eq!(app.leader.tasks(), ["tlog.checkpointer"]);
72
    assert!(until(Duration::from_secs(5), || app.leader.leading() == ["tlog.checkpointer"]).await, "{:?}", app.leader.leading());
73
    let h = health(&app).await;
74
    assert_eq!(h["roles"], serde_json::json!(["web"]));
75
    assert_eq!(h["leading"], serde_json::json!(["tlog.checkpointer"]));
76
    // It serves pages.
77
    let (status, body) = get(&app, "/explore", None).await;
78
    assert_eq!(status, StatusCode::OK, "{body}");
79
}
80
81
#[tokio::test(flavor = "multi_thread")]
82
async fn a_worker_only_hub_serves_its_health_and_metrics_and_no_pages() {
83
    hermetic();
84
    let dir = tempfile::tempdir().unwrap();
85
    let mut cfg = config(dir.path(), &["worker"]);
86
    cfg.server.metrics_token = "m".into();
87
    let app = oxygenhub::build(cfg).await.unwrap();
88
    let mut tasks = WORKER_SINGLETONS.to_vec();
89
    tasks.extend(["diffusion.backfill", "images.learn_archs"]);
90
    tasks.sort();
91
    assert_eq!(app.leader.tasks(), tasks, "every worker task, and not the log's checkpointer");
92
    assert!(
93
        until(Duration::from_secs(5), || WORKER_SINGLETONS.iter().all(|n| app.leader.leading().contains(n))).await,
94
        "{:?}",
95
        app.leader.leading()
96
    );
97
    let h = health(&app).await;
98
    assert_eq!(h["roles"], serde_json::json!(["worker"]));
99
    assert!(h["leading"].as_array().unwrap().iter().any(|n| n == "flow.scheduler"), "{h}");
100
    for path in ["/", "/explore", "/login", "/api/v1/version", "/andy/demo", "/log/checkpoint"] {
101
        assert_eq!(get(&app, path, None).await.0, StatusCode::NOT_FOUND, "{path} is the web role's");
102
    }
103
    let (status, body) = get(&app, "/metrics", Some("m")).await;
104
    assert_eq!(status, StatusCode::OK, "{body}");
105
    // It reads the transparency log but never writes it.
106
    assert!(app.tlog.append(&[]).is_err());
107
}
108
109
#[tokio::test(flavor = "multi_thread")]
110
async fn a_hub_takes_every_role_by_default() {
111
    hermetic();
112
    let dir = tempfile::tempdir().unwrap();
113
    let app = oxygenhub::build(config(dir.path(), &[])).await.unwrap();
114
    assert_eq!(app.leader.tasks().len(), 10, "{:?}", app.leader.tasks());
115
    assert_eq!(health(&app).await["roles"], serde_json::json!(["web", "worker"]));
116
    assert_eq!(get(&app, "/explore", None).await.0, StatusCode::OK);
117
    let mut bad = config(dir.path(), &["web", "cron"]);
118
    bad.server.roles.push(String::new());
119
    let e = oxygenhub::build(bad).await.err().expect("an unknown role is refused").to_string();
120
    assert!(e.contains("unknown role \"cron\""), "{e}");
121
}
122
123
/// Two hubs on one PostgreSQL database (one schema): exactly one leads
124
/// each singleton; when its connection drops it stops them, and when it
125
/// stops the other takes every one over within seconds.
126
#[cfg(feature = "postgres")]
127
#[tokio::test(flavor = "multi_thread")]
128
async fn two_hubs_on_one_database_run_each_singleton_once_and_fail_over() {
129
    hermetic();
130
    let (da, dbb) = (tempfile::tempdir().unwrap(), tempfile::tempdir().unwrap());
131
    let schema = format!("hub_roles_{}", da.path().file_name().unwrap().to_string_lossy());
132
    let pg = |dir: &Path| {
133
        let mut c = config(dir, &[]);
134
        c.database.schema = schema.clone();
135
        c
136
    };
137
    let a = oxygenhub::build(pg(da.path())).await.unwrap();
138
    let b = oxygenhub::build(pg(dbb.path())).await.unwrap();
139
    let mut singletons = WORKER_SINGLETONS.to_vec();
140
    singletons.push("tlog.checkpointer");
141
    let once = |a: &Shared, b: &Shared| {
142
        let (la, lb) = (a.leader.leading(), b.leader.leading());
143
        singletons.iter().all(|n| la.contains(n) != lb.contains(n))
144
    };
145
    assert!(until(Duration::from_secs(10), || once(&a, &b)).await, "a {:?}, b {:?}", a.leader.leading(), b.leader.leading());
146
147
    // A's connection is ended, as when the network drops: its locks free,
148
    // it stops what it led, and between them each runs once again.
149
    let pid = a.leader.session_pid().await.expect("a has a session");
150
    let ended: bool = sqlx::query_scalar("SELECT pg_terminate_backend($1)").bind(pid).fetch_one(&b.db).await.unwrap();
151
    assert!(ended);
152
    let start = Instant::now();
153
    loop {
154
        let reconnected = a.leader.session_pid().await.is_some_and(|p| p != pid);
155
        if reconnected && once(&a, &b) {
156
            break;
157
        }
158
        assert!(start.elapsed() < Duration::from_secs(15), "a {:?}, b {:?}", a.leader.leading(), b.leader.leading());
159
        tokio::time::sleep(Duration::from_millis(100)).await;
160
    }
161
162
    // The hub leading the scheduler stops; the other takes everything over.
163
    let (gone, left) = if a.leader.leading().contains(&"flow.scheduler") { (&a, &b) } else { (&b, &a) };
164
    gone.leader.resign().await;
165
    let start = Instant::now();
166
    assert!(
167
        until(Duration::from_secs(10), || singletons.iter().all(|n| left.leader.leading().contains(n))).await,
168
        "left {:?}",
169
        left.leader.leading()
170
    );
171
    eprintln!("took over in {:.1} s", start.elapsed().as_secs_f64());
172
    assert!(gone.leader.leading().is_empty(), "{:?}", gone.leader.leading());
173
    assert_eq!(health(left).await["leading"].as_array().unwrap().len(), left.leader.leading().len());
174
}