NRG: Decouple Raft transport layer

The Raft implementation was tightly coupled to the server's
internal client and send queue for the RPC communication.
This makes it difficult to test scenarios like network
partitions in deterministic manner.
The primary benefit of this change is improved testability. A
new mockTransport is introduced for testing, which allows for
simulating network partitions and for injecting behavior after
a message is sent.

Signed-off-by: Daniele Sciascia <[email protected]>

Author: sciascid

server/raft.go

@@ -184,7 +184,6 @@ type raft struct {
 	vote   string // Our current vote state
 
 	s  *Server    // Reference to top-level server
-	c  *client    // Internal client for subscriptions
 	js *jetStream // JetStream, if running, to see if we are out of resources
 
 	hasleader atomic.Bool // Is there a group leader right now?
@@ -203,7 +202,7 @@ type raft struct {
 	asubj  string // Append entries subject
 	areply string // Append entries responses subject
 
-	sq    *sendq        // Send queue for outbound RPC messages
+	t     raftTransport // Transport that handles Raft messaging
 	aesub *subscription // Subscription for handleAppendEntry callbacks
 
 	wtv []byte // Term and vote to be written
@@ -296,6 +295,8 @@ type RaftConfig struct {
 	// We need to protect against losing state due to the new peers starting with an empty log.
 	// Therefore, these empty servers can't try to become leader until they at least have _some_ state.
 	ScaleUp bool
+
+	Transport raftTransportFactory
 }
 
 var (
@@ -437,6 +438,12 @@ func (s *Server) initRaftNode(accName string, cfg *RaftConfig, labels pprofLabel
 		extSt:    ps.domainExt,
 	}
 
+	factory := cfg.Transport
+	if factory == nil {
+		factory = defaultRaftTransport
+	}
+	n.t = factory(s, n)
+
 	// Setup our internal subscriptions for proposals, votes and append entries.
 	// If we fail to do this for some reason then this is fatal — we cannot
 	// continue setting up or the Raft node may be partially/totally isolated.
@@ -681,33 +688,11 @@ func (n *raft) recreateInternalSubsLocked() error {
 	// the next step...
 	n.cancelCatchup()
 
-	// If we have an existing client then tear down any existing
-	// subscriptions and close the internal client.
-	if c := n.c; c != nil {
-		c.mu.Lock()
-		subs := make([]*subscription, 0, len(c.subs))
-		for _, sub := range c.subs {
-			subs = append(subs, sub)
-		}
-		c.mu.Unlock()
-		for _, sub := range subs {
-			n.unsubscribe(sub)
-		}
-		c.closeConnection(InternalClient)
-	}
-
 	if n.acc != nrgAcc {
 		n.debug("Subscribing in '%s'", nrgAcc.GetName())
 	}
 
-	c := n.s.createInternalSystemClient()
-	c.registerWithAccount(nrgAcc)
-	if nrgAcc.sq == nil {
-		nrgAcc.sq = n.s.newSendQ(nrgAcc)
-	}
-	n.c = c
-	n.sq = nrgAcc.sq
-	n.acc = nrgAcc
+	n.t.Reset(nrgAcc)
 
 	// Recreate any internal subscriptions for voting, append
 	// entries etc in the new account.
@@ -1968,17 +1953,12 @@ func (n *raft) newInbox() string {
 // Our internal subscribe.
 // Lock should be held.
 func (n *raft) subscribe(subject string, cb msgHandler) (*subscription, error) {
-	if n.c == nil {
-		return nil, errNoInternalClient
-	}
-	return n.s.systemSubscribe(subject, _EMPTY_, false, n.c, cb)
+	return n.t.Subscribe(subject, cb)
 }
 
 // Lock should be held.
 func (n *raft) unsubscribe(sub *subscription) {
-	if n.c != nil && sub != nil {
-		n.c.processUnsub(sub.sid)
-	}
+	n.t.Unsubscribe(sub)
 }
 
 // Lock should be held.
@@ -2103,19 +2083,7 @@ runner:
 	n.Lock()
 	defer n.Unlock()
 
-	if c := n.c; c != nil {
-		var subs []*subscription
-		c.mu.Lock()
-		for _, sub := range c.subs {
-			subs = append(subs, sub)
-		}
-		c.mu.Unlock()
-		for _, sub := range subs {
-			n.unsubscribe(sub)
-		}
-		c.closeConnection(InternalClient)
-		n.c = nil
-	}
+	n.t.Close()
 
 	// Unregistering ipQueues do not prevent them from push/pop
 	// just will remove them from the central monitoring map
@@ -4640,15 +4608,11 @@ func (n *raft) requestVote() {
 }
 
 func (n *raft) sendRPC(subject, reply string, msg []byte) {
-	if n.sq != nil {
-		n.sq.send(subject, reply, nil, msg)
-	}
+	n.t.Publish(subject, reply, msg)
 }
 
 func (n *raft) sendReply(subject string, msg []byte) {
-	if n.sq != nil {
-		n.sq.send(subject, _EMPTY_, nil, msg)
-	}
+	n.t.Publish(subject, _EMPTY_, msg)
 }
 
 func (n *raft) wonElection(votes int) bool {

server/raft_helpers_test.go

@@ -126,21 +126,26 @@ func (sg smGroup) lockFollowers() []stateMachine {
 // Create a raft group and place on numMembers servers at random.
 // Filestore based.
 func (c *cluster) createRaftGroup(name string, numMembers int, smf smFactory) smGroup {
-	return c.createRaftGroupEx(name, numMembers, smf, FileStorage)
+	return c.createRaftGroupEx(name, numMembers, smf, defaultRaftTransport, FileStorage)
 }
 
 func (c *cluster) createMemRaftGroup(name string, numMembers int, smf smFactory) smGroup {
-	return c.createRaftGroupEx(name, numMembers, smf, MemoryStorage)
+	return c.createRaftGroupEx(name, numMembers, smf, defaultRaftTransport, MemoryStorage)
 }
 
-func (c *cluster) createRaftGroupEx(name string, numMembers int, smf smFactory, st StorageType) smGroup {
+func (c *cluster) createMockMemRaftGroup(name string, members int, smf smFactory) (*raftTransportHub, raftTransportFactory, smGroup) {
+	hub, rtf := mockTransportFactory()
+	return hub, rtf, c.createRaftGroupEx(name, members, smf, rtf, MemoryStorage)
+}
+
+func (c *cluster) createRaftGroupEx(name string, numMembers int, smf smFactory, rtf raftTransportFactory, st StorageType) smGroup {
 	c.t.Helper()
 	if numMembers > len(c.servers) {
 		c.t.Fatalf("Members > Peers: %d vs  %d", numMembers, len(c.servers))
 	}
 	servers := append([]*Server{}, c.servers...)
 	rand.Shuffle(len(servers), func(i, j int) { servers[i], servers[j] = servers[j], servers[i] })
-	return c.createRaftGroupWithPeers(name, servers[:numMembers], smf, st)
+	return c.createRaftGroupWithPeers(name, servers[:numMembers], smf, rtf, st)
 }
 
 func (c *cluster) createWAL(name string, st StorageType) WAL {
@@ -189,42 +194,48 @@ func (c *cluster) createStateMachine(s *Server, cfg *RaftConfig, peers []string,
 	return sm
 }
 
-func (c *cluster) createRaftGroupWithPeers(name string, servers []*Server, smf smFactory, st StorageType) smGroup {
+func (c *cluster) createRaftGroupWithPeers(name string, servers []*Server, smf smFactory, rtf raftTransportFactory, st StorageType) smGroup {
 	c.t.Helper()
 
 	var sg smGroup
 	peers := serverPeerNames(servers)
 
 	for _, s := range servers {
 		cfg := &RaftConfig{
-			Name:  name,
-			Store: c.t.TempDir(),
-			Log:   c.createWAL(name, st)}
+			Name:      name,
+			Store:     c.t.TempDir(),
+			Log:       c.createWAL(name, st),
+			Transport: rtf}
 		sg = append(sg, c.createStateMachine(s, cfg, peers, smf))
 	}
 	return sg
 }
 
-func (c *cluster) addNodeEx(name string, smf smFactory, st StorageType) stateMachine {
+func (c *cluster) addNodeEx(name string, smf smFactory, rtf raftTransportFactory, st StorageType) stateMachine {
 	c.t.Helper()
 
 	server := c.addInNewServer()
 
 	cfg := &RaftConfig{
-		Name:  name,
-		Store: c.t.TempDir(),
-		Log:   c.createWAL(name, st)}
+		Name:      name,
+		Store:     c.t.TempDir(),
+		Log:       c.createWAL(name, st),
+		Transport: rtf}
 
 	peers := serverPeerNames(c.servers)
 	return c.createStateMachine(server, cfg, peers, smf)
 }
 
 func (c *cluster) addRaftNode(name string, smf smFactory) stateMachine {
-	return c.addNodeEx(name, smf, FileStorage)
+	return c.addNodeEx(name, smf, defaultRaftTransport, FileStorage)
 }
 
 func (c *cluster) addMemRaftNode(name string, smf smFactory) stateMachine {
-	return c.addNodeEx(name, smf, MemoryStorage)
+	return c.addNodeEx(name, smf, defaultRaftTransport, MemoryStorage)
+}
+
+func (c *cluster) addMockMemRaftNode(name string, rtf raftTransportFactory, smf smFactory) stateMachine {
+	return c.addNodeEx(name, smf, rtf, MemoryStorage)
 }
 
 // Driver program for the state machine.

server/raft_test.go

@@ -4627,3 +4627,110 @@ func TestNRGSingleNodeElection(t *testing.T) {
 	require_Equal(t, newLeader.node().ClusterSize(), 3)
 	require_False(t, newLeader.node().MembershipChangeInProgress())
 }
+
+func TestNRGPartitionedPeerRemove(t *testing.T) {
+	c := createJetStreamClusterExplicit(t, "R2S", 2)
+	defer c.shutdown()
+
+	hub, _, rg := c.createMockMemRaftGroup("MOCK", 2, newStateAdder)
+	defer hub.healPartitions()
+
+	leader := rg.waitOnLeader()
+	followers := rg.followers()
+	require_Equal(t, len(followers), 1)
+	require_Equal(t, leader.node().ClusterSize(), 2)
+
+	// Remove the follower while the leader is partitioned away
+	hub.partition(leader.node().ID(), 1)
+	leader.node().ProposeRemovePeer(followers[0].node().ID())
+
+	// Follower should can't get elected, but let's try anyway
+	followers[0].node().CampaignImmediately()
+
+	// Expect progress on the leader side
+	checkFor(t, 1*time.Second, 10*time.Millisecond, func() error {
+		if leader.node().ClusterSize() != 1 {
+			return errors.New("node removal still in progress")
+		}
+		return nil
+	})
+
+	checkFor(t, 1*time.Second, 10*time.Millisecond, func() error {
+		if leader.node().MembershipChangeInProgress() {
+			return errors.New("membership still in progress")
+		}
+		return nil
+	})
+
+	require_Equal(t, leader.node().ClusterSize(), 1)
+	require_False(t, leader.node().MembershipChangeInProgress())
+
+	// Follower has not changed
+	require_Equal(t, followers[0].node().State(), Follower)
+	require_Equal(t, followers[0].node().ClusterSize(), 2)
+	require_False(t, followers[0].node().MembershipChangeInProgress())
+
+	// Heal the partition, and expect the follower to get the bad news...
+	hub.heal(leader.node().ID())
+
+	checkFor(t, 1*time.Second, 10*time.Millisecond, func() error {
+		if followers[0].node().ClusterSize() != 1 {
+			return errors.New("node removal still in progress")
+		}
+		return nil
+	})
+}
+
+func TestNRGLeaderWithoutQuorumAfterPeerAdd(t *testing.T) {
+	c := createJetStreamClusterExplicit(t, "R3S", 3)
+	defer c.shutdown()
+
+	hub, rtf, rg := c.createMockMemRaftGroup("MOCK", 3, newStateAdder)
+	defer hub.healPartitions()
+
+	leader := rg.waitOnLeader()
+	followers := rg.followers()
+
+	// Setup a after message hook to create a partition as soon as
+	// the leader publishes a EntryAddPeer. The partition will
+	// prevent committing the entry.
+	hub.setAfterMsgHook(func(subject, reply string, msg []byte) {
+		if subject != "$NRG.AE.MOCK" {
+			return
+		}
+		ae, _ := decodeAppendEntry(msg, nil, reply)
+		if ae == nil || len(ae.entries) != 1 {
+			return
+		}
+		if ae.leader != leader.node().ID() {
+			return
+		}
+		require_Equal(t, ae.entries[0].Type, EntryAddPeer)
+
+		// After EntryAddPeer is published, partition the
+		// leader and one of the followers. This partition
+		// can't commit the entry.
+		hub.partition(leader.node().ID(), 1)
+		hub.partition(followers[0].node().ID(), 1)
+	})
+
+	newNode := c.addMockMemRaftNode("MOCK", rtf, newStateAdder)
+
+	// At some point here the cluster gets partitioned:
+	// {leader, followers[0]} and {newNode, followers[1]}
+	// neither side should be able to make progress.
+	newGroup := smGroup{newNode, followers[1]}
+	newLeader := newGroup.waitOnLeader()
+
+	// If the bug is present: we managed to elect a new leader,
+	// in a 4 node cluster, with only two nodes in the partition!
+	// This is because of the following sequence of events:
+	// 1) the follower has received the EntryPeerAdd
+	// 2) the leader and the other follower have partitioned away
+	// 3) the entry is uncommitted, however the follower has added
+	///   the new peer to its peer set, but won't adjust cluster
+	//    size and quorum until after the entry is committed.
+	// 4) follower becomes a canditate and will become leader with
+	//    with a single vote from the new node
+	require_Equal(t, newLeader, nil)
+}