photoncloud-monorepo/flaredb/crates/flaredb-server/src/raft_service.rs

use crate::config::NamespaceManager;
use crate::store::Store;
use flaredb_proto::raft_server::raft_service_server::RaftService;
use flaredb_proto::raft_server::{
    FetchRangeRequest, FetchRangeResponse, ForwardEventualRequest, GetMerkleRequest,
    GetMerkleResponse, GetModeRequest, GetModeResponse, ListNamespaceModesRequest,
    ListNamespaceModesResponse, MerkleRange, NamespaceMode, OpenRaftAppendEntriesRequest,
    OpenRaftAppendEntriesResponse, OpenRaftSnapshotRequest, OpenRaftSnapshotResponse,
    OpenRaftVoteRequest, OpenRaftVoteResponse, RaftMessage, RaftResponse,
    UpdateNamespaceModeRequest, UpdateNamespaceModeResponse,
};
use flaredb_raft::{FlareNode, FlareNodeId, FlareTypeConfig};
use openraft::raft::{AppendEntriesRequest, VoteRequest};
use std::sync::Arc;
use tonic::{Request, Response, Status};

pub struct RaftServiceImpl {
    store: Arc<Store>,
    #[allow(dead_code)]
    local_id: u64,
}

impl RaftServiceImpl {
    pub fn new(store: Arc<Store>, local_id: u64) -> Self {
        Self { store, local_id }
    }
}

#[tonic::async_trait]
impl RaftService for RaftServiceImpl {
    async fn send(&self, _request: Request<RaftMessage>) -> Result<Response<RaftResponse>, Status> {
        // Legacy raft-rs RPC - no longer supported
        Err(Status::unimplemented(
            "Legacy raft-rs RPC not supported. Use OpenRaft RPCs (vote_v2, append_entries_v2).",
        ))
    }

    async fn get_mode(
        &self,
        request: Request<GetModeRequest>,
    ) -> Result<Response<GetModeResponse>, Status> {
        let req = request.into_inner();
        let ns = if req.namespace.is_empty() {
            "default".to_string()
        } else {
            req.namespace
        };
        let mode = self.store.namespace_manager.get_namespace(&ns).mode; // Use namespace_manager
        let mode_str = NamespaceManager::mode_as_str(&mode).to_string();
        Ok(Response::new(GetModeResponse { mode: mode_str }))
    }

    async fn update_namespace_mode(
        &self,
        request: Request<UpdateNamespaceModeRequest>,
    ) -> Result<Response<UpdateNamespaceModeResponse>, Status> {
        let req = request.into_inner();
        if req.namespace.is_empty() {
            return Err(Status::invalid_argument("namespace is required"));
        }
        let mode = match req.mode.to_lowercase().as_str() {
            "strong" => crate::config::ConsistencyMode::Strong,
            "eventual" => crate::config::ConsistencyMode::Eventual,
            _ => {
                return Err(Status::invalid_argument(
                    "mode must be 'strong' or 'eventual'",
                ))
            }
        };
        let cfg = self
            .store
            .namespace_manager // Use namespace_manager
            .set_namespace_mode(&req.namespace, mode)
            .map_err(Status::failed_precondition)?;
        let mode_str = NamespaceManager::mode_as_str(&cfg.mode).to_string();
        let ns_mode = NamespaceMode {
            namespace: cfg.name.clone(),
            id: cfg.id,
            mode: mode_str,
            from_default: !cfg.explicit,
        };
        Ok(Response::new(UpdateNamespaceModeResponse {
            mode: Some(ns_mode),
        }))
    }

    async fn forward_eventual(
        &self,
        request: Request<ForwardEventualRequest>,
    ) -> Result<Response<RaftResponse>, Status> {
        let req = request.into_inner();
        if let Some(node) = self.store.get_raft_node(req.region_id).await {
            if !node.is_leader().await {
                return Err(Status::failed_precondition("not leader"));
            }
            node.write_kv(req.namespace_id, req.key, req.value, req.ts)
                .await
                .map_err(|e| Status::internal(format!("raft write failed: {e}")))?;
            return Ok(Response::new(RaftResponse {}));
        }
        Err(Status::failed_precondition("region not found"))
    }

    async fn list_namespace_modes(
        &self,
        _request: Request<ListNamespaceModesRequest>,
    ) -> Result<Response<ListNamespaceModesResponse>, Status> {
        let list = self
            .store
            .namespace_manager
            .list_namespaces()
            .into_iter()
            .map(|cfg| NamespaceMode {
                namespace: cfg.name,
                id: cfg.id,
                mode: NamespaceManager::mode_as_str(&cfg.mode).to_string(),
                from_default: !cfg.explicit,
            })
            .collect();
        Ok(Response::new(ListNamespaceModesResponse {
            namespaces: list,
        }))
    }

