1mod 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
42pub fn enabled() -> bool {
44 crate::config::get("cluster.secret").is_ok_and(|s| !s.is_empty())
45}
46
47pub 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#[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
148pub const SINGLE_NODE_ID: NodeId = 1;
150
151pub 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#[derive(Clone)]
170pub struct Node {
171 pub raft: Raft,
172 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 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
233pub 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
256pub 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 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
280pub 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
291pub 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 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 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
356struct 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;