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]>
nats-io · Jun 27, 2024 · f6b2099 · f6b2099
1 parent 1761eac
commit f6b2099
Show file tree

Hide file tree

Showing 10 changed files with 291 additions and 29 deletions.
diff --git a/server/accounts.go b/server/accounts.go
@@ -61,6 +61,7 @@ type Account struct {
 	sqmu         sync.Mutex
 	sl           *Sublist
 	ic           *client
+	sq           *sendq
 	isid         uint64
 	etmr         *time.Timer
 	ctmr         *time.Timer

diff --git a/server/events.go b/server/events.go
@@ -263,6 +263,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.
@@ -288,6 +289,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"`
@@ -486,6 +498,7 @@ RESET:
 						// New capability based flags.
 						si.SetJetStreamEnabled()
 						si.SetBinaryStreamSnapshot()
+						si.SetAccountNRG()
 					}
 				}
 				var b []byte
@@ -912,6 +925,7 @@ func (s *Server) sendStatsz(subj string) {
 	// JetStream
 	if js := s.js.Load(); js != nil {
 		jStat := &JetStreamVarz{}
+		jStat.AccountNRGActive = s.accountNRG.Load()
 		s.mu.RUnlock()
 		js.mu.RLock()
 		c := js.config
@@ -1589,7 +1603,9 @@ func (s *Server) remoteServerUpdate(sub *subscription, c *client, _ *Account, su
 		false,
 		si.JetStreamEnabled(),
 		si.BinaryStreamSnapshot(),
+		si.AccountNRG(),
 	})
+	s.updateNRGAccountStatus()
 }
 
 // updateRemoteServer is called when we have an update from a remote server.
@@ -1636,14 +1652,56 @@ 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() {
+	var raftNodes []RaftNode
+	s.optsMu.RLock()
+	supported := s.opts.JetStreamAccountNRG
+	s.optsMu.RUnlock()
+	if supported {
+		s.rnMu.Lock()
+		raftNodes = make([]RaftNode, 0, len(s.raftNodes))
+		for _, n := range s.raftNodes {
+			raftNodes = append(raftNodes, n)
+		}
+		s.rnMu.Unlock()
+		s.mu.Lock()
+		s.nodeToInfo.Range(func(key, value any) bool {
+			si := value.(nodeInfo)
+			if !s.sameDomain(si.domain) {
+				return true
+			}
+			if supported = supported && si.accountNRG; !supported {
+				return false
+			}
+			return true
+		})
+		s.mu.Unlock()
+	}
+	if s.accountNRG.CompareAndSwap(!supported, supported) {
+		if supported {
+			s.Noticef("Moving NRG traffic into asset accounts")
+		} else {
+			s.Warnf("Moving NRG traffic back into system account due to old nodes coming online")
+		}
+		for _, n := range raftNodes {
+			n.RecreateInternalSubs(supported)
+		}
+	}
+}
+
 // 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.

diff --git a/server/jetstream.go b/server/jetstream.go
@@ -48,6 +48,7 @@ type JetStreamConfig struct {
 	Domain       string        `json:"domain,omitempty"`
 	CompressOK   bool          `json:"compress_ok,omitempty"`
 	UniqueTag    string        `json:"unique_tag,omitempty"`
+	AccountNRG   bool          `json:"account_nrg_enabled,omitempty"`
 }
 
 // Statistics about JetStream for this server.
@@ -544,6 +545,7 @@ func (s *Server) restartJetStream() error {
 		MaxMemory:    opts.JetStreamMaxMemory,
 		MaxStore:     opts.JetStreamMaxStore,
 		Domain:       opts.JetStreamDomain,
+		AccountNRG:   opts.JetStreamAccountNRG,
 	}
 	s.Noticef("Restarting JetStream")
 	err := s.EnableJetStream(&cfg)
@@ -2513,6 +2515,7 @@ func (s *Server) dynJetStreamConfig(storeDir string, maxStore, maxMem int64) *Je
 	}
 
 	opts := s.getOpts()
+	jsc.AccountNRG = opts.JetStreamAccountNRG
 
 	// Sync options.
 	jsc.SyncInterval = opts.SyncInterval

