nodedb-raft 0.0.0

//! RPC handlers for incoming Raft messages.

use tracing::{debug, info, warn};

use crate::message::{
    AppendEntriesRequest, AppendEntriesResponse, InstallSnapshotRequest, InstallSnapshotResponse,
    RequestVoteRequest, RequestVoteResponse,
};
use crate::state::NodeRole;
use crate::storage::LogStorage;

use super::core::RaftNode;

impl<S: LogStorage> RaftNode<S> {
    /// Handle incoming AppendEntries RPC.
    pub fn handle_append_entries(&mut self, req: &AppendEntriesRequest) -> AppendEntriesResponse {
        if req.term < self.hard_state.current_term {
            return AppendEntriesResponse {
                term: self.hard_state.current_term,
                success: false,
                last_log_index: self.log.last_index(),
            };
        }

        if req.term > self.hard_state.current_term || self.role == NodeRole::Candidate {
            self.become_follower(req.term);
        }

        self.leader_id = req.leader_id;
        self.reset_election_timeout();

        // Check prev_log consistency.
        if req.prev_log_index > 0 {
            match self.log.term_at(req.prev_log_index) {
                Some(term) if term == req.prev_log_term => {}
                _ => {
                    return AppendEntriesResponse {
                        term: self.hard_state.current_term,
                        success: false,
                        last_log_index: self.log.last_index(),
                    };
                }
            }
        }

        if let Err(e) = self.log.append_entries(req.prev_log_index, &req.entries) {
            warn!(group = self.config.group_id, error = %e, "append_entries failed");
            return AppendEntriesResponse {
                term: self.hard_state.current_term,
                success: false,
                last_log_index: self.log.last_index(),
            };
        }

        if req.leader_commit > self.volatile.commit_index {
            self.volatile.commit_index = req.leader_commit.min(self.log.last_index());
            self.collect_committed_entries();
        }

        AppendEntriesResponse {
            term: self.hard_state.current_term,
            success: true,
            last_log_index: self.log.last_index(),
        }
    }

    /// Handle incoming RequestVote RPC.
    pub fn handle_request_vote(&mut self, req: &RequestVoteRequest) -> RequestVoteResponse {
        if req.term > self.hard_state.current_term {
            self.become_follower(req.term);
        }

        if req.term < self.hard_state.current_term {
            return RequestVoteResponse {
                term: self.hard_state.current_term,
                vote_granted: false,
            };
        }

        let voted_for = self.hard_state.voted_for;
        let can_vote = voted_for == 0 || voted_for == req.candidate_id;

        let log_ok = req.last_log_term > self.log.last_term()
            || (req.last_log_term == self.log.last_term()
                && req.last_log_index >= self.log.last_index());

        if can_vote && log_ok {
            self.hard_state.voted_for = req.candidate_id;
            self.persist_hard_state();
            self.reset_election_timeout();

            debug!(
                node = self.config.node_id,
                group = self.config.group_id,
                candidate = req.candidate_id,
                term = req.term,
                "granted vote"
            );

            RequestVoteResponse {
                term: self.hard_state.current_term,
                vote_granted: true,
            }
        } else {
            RequestVoteResponse {
                term: self.hard_state.current_term,
                vote_granted: false,
            }
        }
    }

    /// Handle AppendEntries response from a peer (leader only).
    pub fn handle_append_entries_response(&mut self, peer: u64, resp: &AppendEntriesResponse) {
        if resp.term > self.hard_state.current_term {
            self.become_follower(resp.term);
            return;
        }

        if self.role != NodeRole::Leader {
            return;
        }

        let leader = match self.leader_state.as_mut() {
            Some(ls) => ls,
            None => return,
        };

        if resp.success {
            let new_match = resp.last_log_index;
            if new_match > leader.match_index_for(peer) {
                leader.set_match_index(peer, new_match);
                leader.set_next_index(peer, new_match + 1);
            }
            self.try_advance_commit_index();
        } else {
            let new_next = resp.last_log_index + 1;
            let current_next = leader.next_index_for(peer);
            if new_next < current_next {
                leader.set_next_index(peer, new_next.max(1));
            } else {
                leader.set_next_index(peer, current_next.saturating_sub(1).max(1));
            }
            self.send_append_entries(peer);
        }
    }

