dcache/src/protobuf.rs

use std::sync::Arc;

use openraft::BasicNode;
use serde::de::DeserializeOwned;
use serde::Serialize;
use tonic::Response;

use dcache::dcache_service_server::DcacheService;
use dcache::{Learner, RaftReply, RaftRequest};

use crate::app::DcacheApp;

pub mod dcache {
    tonic::include_proto!("dcache"); // The string specified here must match the proto package name
}

#[derive(Clone)]
pub struct MyDcacheImpl {
    app: Arc<DcacheApp>,
}

impl MyDcacheImpl {
    pub fn new(app: Arc<DcacheApp>) -> Self {
        Self { app }
    }
}

#[tonic::async_trait]
impl DcacheService for MyDcacheImpl {
    async fn add_learner(
        &self,
        request: tonic::Request<Learner>,
    ) -> std::result::Result<tonic::Response<RaftReply>, tonic::Status> {
        let req = request.into_inner();
        let node_id = req.id;
        let node = BasicNode {
            addr: req.addr.clone(),
        };
        println!("Learner added: {:?}", &req.addr);
        let res = self.app.raft.add_learner(node_id, node, true).await;
        Ok(Response::new(res.into()))
    }

    async fn write(
        &self,
        request: tonic::Request<RaftRequest>,
    ) -> std::result::Result<tonic::Response<RaftReply>, tonic::Status> {
        let req = request.into_inner();
        let req = serde_json::from_str(&req.data).unwrap();
        let res = self.app.raft.client_write(req).await;
        Ok(Response::new(res.into()))
    }
    /// / Forward a request to other
    async fn forward(
        &self,
        _request: tonic::Request<RaftRequest>,
    ) -> std::result::Result<tonic::Response<RaftReply>, tonic::Status> {
        unimplemented!();
    }
    async fn append_entries(
        &self,
        request: tonic::Request<RaftRequest>,
    ) -> std::result::Result<tonic::Response<RaftReply>, tonic::Status> {
        let req = request.into_inner();
        let req = serde_json::from_str(&req.data).unwrap();
        let res = self.app.raft.append_entries(req).await;
        Ok(Response::new(res.into()))
    }
    async fn install_snapshot(
        &self,
        request: tonic::Request<RaftRequest>,
    ) -> std::result::Result<tonic::Response<RaftReply>, tonic::Status> {
        let req = request.into_inner();
        let req = serde_json::from_str(&req.data).unwrap();
        let res = self.app.raft.install_snapshot(req).await;
        Ok(Response::new(res.into()))
    }
    async fn vote(
        &self,
        request: tonic::Request<RaftRequest>,
    ) -> std::result::Result<tonic::Response<RaftReply>, tonic::Status> {
        let req = request.into_inner();
        let req = serde_json::from_str(&req.data).unwrap();
        let res = self.app.raft.vote(req).await;
        Ok(Response::new(res.into()))
    }
}

impl<T, E> From<RaftReply> for Result<T, E>
where
    T: DeserializeOwned,
    E: DeserializeOwned,
{
    fn from(msg: RaftReply) -> Self {
        if !msg.data.is_empty() {
            let resp: T = serde_json::from_str(&msg.data).expect("fail to deserialize");
            Ok(resp)
        } else {
            let err: E = serde_json::from_str(&msg.error).expect("fail to deserialize");
            Err(err)
        }
    }
}

impl<T, E> From<Result<T, E>> for RaftReply
where
    T: Serialize,
    E: Serialize,
{
    fn from(r: Result<T, E>) -> Self {
        match r {
            Ok(x) => {
                let data = serde_json::to_string(&x).expect("fail to serialize");
                RaftReply {
                    data,
                    error: Default::default(),
                }
            }
            Err(e) => {
                let error = serde_json::to_string(&e).expect("fail to serialize");
                RaftReply {
                    data: Default::default(),
                    error,
                }
            }
        }
    }
}