MIT 6.5840 Lab 3A: Raft - leader election (2025)

const (
	// Because the tester limits you tens of heartbeats per second, you will
	// have to use an election timeout larger than the paper's 150 to 300
	// milliseconds, but not too large, because then you may fail to elect
	// a leader within five seconds.
	ELECTIONTIMEOUTLEAST = 700 * time.Millisecond
	ELECTIONTIMEOUTMOST  = 1200 * time.Millisecond
)

func getElectionTimeout() time.Duration {
	return ELECTIONTIMEOUTLEAST + (time.Duration(rand.Int63()) % (ELECTIONTIMEOUTMOST - ELECTIONTIMEOUTLEAST))
}

type Raft struct {
	...
	startTime time.Time
	...
}

func Make(peers []*labrpc.ClientEnd, me int,
	persister *tester.Persister, applyCh chan raftapi.ApplyMsg) raftapi.Raft {
	rf := &Raft{}
	...
	rf.startTime = time.Now()
	...
	return rf
}

electionTimeout := getElectionTimeout()
rf.mu.Lock()
shouldBeginElection := time.Since(rf.startTime) > electionTimeout
rf.mu.Unlock()

rf.mu.Lock()
rf.startTime = time.Now()
rf.mu.Unlock()

func Make(peers []*labrpc.ClientEnd, me int,
	persister *tester.Persister, applyCh chan raftapi.ApplyMsg) raftapi.Raft {
	rf := &Raft{}
	...
	rf.state = FOLLOWER
	...

	// start ticker goroutine to start elections
	go rf.ticker()

	return rf
}

func (rf *Raft) ticker() {
	for rf.killed() == false {

		// Your code here (3A)
		// Check if a leader election should be started.
		electionTimeout := getElectionTimeout()
		rf.mu.Lock()
		shouldBeginElection := time.Since(rf.startTime) > electionTimeout
		rf.mu.Unlock()

		if shouldBeginElection {
			isLeader := rf.beginElection()
			if isLeader {
				rf.sendHeartbeats()
			}
		}

		// pause for a random amount of time between 50 and 350
		// milliseconds.
		ms := 50 + (rand.Int63() % 300)
		time.Sleep(time.Duration(ms) * time.Millisecond)
	}
}

const (
	FOLLOWER  = "FOLLOW" // 便于调试，你也可以直接用 iota
	CANDIDATE = "CANDID"
	LEADER    = "LEADER"
)

type LogEntry struct {
	Command interface{}
	Term    int
}

type Raft struct {
	...
	// Persistent state on all servers
	currentTerm int         // latest term server has seen
	votedFor    int         // candidateId that received vote in current term (or -1 if none)
	log         []*LogEntry // log entries

	// Volatile state on all servers
	state           string             // the server's state
	startTime       time.Time          // the time an election timeout starts
	cancelElection  context.CancelFunc // can be used to cancel the election immediately
	cancelHeartbeat context.CancelFunc // can be used to cancel the heartbeat immediately
	...
}

func Make(peers []*labrpc.ClientEnd, me int,
	persister *tester.Persister, applyCh chan raftapi.ApplyMsg) raftapi.Raft {
	rf := &Raft{}
	rf.peers = peers
	rf.persister = persister
	rf.me = me

	// Your initialization code here (3A, 3B, 3C).
	rf.votedFor = -1
	rf.state = FOLLOWER
	rf.startTime = time.Now()

	// initialize from state persisted before a crash
	rf.readPersist(persister.ReadRaftState())

	// start ticker goroutine to start elections
	go rf.ticker()

	return rf
}

// return currentTerm and whether this server
// believes it is the leader.
func (rf *Raft) GetState() (int, bool) {

	var term int
	var isleader bool
	// Your code here (3A).
	rf.mu.Lock()
	defer rf.mu.Unlock()

	term = rf.currentTerm
	isleader = (rf.state == LEADER)

	return term, isleader
}

