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::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
53pub fn enabled() -> bool {
55 crate_config::get("cluster.secret").is_ok_and(|s| !s.is_empty())
56}
57
58pub 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#[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
159pub const SINGLE_NODE_ID: NodeId = 1;
161
162pub 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#[derive(Clone)]
181pub struct Node {
182 pub raft: Raft,
183 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 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
244pub 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
267pub 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 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
291pub 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
302pub 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 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 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
362struct 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;