    async fn get_merkle(
        &self,
        request: Request<GetMerkleRequest>,
    ) -> Result<Response<GetMerkleResponse>, Status> {
        let req = request.into_inner();
        let chunk = if req.chunk_size == 0 {
            256
        } else {
            req.chunk_size as usize
        };
        let (root, leaves) =
            crate::merkle::build_merkle(self.store.engine(), req.namespace_id, chunk)
                .map_err(|e| Status::internal(format!("merkle: {e}")))?;
        let ranges = leaves
            .into_iter()
            .map(|l| MerkleRange {
                start_key: l.start,
                end_key: l.end,
                hash: l.hash,
            })
            .collect();
        Ok(Response::new(GetMerkleResponse {
            root,
            leaves: ranges,
        }))
    }

    async fn fetch_range(
        &self,
        request: Request<FetchRangeRequest>,
    ) -> Result<Response<FetchRangeResponse>, Status> {
        let req = request.into_inner();
        let (keys, values) = crate::merkle::fetch_range(
            self.store.engine(),
            req.namespace_id,
            &req.start_key,
            &req.end_key,
        )
        .await
        .map_err(|e| Status::internal(format!("fetch_range: {e}")))?;
        Ok(Response::new(FetchRangeResponse { keys, values }))
    }

    // OpenRaft RPCs

    async fn vote_v2(
        &self,
        request: Request<OpenRaftVoteRequest>,
    ) -> Result<Response<OpenRaftVoteResponse>, Status> {
        let req = request.into_inner();
        let node = self
            .store
            .get_raft_node(req.region_id)
            .await
            .ok_or_else(|| Status::failed_precondition("region not found"))?;

        let vote_req: VoteRequest<FlareNodeId> = serde_json::from_slice(&req.data)
            .map_err(|e| Status::invalid_argument(format!("invalid vote request: {}", e)))?;

        let resp = node
            .raft
            .vote(vote_req)
            .await
            .map_err(|e| Status::internal(format!("vote failed: {}", e)))?;

        let data = serde_json::to_vec(&resp)
            .map_err(|e| Status::internal(format!("serialize response: {}", e)))?;

        Ok(Response::new(OpenRaftVoteResponse { data }))
    }

    async fn append_entries_v2(
        &self,
        request: Request<OpenRaftAppendEntriesRequest>,
    ) -> Result<Response<OpenRaftAppendEntriesResponse>, Status> {
        let req = request.into_inner();
        let node = self
            .store
            .get_raft_node(req.region_id)
            .await
            .ok_or_else(|| Status::failed_precondition("region not found"))?;

        let append_req: AppendEntriesRequest<FlareTypeConfig> = serde_json::from_slice(&req.data)
            .map_err(|e| {
            Status::invalid_argument(format!("invalid append_entries request: {}", e))
        })?;

        let resp = node
            .raft
            .append_entries(append_req)
            .await
            .map_err(|e| Status::internal(format!("append_entries failed: {}", e)))?;

        let data = serde_json::to_vec(&resp)
            .map_err(|e| Status::internal(format!("serialize response: {}", e)))?;

        Ok(Response::new(OpenRaftAppendEntriesResponse { data }))
    }

    async fn install_snapshot_v2(
        &self,
        request: Request<OpenRaftSnapshotRequest>,
    ) -> Result<Response<OpenRaftSnapshotResponse>, Status> {
        let req = request.into_inner();
        let node = self
            .store
            .get_raft_node(req.region_id)
            .await
            .ok_or_else(|| Status::failed_precondition("region not found"))?;

        let vote: openraft::Vote<FlareNodeId> = serde_json::from_slice(&req.vote)
            .map_err(|e| Status::invalid_argument(format!("invalid vote: {}", e)))?;

        let meta: openraft::SnapshotMeta<FlareNodeId, FlareNode> =
            serde_json::from_slice(&req.meta)
                .map_err(|e| Status::invalid_argument(format!("invalid meta: {}", e)))?;

        let snapshot = openraft::Snapshot {
            meta,
            snapshot: Box::new(std::io::Cursor::new(req.data)),
        };

        let resp = node
            .raft
            .install_full_snapshot(vote, snapshot)
            .await
            .map_err(|e| Status::internal(format!("install_snapshot failed: {}", e)))?;

        let data = serde_json::to_vec(&resp)
            .map_err(|e| Status::internal(format!("serialize response: {}", e)))?;

        Ok(Response::new(OpenRaftSnapshotResponse { data }))
    }
}