Lab 01 — Raft Leader Election Simulation

Phase: 07 — Production Systems
Difficulty: ⭐⭐⭐⭐☆ | Estimated Time: 2–2.5 hours
Artifacts: RaftNode.java, RaftCluster.java, RaftElectionTest.java
Command: mvn test


Background

Raft is the consensus algorithm used by:

  • Apache ZooKeeper (leader election for coordination)
  • Kafka's KRaft mode (replaces ZooKeeper as of Kafka 3.3)
  • etcd (Kubernetes control plane)
  • CockroachDB, TiKV, and most modern distributed databases

This lab implements the leader election portion of Raft — the part that decides which node is in charge. Log replication (how writes are propagated) is out of scope.


Setup

cd phase-07-production-systems/lab-01-raft-election
mvn test

All 5 tests should pass.


Step 1 — Understand the State Machine

Read RaftNode.java and trace through what happens when a FOLLOWER calls startElection():

  1. What fields change?
  2. What is the new value of currentTerm?
  3. How many votes does the node have before asking any peers?

Step 2 — Trace a 3-Node Election

Read RaftCluster.runElection() and trace through a 3-node cluster where node-1 starts the election. Draw the sequence of method calls and state transitions:

node-1: startElection()       → state=CANDIDATE, term=1, votes=1
node-1 → node-2: requestVote(term=1, candidateId="node-1") → ?
node-1 → node-3: requestVote(term=1, candidateId="node-1") → ?
...

What does node-2 do if it is already at term=2 from a previous aborted election?


Step 3 — Run the Tests

mvn test

Expected output:

[INFO] Tests run: 5, Failures: 0, Errors: 0, Skipped: 0
[INFO] BUILD SUCCESS

Examine each test and make sure you can explain why it passes, not just that it does.


Step 4 — Extend: Make It Fail, Then Fix It

