"""Minimal Raft consensus for tracker shard assignments. Only shard-assignment commands (register/deregister) go through the log. Node liveness (heartbeats) is handled separately via CRDT gossip — these are high-frequency writes where eventual consistency is fine. Election timeout: random 150–300 ms (tight, suits in-process tests). Leader heartbeat interval: 50 ms. """ from __future__ import annotations import json import random import threading import time import urllib.error import urllib.request from dataclasses import dataclass from typing import Callable @dataclass class LogEntry: term: int command: str # "register" | "deregister" payload: dict class RaftNode: """Single Raft participant. ``apply_fn(command, payload)`` is called (under no external lock) when an entry is committed. Implementors must apply the command atomically. """ ELECTION_MIN = 0.15 # seconds ELECTION_MAX = 0.30 HB_INTERVAL = 0.05 # leader heartbeat interval # Role constants FOLLOWER = "follower" CANDIDATE = "candidate" LEADER = "leader" def __init__( self, self_url: str, peers: list[str], apply_fn: Callable[[str, dict], None], ) -> None: self.self_url = self_url self.peers = list(peers) self._apply_fn = apply_fn self._lock = threading.Lock() self.role: str = self.FOLLOWER self.current_term: int = 0 self.voted_for: str | None = None self.log: list[LogEntry] = [] self.commit_index: int = -1 self.last_applied: int = -1 self.leader_url: str | None = None # Leader bookkeeping per peer self._next_index: dict[str, int] = {} self._match_index: dict[str, int] = {} self._election_deadline: float = self._fresh_deadline() self._running = False # ------------------------------------------------------------------ start/stop def start(self) -> None: self._running = True threading.Thread(target=self._tick_loop, daemon=True, name="raft-tick").start() def stop(self) -> None: self._running = False # ------------------------------------------------------------------ public API @property def is_leader(self) -> bool: with self._lock: return self.role == self.LEADER def leader(self) -> str | None: with self._lock: return self.leader_url def status(self) -> dict: with self._lock: return { "role": self.role, "term": self.current_term, "leader": self.leader_url, "log_length": len(self.log), "commit_index": self.commit_index, } def propose(self, command: str, payload: dict) -> bool: """Leader: append and replicate an entry. Returns True when committed. Blocks until majority replication or failure. Must be called only on the leader; returns False immediately if this node is not the leader. """ with self._lock: if self.role != self.LEADER: return False entry = LogEntry(self.current_term, command, payload) self.log.append(entry) entry_index = len(self.log) - 1 term = self.current_term self._replicate_to_peers(term) with self._lock: return self.commit_index >= entry_index # ------------------------------------------------------------------ RPC handlers def handle_request_vote(self, req: dict) -> dict: with self._lock: term = int(req["term"]) candidate = req["candidate_url"] last_li = int(req["last_log_index"]) last_lt = int(req["last_log_term"]) if term > self.current_term: self._step_down(term) if term < self.current_term: return {"term": self.current_term, "vote_granted": False} my_li = len(self.log) - 1 my_lt = self.log[-1].term if self.log else 0 log_ok = (last_lt > my_lt) or (last_lt == my_lt and last_li >= my_li) if (self.voted_for is None or self.voted_for == candidate) and log_ok: self.voted_for = candidate self._reset_deadline() return {"term": self.current_term, "vote_granted": True} return {"term": self.current_term, "vote_granted": False} def handle_append_entries(self, req: dict) -> dict: with self._lock: term = int(req["term"]) leader_url = req["leader_url"] prev_li = int(req["prev_log_index"]) prev_lt = int(req["prev_log_term"]) entries_raw = req.get("entries", []) ldr_commit = int(req["leader_commit"]) if term > self.current_term: self._step_down(term) if term < self.current_term: return {"term": self.current_term, "success": False} # Valid AppendEntries from current leader self._reset_deadline() self.leader_url = leader_url if self.role == self.CANDIDATE: self.role = self.FOLLOWER # Consistency check on prev entry if prev_li >= 0: if prev_li >= len(self.log): return {"term": self.current_term, "success": False} if self.log[prev_li].term != prev_lt: self.log = self.log[:prev_li] return {"term": self.current_term, "success": False} # Append new entries (detect and overwrite conflicts) for i, raw in enumerate(entries_raw): idx = prev_li + 1 + i entry = LogEntry(int(raw["term"]), raw["command"], raw["payload"]) if idx < len(self.log): if self.log[idx].term != entry.term: self.log = self.log[:idx] self.log.append(entry) else: self.log.append(entry) if ldr_commit > self.commit_index: self.commit_index = min(ldr_commit, len(self.log) - 1) self._apply_up_to_commit() return {"term": self.current_term, "success": True} # ------------------------------------------------------------------ internals def _tick_loop(self) -> None: while self._running: time.sleep(0.01) with self._lock: role = self.role deadline = self._election_deadline if role == self.LEADER: self._send_heartbeats() time.sleep(self.HB_INTERVAL) elif time.monotonic() > deadline: self._start_election() def _send_heartbeats(self) -> None: with self._lock: term = self.current_term ldr_commit = self.commit_index log_snapshot = list(self.log) for peer in self.peers: try: self._send_append_entries(peer, term, ldr_commit, log_snapshot) except Exception: pass def _replicate_to_peers(self, term: int) -> None: """Send AppendEntries to all peers and update commit_index on majority ack.""" with self._lock: ldr_commit = self.commit_index log_snapshot = list(self.log) acks = 1 # leader counts as 1 for peer in self.peers: try: ok = self._send_append_entries(peer, term, ldr_commit, log_snapshot) if ok: acks += 1 except Exception: pass # Advance commit if majority replicated with self._lock: if self.role != self.LEADER or self.current_term != term: return majority_index = len(self.log) - 1 while majority_index > self.commit_index: if self.log[majority_index].term == self.current_term: count = 1 + sum( 1 for p in self.peers if self._match_index.get(p, -1) >= majority_index ) if count > (len(self.peers) + 1) / 2: self.commit_index = majority_index self._apply_up_to_commit() break majority_index -= 1 def _send_append_entries( self, peer: str, term: int, ldr_commit: int, log_snapshot: list[LogEntry] ) -> bool: with self._lock: next_idx = self._next_index.get(peer, len(log_snapshot)) prev_li = next_idx - 1 prev_lt = log_snapshot[prev_li].term if 0 <= prev_li < len(log_snapshot) else 0 entries = [ {"term": e.term, "command": e.command, "payload": e.payload} for e in log_snapshot[next_idx:] ] body = json.dumps({ "term": term, "leader_url": self.self_url, "prev_log_index": prev_li, "prev_log_term": prev_lt, "entries": entries, "leader_commit": ldr_commit, }).encode() req = urllib.request.Request( f"{peer}/v1/raft/append", data=body, headers={"Content-Type": "application/json"}, method="POST", ) with urllib.request.urlopen(req, timeout=0.15) as r: resp = json.loads(r.read()) with self._lock: if resp.get("success"): new_match = len(log_snapshot) - 1 self._match_index[peer] = max(self._match_index.get(peer, -1), new_match) self._next_index[peer] = new_match + 1 return True else: if resp.get("term", 0) > self.current_term: self._step_down(resp["term"]) else: self._next_index[peer] = max(0, self._next_index.get(peer, 1) - 1) return False def _start_election(self) -> None: with self._lock: self.current_term += 1 self.role = self.CANDIDATE self.voted_for = self.self_url term = self.current_term my_li = len(self.log) - 1 my_lt = self.log[-1].term if self.log else 0 self._reset_deadline() votes = 1 for peer in self.peers: try: body = json.dumps({ "term": term, "candidate_url": self.self_url, "last_log_index": my_li, "last_log_term": my_lt, }).encode() req = urllib.request.Request( f"{peer}/v1/raft/vote", data=body, headers={"Content-Type": "application/json"}, method="POST", ) with urllib.request.urlopen(req, timeout=0.1) as r: resp = json.loads(r.read()) if resp.get("vote_granted"): votes += 1 elif resp.get("term", 0) > term: with self._lock: self._step_down(resp["term"]) return except Exception: pass with self._lock: if self.role == self.CANDIDATE and self.current_term == term: total = len(self.peers) + 1 if votes > total / 2: self._become_leader() else: self.role = self.FOLLOWER self._reset_deadline() def _become_leader(self) -> None: """Must be called with _lock held.""" self.role = self.LEADER self.leader_url = self.self_url for peer in self.peers: self._next_index[peer] = len(self.log) self._match_index[peer] = -1 print(f"[raft] {self.self_url} became leader (term {self.current_term})", flush=True) def _step_down(self, new_term: int) -> None: """Must be called with _lock held.""" self.current_term = new_term self.role = self.FOLLOWER self.voted_for = None self._reset_deadline() def _apply_up_to_commit(self) -> None: """Must be called with _lock held.""" while self.last_applied < self.commit_index: self.last_applied += 1 entry = self.log[self.last_applied] try: self._apply_fn(entry.command, entry.payload) except Exception: pass def _reset_deadline(self) -> None: self._election_deadline = self._fresh_deadline() @staticmethod def _fresh_deadline() -> float: return time.monotonic() + random.uniform( RaftNode.ELECTION_MIN, RaftNode.ELECTION_MAX )