Skip to main content

orch/cluster/
mod.rs

1//! Raft between orch nodes. Only coordination state goes through the log, the design is in
2//! `docs/specs/2026-10-26-multi-machine-cluster-design.md`.
3
4mod command;
5pub mod config;
6pub mod discovery;
7pub mod local;
8pub mod network;
9pub mod progress;
10mod storage;
11
12use std::collections::BTreeMap;
13use std::io::Cursor;
14use std::path::Path;
15use std::sync::Arc;
16use std::time::Duration;
17
18use openraft::error::{ClientWriteError, InstallSnapshotError, RPCError, RaftError, Unreachable};
19use openraft::network::{RPCOption, RaftNetwork, RaftNetworkFactory};
20use openraft::raft::{
21    AppendEntriesRequest, AppendEntriesResponse, InstallSnapshotRequest, InstallSnapshotResponse,
22    VoteRequest, VoteResponse,
23};
24use openraft::storage::Adaptor;
25use openraft::{BasicNode, SnapshotPolicy};
26
27use crate::store::TaskStore;
28
29pub use command::{Command, Reply, Request};
30pub use storage::RaftStore;
31
32pub type NodeId = u64;
33
34openraft::declare_raft_types!(
35    pub TypeConfig:
36        D = Request,
37        R = Reply,
38);
39
40pub type Raft = openraft::Raft<TypeConfig>;
41
42/// A `cluster:` section turns the cluster on. Without one, orch writes straight to SQLite as it always has
43pub fn enabled() -> bool {
44    crate::config::get("cluster.secret").is_ok_and(|s| !s.is_empty())
45}
46
47/// Start this machine's node from the `cluster:` section and the tailnet. A learner keeps asking
48/// the cluster to add it until a member proves the secret and takes it
49pub async fn start(store: &TaskStore, raft_db: &Path) -> anyhow::Result<Node> {
50    let cfg = config::ClusterConfig::load()?
51        .ok_or_else(|| anyhow::anyhow!("no cluster.secret in the config"))?;
52    let tailnet = discovery::status().await?;
53    let plan = config::plan(&cfg, &tailnet)?;
54    tracing::info!(
55        name = %plan.name,
56        addr = %plan.addr,
57        elect = plan.elect,
58        learner = plan.learner,
59        "starting cluster node"
60    );
61    let node = start_node(
62        store,
63        raft_db,
64        NodeConfig {
65            name: plan.name.clone(),
66            id: plan.id,
67            addr: plan.addr.clone(),
68            secret: cfg.secret.clone(),
69            voters: plan.voters.clone(),
70            elect: plan.elect,
71        },
72    )
73    .await?;
74    if plan.learner {
75        tokio::spawn(join_loop(node.clone(), plan, cfg));
76    }
77
78    Ok(node)
79}
80
81async fn join_loop(node: Node, plan: config::Plan, cfg: config::ClusterConfig) {
82    loop {
83        let member = node
84            .raft
85            .metrics()
86            .borrow()
87            .membership_config
88            .membership()
89            .nodes()
90            .any(|(id, _)| *id == plan.id);
91        if member {
92            return;
93        }
94
95        if let Ok(tailnet) = discovery::status().await {
96            for peer in tailnet.peers.iter().filter(|p| p.online) {
97                let addr = format!("{}:{}", peer.ip, cfg.port);
98                if network::probe(&addr, &cfg.secret, Some(&peer.name))
99                    .await
100                    .is_none()
101                {
102                    continue;
103                }
104
105                let join = network::Join {
106                    id: plan.id,
107                    addr: plan.addr.clone(),
108                };
109                match network::send_join(&addr, &cfg.secret, &join).await {
110                    Ok(()) => tracing::info!(via = %peer.name, "joined the cluster as a learner"),
111                    Err(e) => tracing::warn!(via = %peer.name, err = %e, "join refused"),
112                }
113
114                break;
115            }
116        }
117
118        tokio::time::sleep(Duration::from_secs(30)).await;
119    }
120}
121
122/// One member as the leader sees it, for `orch nodes`
123#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
124pub struct Member {
125    pub id: NodeId,
126    pub addr: String,
127    pub voter: bool,
128    pub leader: bool,
129    pub me: bool,
130}
131
132pub fn members(node: &Node) -> Vec<Member> {
133    let m = node.raft.metrics().borrow().clone();
134    let membership = m.membership_config.membership();
135    let voters: std::collections::BTreeSet<NodeId> = membership.voter_ids().collect();
136    membership
137        .nodes()
138        .map(|(id, n)| Member {
139            id: *id,
140            addr: n.addr.clone(),
141            voter: voters.contains(id),
142            leader: m.current_leader == Some(*id),
143            me: *id == m.id,
144        })
145        .collect()
146}
147
148/// Id of the node when orch runs alone
149pub const SINGLE_NODE_ID: NodeId = 1;
150
151/// The proposer's clock, in the format the store writes timestamps in
152pub fn now_stamp() -> String {
153    chrono::Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string()
154}
155
156fn config() -> anyhow::Result<openraft::Config> {
157    Ok(openraft::Config {
158        cluster_name: "orch".into(),
159        heartbeat_interval: 250,
160        election_timeout_min: 1000,
161        election_timeout_max: 2000,
162        snapshot_policy: SnapshotPolicy::LogsSinceLast(1000),
163        ..Default::default()
164    }
165    .validate()?)
166}
167
168/// A running node. The store holds one and sends every replicated write through it
169#[derive(Clone)]
170pub struct Node {
171    pub raft: Raft,
172    /// None on a single node, which has no peers to forward to
173    secret: Option<Arc<str>>,
174    server: Option<Arc<tokio::sync::Notify>>,
175}
176
177impl std::fmt::Debug for Node {
178    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
179        f.debug_struct("Node")
180            .field("server", &self.server.is_some())
181            .finish_non_exhaustive()
182    }
183}
184
185impl Node {
186    /// Propose on this node, or forward to the leader when this node follows. During an election
187    /// the leader can be unknown or stale, so it waits and tries again, but only a write that was
188    /// certainly not applied is sent again
189    pub async fn write(&self, req: Request) -> anyhow::Result<Reply> {
190        let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
191        loop {
192            let err = match self.raft.client_write(req.clone()).await {
193                Ok(r) => return Ok(r.data),
194                Err(e) => e,
195            };
196            let RaftError::APIError(ClientWriteError::ForwardToLeader(f)) = &err else {
197                return Err(err.into());
198            };
199            let Some(secret) = &self.secret else {
200                return Err(err.into());
201            };
202            let reason = match &f.leader_node {
203                Some(leader) => match network::forward(&leader.addr, secret, &req).await {
204                    Ok(reply) => return Ok(reply),
205                    Err(network::ForwardError::Unknown(e)) => return Err(e),
206                    Err(network::ForwardError::NotApplied(why)) => why,
207                },
208                None => "no leader known".to_string(),
209            };
210            if tokio::time::Instant::now() >= deadline {
211                anyhow::bail!("no leader took the write: {reason}");
212            }
213
214            tokio::time::sleep(Duration::from_millis(200)).await;
215        }
216    }
217
218    pub fn is_leader(&self) -> bool {
219        let m = self.raft.metrics();
220        let m = m.borrow();
221        m.current_leader == Some(m.id)
222    }
223
224    pub async fn stop(self) {
225        if let Some(server) = self.server {
226            server.notify_one();
227        }
228
229        let _ = self.raft.shutdown().await;
230    }
231}
232
233/// How one node of a cluster starts: its id, the address it listens on, the shared secret, the
234/// voters to initialize with, and whether it may lead (a witness may not)
235pub struct NodeConfig {
236    pub name: String,
237    pub id: NodeId,
238    pub addr: String,
239    pub secret: String,
240    pub voters: BTreeMap<NodeId, String>,
241    pub elect: bool,
242}
243
244impl std::fmt::Debug for NodeConfig {
245    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
246        f.debug_struct("NodeConfig")
247            .field("name", &self.name)
248            .field("id", &self.id)
249            .field("addr", &self.addr)
250            .field("voters", &self.voters)
251            .field("elect", &self.elect)
252            .finish_non_exhaustive()
253    }
254}
255
256/// Start a one-node cluster over `store`, with its log in `raft_db`, and route the store's
257/// replicated writes through it
258pub async fn start_single(store: &TaskStore, raft_db: &Path) -> anyhow::Result<Node> {
259    let raft_store = RaftStore::open(raft_db, store.clone()).await?;
260    let (log, sm) = Adaptor::new(raft_store);
261    let raft = Raft::new(SINGLE_NODE_ID, Arc::new(config()?), NoNetwork, log, sm).await?;
262    let members = BTreeMap::from([(SINGLE_NODE_ID, BasicNode::default())]);
263    // A restart finds the membership in the log, and initialize refuses a second time
264    if let Err(e) = raft.initialize(members).await {
265        tracing::debug!(err = %e, "raft already initialized");
266    }
267
268    raft.wait(Some(Duration::from_secs(10)))
269        .current_leader(SINGLE_NODE_ID, "single node leads")
270        .await?;
271    let node = Node {
272        raft,
273        secret: None,
274        server: None,
275    };
276    store.attach_node(node.clone());
277    Ok(node)
278}
279
280/// Start one node of a cluster and serve its RPCs on `cfg.addr`. Every voter initializes with
281/// the same membership, which openraft accepts, so no voter is special
282pub async fn start_node(
283    store: &TaskStore,
284    raft_db: &Path,
285    cfg: NodeConfig,
286) -> anyhow::Result<Node> {
287    let listener = tokio::net::TcpListener::bind(&cfg.addr).await?;
288    start_node_on(store, raft_db, cfg, listener).await
289}
290
291/// Like `start_node`, with a listener the caller already bound. `cfg.addr` is the address peers use
292pub async fn start_node_on(
293    store: &TaskStore,
294    raft_db: &Path,
295    cfg: NodeConfig,
296    listener: tokio::net::TcpListener,
297) -> anyhow::Result<Node> {
298    let secret: Arc<str> = Arc::from(cfg.secret);
299    let raft_store = RaftStore::open(raft_db, store.clone()).await?;
300    let (log, sm) = Adaptor::new(raft_store);
301    let config = openraft::Config {
302        enable_elect: cfg.elect,
303        ..config()?
304    }
305    .validate()?;
306    let network = network::HttpNetwork {
307        secret: std::sync::Arc::clone(&secret),
308    };
309    let raft = Raft::new(cfg.id, Arc::new(config), network, log, sm).await?;
310    let me = network::Hello {
311        name: cfg.name,
312        id: cfg.id,
313        witness: !cfg.elect,
314        leader: None,
315        term: 0,
316        proof: String::new(),
317    };
318    let app = network::router(
319        raft.clone(),
320        std::sync::Arc::clone(&secret),
321        me,
322        cfg.addr.clone(),
323    );
324    let stop = Arc::new(tokio::sync::Notify::new());
325    let stopped = std::sync::Arc::clone(&stop);
326    tokio::spawn(async move {
327        // Graceful shutdown also closes kept-alive connections, which would keep answering for a stopped node
328        let served = axum::serve(listener, app)
329            .with_graceful_shutdown(async move { stopped.notified().await })
330            .await;
331        if let Err(e) = served {
332            tracing::warn!(err = %e, "cluster listener stopped");
333        }
334    });
335    // initialize starts an election whatever enable_elect says, so a witness waits for the leader
336    if cfg.elect && !cfg.voters.is_empty() {
337        let members: BTreeMap<NodeId, BasicNode> = cfg
338            .voters
339            .into_iter()
340            .map(|(id, addr)| (id, BasicNode::new(addr)))
341            .collect();
342        if let Err(e) = raft.initialize(members).await {
343            tracing::debug!(err = %e, "raft already initialized");
344        }
345    }
346
347    let node = Node {
348        raft,
349        secret: Some(secret),
350        server: Some(stop),
351    };
352    store.attach_node(node.clone());
353    Ok(node)
354}
355
356/// A single node never sends an RPC. The tailnet transport replaces this
357struct NoNetwork;
358
359impl RaftNetworkFactory<TypeConfig> for NoNetwork {
360    type Network = NoNetwork;
361
362    async fn new_client(&mut self, _target: NodeId, _node: &BasicNode) -> NoNetwork {
363        NoNetwork
364    }
365}
366
367fn unreachable<E: std::error::Error>() -> RPCError<NodeId, BasicNode, E> {
368    RPCError::Unreachable(Unreachable::new(&std::io::Error::other(
369        "no network on a single node",
370    )))
371}
372
373impl RaftNetwork<TypeConfig> for NoNetwork {
374    async fn append_entries(
375        &mut self,
376        _req: AppendEntriesRequest<TypeConfig>,
377        _option: RPCOption,
378    ) -> Result<AppendEntriesResponse<NodeId>, RPCError<NodeId, BasicNode, RaftError<NodeId>>> {
379        Err(unreachable())
380    }
381
382    async fn install_snapshot(
383        &mut self,
384        _req: InstallSnapshotRequest<TypeConfig>,
385        _option: RPCOption,
386    ) -> Result<
387        InstallSnapshotResponse<NodeId>,
388        RPCError<NodeId, BasicNode, RaftError<NodeId, InstallSnapshotError>>,
389    > {
390        Err(unreachable())
391    }
392
393    async fn vote(
394        &mut self,
395        _req: VoteRequest<NodeId>,
396        _option: RPCOption,
397    ) -> Result<VoteResponse<NodeId>, RPCError<NodeId, BasicNode, RaftError<NodeId>>> {
398        Err(unreachable())
399    }
400}
401
402#[cfg(test)]
403mod tests;