NRG: Decouple Raft transport layer

server/jetstream_jwt_test.go

@@ -1810,7 +1810,7 @@ func TestJetStreamJWTClusterAccountNRG(t *testing.T) {
 			// in-account or not.
 			for _, rg := range raftNodes {
 				rg.Lock()
-				rgAcc := rg.acc
+				rgAcc := rg.t.Account()
 				rg.Unlock()
 				switch state {
 				case "system":
@@ -1911,7 +1911,7 @@ func TestJetStreamJWTClusterAccountNRGPersistsAfterRestart(t *testing.T) {
 
 		for _, rg := range raftNodes {
 			rg.Lock()
-			rgAcc := rg.acc
+			rgAcc := rg.t.Account()
 			rg.Unlock()
 			require_Equal(t, rgAcc.Name, aExpPub)
 			require_Equal(t, rza[rg.group].SystemAcc, false)

server/monitor.go

@@ -4257,7 +4257,7 @@ func (s *Server) Raftz(opts *RaftzOptions) *RaftzStatus {
 			PTerm:         n.pterm,
 			PIndex:        n.pindex,
 			SystemAcc:     n.IsSystemAccount(),
-			TrafficAcc:    n.acc.GetName(),
+			TrafficAcc:    n.t.Account().GetName(),
 			IPQPropLen:    n.prop.len(),
 			IPQEntryLen:   n.entry.len(),
 			IPQRespLen:    n.resp.len(),

server/raft.go

@@ -155,7 +155,6 @@ type raft struct {
 
 	created time.Time      // Time that the group was created
 	accName string         // Account name of the asset this raft group is for
-	acc     *Account       // Account that NRG traffic will be sent/received in
 	group   string         // Raft group
 	sd      string         // Store directory
 	id      string         // Node ID
@@ -202,7 +201,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?
@@ -221,7 +219,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
@@ -319,6 +317,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 (
@@ -464,6 +464,12 @@ func (s *Server) initRaftNode(accName string, cfg *RaftConfig, labels pprofLabel
 		observer: cfg.Observer,
 	}
 
+	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.
@@ -658,7 +664,10 @@ func (n *raft) IsSystemAccount() bool {
 func (n *raft) GetTrafficAccountName() string {
 	n.RLock()
 	defer n.RUnlock()
-	return n.acc.GetName()
+	if n.t == nil {
+		return (*Account)(nil).GetName()
+	}
+	return n.t.Account().GetName()
 }
 
 func (n *raft) RecreateInternalSubs() error {
@@ -708,7 +717,7 @@ func (n *raft) recreateInternalSubsLocked() error {
 			}
 		}
 	}
-	if n.aesub != nil && n.acc == nrgAcc {
+	if n.aesub != nil && n.t.Account() == nrgAcc {
 		// Subscriptions already exist and the account NRG state
 		// hasn't changed.
 		return nil
@@ -719,33 +728,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 {
+	if n.t.Account() != 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.
@@ -2325,17 +2312,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.
@@ -2460,19 +2442,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
@@ -5131,15 +5101,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

@@ -26,6 +26,7 @@ import (
 	"reflect"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -6221,3 +6222,62 @@ func TestNRGReset(t *testing.T) {
 	// Append entry cache cleared.
 	require_Equal(t, len(n.pae), 0)
 }
+
+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().node()
+	require_Equal(t, leader.ClusterSize(), 2)
+	require_Equal(t, len(rg.followers()), 1)
+	follower := rg.followers()[0].node()
+
+	var hookCalls atomic.Int64
+	hub.setAfterMsgHook(func(subject, reply string, msg []byte) {
+		hookCalls.Add(1)
+	})
+
+	// Remove the follower while the leader is partitioned away
+	hub.partition(leader.ID(), 1)
+	leader.ProposeRemovePeer(follower.ID())
+
+	// Follower can't get elected as leader, but let's try anyway
+	follower.CampaignImmediately()
+
+	// Expect progress on the leader side
+	checkFor(t, 1*time.Second, 10*time.Millisecond, func() error {
+		if leader.ClusterSize() != 1 {
+			return errors.New("node removal still in progress")
+		}
+		return nil
+	})
+
+	checkFor(t, 1*time.Second, 10*time.Millisecond, func() error {
+		if leader.MembershipChangeInProgress() {
+			return errors.New("membership still in progress")
+		}
+		return nil
+	})
+
+	require_Equal(t, leader.ClusterSize(), 1)
+	require_False(t, leader.MembershipChangeInProgress())
+	require_True(t, hookCalls.Load() > 0)
+
+	// Follower has not changed
+	require_Equal(t, follower.State(), Follower)
+	require_Equal(t, follower.ClusterSize(), 2)
+	require_False(t, follower.MembershipChangeInProgress())
+
+	// Heal the partition, and expect the follower to get the bad news...
+	hub.heal(leader.ID())
+
+	checkFor(t, 1*time.Second, 10*time.Millisecond, func() error {
+		if follower.ClusterSize() != 1 {
+			return errors.New("node removal still in progress")
+		}
+		return nil
+	})
+}