// example RequestVote RPC arguments structure.
// field names must start with capital letters!
type RequestVoteArgs struct {
	// Your data here (3A, 3B).
	Term         int // candidate's term
	CandidateId  int // the candidate who is requesting vote
	LastLogIndex int // index of candidate's last log entry
	LastLogTerm  int // term of candidate's last log entry
}

// example RequestVote RPC reply structure.
// field names must start with capital letters!
type RequestVoteReply struct {
	// Your data here (3A).
	Term        int  // the receiver's currentTerm, for candidate to update itself
	VoteGranted bool // true means candidate received vote
}

// example RequestVote RPC handler.
func (rf *Raft) RequestVote(args *RequestVoteArgs, reply *RequestVoteReply) {
	// Your code here (3A, 3B).
	rf.mu.Lock()
	defer rf.mu.Unlock()

	// "If a server receives a request with a stale term number, it rejects the
	// request."
	if rf.currentTerm > args.Term {
		reply.Term = rf.currentTerm
		reply.VoteGranted = false
		return
	}

	// "If one server's current term is smaller than the other's, then it updates
	// its current term to the larger value. If a candidate or leader discovers
	// that its term is out of date, it immediately reverts to follower state."
	if rf.currentTerm < args.Term {
		rf.currentTerm = args.Term
		rf.votedFor = -1
		rf.state = FOLLOWER
		// reset election timer if you discovers that your term is out of date
		rf.startTime = time.Now()

		if rf.cancelElection != nil {
			rf.cancelElection()
		}

		if rf.cancelHeartbeat != nil {
			rf.cancelHeartbeat()
		}
	}

	// "Each server will vote for at most one candidate in a given term, on a
	// first-come-first-served basis."
	if rf.votedFor != -1 && rf.votedFor != args.CandidateId {
		reply.Term = rf.currentTerm
		reply.VoteGranted = false
		return
	}

	// "The RPC includes information about the candidate's log, and the voter
	// denies its vote if its own log is more up-to-date than that of the candidate."
	lastLogIndex := len(rf.log) - 1
	var lastLogTerm int
	if lastLogIndex == -1 {
		lastLogTerm = -1
	} else {
		lastLogTerm = rf.log[lastLogIndex].Term
	}
	if lastLogTerm > args.LastLogTerm ||
		(lastLogTerm == args.LastLogTerm && lastLogIndex > args.LastLogIndex) {
		reply.Term = rf.currentTerm
		reply.VoteGranted = false
		return
	}

	reply.Term = rf.currentTerm
	reply.VoteGranted = true
	rf.votedFor = args.CandidateId
	// reset election timer if you grant a vote to another peer
	rf.startTime = time.Now()

	// "Updated on stable storage before responding to RPCs."
	rf.persist()
}

type AppendEntriesArgs struct {
	Term         int // leader's term
	LeaderId     int // so follower can redirect clients
	PrevLogIndex int // index of log entry immediately preceding new ones
	PrevLogTerm  int // term of PrevLogIndex entry
}

type AppendEntriesReply struct {
	Term int // currentTerm, for leader to update itself
}

func (rf *Raft) AppendEntries(args *AppendEntriesArgs, reply *AppendEntriesReply) {
	rf.mu.Lock()
	defer rf.mu.Unlock()

	// "If a server receives a request with a stale term number, it rejects the
	// request."
	if rf.currentTerm > args.Term {
		reply.Term = rf.currentTerm
		reply.Success = false
		return
	}

	// "If one server's current term is smaller than the other's, then it updates
	// its current term to the larger value. If a candidate or leader discovers
	// that its term is out of date, it immediately reverts to follower state."
	if rf.currentTerm < args.Term {
		rf.currentTerm = args.Term
		rf.votedFor = -1
		rf.state = FOLLOWER
		// reset election timer if you discovers that your term is out of date
		rf.startTime = time.Now()
	}

	reply.Term = rf.currentTerm

    // reset election timer if you get an AppendEntries RPC from the current leader
	rf.startTime = time.Now()

	if rf.cancelElection != nil {
		rf.cancelElection()
	}

	if rf.cancelHeartbeat != nil {
		rf.cancelHeartbeat()
	}

	// "Updated on stable storage before responding to RPCs."
	rf.persist()
}

