tatara_engine/cluster/
network.rs1use 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#[derive(Clone)]
17pub struct RaftHttpNetwork {
18 client: reqwest::Client,
19 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
44pub 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 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}