1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
// buildOutbound constructs either snapshot or append args for a given follower.
// withEntries controls whether to include log entries (append) or send an empty heartbeat.
// Returns: needSnapshot, snapshotArgs, appendArgs, valid (still leader in term when built).
func (rf *Raft) buildOutbound(server int, term int, withEntries bool) (bool, InstallSnapshotArgs, AppendEntriesArgs, bool) {
	rf.mu.Lock()
	defer rf.mu.Unlock()

	if rf.state != Leader || rf.currentTerm != term {
		return false, InstallSnapshotArgs{}, AppendEntriesArgs{}, false
	}

	first := rf.getFirstLog()
	prevLogIndex := rf.nextIndex[server] - 1
	// Need snapshot
	if prevLogIndex < first.Index {
		snap := InstallSnapshotArgs{
			Term:             term,
			LeaderId:         rf.me,
			LastIncludeIndex: first.Index,
			LastIncludeTerm:  first.Term,
			Data:             rf.persister.ReadSnapshot(),
		}
		return true, snap, AppendEntriesArgs{}, true
	}

	prevLogTerm := rf.log[prevLogIndex-first.Index].Term
	leaderCommit := rf.commitIndex

	args := AppendEntriesArgs{
		Term:         term,
		LeaderID:     rf.me,
		PrevLogIndex: prevLogIndex,
		PrevLogTerm:  prevLogTerm,
		LeaderCommit: leaderCommit,
	}
	if withEntries && rf.nextIndex[server] <= rf.getLastIndex() {
		args.Entries = rf.log[rf.nextIndex[server]-first.Index:]
	} else {
		args.Entries = []LogEntry{}
	}
	return false, InstallSnapshotArgs{}, args, true
}

// handleAppendReply updates follower progress and handles term changes.
// Returns (done) whether replication to this follower can stop.
func (rf *Raft) handleAppendReply(server int, term int, args AppendEntriesArgs, reply AppendEntriesReply) (done bool) {
	rf.mu.Lock()
	defer rf.mu.Unlock()

	if rf.state != Leader || rf.currentTerm != term {
		return true
	}

	if reply.Term > rf.currentTerm {
		rf.currentTerm = reply.Term
		rf.state = Follower
		rf.voteFor = -1
		rf.persist(false, nil)
		go rf.sendHeartbeat(false)
		return true
	}

	if reply.Success {
		if args.PrevLogIndex != rf.nextIndex[server]-1 {
			return true
		}
		if len(args.Entries) > 0 {
			rf.nextIndex[server] = args.PrevLogIndex + len(args.Entries) + 1
			rf.matchIndex[server] = rf.nextIndex[server] - 1
		}
		return true
	}

	// conflict resolution
	if reply.ConflictIndex <= rf.getFirstLog().Index {
		rf.nextIndex[server] = rf.getFirstLog().Index
	} else {
		rf.nextIndex[server] = reply.ConflictIndex
	}
	return false
}

// replicateToPeer keeps trying to bring one follower up to date.
// withEntries=true sends entries; false sends heartbeat-only (but may still install snapshot if needed).
func (rf *Raft) replicateToPeer(server int, withEntries bool, term int) {
	localWithEntries := withEntries
	for !rf.killed() {
		needSnap, snapArgs, appArgs, ok := rf.buildOutbound(server, term, localWithEntries)
		if !ok {
			return
		}

		if needSnap {
			reply := InstallSnapshotReply{}
			if rf.sendInstallSnapshot(server, &snapArgs, &reply) {
				rf.mu.Lock()
				if rf.state != Leader || rf.currentTerm != term {
					rf.mu.Unlock()
					return
				}
				if reply.Term > rf.currentTerm {
					rf.currentTerm = reply.Term
					rf.state = Follower
					rf.voteFor = -1
					rf.persist(false, nil)
					rf.mu.Unlock()
					rf.sendHeartbeat(false)
					return
				}
				// successful install snapshot advances indices
				rf.nextIndex[server] = snapArgs.LastIncludeIndex + 1
				rf.matchIndex[server] = snapArgs.LastIncludeIndex
				rf.mu.Unlock()
				// after snapshot, try append entries
				localWithEntries = true
				continue
			}
			time.Sleep(50 * time.Millisecond)
			continue
		}

		// AppendEntries path
		reply := AppendEntriesReply{}
		if rf.sendAppendEntries(server, &appArgs, &reply) {
			done := rf.handleAppendReply(server, term, appArgs, reply)
			if done {
				return
			}
			// If this was a heartbeat and we hit a conflict, escalate to append
			localWithEntries = true
			continue
		}
		// RPC failed; retry
		time.Sleep(50 * time.Millisecond)
		continue
	}
}

