orbit_runtime/
cluster.rs

1//! Multi-server operation: one active leader, any number of standbys, all on
2//! the same PostgreSQL database (`orbit server --cluster`).
3//!
4//! * **Election.** Every node keeps a dedicated PostgreSQL session (separate
5//!   from the store's connection) and tries `pg_try_advisory_lock` on a fixed
6//!   key once per check interval. The session that holds the lock is the
7//!   leader. PostgreSQL releases the lock when that session ends, so a crashed
8//!   leader is replaced within about one interval.
9//! * **Fencing.** On election the leader increments `orbit_cluster.term` and
10//!   opens its store in [`PgMode::Fenced`] mode: every history append checks
11//!   the term inside its own transaction. A deposed leader that has not yet
12//!   noticed cannot write anything.
13//! * **Liveness.** The leader re-checks its session, its lock and the term
14//!   every interval. It only serves while the last successful check is younger
15//!   than the lease; when a check fails it steps down to standby in-process
16//!   (drops its engine and its lock session) and rejoins the election.
17//! * **Standby.** A standby serves read-only routes from the database through
18//!   an engine opened on a [`PgMode::ReadOnly`] store, and answers every write
19//!   and worker route with `503 NOT_LEADER` and the leader's address.
20
21use crate::engine::{Engine, EngineError};
22use crate::server::{engine, Shared};
23use crate::store::{session_client, PgMode, PgStore, Store, CLUSTER_SCHEMA};
24use axum::body::Body;
25use axum::extract::{Request, State};
26use axum::http::{HeaderValue, Method, StatusCode};
27use axum::middleware::{self, Next};
28use axum::response::{IntoResponse, Response};
29use axum::{Json, Router};
30use serde_json::{json, Value};
31use std::sync::atomic::{AtomicBool, Ordering};
32use std::sync::{Arc, Mutex};
33use std::time::{Duration, Instant};
34
35/// Advisory lock key for leader election (ASCII "orbitldr").
36pub const LEADER_LOCK_KEY: i64 = 0x6f72_6269_746c_6472;
37/// Advisory lock key serialising creation of the cluster table.
38const SCHEMA_LOCK_KEY: i64 = 0x6f72_6269_7473_6368;
39
40/// Builds an engine on a store. The CLI passes one that applies the server's
41/// engine configuration and telemetry, so leader and standby engines are
42/// configured exactly like a single-node server's.
43pub type EngineFactory =
44    Box<dyn Fn(Store) -> std::result::Result<Engine, EngineError> + Send + Sync>;
45
46#[derive(Debug, Clone)]
47pub struct ClusterConfig {
48    /// PostgreSQL URL shared by every node.
49    pub url: String,
50    /// This node's name in health output, logs and `orbit_cluster`.
51    pub node_id: String,
52    /// URL clients and workers should use to reach this node directly.
53    pub advertise_addr: Option<String>,
54    /// How often a standby tries to take the lock and the leader re-checks it.
55    pub check_interval: Duration,
56    /// The leader stops serving when its last successful check is older than
57    /// this. Must be shorter than the server-side session timeout below, so a
58    /// partitioned leader stops before PostgreSQL releases its lock.
59    pub lease: Duration,
60    /// `idle_session_timeout` for the election session (PostgreSQL 14+): the
61    /// server ends a leader session that stopped checking in, which releases
62    /// the lock even when the leader host is unreachable.
63    pub session_timeout: Duration,
64}
65
66impl ClusterConfig {
67    pub fn new(url: &str, node_id: &str) -> ClusterConfig {
68        ClusterConfig {
69            url: url.to_string(),
70            node_id: node_id.to_string(),
71            advertise_addr: None,
72            check_interval: Duration::from_secs(1),
73            lease: Duration::from_secs(3),
74            session_timeout: Duration::from_secs(5),
75        }
76    }
77}
78
79/// A default node id: `$HOSTNAME`, `/etc/hostname`, or a random id.
80pub fn default_node_id() -> String {
81    std::env::var("HOSTNAME")
82        .ok()
83        .or_else(|| std::fs::read_to_string("/etc/hostname").ok())
84        .map(|h| h.trim().to_string())
85        .filter(|h| !h.is_empty())
86        .unwrap_or_else(|| format!("node-{}", &uuid::Uuid::new_v4().to_string()[..8]))
87}
88
89#[derive(Debug, Clone, Copy, PartialEq, Eq)]
90pub enum Role {
91    Standby,
92    Leader,
93}
94
95#[derive(Debug, Clone)]
96struct Status {
97    role: Role,
98    /// The current term as last seen in `orbit_cluster`.
99    term: i64,
100    leader_node: Option<String>,
101    leader_addr: Option<String>,
102    /// Last successful leadership check (leader only).
103    last_ok: Option<Instant>,
104}
105
106/// One cluster member. Created by the CLI, driven by [`start`].
107pub struct Node {
108    config: ClusterConfig,
109    factory: EngineFactory,
110    status: Mutex<Status>,
111    stop: AtomicBool,
112}
113
114impl Node {
115    pub fn new(config: ClusterConfig, factory: EngineFactory) -> Arc<Node> {
116        Arc::new(Node {
117            config,
118            factory,
119            status: Mutex::new(Status {
120                role: Role::Standby,
121                term: 0,
122                leader_node: None,
123                leader_addr: None,
124                last_ok: None,
125            }),
126            stop: AtomicBool::new(false),
127        })
128    }
129
130    pub fn config(&self) -> &ClusterConfig {
131        &self.config
132    }
133
134    fn status(&self) -> std::sync::MutexGuard<'_, Status> {
135        self.status.lock().unwrap_or_else(|p| p.into_inner())
136    }
137
138    /// True while this node leads and its leadership was confirmed within the
139    /// lease. Checked on every request.
140    pub fn is_leader(&self) -> bool {
141        let s = self.status();
142        s.role == Role::Leader && s.last_ok.is_some_and(|t| t.elapsed() < self.config.lease)
143    }
144
145    pub fn role(&self) -> Role {
146        if self.is_leader() {
147            Role::Leader
148        } else {
149            Role::Standby
150        }
151    }
152
153    pub fn term(&self) -> i64 {
154        self.status().term
155    }
156
157    /// The engine a standby serves reads from: recovers nothing, writes nothing.
158    pub fn standby_engine(&self) -> std::result::Result<Engine, EngineError> {
159        let store = PgStore::connect_with(&self.config.url, PgMode::ReadOnly)
160            .map_err(|e| EngineError::new(500, "STORAGE_ERROR", e.to_string()))?;
161        (self.factory)(Store::Postgres(store))
162    }
163
164    /// Stop taking part in the election (tests, embedding). A leader steps
165    /// down at its next check and releases the lock.
166    pub fn stop(&self) {
167        self.stop.store(true, Ordering::SeqCst);
168    }
169
170    fn stopped(&self) -> bool {
171        self.stop.load(Ordering::SeqCst)
172    }
173
174    fn leader_json(&self) -> Value {
175        let s = self.status();
176        match &s.leader_node {
177            Some(n) => json!({ "node_id": n, "addr": s.leader_addr }),
178            None => Value::Null,
179        }
180    }
181
182    /// Fields added to `/v1/health` and `/v1/ready`.
183    fn health_fields(&self) -> Value {
184        json!({
185            "role": if self.is_leader() { "leader" } else { "standby" },
186            "node_id": self.config.node_id,
187            "term": self.term(),
188            "leader": self.leader_json(),
189        })
190    }
191
192    fn not_leader(&self) -> Response {
193        let (node, addr) = {
194            let s = self.status();
195            (s.leader_node.clone(), s.leader_addr.clone())
196        };
197        let message = match (&node, &addr) {
198            (Some(n), Some(a)) if *n != self.config.node_id => {
199                format!("this server is a standby; the leader is `{n}` at {a}")
200            }
201            (Some(n), None) if *n != self.config.node_id => {
202                format!("this server is a standby; the leader is `{n}`")
203            }
204            _ => "this server is a standby and no leader is serving yet; retry shortly".into(),
205        };
206        let mut resp = (
207            StatusCode::SERVICE_UNAVAILABLE,
208            Json(json!({ "error": {
209                "code": "NOT_LEADER",
210                "message": message,
211                "leader_node": node,
212                "leader_addr": addr,
213            }})),
214        )
215            .into_response();
216        let h = resp.headers_mut();
217        h.insert("retry-after", HeaderValue::from_static("1"));
218        if let Some(v) = addr.and_then(|a| HeaderValue::from_str(&a).ok()) {
219            h.insert("x-orbit-leader", v);
220        }
221        resp
222    }
223}
224
225// --------------------------------------------------------------- HTTP gate
226
227#[derive(Clone)]
228struct Gate {
229    node: Option<Arc<Node>>,
230    state: Shared,
231}
232
233/// Wrap the API router: adds `/v1/ready`, the cluster fields of
234/// `/v1/health`, and (in cluster mode) refuses writes on a standby. Runs
235/// before authentication, so load balancers can probe `/v1/ready` without
236/// credentials.
237pub fn layer(app: Router, node: Option<Arc<Node>>, state: Shared) -> Router {
238    app.layer(middleware::from_fn_with_state(Gate { node, state }, gate))
239}
240
241/// Requests a standby answers itself (from the database).
242fn is_read(method: &Method, path: &str) -> bool {
243    (method == Method::GET || method == Method::HEAD) && path.starts_with("/v1/")
244}
245
246async fn gate(State(g): State<Gate>, req: Request, next: Next) -> Response {
247    let path = req.uri().path().to_string();
248    if path == "/v1/ready" {
249        let (ready, mut body) = match &g.node {
250            None => (true, json!({ "role": "single" })),
251            Some(n) => (n.is_leader(), n.health_fields()),
252        };
253        body["ready"] = json!(ready);
254        let status = if ready {
255            StatusCode::OK
256        } else {
257            StatusCode::SERVICE_UNAVAILABLE
258        };
259        return (status, Json(body)).into_response();
260    }
261    if path == "/v1/health" {
262        let resp = next.run(req).await;
263        let fields = match &g.node {
264            None => json!({ "role": "single", "node_id": null, "term": null, "leader": null }),
265            Some(n) => n.health_fields(),
266        };
267        return merge_json(resp, fields).await;
268    }
269    let Some(node) = &g.node else {
270        return next.run(req).await;
271    };
272    if node.is_leader() {
273        return next.run(req).await;
274    }
275    if is_read(req.method(), &path) {
276        // Definitions registered on the leader since this standby opened.
277        if path.starts_with("/v1/workflows") || path.contains("/queries/") {
278            let st = g.state.clone();
279            let _ = tokio::task::spawn_blocking(move || engine(&st).refresh_definitions()).await;
280        }
281        return next.run(req).await;
282    }
283    node.not_leader()
284}
285
286/// Add `fields` to a JSON object response.
287pub(crate) async fn merge_json(resp: Response, fields: Value) -> Response {
288    if resp.status() != StatusCode::OK {
289        return resp;
290    }
291    let (mut parts, body) = resp.into_parts();
292    let Ok(bytes) = axum::body::to_bytes(body, 1 << 20).await else {
293        return Response::from_parts(parts, Body::empty());
294    };
295    let mut v: Value = match serde_json::from_slice(&bytes) {
296        Ok(v @ Value::Object(_)) => v,
297        _ => return Response::from_parts(parts, Body::from(bytes)),
298    };
299    if let (Some(obj), Value::Object(extra)) = (v.as_object_mut(), fields) {
300        obj.extend(extra);
301    }
302    let out = serde_json::to_vec(&v).unwrap_or_default();
303    parts.headers.remove("content-length");
304    Response::from_parts(parts, Body::from(out))
305}
306
307/// gRPC interceptor: on a standby every worker call fails with `UNAVAILABLE`
308/// and `orbit-error-code: NOT_LEADER` (SDKs retry it like a connection error).
309pub fn grpc_gate(
310    node: Arc<Node>,
311) -> impl FnMut(tonic::Request<()>) -> Result<tonic::Request<()>, tonic::Status> + Clone {
312    move |req| {
313        if node.is_leader() {
314            return Ok(req);
315        }
316        let addr = node.status().leader_addr.clone();
317        let mut s = tonic::Status::unavailable(match &addr {
318            Some(a) => format!("NOT_LEADER: this server is a standby; the leader is at {a}"),
319            None => "NOT_LEADER: this server is a standby".to_string(),
320        });
321        s.metadata_mut()
322            .insert("orbit-error-code", "NOT_LEADER".parse().expect("ascii"));
323        if let Some(v) = addr.and_then(|a| a.parse().ok()) {
324            s.metadata_mut().insert("x-orbit-leader", v);
325        }
326        Err(s)
327    }
328}
329
330// ---------------------------------------------------------------- election
331
332/// Start the election thread for `node`, swapping engines in `state` as the
333/// node gains and loses leadership. `state` must hold a standby engine
334/// ([`Node::standby_engine`]) to begin with.
335pub fn start(node: Arc<Node>, state: Shared) -> std::io::Result<std::thread::JoinHandle<()>> {
336    tracing::info!(
337        node_id = %node.config.node_id,
338        advertise_addr = node.config.advertise_addr.as_deref().unwrap_or(""),
339        "cluster mode: standby until elected"
340    );
341    std::thread::Builder::new()
342        .name("orbit-cluster".into())
343        .spawn(move || election_loop(&node, &state))
344}
345
346fn now_ms() -> i64 {
347    std::time::SystemTime::now()
348        .duration_since(std::time::UNIX_EPOCH)
349        .map(|d| d.as_millis() as i64)
350        .unwrap_or(0)
351}
352
353fn connect(config: &ClusterConfig) -> Result<postgres::Client, String> {
354    let mut c = session_client(&config.url, &format!("orbit-cluster:{}", config.node_id))
355        .map_err(|e| e.to_string())?;
356    // Older servers lack idle_session_timeout; TCP keepalives still apply.
357    let _ = c.batch_execute(&format!(
358        "SET idle_session_timeout = '{}ms'",
359        config.session_timeout.as_millis()
360    ));
361    c.batch_execute(&format!(
362        "SET statement_timeout = '{}ms'",
363        config.lease.as_millis()
364    ))
365    .map_err(|e| e.to_string())?;
366    let mut tx = c.transaction().map_err(|e| e.to_string())?;
367    tx.execute("SELECT pg_advisory_xact_lock($1)", &[&SCHEMA_LOCK_KEY])
368        .map_err(|e| e.to_string())?;
369    let exists: bool = tx
370        .query_one("SELECT to_regclass('orbit_cluster') IS NOT NULL", &[])
371        .map_err(|e| e.to_string())?
372        .get(0);
373    if !exists {
374        tx.batch_execute(CLUSTER_SCHEMA)
375            .map_err(|e| e.to_string())?;
376    }
377    tx.commit().map_err(|e| e.to_string())?;
378    Ok(c)
379}
380
381fn sleep_interval(node: &Node) {
382    std::thread::sleep(node.config.check_interval);
383}
384
385fn election_loop(node: &Node, state: &Shared) {
386    let mut reader_ok = true;
387    let mut followed: Option<(i64, Option<String>)> = None;
388    while !node.stopped() {
389        let mut client = match connect(&node.config) {
390            Ok(c) => c,
391            Err(e) => {
392                tracing::warn!(error = %e, "cluster: cannot reach PostgreSQL for leader election; retrying");
393                sleep_interval(node);
394                continue;
395            }
396        };
397        // Standby: try to take the lock, otherwise follow the current leader.
398        let elected = loop {
399            if node.stopped() {
400                return;
401            }
402            if !reader_ok {
403                reader_ok = install_standby_engine(node, state);
404            }
405            match client.query_one("SELECT pg_try_advisory_lock($1)", &[&LEADER_LOCK_KEY]) {
406                Ok(r) if r.get::<_, bool>(0) => break true,
407                Ok(_) => {}
408                Err(e) => {
409                    tracing::warn!(error = %e, "cluster: election session failed; reconnecting");
410                    break false;
411                }
412            }
413            match client.query_one(
414                "SELECT term, leader_node, leader_addr FROM orbit_cluster WHERE id = 1",
415                &[],
416            ) {
417                Ok(r) => {
418                    let (term, leader, addr): (i64, Option<String>, Option<String>) =
419                        (r.get(0), r.get(1), r.get(2));
420                    if followed.as_ref() != Some(&(term, leader.clone())) {
421                        tracing::info!(
422                            term,
423                            leader = leader.as_deref().unwrap_or("none"),
424                            "cluster: standby; following the leader"
425                        );
426                        followed = Some((term, leader.clone()));
427                    }
428                    let mut s = node.status();
429                    s.term = term;
430                    s.leader_node = leader;
431                    s.leader_addr = addr;
432                }
433                Err(e) => {
434                    tracing::warn!(error = %e, "cluster: election session failed; reconnecting");
435                    break false;
436                }
437            }
438            // Long-polls (`/wait`) on a standby re-read the database.
439            state.changed.notify_waiters();
440            sleep_interval(node);
441        };
442        if !elected {
443            sleep_interval(node);
444            continue;
445        }
446        let reason = lead(node, state, &mut client);
447        // Step down: stop serving first, then drop the leader engine.
448        {
449            let mut s = node.status();
450            s.role = Role::Standby;
451            s.last_ok = None;
452        }
453        tracing::error!(node_id = %node.config.node_id, reason = %reason, "cluster: lost leadership; stepped down to standby");
454        drop(client); // ends the session, releasing the lock if still held
455        reader_ok = install_standby_engine(node, state);
456        followed = None;
457        if !node.stopped() {
458            sleep_interval(node);
459        }
460    }
461}
462
463/// Replace the engine with a standby engine. If the database is unreachable,
464/// install an empty in-memory engine so nothing of the old leader keeps
465/// running, and report failure so the caller retries.
466fn install_standby_engine(node: &Node, state: &Shared) -> bool {
467    let (engine_new, ok) = match node.standby_engine() {
468        Ok(e) => (e, true),
469        Err(e) => {
470            tracing::warn!(error = %e, "cluster: cannot open the standby store; retrying");
471            match Store::in_memory()
472                .map_err(|e| e.to_string())
473                .and_then(|s| (node.factory)(s).map_err(|e| e.to_string()))
474            {
475                Ok(e) => (e, false),
476                Err(_) => return false,
477            }
478        }
479    };
480    swap_engine(state, engine_new);
481    ok
482}
483
484fn swap_engine(state: &Shared, mut new: Engine) {
485    new.changed = state.changed.clone();
486    let old = std::mem::replace(&mut *engine(state), new);
487    drop(old);
488    state.changed.notify_waiters();
489}
490
491/// Take office, then hold it until a check fails. Returns why it ended.
492fn lead(node: &Node, state: &Shared, client: &mut postgres::Client) -> String {
493    let addr = node.config.advertise_addr.clone();
494    let term: i64 = match client.query_one(
495        "UPDATE orbit_cluster SET term = term + 1, leader_node = $1, leader_addr = $2, since_ms = $3 WHERE id = 1 RETURNING term",
496        &[&node.config.node_id, &addr, &now_ms()],
497    ) {
498        Ok(r) => r.get(0),
499        Err(e) => return format!("cannot record the new term: {e}"),
500    };
501    {
502        let mut s = node.status();
503        s.term = term;
504        s.leader_node = Some(node.config.node_id.clone());
505        s.leader_addr = addr.clone();
506    }
507    tracing::info!(node_id = %node.config.node_id, term, "cluster: elected leader; recovering executions from history");
508    let started = Instant::now();
509    // Recovery may take longer than the idle timeout of this (idle) session;
510    // suspend it meanwhile so the lock is not released under our feet.
511    let _ = client.batch_execute("SET idle_session_timeout = 0");
512    let opened = PgStore::connect_with(&node.config.url, PgMode::Fenced(term))
513        .map_err(|e| e.to_string())
514        .and_then(|s| (node.factory)(Store::Postgres(s)).map_err(|e| e.to_string()));
515    let _ = client.batch_execute(&format!(
516        "SET idle_session_timeout = '{}ms'",
517        node.config.session_timeout.as_millis()
518    ));
519    let leader_engine = match opened {
520        Ok(e) => e,
521        Err(e) => return format!("cannot open the engine: {e}"),
522    };
523    let active = leader_engine.active_count();
524    swap_engine(state, leader_engine);
525    {
526        let mut s = node.status();
527        s.role = Role::Leader;
528        s.last_ok = Some(Instant::now());
529    }
530    tracing::info!(node_id = %node.config.node_id, term, active_executions = active, recovery_ms = started.elapsed().as_millis() as u64, "cluster: leader is serving");
531    loop {
532        sleep_interval(node);
533        if node.stopped() {
534            return "node is shutting down".into();
535        }
536        let check = client.query_one(
537            "SELECT (SELECT term FROM orbit_cluster WHERE id = 1),
538                    EXISTS (SELECT 1 FROM pg_locks WHERE locktype = 'advisory' AND pid = pg_backend_pid() AND granted)",
539            &[],
540        );
541        match check {
542            Ok(r) => {
543                let (current, holds): (Option<i64>, bool) = (r.get(0), r.get(1));
544                if !holds {
545                    return "the election lock is no longer held".into();
546                }
547                if current != Some(term) {
548                    return format!(
549                        "the cluster moved to term {} (this node led term {term})",
550                        current.unwrap_or(-1)
551                    );
552                }
553                // A slow check only pauses serving (the lease lapsed); the
554                // lock and the term prove nobody else took over meanwhile.
555                node.status().last_ok = Some(Instant::now());
556            }
557            Err(e) => return format!("the election session failed: {e}"),
558        }
559    }
560}
orbit_runtime/cluster.rs

orbit_runtime/
cluster.rs
