S073 · Consensus — Paxos & Raft Intuition

#!/usr/bin/env python3
"""raft_election.py — a toy of Raft's leader election in one file.
 
Run:  python raft_election.py
No deps. Simulates 5 nodes, a network with adjustable drop rate,
and prints a timeline of terms + votes + wins.
 
Not production Raft — no log replication, no persistence, no snapshots.
Enough to feel the protocol.
"""
from __future__ import annotations
 
import random
import threading
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional
 
random.seed(7)
 
 
class Role(Enum):
    FOLLOWER = "F"
    CANDIDATE = "C"
    LEADER = "L"
 
 
@dataclass
class VoteReq:
    term: int
    candidate: str
 
 
@dataclass
class VoteResp:
    term: int
    granted: bool
    voter: str
 
 
@dataclass
class Heartbeat:
    term: int
    leader: str
 
 
class Network:
    """Simulated network with a drop rate + latency range."""
    def __init__(self, drop: float = 0.0, latency: tuple[float, float] = (0.005, 0.02)):
        self.drop = drop
        self.latency = latency
        self.nodes: dict[str, "Node"] = {}
        self.partitions: set[frozenset[str]] = set()
 
    def register(self, node: "Node") -> None:
        self.nodes[node.name] = node
 
    def partition(self, group_a: set[str], group_b: set[str]) -> None:
        """Prevent any message between group_a and group_b."""
        for a in group_a:
            for b in group_b:
                self.partitions.add(frozenset([a, b]))
 
    def heal(self) -> None:
        self.partitions.clear()
 
    def send(self, src: str, dst: str, msg: object) -> None:
        if frozenset([src, dst]) in self.partitions:
            return
        if random.random() < self.drop:
            return
        delay = random.uniform(*self.latency)
        def deliver() -> None:
            self.nodes[dst].receive(src, msg)
        threading.Timer(delay, deliver).start()
 
 
class Node:
    def __init__(self, name: str, peers: list[str], net: Network):
        self.name = name
        self.peers = peers
        self.net = net
        self.role = Role.FOLLOWER
        self.term = 0
        self.voted_for: Optional[str] = None
        self.leader: Optional[str] = None
        self.votes_received: set[str] = set()
        self.lock = threading.Lock()
        self.stopped = False
        self._reset_election_timer()
        # Start the tick loop
        threading.Thread(target=self._tick, daemon=True).start()
 
    def _reset_election_timer(self) -> None:
        self.election_deadline = time.time() + random.uniform(0.15, 0.30)
 
    def _tick(self) -> None:
        while not self.stopped:
            time.sleep(0.01)
            with self.lock:
                if self.role == Role.LEADER:
                    self._send_heartbeats()
                elif time.time() > self.election_deadline:
                    self._start_election()
 
    def _start_election(self) -> None:
        self.role = Role.CANDIDATE
        self.term += 1
        self.voted_for = self.name
        self.votes_received = {self.name}
        self._reset_election_timer()
        self._log(f"election started, term={self.term}")
        for p in self.peers:
            self.net.send(self.name, p, VoteReq(self.term, self.name))
 
    def _send_heartbeats(self) -> None:
        for p in self.peers:
            self.net.send(self.name, p, Heartbeat(self.term, self.name))
 
    def receive(self, src: str, msg: object) -> None:
        with self.lock:
            if isinstance(msg, VoteReq):
                self._handle_vote_req(msg)
            elif isinstance(msg, VoteResp):
                self._handle_vote_resp(msg)
            elif isinstance(msg, Heartbeat):
                self._handle_heartbeat(msg)
 
    def _handle_vote_req(self, m: VoteReq) -> None:
        if m.term > self.term:
            self.term = m.term
            self.voted_for = None
            self.role = Role.FOLLOWER
        grant = (m.term == self.term
                 and (self.voted_for is None or self.voted_for == m.candidate))
        if grant:
            self.voted_for = m.candidate
            self._reset_election_timer()
        self.net.send(self.name, m.candidate, VoteResp(self.term, grant, self.name))
 
    def _handle_vote_resp(self, m: VoteResp) -> None:
        if self.role != Role.CANDIDATE or m.term != self.term:
            return
        if m.granted:
            self.votes_received.add(m.voter)
            majority = (len(self.peers) + 1) // 2 + 1
            if len(self.votes_received) >= majority:
                self.role = Role.LEADER
                self.leader = self.name
                self._log(f"WON election, term={self.term}, votes={self.votes_received}")
 
    def _handle_heartbeat(self, m: Heartbeat) -> None:
        if m.term >= self.term:
            self.term = m.term
            self.role = Role.FOLLOWER
            self.leader = m.leader
            self.voted_for = None
            self._reset_election_timer()
 
    def _log(self, msg: str) -> None:
        print(f"[{time.time()%100:6.2f}] {self.name}: {msg}")
 
    def stop(self) -> None:
        self.stopped = True
 
 
def main() -> None:
    net = Network(drop=0.0)
    names = ["n1", "n2", "n3", "n4", "n5"]
    nodes = {n: Node(n, [p for p in names if p != n], net) for n in names}
    for n in nodes.values():
        net.register(n)
 
    print("--- start: watch a leader emerge ---")
    time.sleep(1.5)
 
    # Simulate a partition isolating the current leader
    leader = next((n for n in nodes.values() if n.role == Role.LEADER), None)
    if leader:
        others = {n.name for n in nodes.values() if n.name != leader.name}
        print(f"\n--- partition {leader.name} away from the rest ---")
        net.partition({leader.name}, others)
        time.sleep(1.5)
        print("\n--- heal the partition ---")
        net.heal()
        time.sleep(1.0)
 
    for n in nodes.values():
        n.stop()
 
 
if __name__ == "__main__":
    main()