Skip to main content

tatara_engine/cluster/
network.rs

1use openraft::error::{InstallSnapshotError, RPCError, RaftError};
2use openraft::network::{RPCOption, RaftNetwork, RaftNetworkFactory};
3use openraft::raft::{
4    AppendEntriesRequest, AppendEntriesResponse, InstallSnapshotRequest, InstallSnapshotResponse,
5    VoteRequest, VoteResponse,
6};
7use openraft::BasicNode;
8use std::collections::HashMap;
9use std::sync::Arc;
10use tokio::sync::RwLock;
11
12use super::raft_sm::TypeConfig;
13use tatara_core::cluster::types::NodeId;
14
15/// Network layer for Raft — HTTP-based RPCs between nodes.
16#[derive(Clone)]
17pub struct RaftHttpNetwork {
18    client: reqwest::Client,
19    /// Map of node_id → HTTP base address (e.g., "http://10.0.0.1:4648")
20    peers: Arc<RwLock<HashMap<NodeId, String>>>,
21}
22
23impl RaftHttpNetwork {
24    pub fn new() -> Self {
25        Self {
26            client: reqwest::Client::new(),
27            peers: Arc::new(RwLock::new(HashMap::new())),
28        }
29    }
30
31    pub async fn update_peer(&self, node_id: NodeId, addr: String) {
32        self.peers.write().await.insert(node_id, addr);
33    }
34
35    pub async fn remove_peer(&self, node_id: &NodeId) {
36        self.peers.write().await.remove(node_id);
37    }
38
39    async fn peer_addr(&self, node_id: &NodeId) -> Option<String> {
40        self.peers.read().await.get(node_id).cloned()
41    }
42}
43
44/// A connection to a single Raft peer.
45pub struct RaftHttpConnection {
46    client: reqwest::Client,
47    target_addr: String,
48    target_id: NodeId,
49}
50
51impl RaftNetworkFactory<TypeConfig> for RaftHttpNetwork {
52    type Network = RaftHttpConnection;
53
54    async fn new_client(&mut self, target: NodeId, node: &BasicNode) -> Self::Network {
55        // Use node.addr from openraft's BasicNode, or fall back to our peer map
56        let addr = if !node.addr.is_empty() {
57            node.addr.clone()
58        } else {
59            self.peer_addr(&target)
60                .await
61                .unwrap_or_else(|| "http://127.0.0.1:4648".to_string())
62        };
63
64        RaftHttpConnection {
65            client: self.client.clone(),
66            target_addr: addr,
67            target_id: target,
68        }
69    }
70}
71
72impl RaftNetwork<TypeConfig> for RaftHttpConnection {
73    async fn append_entries(
74        &mut self,
75        rpc: AppendEntriesRequest<TypeConfig>,
76        _option: RPCOption,
77    ) -> Result<AppendEntriesResponse<NodeId>, RPCError<NodeId, BasicNode, RaftError<NodeId>>> {
78        let url = format!("{}/raft/append", self.target_addr);
79        let resp = self
80            .client
81            .post(&url)
82            .json(&rpc)
83            .send()
84            .await
85            .map_err(|e| new_rpc_error(self.target_id, &e))?;
86
87        let result: AppendEntriesResponse<NodeId> = resp
88            .json()
89            .await
90            .map_err(|e| new_rpc_error(self.target_id, &e))?;
91
92        Ok(result)
93    }
94
95    async fn install_snapshot(
96        &mut self,
97        rpc: InstallSnapshotRequest<TypeConfig>,
98        _option: RPCOption,
99    ) -> Result<
100        InstallSnapshotResponse<NodeId>,
101        RPCError<NodeId, BasicNode, RaftError<NodeId, InstallSnapshotError>>,
102    > {
103        let url = format!("{}/raft/snapshot", self.target_addr);
104        let resp = self
105            .client
106            .post(&url)
107            .json(&rpc)
108            .send()
109            .await
110            .map_err(|e| new_rpc_error_snap(self.target_id, &e))?;
111
112        let result: InstallSnapshotResponse<NodeId> = resp
113            .json()
114            .await
115            .map_err(|e| new_rpc_error_snap(self.target_id, &e))?;
116
117        Ok(result)
118    }
119
120    async fn vote(
121        &mut self,
122        rpc: VoteRequest<NodeId>,
123        _option: RPCOption,
124    ) -> Result<VoteResponse<NodeId>, RPCError<NodeId, BasicNode, RaftError<NodeId>>> {
125        let url = format!("{}/raft/vote", self.target_addr);
126        let resp = self
127            .client
128            .post(&url)
129            .json(&rpc)
130            .send()
131            .await
132            .map_err(|e| new_rpc_error(self.target_id, &e))?;
133
134        let result: VoteResponse<NodeId> = resp
135            .json()
136            .await
137            .map_err(|e| new_rpc_error(self.target_id, &e))?;
138
139        Ok(result)
140    }
141}
142
143fn new_rpc_error(
144    target: NodeId,
145    e: &reqwest::Error,
146) -> RPCError<NodeId, BasicNode, RaftError<NodeId>> {
147    let io_err = std::io::Error::new(
148        std::io::ErrorKind::ConnectionRefused,
149        format!("Node {} unreachable: {}", target, e),
150    );
151    RPCError::Unreachable(openraft::error::Unreachable::new(&io_err))
152}
153
154fn new_rpc_error_snap(
155    target: NodeId,
156    e: &reqwest::Error,
157) -> RPCError<NodeId, BasicNode, RaftError<NodeId, InstallSnapshotError>> {
158    let io_err = std::io::Error::new(
159        std::io::ErrorKind::ConnectionRefused,
160        format!("Node {} unreachable: {}", target, e),
161    );
162    RPCError::Unreachable(openraft::error::Unreachable::new(&io_err))
163}