    /// Handle RequestVote response (candidate only).
    pub fn handle_request_vote_response(&mut self, _peer: u64, resp: &RequestVoteResponse) {
        if resp.term > self.hard_state.current_term {
            self.become_follower(resp.term);
            return;
        }

        if self.role != NodeRole::Candidate {
            return;
        }

        if resp.vote_granted {
            self.votes_received.push(resp.term);
            let vote_count = self.votes_received.len() + 1; // +1 for self-vote

            if vote_count >= self.config.quorum() {
                self.become_leader();
            }
        }
    }

    /// Handle incoming InstallSnapshot RPC (Raft paper Figure 13).
    ///
    /// Called on followers that are too far behind for log-based catch-up.
    /// The leader sends its snapshot; the follower replaces its log and state.
    pub fn handle_install_snapshot(
        &mut self,
        req: &InstallSnapshotRequest,
    ) -> InstallSnapshotResponse {
        if req.term < self.hard_state.current_term {
            return InstallSnapshotResponse {
                term: self.hard_state.current_term,
            };
        }

        if req.term > self.hard_state.current_term {
            self.become_follower(req.term);
        }

        self.leader_id = req.leader_id;
        self.reset_election_timeout();

        if req.done && req.last_included_index > self.log.snapshot_index() {
            info!(
                node = self.config.node_id,
                group = self.config.group_id,
                snapshot_index = req.last_included_index,
                snapshot_term = req.last_included_term,
                "applying installed snapshot"
            );

            self.log
                .apply_snapshot(req.last_included_index, req.last_included_term);

            if self.volatile.commit_index < req.last_included_index {
                self.volatile.commit_index = req.last_included_index;
            }
            if self.volatile.last_applied < req.last_included_index {
                self.volatile.last_applied = req.last_included_index;
            }
        }

        InstallSnapshotResponse {
            term: self.hard_state.current_term,
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::{Duration, Instant};

    use crate::message::{AppendEntriesRequest, LogEntry, RequestVoteRequest, RequestVoteResponse};
    use crate::state::NodeRole;
    use crate::storage::MemStorage;

    use super::super::core::RaftConfig;
    use super::*;

    fn test_config(node_id: u64, peers: Vec<u64>) -> RaftConfig {
        RaftConfig {
            node_id,
            group_id: 1,
            peers,
            election_timeout_min: Duration::from_millis(150),
            election_timeout_max: Duration::from_millis(300),
            heartbeat_interval: Duration::from_millis(50),
        }
    }

    #[test]
    fn follower_rejects_old_term() {
        let config = test_config(1, vec![2, 3]);
        let mut node = RaftNode::new(config, MemStorage::new());
        node.hard_state.current_term = 5;

        let req = AppendEntriesRequest {
            term: 3,
            leader_id: 2,
            prev_log_index: 0,
            prev_log_term: 0,
            entries: vec![],
            leader_commit: 0,
            group_id: 1,
        };

        let resp = node.handle_append_entries(&req);
        assert!(!resp.success);
        assert_eq!(resp.term, 5);
    }

    #[test]
    fn follower_accepts_valid_append() {
        let config = test_config(1, vec![2, 3]);
        let mut node = RaftNode::new(config, MemStorage::new());

        let req = AppendEntriesRequest {
            term: 1,
            leader_id: 2,
            prev_log_index: 0,
            prev_log_term: 0,
            entries: vec![
                LogEntry {
                    term: 1,
                    index: 1,
                    data: b"a".to_vec(),
                },
                LogEntry {
                    term: 1,
                    index: 2,
                    data: b"b".to_vec(),
                },
            ],
            leader_commit: 1,
            group_id: 1,
        };

        let resp = node.handle_append_entries(&req);
        assert!(resp.success);
        assert_eq!(resp.last_log_index, 2);
        assert_eq!(node.commit_index(), 1);
        assert_eq!(node.leader_id(), 2);
    }

    #[test]
    fn vote_grant_and_reject() {
        let config = test_config(1, vec![2, 3]);
        let mut node = RaftNode::new(config, MemStorage::new());

        let req = RequestVoteRequest {
            term: 1,
            candidate_id: 2,
            last_log_index: 0,
            last_log_term: 0,
            group_id: 1,
        };
        let resp = node.handle_request_vote(&req);
        assert!(resp.vote_granted);

        let req2 = RequestVoteRequest {
            term: 1,
            candidate_id: 3,
            last_log_index: 0,
            last_log_term: 0,
            group_id: 1,
        };
        let resp2 = node.handle_request_vote(&req2);
        assert!(!resp2.vote_granted);
    }

    #[test]
    fn three_node_election() {
        let config1 = test_config(1, vec![2, 3]);
        let config2 = test_config(2, vec![1, 3]);
        let config3 = test_config(3, vec![1, 2]);

        let mut node1 = RaftNode::new(config1, MemStorage::new());
        let mut node2 = RaftNode::new(config2, MemStorage::new());
        let mut node3 = RaftNode::new(config3, MemStorage::new());

        node1.election_deadline = Instant::now() - Duration::from_millis(1);
        node1.tick();
        assert_eq!(node1.role(), NodeRole::Candidate);

        let ready = node1.take_ready();
        assert_eq!(ready.vote_requests.len(), 2);

        let resp2 = node2.handle_request_vote(&ready.vote_requests[0].1);
        let resp3 = node3.handle_request_vote(&ready.vote_requests[1].1);
        assert!(resp2.vote_granted);
        assert!(resp3.vote_granted);

        node1.handle_request_vote_response(2, &resp2);
        assert_eq!(node1.role(), NodeRole::Leader);
    }

    #[test]
    fn three_node_replication() {
        let config1 = test_config(1, vec![2, 3]);
        let config2 = test_config(2, vec![1, 3]);

        let mut node1 = RaftNode::new(config1, MemStorage::new());
        let mut node2 = RaftNode::new(config2, MemStorage::new());

        node1.election_deadline = Instant::now() - Duration::from_millis(1);
        node1.tick();
        let ready = node1.take_ready();
        let resp2 = node2.handle_request_vote(&ready.vote_requests[0].1);
        node1.handle_request_vote_response(2, &resp2);
        assert_eq!(node1.role(), NodeRole::Leader);

        let heartbeat_ready = node1.take_ready();
        for (peer_id, msg) in &heartbeat_ready.messages {
            if *peer_id == 2 {
                let resp = node2.handle_append_entries(msg);
                node1.handle_append_entries_response(2, &resp);
            }
        }

        let idx = node1.propose(b"cmd1".to_vec()).unwrap();
        assert_eq!(idx, 2);

        let ready = node1.take_ready();
        for (peer_id, msg) in &ready.messages {
            if *peer_id == 2 {
                let resp = node2.handle_append_entries(msg);
                assert!(resp.success);
                node1.handle_append_entries_response(2, &resp);
            }
        }

        let ready = node1.take_ready();
        let committed: Vec<_> = ready
            .committed_entries
            .iter()
            .filter(|e| !e.data.is_empty())
            .collect();
        assert_eq!(committed.len(), 1);
        assert_eq!(committed[0].data, b"cmd1");
    }

    #[test]
    fn leader_steps_down_on_higher_term() {
        let config = test_config(1, vec![2, 3]);
        let mut node = RaftNode::new(config, MemStorage::new());

        node.election_deadline = Instant::now() - Duration::from_millis(1);
        node.tick();
        let _ready = node.take_ready();
        let resp = RequestVoteResponse {
            term: 1,
            vote_granted: true,
        };
        node.handle_request_vote_response(2, &resp);
        assert_eq!(node.role(), NodeRole::Leader);

        let req = AppendEntriesRequest {
            term: 5,
            leader_id: 2,
            prev_log_index: 0,
            prev_log_term: 0,
            entries: vec![],
            leader_commit: 0,
            group_id: 1,
        };
        node.handle_append_entries(&req);
        assert_eq!(node.role(), NodeRole::Follower);
        assert_eq!(node.current_term(), 5);
        assert_eq!(node.leader_id(), 2);
    }
}