dcache/src/network/raft_network_impl.rs

224 lines
6.5 KiB
Rust
Raw Normal View History

2023-05-26 00:42:35 +05:30
/*
* mCaptcha - A proof of work based DoS protection system
* Copyright © 2023 Aravinth Manivannan <realravinth@batsense.net>
2023-05-26 00:42:35 +05:30
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU Affero General Public License as
* published by the Free Software Foundation, either version 3 of the
* License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU Affero General Public License for more details.
*
* You should have received a copy of the GNU Affero General Public License
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
use std::sync::Arc;
use std::time::Duration;
use super::management::HealthStatus;
use crate::DcacheNodeId;
use crate::DcacheTypeConfig;
2023-05-24 21:22:14 +05:30
use async_trait::async_trait;
use openraft::error::InstallSnapshotError;
use openraft::error::NetworkError;
use openraft::error::RPCError;
use openraft::error::RaftError;
use openraft::raft::AppendEntriesRequest;
use openraft::raft::AppendEntriesResponse;
use openraft::raft::InstallSnapshotRequest;
use openraft::raft::InstallSnapshotResponse;
use openraft::raft::VoteRequest;
use openraft::raft::VoteResponse;
use openraft::BasicNode;
use openraft::RaftNetwork;
use openraft::RaftNetworkFactory;
use serde::de::DeserializeOwned;
use serde::Serialize;
use tokio::sync::mpsc::Sender;
use tonic::transport::channel::Channel;
use tower_service::Service;
2023-05-24 21:22:14 +05:30
use crate::pool::*;
2023-12-26 15:13:41 +05:30
use crate::protobuf::dcache::dcache_service_client::DcacheServiceClient;
use crate::protobuf::dcache::RaftRequest;
2023-12-26 14:58:55 +05:30
#[derive(Debug)]
struct ChannelManager {}
#[async_trait]
impl ItemManager for ChannelManager {
type Key = String;
type Item = Channel;
type Error = tonic::transport::Error;
async fn build(&self, addr: &Self::Key) -> Result<Channel, tonic::transport::Error> {
tonic::transport::Endpoint::new(addr.clone())?
.connect()
.await
}
async fn check(&self, mut ch: Channel) -> Result<Channel, tonic::transport::Error> {
futures::future::poll_fn(|cx| (&mut ch).poll_ready(cx)).await?;
Ok(ch)
}
}
pub struct DcacheNetwork {
pub signal: Sender<HealthStatus>,
conn_pool: Pool<ChannelManager>,
}
2023-05-24 21:22:14 +05:30
2023-12-26 14:58:55 +05:30
pub enum RPCType {
Vote,
Snapshot,
Append,
}
2023-05-27 10:28:52 +05:30
impl DcacheNetwork {
2023-12-26 15:13:41 +05:30
pub fn new(signal: Sender<HealthStatus>) -> Self {
let mgr = ChannelManager {};
Self {
signal,
conn_pool: Pool::new(mgr, Duration::from_millis(50)),
}
}
2023-05-24 21:22:14 +05:30
pub async fn send_rpc<Req, Resp, Err>(
&self,
2023-05-27 10:28:52 +05:30
target: DcacheNodeId,
2023-05-24 21:22:14 +05:30
target_node: &BasicNode,
req: Req,
2023-12-26 14:58:55 +05:30
event: RPCType,
2023-05-27 10:28:52 +05:30
) -> Result<Resp, RPCError<DcacheNodeId, BasicNode, Err>>
2023-05-24 21:22:14 +05:30
where
Req: Serialize,
Err: std::error::Error + DeserializeOwned,
Resp: DeserializeOwned,
{
let mut client = self.make_client(&target, target_node).await;
2023-05-24 21:22:14 +05:30
2023-12-26 14:58:55 +05:30
let res = match event {
RPCType::Vote => {
client
.vote(RaftRequest {
data: serde_json::to_string(&req).unwrap(),
})
.await
}
2023-05-24 21:22:14 +05:30
2023-12-26 14:58:55 +05:30
RPCType::Snapshot => {
client
.install_snapshot(RaftRequest {
data: serde_json::to_string(&req).unwrap(),
})
.await
}
2023-05-24 21:22:14 +05:30
2023-12-26 14:58:55 +05:30
RPCType::Append => {
client
.append_entries(RaftRequest {
data: serde_json::to_string(&req).unwrap(),
})
.await
}
};
match res {
Ok(res) => {
let signal2 = self.signal.clone();
let fut = async move {
let _ = signal2.send(HealthStatus::Healthy(target)).await;
};
tokio::spawn(fut);
let res = res.into_inner();
Ok(serde_json::from_str(&res.data).unwrap())
}
Err(e) => {
let _ = self.signal.send(HealthStatus::Down(target)).await;
Err(RPCError::Network(NetworkError::new(&e)))
}
}
2023-05-24 21:22:14 +05:30
}
pub async fn make_client(
&self,
target: &DcacheNodeId,
target_node: &BasicNode,
) -> DcacheServiceClient<Channel> {
let addr = format!("http://{}", &target_node.addr);
tracing::debug!("connect: target={}: {}", target, addr);
let channel = self.conn_pool.get(&addr).await.unwrap();
let client = DcacheServiceClient::new(channel);
tracing::info!("connected: target={}: {}", target, addr);
client
}
2023-05-24 21:22:14 +05:30
}
2023-05-27 10:28:52 +05:30
// NOTE: This could be implemented also on `Arc<DcacheNetwork>`, but since it's empty, implemented
2023-05-24 21:22:14 +05:30
// directly.
#[async_trait]
impl RaftNetworkFactory<DcacheTypeConfig> for Arc<DcacheNetwork> {
2023-05-27 10:28:52 +05:30
type Network = DcacheNetworkConnection;
2023-05-24 21:22:14 +05:30
2023-05-27 10:28:52 +05:30
async fn new_client(&mut self, target: DcacheNodeId, node: &BasicNode) -> Self::Network {
DcacheNetworkConnection {
owner: self.clone(),
2023-05-24 21:22:14 +05:30
target,
target_node: node.clone(),
}
}
}
2023-05-27 10:28:52 +05:30
pub struct DcacheNetworkConnection {
owner: Arc<DcacheNetwork>,
2023-05-27 10:28:52 +05:30
target: DcacheNodeId,
2023-05-24 21:22:14 +05:30
target_node: BasicNode,
}
#[async_trait]
2023-05-27 10:28:52 +05:30
impl RaftNetwork<DcacheTypeConfig> for DcacheNetworkConnection {
2023-05-24 21:22:14 +05:30
async fn send_append_entries(
&mut self,
2023-05-27 10:28:52 +05:30
req: AppendEntriesRequest<DcacheTypeConfig>,
2023-05-24 21:22:14 +05:30
) -> Result<
2023-05-27 10:28:52 +05:30
AppendEntriesResponse<DcacheNodeId>,
RPCError<DcacheNodeId, BasicNode, RaftError<DcacheNodeId>>,
2023-05-24 21:22:14 +05:30
> {
self.owner
2023-12-26 15:13:41 +05:30
.send_rpc(self.target, &self.target_node, req, RPCType::Append)
2023-05-24 21:22:14 +05:30
.await
}
async fn send_install_snapshot(
&mut self,
2023-05-27 10:28:52 +05:30
req: InstallSnapshotRequest<DcacheTypeConfig>,
2023-05-24 21:22:14 +05:30
) -> Result<
2023-05-27 10:28:52 +05:30
InstallSnapshotResponse<DcacheNodeId>,
RPCError<DcacheNodeId, BasicNode, RaftError<DcacheNodeId, InstallSnapshotError>>,
2023-05-24 21:22:14 +05:30
> {
self.owner
2023-12-26 15:13:41 +05:30
.send_rpc(self.target, &self.target_node, req, RPCType::Append)
2023-05-24 21:22:14 +05:30
.await
}
async fn send_vote(
&mut self,
2023-05-27 10:28:52 +05:30
req: VoteRequest<DcacheNodeId>,
2023-05-24 21:22:14 +05:30
) -> Result<
2023-05-27 10:28:52 +05:30
VoteResponse<DcacheNodeId>,
RPCError<DcacheNodeId, BasicNode, RaftError<DcacheNodeId>>,
2023-05-24 21:22:14 +05:30
> {
self.owner
2023-12-26 15:13:41 +05:30
.send_rpc(self.target, &self.target_node, req, RPCType::Vote)
2023-05-24 21:22:14 +05:30
.await
}
}