NRG: Truncate down to committed on step down

Signed-off-by: Maurice van Veen <[email protected]>

Author: MauriceVanVeen

server/jetstream_cluster_1_test.go

@@ -8495,6 +8495,9 @@ func TestJetStreamClusterSnapshotStreamAssetOnShutdown(t *testing.T) {
 	// Publish, so we have something new to snapshot.
 	_, err = js.Publish("foo", nil)
 	require_NoError(t, err)
+	checkFor(t, 2*time.Second, 200*time.Millisecond, func() error {
+		return checkState(t, c, globalAccountName, "TEST")
+	})
 
 	// Shutdown servers, and check if all made stream snapshots.
 	for _, s := range c.servers {

server/raft.go

@@ -169,6 +169,8 @@ type raft struct {
 	pindex  uint64 // Previous index from the last snapshot
 	commit  uint64 // Index of the most recent commit
 	applied uint64 // Index of the most recently applied commit
+	lterm   uint64 // Term of the entry before this leader's term started.
+	lindex  uint64 // Index of the entry before this leader's term started
 
 	aflr uint64 // Index when to signal initial messages have been applied after becoming leader. 0 means signaling is disabled.
 
@@ -188,7 +190,7 @@ type raft struct {
 	isSysAcc    atomic.Bool // Are we utilizing the system account?
 	maybeLeader bool        // The group had a preferred leader. And is maybe already acting as leader prior to scale up.
 
-	observer bool // The node is observing, i.e. not participating in voting
+	observer bool // The node is observing, i.e. not able to become leader
 
 	extSt extensionState // Extension state
 
@@ -3303,6 +3305,16 @@ func (n *raft) truncateWAL(term, index uint64) {
 	n.pterm, n.pindex = term, index
 }
 
+// Truncate our WAL to remove all entries from the current leader's term if
+// none were committed (no quorum). Keep all of them if at least one had quorum.
+// Lock should be held.
+func (n *raft) truncateWALForLeaderTerm() {
+	// Had at least one entry in the current leader's term and none were committed.
+	if n.pindex > n.lindex && n.lindex >= n.commit {
+		n.truncateWAL(n.lterm, n.lindex)
+	}
+}
+
 // Reset our WAL. This is equivalent to truncating all data from the log.
 // Lock should be held.
 func (n *raft) resetWAL() {
@@ -4290,6 +4302,8 @@ retry:
 	var leadChange bool
 	if pstate == Leader && state != Leader {
 		leadChange = true
+		// Truncate entries from this leader's term if none had quorum.
+		n.truncateWALForLeaderTerm()
 		n.updateLeadChange(false)
 		// Drain the append entry response and proposal queues.
 		n.resp.drain()
@@ -4383,6 +4397,8 @@ func (n *raft) switchToLeader() {
 	sendHB := state.LastSeq > n.commit
 
 	n.lxfer = false
+	n.lterm = n.pterm
+	n.lindex = n.pindex
 	n.updateLeader(n.id)
 	leadChange := n.switchState(Leader)
 

server/raft_helpers_test.go

@@ -18,6 +18,7 @@ package server
 
 import (
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"math/rand"
 	"sync"
@@ -329,14 +330,15 @@ func (a *stateAdder) snapshot(t *testing.T) {
 func (rg smGroup) waitOnTotal(t *testing.T, expected int64) {
 	t.Helper()
 	checkFor(t, 5*time.Second, 200*time.Millisecond, func() error {
+		var err error
 		for _, sm := range rg {
 			asm := sm.(*stateAdder)
 			if total := asm.total(); total != expected {
-				return fmt.Errorf("Adder on %v has wrong total: %d vs %d",
-					asm.server(), total, expected)
+				err = errors.Join(err, fmt.Errorf("Adder on %v has wrong total: %d vs %d",
+					asm.server(), total, expected))
 			}
 		}
-		return nil
+		return err
 	})
 }
 

server/raft_test.go

@@ -913,22 +913,29 @@ func TestNRGNoResetOnAppendEntryResponse(t *testing.T) {
 	follower := rg.nonLeader().node().(*raft)
 	lsm := rg.leader().(*stateAdder)
 
+	// Ensure all servers can get quorum on an initial set.
+	c.waitOnAllCurrent()
+	lsm.proposeDelta(5)
+	rg.waitOnTotal(t, 5)
+
 	// Subscribe for append entries that aren't heartbeats and respond to
 	// each of them as though it's a non-success and with a higher term.
 	// The higher term in this case is what would cause the leader previously
 	// to reset the entire log which it shouldn't do.
+	var once bool
 	_, err := nc.Subscribe(fmt.Sprintf(raftAppendSubj, "TEST"), func(msg *nats.Msg) {
-		if ae, err := follower.decodeAppendEntry(msg.Data, nil, msg.Reply); err == nil && len(ae.entries) > 0 {
+		if ae, err := follower.decodeAppendEntry(msg.Data, nil, msg.Reply); err == nil && len(ae.entries) > 0 && !once {
 			ar := newAppendEntryResponse(ae.term+1, ae.commit, follower.id, false)
 			require_NoError(t, msg.Respond(ar.encode(nil)))
+			once = true
 		}
 	})
 	require_NoError(t, err)
 
 	// Generate an append entry that the subscriber above can respond to.
 	c.waitOnAllCurrent()
 	lsm.proposeDelta(5)
-	rg.waitOnTotal(t, 5)
+	rg.waitOnTotal(t, 10)
 
 	// The was-leader should now have stepped down, make sure that it
 	// didn't blow away its log in the process.
@@ -2363,6 +2370,105 @@ func TestNRGSignalLeadChangeFalseIfCampaignImmediately(t *testing.T) {
 	}
 }
 
+func TestNRGClearWALOnStepDownWhenZeroCommits(t *testing.T) {
+	n, cleanup := initSingleMemRaftNode(t)
+	defer cleanup()
+
+	// Create a sample entry, the content doesn't matter, just that it's stored.
+	esm := encodeStreamMsgAllowCompress("foo", "_INBOX.foo", nil, nil, 0, 0, true)
+	entries := []*Entry{newEntry(EntryNormal, esm)}
+
+	// Switch this node to leader, and send one entry. It will not get quorum.
+	n.term++
+	n.switchToLeader()
+	require_Equal(t, n.term, 1)
+	require_Equal(t, n.pindex, 0)
+	n.sendAppendEntry(entries)
+	require_Equal(t, n.pindex, 1)
+	require_Equal(t, n.pterm, 1)
+
+	// Shouldn't have any commits up to this point.
+	require_Equal(t, n.commit, 0)
+
+	// If we now step down, we clear our WAL, we had no commits.
+	n.stepdown(noLeader)
+	require_Equal(t, n.commit, 0)
+	require_Equal(t, n.pindex, 0)
+	require_Equal(t, n.pterm, 0)
+}
+
+func TestNRGTruncateDownToPreviousTermOnStepDown(t *testing.T) {
+	tests := []struct {
+		title  string
+		commit uint64
+	}{
+		// The current leader's term is fully removed. None of them had quorum.
+		{title: "no-commits", commit: 0},
+		{title: "half-previous-term", commit: 1},
+		{title: "full-previous-term", commit: 2},
+		// At least one entry in the leader's term had quorum, need to preserve all of them.
+		{title: "half-leader-term", commit: 3},
+	}
+	for _, test := range tests {
+		t.Run(test.title, func(t *testing.T) {
+			n, cleanup := initSingleMemRaftNode(t)
+			defer cleanup()
+
+			// Create a sample entry, the content doesn't matter, just that it's stored.
+			esm := encodeStreamMsgAllowCompress("foo", "_INBOX.foo", nil, nil, 0, 0, true)
+			entries := []*Entry{newEntry(EntryNormal, esm)}
+
+			// Switch this node to leader, and send two entries. The first will get quorum, the second will not.
+			n.term++
+			n.switchToLeader()
+			require_Equal(t, n.term, 1)
+			require_Equal(t, n.pindex, 0)
+			n.sendAppendEntry(entries)
+			n.sendAppendEntry(entries)
+			require_Equal(t, n.pindex, 2)
+			require_Equal(t, n.pterm, 1)
+
+			// Shouldn't have any commits up to this point.
+			require_Equal(t, n.commit, 0)
+
+			// We get quorum on the first entry.
+			for i := uint64(1); i <= test.commit && i <= 2; i++ {
+				require_NoError(t, n.applyCommit(i))
+				require_Equal(t, n.commit, i)
+			}
+
+			// Act like the next entry is sent under a new term.
+			// That will verify we can revert back to the pterm at time of last commit.
+			n.term++
+			n.switchToLeader()
+			n.sendAppendEntry(entries)
+			n.sendAppendEntry(entries)
+			require_Equal(t, n.pindex, 4)
+			require_Equal(t, n.pterm, 2)
+
+			if test.commit == 3 {
+				require_NoError(t, n.applyCommit(3))
+				require_Equal(t, n.commit, 3)
+			}
+
+			// If we now step down, we should remove all entries from the leader's term if none had quorum.
+			// This is required, otherwise we could try to become leader again and trick a follower into
+			// voting for us based on a pterm without quorum. That could result in overwriting an entry
+			// at a pindex where a previous leader knew we had quorum on, but the follower didn't.
+			n.stepdown(noLeader)
+
+			require_Equal(t, n.commit, test.commit)
+			if test.commit < 3 {
+				require_Equal(t, n.pindex, 2)
+				require_Equal(t, n.pterm, 1)
+			} else {
+				require_Equal(t, n.pindex, 4)
+				require_Equal(t, n.pterm, 2)
+			}
+		})
+	}
+}
+
 // This is a RaftChainOfBlocks test where a block is proposed and then we wait for all replicas to apply it before
 // proposing the next one.
 // The test may fail if: