tokioraft/

lib.rs

//! `tokioraft` is a small Tokio-oriented leader election library.
//!
//! The crate implements a minimal Raft-inspired election state machine.
//! It reacts to incoming peer messages, drives local election timers, and
//! emits outgoing election messages together with local role changes.
//!
//! `tokioraft` is intentionally narrow in scope:
//!
//! - it does leader election only
//! - it does not implement a replicated log
//! - it does not provide networking or node discovery
//! - it does not manage cluster membership automatically
//!
//! Typical integration flow:
//!
//! 1. Create one [`TinyRaft`] instance per node
//! 2. Read events from [`TinyRaft::get_receiver_from_raft`]
//! 3. Deliver [`send::MsgRouter::ToNode`] messages using your own transport
//! 4. Feed received peer messages back through [`TinyRaft::on_receive`]
//! 5. Apply membership changes explicitly through [`TinyRaft::set_nodes`]
//!
//! The library is intended for Tokio-based applications and expects the host
//! application not to block Tokio worker threads for long periods of time.
//!
use std::{collections::HashMap, time::Duration};

use async_channel::{Receiver, Sender};
use async_timeouts::Timeout;
use error::{Result, ResultExt, StateError};
use inner::TinyRaftInner;
pub use node::{NodeId, NodeType};
use send::{Message, MessageSender, MessageType, MsgRouter, NodeState, SignalSender};
use tokio::{task, time};
use tracing::{error, info};

pub mod error;
pub mod inner;
mod node;
pub mod send;

type OnReceive = (Sender<(NodeId, Message)>, Receiver<(NodeId, Message)>);

/// A running leader-election node.
///
/// A `TinyRaft` instance owns the local election state machine for a single
/// node. It communicates with the outside world through channels:
///
/// - it emits outgoing peer messages and local state changes
/// - it accepts incoming peer messages from the host application
/// - it accepts explicit cluster membership updates
#[derive(Clone)]
pub struct TinyRaft {
    stop_tinyraft_tx: Sender<()>,
    on_receive: OnReceive,
    raft_change_rx: Receiver<MsgRouter>,
    set_nodes: (Sender<Vec<NodeId>>, Receiver<Vec<NodeId>>),
}

impl TinyRaft {
    /// Starts a new election state machine for a single node.
    ///
    /// `nodes` must contain `node_id`. The returned instance immediately starts
    /// its internal Tokio task and election timers.
    ///
    /// Parameters:
    ///
    /// - `nodes` — current cluster membership known to this node
    /// - `node_id` — identifier of the local node
    /// - `election_timeout` — base timeout before a node starts a new election
    /// - `heartbeat_timeout` — interval used by a leader to emit heartbeat messages
    /// - `random_timeout` — additional random delay added to `election_timeout`
    /// - `leadership_timeout` — optional timeout used to expire leadership without quorum confirmation
    /// - `leader_priority` — optional tie-breaking priority used when handling higher-term vote requests
    /// - `term` — initial local term
    ///
    /// More detail on timing-related parameters:
    ///
    /// - `election_timeout` controls how quickly a node tries to become a
    ///   candidate when it does not observe a stable leader
    /// - `random_timeout` is added on top of `election_timeout` to reduce
    ///   simultaneous elections; if `None`, it defaults to `election_timeout`
    /// - `heartbeat_timeout` controls how often the current leader broadcasts
    ///   `Ping` messages
    /// - `leadership_timeout`, when non-zero, enables a stricter leadership
    ///   confirmation mode where a leader must keep receiving enough `Pong`
    ///   confirmations; if it stops receiving them, leadership eventually expires
    ///
    /// `leader_priority` is optional. When set, it affects how a node reacts to
    /// a higher-term `VoteRequest`: a node with a higher local priority may force
    /// a local restart of the election instead of immediately accepting the remote
    /// candidate.
    ///
    /// `term` normally starts at `0` for fresh in-memory state. Supplying a
    /// non-zero value is useful only if the surrounding application restores term
    /// information from its own state.
    ///
    /// The caller is responsible for:
    ///
    /// - reading events from [`TinyRaft::get_receiver_from_raft`]
    /// - transporting outgoing messages to remote nodes
    /// - passing received peer messages back into [`TinyRaft::on_receive`]
    /// - updating cluster membership through [`TinyRaft::set_nodes`]
    ///
    /// # Panics
    ///
    /// Panics if the provided node list is invalid or does not contain `node_id`.
    #[allow(clippy::too_many_arguments)]
    pub async fn start(
        nodes: impl IntoIterator<Item = impl Into<NodeId>>,
        node_id: impl Into<NodeId>,
        election_timeout: Duration,
        heartbeat_timeout: Duration,
        random_timeout: Option<Duration>,
        leadership_timeout: Option<Duration>,
        leader_priority: Option<u64>,
        term: u64,
    ) -> Self {
        let nodes = match node::check_node_ids(nodes) {
            Ok(nodes) => nodes,
            Err(e) => panic!("{e:?}"),
        };

        let node_id = node_id.into();

        info!(node = %node_id, nodes = ?nodes, term, "starting tinyraft node");

        if !nodes.contains(&node_id) {
            panic!("{}", StateError::SetNodesMissingSelf);
        }

        let leadership_timeout = leadership_timeout.unwrap_or_default();
        let random_timeout = random_timeout.unwrap_or(election_timeout);

        let (stop_tx, stop_rx) = async_channel::unbounded();
        let (start_tx, start_rx) = async_channel::unbounded();
        let (on_receive_tx, on_receive_rx) = async_channel::unbounded();
        let (raft_change_tx, raft_change_rx) = async_channel::unbounded();
        let (set_nodes_tx, set_nodes_rx) = async_channel::unbounded();
        let (heartbeat_tx, heartbeat_rx) = async_channel::unbounded();

        let set_nodes_rx_clone = set_nodes_rx.clone();
        let on_receive_rx_clone = on_receive_rx.clone();

        let heartbeat_sender = SignalSender::new(heartbeat_tx.clone());

        let start_sender = SignalSender::new(start_tx.clone());
        let message_sender = MessageSender::new(raft_change_tx);

        let _ = start_tx.send(()).await;
        let _ = heartbeat_tx.send(()).await;

        task::spawn(async move {
            let mut tinyraft = TinyRaftInner {
                node_type: NodeType::Candidate,
                confirmed: HashMap::new(),
                leader_priority,
                term,
                leader: node_id.clone(),
                followers: vec![],
                votes: HashMap::new(),
                node_id,
                nodes,
                election_timeout,
                heartbeat_timeout,
                start_sender: start_sender.clone(),
                next_term: Timeout::default(),
                heartbeat_sender: heartbeat_sender.clone(),
                message_sender,
                leadership_timeout,
                random_timeout,
            };

            loop {
                tokio::select! {
                    _ = stop_rx.recv() => {
                        info!(node = %tinyraft.node_id, term = tinyraft.term, "stopping tinyraft worker");
                        break
                    },

                    _ = start_rx.recv() => {
                        tinyraft.reset_election_state().await;
                        tinyraft.broadcast_vote_requests().await;
                        tinyraft.publish_candidate_state().await;
                        tinyraft.start_timeout(start_sender.clone()).await;
                    }

                    _ = heartbeat_rx.recv() => {
                        if let NodeType::Leader = tinyraft.node_type {
                            tinyraft.broadcast_heartbeat().await;
                        }

                        let heartbeat_sender = heartbeat_sender.clone();

                        task::spawn(async move {
                            time::sleep(heartbeat_timeout).await;

                            if let Err(e) = heartbeat_sender.send().await {
                                info!(error = ?e, "heartbeat loop stopped because node is shutting down");
                            }
                        });
                    },

                    Ok(nodes) = set_nodes_rx_clone.recv() => {
                        if let Err(e) = tinyraft.set_nodes(nodes).await {
                            error!(node = %tinyraft.node_id, error = ?e, "failed to update cluster nodes");
                        }
                    },

                    Ok((from, message)) = on_receive_rx_clone.recv() => {
                        if let Err(e) = tinyraft.on_receive(from, message).await {
                            error!(node = %tinyraft.node_id, error = ?e, "failed to handle incoming message");
                            panic!("{e:?}")
                        }
                    }
                }
            }
        });

        Self {
            stop_tinyraft_tx: stop_tx,
            on_receive: (on_receive_tx, on_receive_rx),
            raft_change_rx,
            set_nodes: (set_nodes_tx, set_nodes_rx),
        }
    }

