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