func (rf *Raft) sendAppendEntries(server int, args *AppendEntriesArgs, reply *AppendEntriesReply) bool {
	ok := rf.peers[server].Call("Raft.AppendEntries", args, reply)
	return ok
}

ctx, cancel := context.WithCancel(context.Background())
countCh := make(chan bool)
for i := range rf.peers {
	if i != rf.me {
		go func(peer int) {
			...
			doneCh := make(chan bool, 1)
			go func() {
				ok := rf.sendRPC(peer, &args, &reply)
				if ok {
					if reply.xxx {
						doneCh <- true
					} else {
						doneCh <- false
					}
				} else {
					doneCh <- false
				}
			}()

			select {
			case <-ctx.Done():
				return
			case result := <-doneCh:
				countCh <- result
				return
			}
		}(i)
	}
}

timeout := time.NewTimer(xxx)
for {
	select {
	case <-timeout.C:
		return
	case <-ctx.Done():
		return
	case result := <-countCh:
		...
	}
}

ok := rf.sendRPC(peer, &args, &reply)
if ok {
	select {
	case <-ctx.Done():
	case countCh <- true:
	}
} else {
	select {
	case <-ctx.Done():
	case countCh <- false:
	}
}

type Raft struct {
	...
	cancelElection context.CancelFunc // can be used to cancel the election immediately
}

func (rf *Raft) beginElection() bool {
	rf.mu.Lock()
	...
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	rf.cancelElection = cancel
	...
	rf.mu.Unlock()
	...
}

if rf.cancelElection != nil {
	rf.cancelElection()
}

func (rf *Raft) beginElection() bool {
	rf.mu.Lock()
	rf.state = CANDIDATE
	rf.currentTerm++
	rf.votedFor = rf.me
	rf.startTime = time.Now() // reset election timer
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	rf.cancelElection = cancel
	lastLogIndex := len(rf.log) - 1
	var lastLogTerm int
	if lastLogIndex == -1 {
		lastLogTerm = -1
	} else {
		lastLogTerm = rf.log[lastLogIndex].Term
	}
	currentTerm := rf.currentTerm
	rf.mu.Unlock()

	ticket := 1
	count := 0
	voteCh := make(chan bool)

	for i := range rf.peers {
		if i != rf.me {
			go func(peer int) {
				args := RequestVoteArgs{
					Term:         currentTerm,
					CandidateId:  rf.me,
					LastLogIndex: lastLogIndex,
					LastLogTerm:  lastLogTerm,
				}
				var reply RequestVoteReply

				doneCh := make(chan bool, 1) // use buffered channel to avoid goroutine leak
				go func() {
					ok := rf.sendRequestVote(peer, &args, &reply)
					if ok {
						if reply.Term > args.Term {
							rf.mu.Lock()
							rf.currentTerm = reply.Term
							rf.mu.Unlock()
							cancel()
							return
						}
						if reply.VoteGranted {
							doneCh <- true
						} else {
							doneCh <- false
						}
					} else {
						doneCh <- false
					}
				}()

				select {
				case <-ctx.Done():
					return
				case result := <-doneCh:
					voteCh <- result
					return
				}
			}(i)
		}
	}

	electionTimeout := getElectionTimeout()

	for {
		select {
		case <-ctx.Done():
			rf.mu.Lock()
			rf.state = FOLLOWER
			rf.mu.Unlock()
			return false // find a larger term or receive heartbeats from leader
		case vote := <-voteCh:
			count++
			if vote {
				ticket++
			}
			if ticket > len(rf.peers)/2 {
				rf.mu.Lock()
				rf.state = LEADER
				rf.mu.Unlock()
				return true // win the election
			}
			if count == len(rf.peers)-1 {
				return false // lose the election
			}
		default:
			rf.mu.Lock()
			shouldBeginAnotherElection := time.Since(rf.startTime) > electionTimeout
			rf.mu.Unlock()
			if shouldBeginAnotherElection {
				return false // use `return rf.beginElection()` instead?
			}
		}
	}
}