    pub async fn stop(&self) -> Result<()> {
        info!("stop requested");
        self.stop_tinyraft_tx.send(()).await.or_request_stop()
    }

    /// Enqueues a message received from another node.
    ///
    /// This should be called by the host transport layer when it receives a
    /// Raft message from a peer.
    pub async fn on_receive(&self, from: NodeId, message: Message) -> Result<()> {
        info!(from = %from, message = ?message, "enqueueing incoming message");
        self.on_receive
            .0
            .send((from.clone(), message.clone()))
            .await
            .or_request_on_receive(from, message)
    }

    /// Returns a receiver with outgoing peer messages and local state updates.
    ///
    /// The receiver yields [`MsgRouter`] events:
    ///
    /// - [`MsgRouter::ToNode`] for outgoing peer messages
    /// - [`MsgRouter::ToSelf`] for local state updates
    pub fn get_receiver_from_raft(&self) -> Receiver<MsgRouter> {
        self.raft_change_rx.clone()
    }

    /// Replaces the cluster membership known to this node.
    ///
    /// Membership changes are explicit in `tokioraft`: the library does not
    /// discover node additions or removals automatically.
    pub async fn set_nodes(
        &self,
        nodes: impl IntoIterator<Item = impl Into<NodeId>>,
    ) -> Result<()> {
        let nodes = node::check_node_ids(nodes)?;
        info!(nodes = ?nodes, "requesting cluster nodes update");

        self.set_nodes
            .0
            .send(nodes.clone())
            .await
            .or_request_set_nodes(nodes)
    }
}

impl TinyRaftInner {
    async fn reset_election_state(&mut self) {
        info!(node = %self.node_id, next_term = self.term + 1, "starting election");
        self.next_term.stop().await;
        self.term += 1;
        self.leader = self.node_id.clone();
        self.node_type = NodeType::Candidate;
        self.followers.clear();
        self.confirmed.clear();
        self.votes = HashMap::from([(self.node_id.clone(), vec![self.node_id.clone()])]);
    }

    async fn broadcast_vote_requests(&self) {
        for node in self.nodes.iter().filter(|node| *node != &self.node_id) {
            if let Err(e) = self
                .message_sender
                .send(MsgRouter::ToNode(
                    node.clone(),
                    Message {
                        msg_type: MessageType::VoteRequest,
                        term: self.term,
                        leader: Some(self.leader.clone()),
                        priority: self.leader_priority,
                    },
                ))
                .await
            {
                error!(node = %self.node_id, target = %node, term = self.term, error = ?e, "failed to send vote request");
                panic!("start error: {e:?}");
            }
        }
    }

    async fn publish_candidate_state(&self) {
        if let Err(e) = self
            .message_sender
            .send(MsgRouter::ToSelf(NodeState {
                node_type: self.node_type,
                term: self.term,
                leader: None,
                followers: vec![],
            }))
            .await
        {
            error!(node = %self.node_id, term = self.term, error = ?e, "failed to publish candidate state");
            panic!("start error: {e:?}");
        }
    }

    async fn broadcast_heartbeat(&mut self) {
        info!(node = %self.node_id, term = self.term, "sending heartbeat to followers");
        self.confirmed.clear();
        for node in self.nodes.iter().filter(|node| *node != &self.node_id) {
            if let Err(e) = self
                .message_sender
                .send(MsgRouter::ToNode(
                    node.clone(),
                    Message {
                        msg_type: MessageType::Ping,
                        term: self.term,
                        leader: None,
                        priority: self.leader_priority,
                    },
                ))
                .await
            {
                error!(node = %self.node_id, target = %node, term = self.term, error = ?e, "failed to send heartbeat");
                panic!("heartbeat error: {e:?}");
            }
        }
    }
}