NRG: Don't drain response or proposal queues when switching to leader

The `runAsCandidate` and `runAsFollower` goroutines both try to drain
the append entry response and proposal queues if they are still populated
from the previous leadership, but we might be blocked in `runAsCandidate`
or `runAsFollower` while switching to leader waiting for the state change
to complete and accidentally drain incoming proposals.

Signed-off-by: Neil Twigg <[email protected]>

Author: neilalexander

server/ipqueue.go

@@ -230,6 +230,21 @@ func (q *ipQueue[T]) recycle(elts *[]T) {
 	q.pool.Put(elts)
 }
 
+// Renotify will re-fill the notification channel if there are pending
+// reads. Useful if you have read the notification channel but do not
+// intend to pop/popOne/drain etc.
+func (q *ipQueue[T]) renotify() {
+	if q == nil {
+		return
+	}
+	q.Lock()
+	defer q.Unlock()
+	if len(q.elts)-q.pos == 0 {
+		return
+	}
+	q.ch <- struct{}{}
+}
+
 // Returns the current length of the queue.
 func (q *ipQueue[T]) len() int {
 	q.Lock()

server/raft.go

@@ -2261,10 +2261,10 @@ func (n *raft) runAsFollower() {
 			n.votes.popOne()
 		case <-n.resp.ch:
 			// Ignore append entry responses received from before the state change.
-			n.resp.drain()
+			n.drainOrRenotify(n.resp)
 		case <-n.prop.ch:
 			// Ignore proposals received from before the state change.
-			n.prop.drain()
+			n.drainOrRenotify(n.prop)
 		case <-n.reqs.ch:
 			// We've just received a vote request from the network.
 			// Because of drain() it is possible that we get nil from popOne().
@@ -2275,6 +2275,24 @@ func (n *raft) runAsFollower() {
 	}
 }
 
+// For runAsFollower and runAsCandidate, to ensure that we don't hang onto
+// things like append entry responses or proposals from a previous leadership
+// if we are staying in this state, but that we don't nuke them if we're about
+// to become a leader.
+func (n *raft) drainOrRenotify(ipq interface {
+	drain() int
+	renotify()
+}) {
+	if n.State() == Leader {
+		// We're in the process of switching to leader but haven't switched
+		// to runAsLeader yet, put the notification back so runAsLeader will
+		// fetch it on the next select{}.
+		ipq.renotify()
+	} else {
+		ipq.drain()
+	}
+}
+
 // Pool for CommittedEntry re-use.
 var cePool = sync.Pool{
 	New: func() any {
@@ -3412,10 +3430,10 @@ func (n *raft) runAsCandidate() {
 			n.processAppendEntries()
 		case <-n.resp.ch:
 			// Ignore append entry responses received from before the state change.
-			n.resp.drain()
+			n.drainOrRenotify(n.resp)
 		case <-n.prop.ch:
 			// Ignore proposals received from before the state change.
-			n.prop.drain()
+			n.drainOrRenotify(n.prop)
 		case <-n.s.quitCh:
 			return
 		case <-n.quit: