373 lines
13 KiB
Python
373 lines
13 KiB
Python
"""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
|
||
)
|