// broadcastHeartbeat sends heartbeats (no log entries) to all followers.
func (rf *Raft) broadcastHeartbeat() {
	rf.mu.Lock()
	term := rf.currentTerm
	rf.mu.Unlock()
	for i := range rf.peers {
		if i == rf.me {
			continue
		}
		go rf.replicateToPeer(i, false, term)
	}
}

// broadcastAppend sends append entries (with log entries when available) to all followers.
func (rf *Raft) broadcastAppend() {
	rf.mu.Lock()
	term := rf.currentTerm
	rf.mu.Unlock()
	for i := range rf.peers {
		if i == rf.me {
			continue
		}
		go rf.replicateToPeer(i, true, term)
	}
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
func (rf *Raft) InstallSnapshot(args *InstallSnapshotArgs, reply *InstallSnapshotReply) {
	rf.mu.Lock()

	// term handling
	if args.Term < rf.currentTerm {
		reply.Term = rf.currentTerm
		rf.mu.Unlock()
		return
	}
	if args.Term > rf.currentTerm {
		rf.currentTerm = args.Term
		rf.voteFor = -1
		rf.state = Follower
		rf.persist(false, nil)
	}
	reply.Term = rf.currentTerm
	// receiving valid RPC resets election timer
	rf.electionTimer.Reset(rf.getNewElectionTimeout())

	// If snapshot is outdated (we already have equal/newer snapshot), ignore
	if args.LastIncludeIndex <= rf.getFirstLog().Index {
		rf.mu.Unlock()
		return
	}

	// Rebuild log with snapshot sentinel and an optional suffix to retain
	firstIdx := rf.getFirstLog().Index
	lastIdx := rf.getLastIndex()

	var newLog []LogEntry
	if args.LastIncludeIndex <= lastIdx {
		// There may be a suffix to keep
		cut := args.LastIncludeIndex - firstIdx
		// Guard: cut could be out of bound if state is inconsistent
		if cut >= 0 && cut < len(rf.log) && rf.log[cut].Term == args.LastIncludeTerm {
			// keep suffix strictly after snapshot index
			suffix := rf.log[cut+1:]
			newLog = make([]LogEntry, 1+len(suffix))
			newLog[0] = LogEntry{Index: args.LastIncludeIndex, Term: args.LastIncludeTerm}
			copy(newLog[1:], suffix)
		} else {
			// terms conflict; drop everything and start fresh from sentinel
			newLog = []LogEntry
		}
	} else {
		// Snapshot goes beyond our last log; drop everything
		newLog = []LogEntry
	}
	rf.log = newLog

	// bump commitIndex/lastApplied to at least snapshot index
	if rf.commitIndex < args.LastIncludeIndex {
		rf.commitIndex = args.LastIncludeIndex
	}
	if rf.lastApplied < args.LastIncludeIndex {
		rf.lastApplied = args.LastIncludeIndex
	}

	// Persist new state and snapshot atomically
	rf.persist(true, &args.Data)

	// deliver snapshot to service outside the lock
	msg := raftapi.ApplyMsg{
		SnapshotValid: true,
		Snapshot:      args.Data,
		SnapshotTerm:  args.LastIncludeTerm,
		SnapshotIndex: args.LastIncludeIndex,
	}
	rf.mu.Unlock()
	rf.applyCh <- msg // send message out the critical section!
}