Farshid: Professional Life, Personal Passions

# Sample Implementation of RAFT
# Farshid Ashouri

import random
import threading
import time


class Node(threading.Thread):
    def __init__(self, node_id, cluster):
        super().__init__()
        self.node_id = node_id
        self.cluster = cluster
        self.state = "follower"
        self.term = 0
        self.voted_for = None
        self.votes_received = 0
        self.leader_id = None

    def run(self):
        while True:
            if self.state == "follower":
                self.follow()
            elif self.state == "candidate":
                self.start_election()
            elif self.state == "leader":
                self.lead()

    def follow(self):
        timeout = random.uniform(1, 2)
        start_time = time.time()
        while time.time() - start_time < timeout:
            # Check for messages efficiently.
            # This message handling is very basic and needs improvement
            # for a real system (e.g., dedicated message queues per node)
            message_to_process = None
            for i, (dest_node_id, msg) in enumerate(self.cluster.messages):
                if dest_node_id == self.node_id:
                    message_to_process = msg
                    # A real system would consume the message
                    # self.cluster.messages.pop(i)
                    break # Process one message per check cycle for simplicity

            if message_to_process and \
               message_to_process["term"] >= self.term:
                if message_to_process["type"] == "heartbeat":
                    # print(f"Node {self.node_id} (T{self.term}) received
                    # heartbeat from {message_to_process['leader_id']}
                    # (T{message_to_process['term']})")
                    start_time = time.time() # Reset timeout
                    self.term = message_to_process["term"]
                    self.leader_id = message_to_process["leader_id"]
                    self.state = "follower" # Ensure state if was candidate
                    self.voted_for = None
                elif message_to_process["type"] == "vote_request" and \
                     (self.voted_for is None or \
                      self.voted_for == message_to_process["candidate_id"]):
                    # Grant vote if term is newer or same and not yet voted
                    if message_to_process["term"] > self.term:
                        self.term = message_to_process["term"]
                        self.voted_for = None # Reset vote for new term
                        self.state = "follower"
                    if self.voted_for is None:
                        # print(f"Node {self.node_id} (T{self.term})
                        # voted for {message_to_process['candidate_id']}")
                        self.cluster.send_direct_message(
                            message_to_process["candidate_id"],
                            {
                                "type": "vote_response",
                                "term": self.term,
                                "vote_granted": True,
                                "voter_id": self.node_id,
                            }
                        )
                        self.voted_for = message_to_process["candidate_id"]
                        start_time = time.time() # Reset timeout after voting
            time.sleep(0.05) # Brief pause to prevent busy-waiting excessively
        # Timeout expired
        # print(f"Node {self.node_id} (T{self.term}) timed out,
        # becoming candidate.")
        self.state = "candidate"
        self.leader_id = None


    def start_election(self):
        self.term += 1
        # print(f"Node {self.node_id} starting election for term {self.term}")
        self.voted_for = self.node_id
        self.votes_received = 1
        self.leader_id = None # No leader known during election
        self.cluster.broadcast_to_others(
            self.node_id,
            {
                "type": "vote_request",
                "term": self.term,
                "candidate_id": self.node_id,
            }
        )
        start_time = time.time()
        # Shorter timeout for election phase compared to follower phase
        timeout = random.uniform(0.5, 1)

        processed_voter_ids = {self.node_id}

        while time.time() - start_time < timeout:
            message = self.cluster.get_direct_message(self.node_id)
            if (
                message
                and message["type"] == "vote_response"
                and message["term"] == self.term
                and message["vote_granted"]
                and message["voter_id"] not in processed_voter_ids
            ):
                self.votes_received += 1
                processed_voter_ids.add(message["voter_id"])
                # print(f"Node {self.node_id} received vote from
                # {message['voter_id']}, total {self.votes_received}")
            if self.votes_received > len(self.cluster.nodes) // 2:
                # print(f"Node {self.node_id} became leader for term
                # {self.term} with {self.votes_received} votes")
                self.state = "leader"
                return
            time.sleep(0.05)

        # print(f"Node {self.node_id} election timed out or failed,
        # returning to follower")
        self.state = "follower" # Election failed or timed out


    def lead(self):
        # print(f"Node {self.node_id} is now leader for term {self.term}.")
        while self.state == "leader":
            # print(f"Leader {self.node_id} (T{self.term}) sending heartbeats")
            self.cluster.broadcast_to_others(
                self.node_id, # Source
                {
                    "type": "heartbeat",
                    "term": self.term,
                    "leader_id": self.node_id,
                }
            )
            time.sleep(0.3) # Heartbeat interval, must be less than follower timeout
            # Check if a new leader has emerged with a higher term
            # This part is simplified; a real leader checks responses
            # or steps down if it sees a higher term.
            # For this simulation, if a leader receives a message with a higher term
            # (e.g. a vote request from a new candidate), it should step down.
            # This logic is partially in follow() but leaders also need to check.
            # For simplicity here, we assume leader remains leader until failure simulation (not implemented)


class Cluster:
    def __init__(self, num_nodes=5):
        self.num_nodes = num_nodes
        self.nodes = [Node(i, self) for i in range(self.num_nodes)]
        self.message_queues = {i: [] for i in range(self.num_nodes)}
        # Using a lock for message queue access, though Python lists are
        # mostly thread-safe for append/pop, explicit locking is safer practice.
        self.message_lock = threading.Lock()

    def broadcast_to_others(self, sender_id, message):
        with self.message_lock:
            for node_id in range(self.num_nodes):
                if node_id != sender_id:
                    self.message_queues[node_id].append(message)

    def send_direct_message(self, recipient_id, message):
        with self.message_lock:
            if recipient_id in self.message_queues:
                self.message_queues[recipient_id].append(message)

    def get_direct_message(self, node_id):
        with self.message_lock:
            if self.message_queues[node_id]:
                return self.message_queues[node_id].pop(0) # FIFO
        return None

    def run(self):
        print(f"Starting cluster with {self.num_nodes} nodes...")
        for node in self.nodes:
            node.start()
        # Keep main thread alive or join threads if simulation has a defined end
        # For continuous simulation:
        try:
            while True:
                time.sleep(1)
                leader_count = sum(1 for n in self.nodes if n.state == "leader")
                # print(f"Cluster status: Leaders found: {leader_count}")
                # if leader_count > 1:
                #    print("!!! Multiple leaders detected! Error in consensus.")
                # elif leader_count == 0 and any(n.state == "candidate"
                #                                for n in self.nodes):
                #    print("--- Election in progress or no leader yet ---")
                # elif leader_count == 1:
                #    current_leader = [n.node_id for n in self.nodes if n.state == "leader"][0]
                #    current_term = [n.term for n in self.nodes if n.state == "leader"][0]
                #    print(f"--- Leader is {current_leader} in term {current_term} ---")

        except KeyboardInterrupt:
            print("Cluster simulation stopped by user.")


if __name__ == "__main__":
    # The message passing in the original was problematic.
    # I've slightly refactored it to use per-node queues for clarity
    # and more reliable message delivery for the simulation.
    # The original get_message would always return the first message
    # in a global list if it matched node_id,
    # which is not how individual message processing works.
    # The original send_message put responses into global list keyed by candidate_id
    # which is also not a general message delivery.

    # The Python code provided in the article is a very simplified simulation
    # and has several issues for robust Raft behavior (e.g., message delivery,
    # handling of terms in messages correctly by all states, proper log management).
    # The version here includes minor corrections to the simulation's
    # message passing to make it slightly more functional for demonstration.
    # However, it remains a highly conceptual illustration.

    cluster = Cluster(num_nodes=5)
    cluster.run()