//! Roles (docs/roles.md, #272): a web-only hub starts none of the worker's //! background tasks, a worker-only hub serves no pages, and on PostgreSQL //! two hubs on one database run each singleton task once between them. use axum::body::Body; use axum::http::{header, Request, StatusCode}; use http_body_util::BodyExt; use oxygenhub::config::Config; use oxygenhub::Shared; use std::path::Path; use std::time::{Duration, Instant}; use tower::ServiceExt; /// The worker's singleton tasks that run for as long as the process does. const WORKER_SINGLETONS: [&str; 7] = ["backups.scheduler", "diffusion.reconciler", "flow.placement", "flow.reaper", "flow.scheduler", "hooks.worker", "maintenance.nightly"]; fn hermetic() { static ONCE: std::sync::Once = std::sync::Once::new(); static DIR: std::sync::OnceLock = std::sync::OnceLock::new(); ONCE.call_once(|| { let d = tempfile::tempdir().unwrap(); std::env::set_var("OXY_HOME", d.path()); DIR.set(d).ok(); }); } fn config(dir: &Path, roles: &[&str]) -> Config { let mut config = Config::load(dir).unwrap(); config.server.roles = roles.iter().map(|r| r.to_string()).collect(); config.registration.setup_code = "test-setup".into(); config } async fn get(app: &Shared, path: &str, token: Option<&str>) -> (StatusCode, String) { let mut req = Request::get(path); if let Some(t) = token { req = req.header(header::AUTHORIZATION, format!("Bearer {t}")); } let res = oxygenhub::router(app.clone()).oneshot(req.body(Body::empty()).unwrap()).await.unwrap(); let status = res.status(); let body = res.into_body().collect().await.unwrap().to_bytes(); (status, String::from_utf8_lossy(&body).into_owned()) } async fn health(app: &Shared) -> serde_json::Value { let (status, body) = get(app, "/health", None).await; assert_eq!(status, StatusCode::OK, "{body}"); serde_json::from_str(&body).unwrap() } /// Wait up to `limit` for `ok`, checking every 100 ms. async fn until(limit: Duration, mut ok: impl FnMut() -> bool) -> bool { let start = Instant::now(); while start.elapsed() < limit { if ok() { return true; } tokio::time::sleep(Duration::from_millis(100)).await; } ok() } #[tokio::test(flavor = "multi_thread")] async fn a_web_only_hub_starts_none_of_the_workers_tasks() { hermetic(); let dir = tempfile::tempdir().unwrap(); let app = oxygenhub::build(config(dir.path(), &["web"])).await.unwrap(); // Only the transparency log's checkpointer, which goes with the log's // appends in the web role. assert_eq!(app.leader.tasks(), ["tlog.checkpointer"]); assert!(until(Duration::from_secs(5), || app.leader.leading() == ["tlog.checkpointer"]).await, "{:?}", app.leader.leading()); let h = health(&app).await; assert_eq!(h["roles"], serde_json::json!(["web"])); assert_eq!(h["leading"], serde_json::json!(["tlog.checkpointer"])); // It serves pages. let (status, body) = get(&app, "/explore", None).await; assert_eq!(status, StatusCode::OK, "{body}"); } #[tokio::test(flavor = "multi_thread")] async fn a_worker_only_hub_serves_its_health_and_metrics_and_no_pages() { hermetic(); let dir = tempfile::tempdir().unwrap(); let mut cfg = config(dir.path(), &["worker"]); cfg.server.metrics_token = "m".into(); let app = oxygenhub::build(cfg).await.unwrap(); let mut tasks = WORKER_SINGLETONS.to_vec(); tasks.extend(["diffusion.backfill", "images.learn_archs"]); tasks.sort(); assert_eq!(app.leader.tasks(), tasks, "every worker task, and not the log's checkpointer"); assert!( until(Duration::from_secs(5), || WORKER_SINGLETONS.iter().all(|n| app.leader.leading().contains(n))).await, "{:?}", app.leader.leading() ); let h = health(&app).await; assert_eq!(h["roles"], serde_json::json!(["worker"])); assert!(h["leading"].as_array().unwrap().iter().any(|n| n == "flow.scheduler"), "{h}"); for path in ["/", "/explore", "/login", "/api/v1/version", "/andy/demo", "/log/checkpoint"] { assert_eq!(get(&app, path, None).await.0, StatusCode::NOT_FOUND, "{path} is the web role's"); } let (status, body) = get(&app, "/metrics", Some("m")).await; assert_eq!(status, StatusCode::OK, "{body}"); // It reads the transparency log but never writes it. assert!(app.tlog.append(&[]).is_err()); } #[tokio::test(flavor = "multi_thread")] async fn a_hub_takes_every_role_by_default() { hermetic(); let dir = tempfile::tempdir().unwrap(); let app = oxygenhub::build(config(dir.path(), &[])).await.unwrap(); assert_eq!(app.leader.tasks().len(), 10, "{:?}", app.leader.tasks()); assert_eq!(health(&app).await["roles"], serde_json::json!(["web", "worker"])); assert_eq!(get(&app, "/explore", None).await.0, StatusCode::OK); let mut bad = config(dir.path(), &["web", "cron"]); bad.server.roles.push(String::new()); let e = oxygenhub::build(bad).await.err().expect("an unknown role is refused").to_string(); assert!(e.contains("unknown role \"cron\""), "{e}"); } /// Two hubs on one PostgreSQL database (one schema): exactly one leads /// each singleton; when its connection drops it stops them, and when it /// stops the other takes every one over within seconds. #[cfg(feature = "postgres")] #[tokio::test(flavor = "multi_thread")] async fn two_hubs_on_one_database_run_each_singleton_once_and_fail_over() { hermetic(); let (da, dbb) = (tempfile::tempdir().unwrap(), tempfile::tempdir().unwrap()); let schema = format!("hub_roles_{}", da.path().file_name().unwrap().to_string_lossy()); let pg = |dir: &Path| { let mut c = config(dir, &[]); c.database.schema = schema.clone(); c }; let a = oxygenhub::build(pg(da.path())).await.unwrap(); let b = oxygenhub::build(pg(dbb.path())).await.unwrap(); let mut singletons = WORKER_SINGLETONS.to_vec(); singletons.push("tlog.checkpointer"); let once = |a: &Shared, b: &Shared| { let (la, lb) = (a.leader.leading(), b.leader.leading()); singletons.iter().all(|n| la.contains(n) != lb.contains(n)) }; assert!(until(Duration::from_secs(10), || once(&a, &b)).await, "a {:?}, b {:?}", a.leader.leading(), b.leader.leading()); // A's connection is ended, as when the network drops: its locks free, // it stops what it led, and between them each runs once again. let pid = a.leader.session_pid().await.expect("a has a session"); let ended: bool = sqlx::query_scalar("SELECT pg_terminate_backend($1)").bind(pid).fetch_one(&b.db).await.unwrap(); assert!(ended); let start = Instant::now(); loop { let reconnected = a.leader.session_pid().await.is_some_and(|p| p != pid); if reconnected && once(&a, &b) { break; } assert!(start.elapsed() < Duration::from_secs(15), "a {:?}, b {:?}", a.leader.leading(), b.leader.leading()); tokio::time::sleep(Duration::from_millis(100)).await; } // The hub leading the scheduler stops; the other takes everything over. let (gone, left) = if a.leader.leading().contains(&"flow.scheduler") { (&a, &b) } else { (&b, &a) }; gone.leader.resign().await; let start = Instant::now(); assert!( until(Duration::from_secs(10), || singletons.iter().all(|n| left.leader.leading().contains(n))).await, "left {:?}", left.leader.leading() ); eprintln!("took over in {:.1} s", start.elapsed().as_secs_f64()); assert!(gone.leader.leading().is_empty(), "{:?}", gone.leader.leading()); assert_eq!(health(left).await["leading"].as_array().unwrap().len(), left.leader.leading().len()); }