Move NRG traffic into asset account

This adds a new account NRG capability into statsz so that we can
detect when all servers in the cluster support moving traffic into
the asset account, instead of all being in the system account.

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

Author: neilalexander

go.mod

@@ -2,6 +2,8 @@ module github.com/nats-io/nats-server/v2
 
 go 1.22.0
 
+replace github.com/nats-io/jwt/v2 => github.com/nats-io/jwt/v2 v2.5.9-0.20240801130136-270cc45c44ee
+
 require (
 	github.com/google/go-tpm v0.9.0
 	github.com/klauspost/compress v1.17.9

go.sum

@@ -6,8 +6,8 @@ github.com/klauspost/compress v1.17.9 h1:6KIumPrER1LHsvBVuDa0r5xaG0Es51mhhB9BQB2
 github.com/klauspost/compress v1.17.9/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
 github.com/minio/highwayhash v1.0.3 h1:kbnuUMoHYyVl7szWjSxJnxw11k2U709jqFPPmIUyD6Q=
 github.com/minio/highwayhash v1.0.3/go.mod h1:GGYsuwP/fPD6Y9hMiXuapVvlIUEhFhMTh0rxU3ik1LQ=
-github.com/nats-io/jwt/v2 v2.5.8 h1:uvdSzwWiEGWGXf+0Q+70qv6AQdvcvxrv9hPM0RiPamE=
-github.com/nats-io/jwt/v2 v2.5.8/go.mod h1:ZdWS1nZa6WMZfFwwgpEaqBV8EPGVgOTDHN/wTbz0Y5A=
+github.com/nats-io/jwt/v2 v2.5.9-0.20240801130136-270cc45c44ee h1:kiTo11kMbk4UMNjdgXAFFSq8+p1RLI3XPcAukFaiw+g=
+github.com/nats-io/jwt/v2 v2.5.9-0.20240801130136-270cc45c44ee/go.mod h1:ZdWS1nZa6WMZfFwwgpEaqBV8EPGVgOTDHN/wTbz0Y5A=
 github.com/nats-io/nats.go v1.36.0 h1:suEUPuWzTSse/XhESwqLxXGuj8vGRuPRoG7MoRN/qyU=
 github.com/nats-io/nats.go v1.36.0/go.mod h1:Ubdu4Nh9exXdSz0RVWRFBbRfrbSxOYd26oF0wkWclB8=
 github.com/nats-io/nkeys v0.4.7 h1:RwNJbbIdYCoClSDNY7QVKZlyb/wfT6ugvFCiKy6vDvI=

server/accounts.go

@@ -62,6 +62,7 @@ type Account struct {
 	sqmu         sync.Mutex
 	sl           *Sublist
 	ic           *client
+	sq           *sendq
 	isid         uint64
 	etmr         *time.Timer
 	ctmr         *time.Timer
@@ -3679,6 +3680,26 @@ func (s *Server) updateAccountClaimsWithRefresh(a *Account, ac *jwt.AccountClaim
 		a.enableAllJetStreamServiceImportsAndMappings()
 	}
 
+	if a.js != nil {
+		// Check whether the account NRG status changed. If it has then we need to notify the
+		// Raft groups running on the system so that they can move their subs if needed.
+		a.mu.Lock()
+		previous := a.js.nrgAccount
+		switch strings.ToLower(ac.NRGAccount) {
+		case "account":
+			a.js.nrgAccount = a.Name
+		case "system":
+			a.js.nrgAccount = ""
+		default:
+			s.Errorf("Account claim for %q has invalid value %q for account NRG status", a.Name, ac.NRGAccount)
+		}
+		changed := a.js.nrgAccount != previous
+		a.mu.Unlock()
+		if changed {
+			s.updateNRGAccountStatus()
+		}
+	}
+
 	for i, c := range clients {
 		a.mu.RLock()
 		exceeded := a.mconns != jwt.NoLimit && i >= int(a.mconns)

server/events.go

@@ -264,6 +264,7 @@ type ServerInfo struct {
 const (
 	JetStreamEnabled     ServerCapability = 1 << iota // Server had JetStream enabled.
 	BinaryStreamSnapshot                              // New stream snapshot capability.
+	AccountNRG                                        // Move NRG traffic out of system account.
 )
 
 // Set JetStream capability.
@@ -289,6 +290,17 @@ func (si *ServerInfo) BinaryStreamSnapshot() bool {
 	return si.Flags&BinaryStreamSnapshot != 0
 }
 
+// Set account NRG capability.
+func (si *ServerInfo) SetAccountNRG() {
+	si.Flags |= AccountNRG
+}
+
+// AccountNRG indicates whether or not we support moving the NRG traffic out of the
+// system account and into the asset account.
+func (si *ServerInfo) AccountNRG() bool {
+	return si.Flags&AccountNRG != 0
+}
+
 // ClientInfo is detailed information about the client forming a connection.
 type ClientInfo struct {
 	Start      *time.Time    `json:"start,omitempty"`
@@ -475,10 +487,14 @@ RESET:
 					si.Version = VERSION
 					si.Time = time.Now().UTC()
 					si.Tags = tags
+					si.Flags = 0
 					if js {
 						// New capability based flags.
 						si.SetJetStreamEnabled()
 						si.SetBinaryStreamSnapshot()
+						if s.accountNRGAllowed.Load() {
+							si.SetAccountNRG()
+						}
 					}
 				}
 				var b []byte
@@ -1616,7 +1632,8 @@ func (s *Server) remoteServerUpdate(sub *subscription, c *client, _ *Account, su
 	}
 
 	node := getHash(si.Name)
-	s.nodeToInfo.Store(node, nodeInfo{
+	accountNRG := si.AccountNRG()
+	oldInfo, _ := s.nodeToInfo.Swap(node, nodeInfo{
 		si.Name,
 		si.Version,
 		si.Cluster,
@@ -1628,7 +1645,14 @@ func (s *Server) remoteServerUpdate(sub *subscription, c *client, _ *Account, su
 		false,
 		si.JetStreamEnabled(),
 		si.BinaryStreamSnapshot(),
+		accountNRG,
 	})
+	if oldInfo == nil || accountNRG != oldInfo.(nodeInfo).accountNRG {
+		// One of the servers we received statsz from changed its mind about
+		// whether or not it supports in-account NRG, so update the groups
+		// with this information.
+		s.updateNRGAccountStatus()
+	}
 }
 
 // updateRemoteServer is called when we have an update from a remote server.
@@ -1675,14 +1699,35 @@ func (s *Server) processNewServer(si *ServerInfo) {
 				false,
 				si.JetStreamEnabled(),
 				si.BinaryStreamSnapshot(),
+				si.AccountNRG(),
 			})
 		}
 	}
+	go s.updateNRGAccountStatus()
 	// Announce ourselves..
 	// Do this in a separate Go routine.
 	go s.sendStatszUpdate()
 }
 
+// Works out whether all nodes support moving the NRG traffic into
+// the account and moves it appropriately.
+// Server lock MUST NOT be held on entry.
+func (s *Server) updateNRGAccountStatus() {
+	s.rnMu.Lock()
+	raftNodes := make([]RaftNode, 0, len(s.raftNodes))
+	for _, n := range s.raftNodes {
+		raftNodes = append(raftNodes, n)
+	}
+	s.rnMu.Unlock()
+	for _, n := range raftNodes {
+		// In the event that the node is happy that all nodes that
+		// it cares about haven't changed, this will be a no-op.
+		if err := n.RecreateInternalSubs(); err != nil {
+			n.Stop()
+		}
+	}
+}
+
 // If GW is enabled on this server and there are any leaf node connections,
 // this function will send a LeafNode connect system event to the super cluster
 // to ensure that the GWs are in interest-only mode for this account.

server/jetstream.go

@@ -174,6 +174,9 @@ type jsAccount struct {
 	updatesSub *subscription
 	lupdate    time.Time
 	utimer     *time.Timer
+
+	// Which account to send NRG traffic into. Empty string is system account.
+	nrgAccount string
 }
 
 // Track general usage for this account.

server/jetstream_cluster_4_test.go

@@ -2505,6 +2505,133 @@ func TestJetStreamClusterConsumerLeak(t *testing.T) {
 	}
 }
 
+func TestJetStreamClusterAccountNRG(t *testing.T) {
+	c := createJetStreamClusterExplicit(t, "R3S", 3)
+	defer c.shutdown()
+
+	nc, js := jsClientConnect(t, c.randomServer())
+	defer nc.Close()
+
+	snc, _ := jsClientConnect(t, c.randomServer(), nats.UserInfo("admin", "s3cr3t!"))
+	defer snc.Close()
+
+	_, err := js.AddStream(&nats.StreamConfig{
+		Name:      "TEST",
+		Subjects:  []string{"foo"},
+		Storage:   nats.MemoryStorage,
+		Retention: nats.WorkQueuePolicy,
+		Replicas:  3,
+	})
+	require_NoError(t, err)
+
+	leader := c.streamLeader(globalAccountName, "TEST")
+	stream, err := leader.gacc.lookupStream("TEST")
+	require_NoError(t, err)
+	rg := stream.node.(*raft)
+
+	t.Run("Disabled", func(t *testing.T) {
+		// Switch off account NRG on all servers in the cluster.
+		for _, s := range c.servers {
+			s.accountNRGAllowed.Store(false)
+			s.sendStatszUpdate()
+		}
+		time.Sleep(time.Millisecond * 100)
+		for _, s := range c.servers {
+			s.GlobalAccount().js.nrgAccount = ""
+			s.updateNRGAccountStatus()
+		}
+
+		// System account should have interest, but the global account
+		// shouldn't.
+		for _, s := range c.servers {
+			require_True(t, s.sys.account.sl.hasInterest(rg.asubj, true))
+			require_False(t, s.gacc.sl.hasInterest(rg.asubj, true))
+		}
+
+		// First of all check that the Raft traffic is in the system
+		// account, as we haven't moved it elsewhere yet.
+		{
+			sub, err := snc.SubscribeSync(rg.asubj)
+			require_NoError(t, err)
+			require_NoError(t, sub.AutoUnsubscribe(1))
+
+			msg, err := sub.NextMsg(time.Second * 3)
+			require_NoError(t, err)
+			require_True(t, msg != nil)
+		}
+	})
+
+	t.Run("Mixed", func(t *testing.T) {
+		// Switch on account NRG on a single server in the cluster and
+		// leave it off on the rest.
+		for i, s := range c.servers {
+			s.accountNRGAllowed.Store(i == 0)
+			s.sendStatszUpdate()
+		}
+		time.Sleep(time.Millisecond * 100)
+		for i, s := range c.servers {
+			if i == 0 {
+				s.GlobalAccount().js.nrgAccount = globalAccountName
+			} else {
+				s.GlobalAccount().js.nrgAccount = ""
+			}
+			s.updateNRGAccountStatus()
+		}
+
+		// System account should have interest, but the global account
+		// shouldn't.
+		for _, s := range c.servers {
+			require_True(t, s.sys.account.sl.hasInterest(rg.asubj, true))
+			require_False(t, s.gacc.sl.hasInterest(rg.asubj, true))
+		}
+
+		// First of all check that the Raft traffic is in the system
+		// account, as we haven't moved it elsewhere yet.
+		{
+			sub, err := snc.SubscribeSync(rg.asubj)
+			require_NoError(t, err)
+			require_NoError(t, sub.AutoUnsubscribe(1))
+
+			msg, err := sub.NextMsg(time.Second * 3)
+			require_NoError(t, err)
+			require_True(t, msg != nil)
+		}
+	})
+
+	t.Run("Enabled", func(t *testing.T) {
+		// Switch on account NRG on all servers in the cluster.
+		for _, s := range c.servers {
+			s.accountNRGAllowed.Store(true)
+			s.sendStatszUpdate()
+		}
+		time.Sleep(time.Millisecond * 100)
+		for _, s := range c.servers {
+			s.GlobalAccount().js.nrgAccount = globalAccountName
+			s.updateNRGAccountStatus()
+		}
+
+		// Now check that the traffic has moved into the asset acc.
+		// In this case the system account should no longer have
+		// subscriptions for those subjects.
+		{
+			sub, err := nc.SubscribeSync(rg.asubj)
+			require_NoError(t, err)
+			require_NoError(t, sub.AutoUnsubscribe(1))
+
+			msg, err := sub.NextMsg(time.Second * 3)
+			require_NoError(t, err)
+			require_True(t, msg != nil)
+		}
+
+		// The global account should now have interest and the
+		// system account shouldn't.
+		for _, s := range c.servers {
+			require_False(t, s.sys.account.sl.hasInterest(rg.asubj, true))
+			require_True(t, s.gacc.sl.hasInterest(rg.asubj, true))
+		}
+	})
+}
+
 func TestJetStreamClusterWQRoundRobinSubjectRetention(t *testing.T) {
 	c := createJetStreamClusterExplicit(t, "R3S", 3)
 	defer c.shutdown()

server/jetstream_jwt_test.go

@@ -1532,3 +1532,93 @@ func TestJetStreamJWTClusteredTiersR3StreamWithR1ConsumersAndAccounting(t *testi
 	require_Equal(t, r3.Streams, 1)
 	require_Equal(t, r3.Consumers, 0)
 }
+
+func TestJetStreamJWTClusterAccountNRG(t *testing.T) {
+	_, syspub := createKey(t)
+	sysJwt := encodeClaim(t, jwt.NewAccountClaims(syspub), syspub)
+
+	_, aExpPub := createKey(t)
+	accClaim := jwt.NewAccountClaims(aExpPub)
+	accClaim.Name = "acc"
+	accClaim.Limits.JetStreamTieredLimits["R1"] = jwt.JetStreamLimits{DiskStorage: 1100, Consumer: 10, Streams: 1}
+	accClaim.Limits.JetStreamTieredLimits["R3"] = jwt.JetStreamLimits{DiskStorage: 1100, Consumer: 1, Streams: 1}
+	accJwt := encodeClaim(t, accClaim, aExpPub)
+
+	tmlp := `
+		listen: 127.0.0.1:-1
+		server_name: %s
+		jetstream: {max_mem_store: 256MB, max_file_store: 2GB, store_dir: '%s'}
+		leaf {
+			listen: 127.0.0.1:-1
+		}
+		cluster {
+			name: %s
+			listen: 127.0.0.1:%d
+			routes = [%s]
+		}
+	` + fmt.Sprintf(`
+		operator: %s
+		system_account: %s
+		resolver = MEMORY
+		resolver_preload = {
+			%s : %s
+			%s : %s
+		}
+	`, ojwt, syspub, syspub, sysJwt, aExpPub, accJwt)
+
+	c := createJetStreamClusterWithTemplate(t, tmlp, "cluster", 3)
+	defer c.shutdown()
+
+	// We'll try flipping the state a few times and then do some sanity
+	// checks to check that it took effect.
+	for _, state := range []string{"account", "system", "account"} {
+		accClaim.NRGAccount = state
+		accJwt = encodeClaim(t, accClaim, aExpPub)
+
+		for _, s := range c.servers {
+			// Find the account.
+			acc, err := s.lookupAccount(aExpPub)
+			require_NoError(t, err)
+
+			// Submit a claim update using the new JWT.
+			require_NoError(t, s.updateAccountWithClaimJWT(acc, accJwt))
+
+			// Check that everything looks like it should.
+			require_True(t, acc != nil)
+			require_True(t, acc.js != nil)
+			switch state {
+			case "account":
+				require_Equal(t, acc.js.nrgAccount, aExpPub)
+			case "system":
+				require_Equal(t, acc.js.nrgAccount, _EMPTY_)
+			}
+
+			// Now get a list of all of the Raft nodes that should
+			// have been updated by now.
+			s.rnMu.Lock()
+			raftNodes := make([]*raft, 0, len(s.raftNodes))
+			for _, n := range s.raftNodes {
+				rg := n.(*raft)
+				if rg.accName != acc.Name {
+					continue
+				}
+				raftNodes = append(raftNodes, rg)
+			}
+			s.rnMu.Unlock()
+
+			// Check whether each of the Raft nodes reports being
+			// in-account or not.
+			for _, rg := range raftNodes {
+				rg.Lock()
+				rgAcc := rg.acc
+				rg.Unlock()
+				switch state {
+				case "account":
+					require_Equal(t, rgAcc.Name, aExpPub)
+				case "system":
+					require_Equal(t, rgAcc.Name, syspub)
+				}
+			}
+		}
+	}
+}