Add a 6th test that creates a 3-node cluster and verifies that a node with a stale term (lower than the cluster's current term) cannot win an election. Your test must:

  1. Advance the cluster term to 3 (simulate 3 elections)
  2. Create a new "stale" candidate at term 1
  3. Assert it does NOT become leader
Hint

Use requestVote(3, "leader-id") on two of the three nodes before your test to advance their terms. A node that calls requestVote(1, "stale-candidate") after that will receive a VoteResponse with voteGranted=false.


Step 5 — Interview Prep

Write a 3-sentence answer to: "How does Raft prevent two leaders from being elected in the same term?"


Source Files

pom.xml

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
    xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
    xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>

    <groupId>org.example.raft</groupId>
    <artifactId>raft-election-lab</artifactId>
    <version>1.0-SNAPSHOT</version>

    <properties>
        <maven.compiler.source>11</maven.compiler.source>
        <maven.compiler.target>11</maven.compiler.target>
        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
    </properties>

    <dependencies>
        <dependency>
            <groupId>org.junit.jupiter</groupId>
            <artifactId>junit-jupiter</artifactId>
            <version>5.10.2</version>
            <scope>test</scope>
        </dependency>
    </dependencies>

    <build>
        <plugins>
            <plugin>
                <groupId>org.apache.maven.plugins</groupId>
                <artifactId>maven-surefire-plugin</artifactId>
                <version>3.2.5</version>
            </plugin>
        </plugins>
    </build>
</project>

NodeState.java

package org.example.raft;

/**
 * The three states a Raft node can be in at any point in time.
 *
 * <p>
 * State transitions:
 * 
 * <pre>
 *   FOLLOWER  ──(election timeout)──▶  CANDIDATE
 *   CANDIDATE ──(majority votes)─────▶  LEADER
 *   CANDIDATE ──(sees higher term)───▶  FOLLOWER
 *   LEADER    ──(sees higher term)───▶  FOLLOWER
 * </pre>
 *
 * <p>
 * All nodes start as FOLLOWER. There is no direct path from FOLLOWER to LEADER
 * without going through CANDIDATE — a node must win a vote, not just claim
 * leadership.
 */
public enum NodeState {
    FOLLOWER,
    CANDIDATE,
    LEADER
}

VoteResponse.java

package org.example.raft;

/**
 * Response to a {@code RequestVote} RPC (Raft Section 5.2).
 *
 * <p>
 * The receiver always includes its {@code currentTerm} in the response.
 * If the sender's term is stale, it sees a higher term in the response
 * and immediately steps down to FOLLOWER.
 */
public final class VoteResponse {

    private final int term; // receiver's currentTerm
    private final boolean voteGranted;

    public VoteResponse(int term, boolean voteGranted) {
        this.term = term;
        this.voteGranted = voteGranted;
    }

    /** The receiver's current term at the time of the response. */
    public int getTerm() {
        return term;
    }

    /** Whether the receiver granted its vote to the candidate. */
    public boolean isVoteGranted() {
        return voteGranted;
    }

    @Override
    public String toString() {
        return "VoteResponse{term=" + term + ", voteGranted=" + voteGranted + "}";
    }
}

RaftNode.java

package org.example.raft;

/**
 * A single node in a Raft cluster.
 *
 * <p>
 * This is a synchronous, single-JVM simulation. In a real distributed system,
 * these methods would be invoked via RPC (e.g., gRPC). Here, they are called
 * directly by {@link RaftCluster} to simulate the message passing.
 *
 * <p>
 * All public state-mutating methods are synchronized to keep the simulation
 * thread-safe (useful if you extend this lab to run elections on background
 * threads).
 *
 * <h3>Raft invariants maintained by this class:</h3>
 * <ul>
 * <li>A node grants at most one vote per term (tracked by
 * {@code votedFor})</li>
 * <li>A node that sees a higher term immediately steps down to FOLLOWER</li>
 * <li>A node only becomes LEADER if it receives votes from a strict
 * majority</li>
 * </ul>
 */
public class RaftNode {

    private final String id;
    private NodeState state;
    private int currentTerm;
    private String votedFor; // node id we voted for in currentTerm; null if we haven't voted
    private String leaderId; // id of the current leader as learned from heartbeats
    private int votesReceived; // only meaningful when state == CANDIDATE

    public RaftNode(String id) {
        this.id = id;
        this.state = NodeState.FOLLOWER;
        this.currentTerm = 0;
        this.votedFor = null;
        this.leaderId = null;
        this.votesReceived = 0;
    }

    // -------------------------------------------------------------------------
    // RPC handlers (called by RaftCluster to simulate message delivery)
    // -------------------------------------------------------------------------

    /**
     * Handles a {@code RequestVote} RPC from a candidate.
     *
     * <p>
     * Raft Section 5.2 grant conditions:
     * <ol>
     * <li>Reject if {@code candidateTerm < currentTerm} (candidate is behind)</li>
     * <li>If {@code candidateTerm > currentTerm}: update term, step down to
     * FOLLOWER,
     * clear {@code votedFor} — we may now vote for this candidate</li>
     * <li>Grant vote if we haven't voted this term OR already voted for this
     * candidate</li>
     * </ol>
     *
     * @param candidateTerm the term the candidate is running in
     * @param candidateId   the candidate requesting the vote
     * @return response including this node's currentTerm and whether the vote was
     *         granted
     */
    public synchronized VoteResponse requestVote(int candidateTerm, String candidateId) {
        if (candidateTerm < currentTerm) {
            // The candidate is behind — tell it our term so it steps down
            return new VoteResponse(currentTerm, false);
        }

        if (candidateTerm > currentTerm) {
            // Higher term: step down regardless of our current state
            currentTerm = candidateTerm;
            state = NodeState.FOLLOWER;
            votedFor = null;
        }

        // Grant the vote if we haven't committed to anyone else this term
        boolean grant = (votedFor == null || votedFor.equals(candidateId));
        if (grant) {
            votedFor = candidateId;
        }
        return new VoteResponse(currentTerm, grant);
    }

    /**
     * Handles an {@code AppendEntries} RPC from a leader (used here as a
     * heartbeat).
     *
     * <p>
     * In full Raft this RPC also carries log entries. Here we only model the
     * heartbeat behavior: if the leader's term is valid, accept it and reset the
     * election timer (simulated by transitioning/remaining FOLLOWER).
     *
     * @param leaderTerm the leader's current term
     * @param leaderId   the leader's node id
     * @return {@code true} if we accept this leader, {@code false} if term is stale
     */
    public synchronized boolean appendEntries(int leaderTerm, String leaderId) {
        if (leaderTerm < currentTerm) {
            return false; // stale leader — reject
        }
        if (leaderTerm > currentTerm) {
            currentTerm = leaderTerm;
            votedFor = null;
        }
        // Accept the leader: transition (back) to FOLLOWER and record who the leader is
        state = NodeState.FOLLOWER;
        this.leaderId = leaderId;
        return true;
    }

    // -------------------------------------------------------------------------
    // Election actions (called on the candidate node itself)
    // -------------------------------------------------------------------------

    /**
     * Starts an election: increments the term, transitions to CANDIDATE, and
     * casts a self-vote.
     *
     * <p>
     * In real Raft this is triggered when the election timeout fires.
     * Here it is called explicitly by {@link RaftCluster#runElection()}.
     */
    public synchronized void startElection() {
        currentTerm++;
        state = NodeState.CANDIDATE;
        votedFor = id; // every candidate votes for itself
        votesReceived = 1; // the self-vote
        leaderId = null;
    }

    /**
     * Records one received vote. If the candidate now holds a majority of votes
     * in a cluster of {@code clusterSize} nodes, it transitions to LEADER.
     *
     * @param clusterSize total number of nodes in the cluster (used to compute
     *                    majority)
     */
    public synchronized void countVote(int clusterSize) {
        if (state != NodeState.CANDIDATE) {
            return; // already elected or stepped down — ignore late votes
        }
        votesReceived++;
        int majority = (clusterSize / 2) + 1;
        if (votesReceived >= majority) {
            state = NodeState.LEADER;
        }
    }

    // -------------------------------------------------------------------------
    // Getters
    // -------------------------------------------------------------------------

    public String getId() {
        return id;
    }

    public synchronized NodeState getState() {
        return state;
    }

    public synchronized int getCurrentTerm() {
        return currentTerm;
    }

    public synchronized String getVotedFor() {
        return votedFor;
    }

    /** The leader id as last communicated via a heartbeat; null if unknown. */
    public synchronized String getLeaderId() {
        return leaderId;
    }

    @Override
    public String toString() {
        return "RaftNode{id='" + id + "', state=" + state
                + ", term=" + currentTerm + ", votedFor='" + votedFor + "'}";
    }
}

RaftCluster.java

package org.example.raft;

import java.util.Arrays;
import java.util.List;
import java.util.Optional;

/**
 * Coordinates leader election across a fixed set of {@link RaftNode}s.
 *
 * <p>
 * This is a synchronous, deterministic simulation of the Raft election
 * protocol.
 * In a real cluster, the timing is driven by election timeouts (random timers
 * per node).
 * Here, elections are triggered explicitly so tests are deterministic.
 *
 * <p>
 * The cluster size is fixed at construction time. Node failures are simulated
 * by excluding specific nodes from an election via
 * {@link #runElectionExcluding(RaftNode)}.
 */
public class RaftCluster {

    private final List<RaftNode> nodes;

    public RaftCluster(RaftNode... nodes) {
        if (nodes == null || nodes.length < 3) {
            throw new IllegalArgumentException(
                    "Raft requires at least 3 nodes to tolerate 1 failure");
        }
        this.nodes = Arrays.asList(nodes);
    }

    /**
     * Simulates one complete election cycle:
     * <ol>
     * <li>The first FOLLOWER (or CANDIDATE) starts an election</li>
     * <li>It broadcasts {@code RequestVote} to all peers</li>
     * <li>If a majority votes for it, it transitions to LEADER</li>
     * <li>The new leader sends heartbeats so all followers learn the leader id</li>
     * </ol>
     *
     * @return the elected leader, or empty if no candidate won a majority (split
     *         vote)
     */
    public Optional<RaftNode> runElection() {
        RaftNode candidate = findCandidate(null);
        if (candidate == null) {
            return getLeader(); // already have a leader
        }

        candidate.startElection();
        broadcastRequestVote(candidate, null);

        if (candidate.getState() == NodeState.LEADER) {
            sendHeartbeats(candidate, null);
        }

        return getLeader();
    }

    /**
     * Simulates a network partition: runs an election excluding the given node.
     *
     * <p>
     * The partitioned node receives no RPCs and sends none. The remaining
     * nodes form a sub-cluster and, if they represent a majority, elect a new
     * leader.
     *
     * <p>
     * With a 3-node cluster and 1 partitioned node, the 2 remaining nodes
     * still constitute a majority (⌊3/2⌋ + 1 = 2).
     *
     * @param partitioned the node to exclude from the election
     * @return the new leader among the non-partitioned nodes, or empty if none
     */
    public Optional<RaftNode> runElectionExcluding(RaftNode partitioned) {
        RaftNode candidate = findCandidate(partitioned);
        if (candidate == null) {
            return Optional.empty();
        }

        candidate.startElection();
        broadcastRequestVote(candidate, partitioned);

        if (candidate.getState() == NodeState.LEADER) {
            sendHeartbeats(candidate, partitioned);
        }

        return nodes.stream()
                .filter(n -> n != partitioned && n.getState() == NodeState.LEADER)
                .findFirst();
    }

    /**
     * Sends heartbeats from {@code leader} to all other (non-excluded) nodes.
     * In real Raft, this is sent periodically; here it is called once after
     * election.
     */
    public void sendHeartbeats(RaftNode leader, RaftNode exclude) {
        for (RaftNode node : nodes) {
            if (node == leader || node == exclude)
                continue;
            node.appendEntries(leader.getCurrentTerm(), leader.getId());
        }
    }

    /** Returns the current leader, if any. */
    public Optional<RaftNode> getLeader() {
        return nodes.stream()
                .filter(n -> n.getState() == NodeState.LEADER)
                .findFirst();
    }

    /**
     * Returns the number of nodes currently in the LEADER state. Must be 0 or 1.
     */
    public long countLeaders() {
        return nodes.stream()
                .filter(n -> n.getState() == NodeState.LEADER)
                .count();
    }

    public List<RaftNode> getNodes() {
        return nodes;
    }

    // -------------------------------------------------------------------------
    // Private helpers
    // -------------------------------------------------------------------------

    /**
     * Finds the first node eligible to start an election (FOLLOWER or CANDIDATE,
     * and not the excluded node).
     */
    private RaftNode findCandidate(RaftNode exclude) {
        return nodes.stream()
                .filter(n -> n != exclude)
                .filter(n -> n.getState() == NodeState.FOLLOWER
                        || n.getState() == NodeState.CANDIDATE)
                .findFirst()
                .orElse(null);
    }

    /**
     * Broadcasts {@code RequestVote} from {@code candidate} to all peers except
     * itself
     * and the excluded node. Calls {@link RaftNode#countVote(int)} on the candidate
     * for each granted vote, using the full cluster size for the majority
     * calculation.
     */
    private void broadcastRequestVote(RaftNode candidate, RaftNode exclude) {
        int term = candidate.getCurrentTerm();
        for (RaftNode peer : nodes) {
            if (peer == candidate || peer == exclude)
                continue;
            VoteResponse response = peer.requestVote(term, candidate.getId());
            if (response.isVoteGranted()) {
                // Use full cluster size — majority is over all nodes, not just available ones
                candidate.countVote(nodes.size());
            }
        }
    }
}

RaftElectionTest.java

package org.example.raft;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.util.Optional;

import static org.junit.jupiter.api.Assertions.*;

class RaftElectionTest {

    private RaftNode node1;
    private RaftNode node2;
    private RaftNode node3;
    private RaftCluster cluster;

    @BeforeEach
    void setUp() {
        node1 = new RaftNode("node-1");
        node2 = new RaftNode("node-2");
        node3 = new RaftNode("node-3");
        cluster = new RaftCluster(node1, node2, node3);
    }

    /**
     * A 3-node cluster must elect exactly one leader.
     * Verifies the fundamental Raft safety property: at most one leader per term.
     */
    @Test
    void testElectionProducesExactlyOneLeader() {
        Optional<RaftNode> leader = cluster.runElection();

        assertTrue(leader.isPresent(), "Election must produce a leader");
        assertEquals(NodeState.LEADER, leader.get().getState());
        assertEquals(1, cluster.countLeaders(), "Exactly one node must be LEADER");
    }

    /**
     * After election + heartbeats, every follower's term must be <= the leader's
     * term.
     * Verifies the Raft term invariant: the leader always has the highest term.
     */
    @Test
    void testLeaderHasHighestTerm() {
        Optional<RaftNode> leader = cluster.runElection();

        assertTrue(leader.isPresent());
        int leaderTerm = leader.get().getCurrentTerm();

        cluster.getNodes().forEach(n -> assertTrue(n.getCurrentTerm() <= leaderTerm,
                n.getId() + " has term=" + n.getCurrentTerm()
                        + " which exceeds leader term=" + leaderTerm));
    }

    /**
     * After the leader sends heartbeats, all followers must know who the leader is.
     * Verifies that AppendEntries (heartbeat) propagates leader identity.
     */
    @Test
    void testFollowersLearnLeaderViaHeartbeat() {
        Optional<RaftNode> leader = cluster.runElection();

        assertTrue(leader.isPresent());
        String leaderId = leader.get().getId();

        cluster.getNodes().stream()
                .filter(n -> n.getState() == NodeState.FOLLOWER)
                .forEach(n -> assertEquals(leaderId, n.getLeaderId(),
                        n.getId() + " should know the leader after receiving a heartbeat"));
    }

    /**
     * A node at term=2 must reject a RequestVote for term=1.
     * Verifies Raft's stale-term rejection rule.
     */
    @Test
    void testStaleCandidateIsRejected() {
        // Advance node2 and node3 to term=2 by having them participate in an election
        node2.startElection(); // node2 → CANDIDATE, term=1
        node2.startElection(); // node2 → CANDIDATE, term=2 (simulates a second aborted attempt)

        // node1 is still at term=0. If node1 tries to run for election at term=1,
        // node2 (at term=2) must reject it.
        node1.startElection(); // node1 → CANDIDATE, term=1

        VoteResponse response = node2.requestVote(1, "node-1");

        assertFalse(response.isVoteGranted(),
                "node2 at term=2 must reject a RequestVote for term=1");
        assertEquals(2, response.getTerm(),
                "Response must include the node's current term so the sender can update");
    }

    /**
     * When the current leader is network-partitioned (excluded from RPCs), the
     * remaining nodes must elect a new leader — a different node.
     * Verifies Raft's liveness property: the cluster recovers from a leader
     * failure.
     */
    @Test
    void testPartitionedLeaderCausesNewElection() {
        // Step 1: elect the initial leader
        Optional<RaftNode> firstLeader = cluster.runElection();
        assertTrue(firstLeader.isPresent(), "Initial election must succeed");

        String firstLeaderId = firstLeader.get().getId();

        // Step 2: simulate leader partition — run a new election without the old leader
        Optional<RaftNode> newLeader = cluster.runElectionExcluding(firstLeader.get());

        assertTrue(newLeader.isPresent(),
                "Cluster must elect a new leader after partitioning the old one");
        assertNotEquals(firstLeaderId, newLeader.get().getId(),
                "The new leader must be a different node than the partitioned leader");
    }
}