diff --git a/server/jetstream_cluster_4_test.go b/server/jetstream_cluster_4_test.go
@@ -2329,3 +2329,115 @@ func TestJetStreamClusterAckFloorBetweenLeaderAndFollowers(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)
+
+	// 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))
+	}
+
+	// Check that varz tells us what we expect.
+	// set a header to make sure request parsing knows to ignore them
+	ri := snc.NewRespInbox()
+	replies, err := snc.SubscribeSync(ri)
+	require_NoError(t, err)
+	require_NoError(t, snc.PublishMsg(&nats.Msg{
+		Subject: "$SYS.REQ.SERVER.PING.VARZ",
+		Reply:   ri,
+	}))
+	for i := 0; i < len(c.servers); i++ {
+		msg, err := replies.NextMsg(time.Second * 5)
+		require_NoError(t, err)
+		var v Varz
+		require_NoError(t, json.Unmarshal(msg.Data, &v))
+		//require_Equal(t, v.JetStream.Config.AccountNRG, false)
+		require_Equal(t, v.JetStream.AccountNRGActive, false)
+	}
+	p, _, _ := replies.Pending()
+	require_Equal(t, p, 0)
+
+	// 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)
+	}
+
+	// Switch on account NRG on all servers in the cluster. Then
+	// we wait, as we will need statsz to be sent for all servers
+	// in the cluster.
+	for _, s := range c.servers {
+		s.optsMu.Lock()
+		s.opts.JetStreamAccountNRG = true
+		s.optsMu.Unlock()
+		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))
+	}
+
+	// Check varz again.
+	require_NoError(t, snc.PublishMsg(&nats.Msg{
+		Subject: "$SYS.REQ.SERVER.PING.VARZ",
+		Reply:   ri,
+	}))
+	for i := 0; i < len(c.servers); i++ {
+		msg, err := replies.NextMsg(time.Second * 5)
+		require_NoError(t, err)
+		var v Varz
+		//Server Varz `json:"server"`
+		//}
+		require_NoError(t, json.Unmarshal(msg.Data, &v))
+		t.Logf("JSON: %s", msg.Data)
+		t.Logf("Varz: %+v", v)
+		//require_Equal(t, v.JetStream.Config.AccountNRG, true)
+		//require_Equal(t, v.JetStream.AccountNRGActive, true)
+	}
+}
diff --git a/server/monitor.go b/server/monitor.go
@@ -1226,9 +1226,10 @@ type Varz struct {
 
 // JetStreamVarz contains basic runtime information about jetstream
 type JetStreamVarz struct {
-	Config *JetStreamConfig `json:"config,omitempty"`
-	Stats  *JetStreamStats  `json:"stats,omitempty"`
-	Meta   *MetaClusterInfo `json:"meta,omitempty"`
+	Config           *JetStreamConfig `json:"config,omitempty"`
+	Stats            *JetStreamStats  `json:"stats,omitempty"`
+	Meta             *MetaClusterInfo `json:"meta,omitempty"`
+	AccountNRGActive bool             `json:"account_nrg_active,omitempty"`
 }
 
 // ClusterOptsVarz contains monitoring cluster information
@@ -1443,6 +1444,7 @@ func (s *Server) HandleRoot(w http.ResponseWriter, r *http.Request) {
 }
 
 func (s *Server) updateJszVarz(js *jetStream, v *JetStreamVarz, doConfig bool) {
+	v.AccountNRGActive = s.accountNRG.Load()
 	if doConfig {
 		js.mu.RLock()
 		// We want to snapshot the config since it will then be available outside
@@ -1821,6 +1823,7 @@ func (s *Server) HandleVarz(w http.ResponseWriter, r *http.Request) {
 		// Now update server's varz
 		s.mu.RLock()
 		sv := &s.varz.JetStream
+		sv.AccountNRGActive = s.accountNRG.Load()
 		if created {
 			sv.Config = v.Config
 		}
@@ -2794,11 +2797,12 @@ type JSInfo struct {
 	Disabled bool            `json:"disabled,omitempty"`
 	Config   JetStreamConfig `json:"config,omitempty"`
 	JetStreamStats
-	Streams   int              `json:"streams"`
-	Consumers int              `json:"consumers"`
-	Messages  uint64           `json:"messages"`
-	Bytes     uint64           `json:"bytes"`
-	Meta      *MetaClusterInfo `json:"meta_cluster,omitempty"`
+	Streams          int              `json:"streams"`
+	Consumers        int              `json:"consumers"`
+	Messages         uint64           `json:"messages"`
+	Bytes            uint64           `json:"bytes"`
+	Meta             *MetaClusterInfo `json:"meta_cluster,omitempty"`
+	AccountNRGActive bool             `json:"account_nrg_active,omitempty"`
 
 	// aggregate raft info
 	AccountDetails []*AccountDetail `json:"account_details,omitempty"`

diff --git a/server/opts.go b/server/opts.go
@@ -317,6 +317,7 @@ type Options struct {
 	JetStreamLimits       JSLimitOpts
 	JetStreamTpm          JSTpmOpts
 	JetStreamMaxCatchup   int64
+	JetStreamAccountNRG   bool
 	StoreDir              string            `json:"-"`
 	SyncInterval          time.Duration     `json:"-"`
 	SyncAlways            bool              `json:"-"`
@@ -2310,6 +2311,8 @@ func parseJetStream(v any, opts *Options, errors *[]error, warnings *[]error) er
 					return &configErr{tk, fmt.Sprintf("%s %s", strings.ToLower(mk), err)}
 				}
 				opts.JetStreamMaxCatchup = s
+			case "account_nrg":
+				opts.JetStreamAccountNRG = mv.(bool)
 			default:
 				if !tk.IsUsedVariable() {
 					err := &unknownConfigFieldErr{