const (
	// Heartbeat RPCs should finish within this time. Cannot be too small,
	// beacuse the leader will convert to follower state easily. Cannot be
	// too large, because the leader may not convert to follower just in time
	// when there's a network partition and the leader become disconnected
	// from the other servers.
	HEARTBEATTIMEOUT = 500 * time.Millisecond
	// The tester requires that the leader send heartbeat RPCs no more than
	// ten times per second.
	HEARTBEATINTERVAL = 100 * time.Millisecond
)

func (rf *Raft) sendHeartbeats() {
	defer func() {
		rf.mu.Lock()
		rf.state = FOLLOWER
		rf.mu.Unlock()
	}()

	rf.mu.Lock()
	currentTerm := rf.currentTerm
	rf.mu.Unlock()

	for {
		countCh := make(chan bool)
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		for i := range rf.peers {
			if i != rf.me {
				go func(peer int) {
					args := AppendEntriesArgs{
						Term:     currentTerm,
						LeaderId: rf.me,
					}
					var reply AppendEntriesReply

					doneCh := make(chan bool, 1)
					go func() {
						// this may spend a lot of time to finish
						ok := rf.sendAppendEntries(peer, &args, &reply)
						if ok {
							if args.Term < reply.Term {
								rf.mu.Lock()
								rf.currentTerm = reply.Term
								rf.mu.Unlock()
								cancel()
							} else {
								doneCh <- true
							}
						} else {
							countCh <- false
						}
					}()

					select {
					case <-ctx.Done():
						return
					case result := <-doneCh:
						countCh <- result
					}
				}(i)
			}
		}

		success := 0
		count := 0
		// each round of heartbeat should finish within this amount of time
		timeout := time.NewTimer(HEARTBEATTIMEOUT)
		defer timeout.Stop()
	OuterLoop:
		for {
			select {
			case <-timeout.C:
				return
			case <-ctx.Done():
				return
			case result := <-countCh:
				count++
				if result {
					success++
				}
				if count == len(rf.peers)-1 {
					if success < len(rf.peers)/2 {
						return
					} else {
						break OuterLoop
					}
				}
			}
		}
		time.Sleep(HEARTBEATINTERVAL)
	}
}

const (
	ELECTIONTIMEOUTLEAST = 500 * time.Millisecond
	ELECTIONTIMEOUTMOST  = 1000 * time.Millisecond
)

type Raft struct {
	...
	electionTimeout *time.Timer
}

func getElectionTimeout() time.Duration {
	return ELECTIONTIMEOUTLEAST + (time.Duration(rand.Int63()) % ELECTIONTIMEOUTMOST)
}

func Make(peers []*labrpc.ClientEnd, me int,
	persister *tester.Persister, applyCh chan raftapi.ApplyMsg) raftapi.Raft {
	rf := &Raft{}
	...
	rf.electionTimeout = time.NewTimer(getElectionTimeout())
	...
	go rf.ticker()
	return rf
}

func (rf *Raft) ticker() {
	for rf.killed() == false {
		<-rf.electionTimeout.C
		isLeader := rf.beginElection()
		if isLeader {
			rf.sendHeartbeats()
		}
	}
}

func (rf *Raft) resetElectionTimeout() {
	rf.electionTimeout.Reset(getElectionTimeout())
}

func (rf *Raft) RequestVote(args *RequestVoteArgs, reply *RequestVoteReply) {
	...
	rf.resetElectionTimeout()
	...
}

func (rf *Raft) AppendEntries(args *AppendEntriesArgs, reply *AppendEntriesReply) {
	...
	rf.resetElectionTimeout()
	...
}

func (rf *Raft) beginElection() bool {
	...
	rf.resetElectionTimeout()
	...
	select {
	case <-rf.electionTimeout.C:
		rf.electionTimeout.Reset(0) // 这边很重要，否则会导致 ticker 中 <-rf.electionTimeout.C 阻塞
		return false
	...
	}
}