Repository: nats-io/nats-server
Pull request #4687 (2023-10-22 05:12:16)

[IMPROVED] WriteFullState Allocations
Avoid additional allocations when doing a writeFullState for large state.
Track the total subject length of all subjects in the PSIM structure.
This is to optimize the writeFullState when the state is large due to a large PSIM with large subjects.

Signed-off-by: Derek Collison <derek@nats.io>



Patch:
diff --git a/server/client.go b/server/client.go
index 4674d354..860fb47b 100644
--- a/server/client.go
+++ b/server/client.go
@@ -5311,10 +5311,10 @@ func (c *client) getAccAndResultFromCache() (*Account, *SublistResult) {
 	// Check our cache.
 	if pac, ok = c.in.pacache[string(c.pa.pacache)]; ok {
 		// Check the genid to see if it's still valid.
-		// Since v2.10.0, the config reload of accounts has been fixed
-		// and an account's sublist pointer should not change, so no need to
-		// lock to access it.
+		// sl could be swapped out on reload so need to lock.
+		pac.acc.mu.RLock()
 		sl := pac.acc.sl
+		pac.acc.mu.RUnlock()
 
 		if genid := atomic.LoadUint64(&sl.genid); genid != pac.genid {
 			ok = false
@@ -5326,15 +5326,15 @@ func (c *client) getAccAndResultFromCache() (*Account, *SublistResult) {
 	}
 
 	if !ok {
-		if c.kind == ROUTER && len(c.route.accName) > 0 {
-			acc = c.acc
-		} else {
-			// Match correct account and sublist.
-			if acc, _ = c.srv.LookupAccount(string(c.pa.account)); acc == nil {
-				return nil, nil
-			}
+		// Match correct account and sublist.
+		if acc, _ = c.srv.LookupAccount(string(c.pa.account)); acc == nil {
+			return nil, nil
 		}
+
+		// sl could be swapped out on reload so need to lock.
+		acc.mu.RLock()
 		sl := acc.sl
+		acc.mu.RUnlock()
 
 		// Match against the account sublist.
 		r = sl.Match(string(c.pa.subject))
diff --git a/server/filestore.go b/server/filestore.go
index 4238ed21..3aadeb56 100644
--- a/server/filestore.go
+++ b/server/filestore.go
@@ -176,6 +176,7 @@ type fileStore struct {
 	blks        []*msgBlock
 	bim         map[uint32]*msgBlock
 	psim        map[string]*psi
+	tsl         int
 	hh          hash.Hash64
 	qch         chan struct{}
 	fch         chan struct{}
@@ -428,7 +429,7 @@ func newFileStoreWithCreated(fcfg FileStoreConfig, cfg StreamConfig, created tim
 		prior := fs.state
 		// Reset anything that could have been set from above.
 		fs.state = StreamState{}
-		fs.psim = make(map[string]*psi)
+		fs.psim, fs.tsl = make(map[string]*psi), 0
 		fs.bim = make(map[uint32]*msgBlock)
 		fs.blks = nil
 		fs.tombs = nil
@@ -1553,7 +1554,7 @@ func (fs *fileStore) recoverFullState() (rerr error) {
 
 	// Check for per subject info.
 	if numSubjects := int(readU64()); numSubjects > 0 {
-		fs.psim = make(map[string]*psi, numSubjects)
+		fs.psim, fs.tsl = make(map[string]*psi, numSubjects), 0
 		for i := 0; i < numSubjects; i++ {
 			if lsubj := int(readU64()); lsubj > 0 {
 				if bi+lsubj > len(buf) {
@@ -1570,6 +1571,7 @@ func (fs *fileStore) recoverFullState() (rerr error) {
 					psi.lblk = psi.fblk
 				}
 				fs.psim[subj] = psi
+				fs.tsl += len(subj)
 			}
 		}
 	}
@@ -2041,7 +2043,7 @@ func (fs *fileStore) expireMsgsOnRecover() {
 			lmb.writeTombstone(last.seq, last.ts)
 		}
 		// Clear any global subject state.
-		fs.psim = make(map[string]*psi)
+		fs.psim, fs.tsl = make(map[string]*psi), 0
 	}
 
 	// If we purged anything, make sure we kick flush state loop.
@@ -2940,6 +2942,7 @@ func (fs *fileStore) storeRawMsg(subj string, hdr, msg []byte, seq uint64, ts in
 			}
 		} else {
 			fs.psim[subj] = &psi{total: 1, fblk: index, lblk: index}
+			fs.tsl += len(subj)
 		}
 	}
 
@@ -3214,7 +3217,7 @@ func (fs *fileStore) enforceMsgPerSubjectLimit(fireCallback bool) {
 	if numMsgs != fs.state.Msgs {
 		fs.warn("Detected skew in subject-based total (%d) vs raw total (%d), rebuilding", numMsgs, fs.state.Msgs)
 		// Clear any global subject state.
-		fs.psim = make(map[string]*psi)
+		fs.psim, fs.tsl = make(map[string]*psi), 0
 		for _, mb := range fs.blks {
 			ld, _, err := mb.rebuildState()
 			if err != nil && ld != nil {
@@ -3323,6 +3326,7 @@ func (fs *fileStore) removePerSubject(subj string) {
 			info.fblk = info.lblk
 		} else if info.total == 0 {
 			delete(fs.psim, subj)
+			fs.tsl -= len(subj)
 		}
 	}
 }
@@ -4268,7 +4272,7 @@ func (fs *fileStore) checkMsgs() *LostStreamData {
 	fs.checkAndFlushAllBlocks()
 
 	// Clear any global subject state.
-	fs.psim = make(map[string]*psi)
+	fs.psim, fs.tsl = make(map[string]*psi), 0
 
 	for _, mb := range fs.blks {
 		// Make sure encryption loaded if needed for the block.
@@ -6141,7 +6145,7 @@ func (fs *fileStore) purge(fseq uint64) (uint64, error) {
 	fs.lmb = nil
 	fs.bim = make(map[uint32]*msgBlock)
 	// Clear any per subject tracking.
-	fs.psim = make(map[string]*psi)
+	fs.psim, fs.tsl = make(map[string]*psi), 0
 	// Mark dirty
 	fs.dirty++
 
@@ -6414,7 +6418,7 @@ func (fs *fileStore) reset() error {
 	fs.blks, fs.lmb = nil, nil
 
 	// Reset subject mappings.
-	fs.psim = make(map[string]*psi)
+	fs.psim, fs.tsl = make(map[string]*psi), 0
 	fs.bim = make(map[uint32]*msgBlock)
 
 	// If we purged anything, make sure we kick flush state loop.
@@ -6701,7 +6705,7 @@ func (mb *msgBlock) recalculateFirstForSubj(subj string, startSeq uint64, ss *Si
 // Lock should be held.
 func (fs *fileStore) resetGlobalPerSubjectInfo() {
 	// Clear any global subject state.
-	fs.psim = make(map[string]*psi)
+	fs.psim, fs.tsl = make(map[string]*psi), 0
 	for _, mb := range fs.blks {
 		fs.populateGlobalPerSubjectInfo(mb)
 	}
@@ -6799,6 +6803,7 @@ func (fs *fileStore) populateGlobalPerSubjectInfo(mb *msgBlock) {
 				}
 			} else {
 				fs.psim[subj] = &psi{total: ss.Msgs, fblk: mb.index, lblk: mb.index}
+				fs.tsl += len(subj)
 			}
 		}
 	}
@@ -6949,10 +6954,19 @@ func (fs *fileStore) writeFullState() error {
 		return nil
 	}
 
-	var _buf [32 * 1024]byte
-	_buf[0], _buf[1] = fullStateMagic, fullStateVersion
-	buf := _buf[:hdrLen]
+	// For calculating size.
+	numSubjects := len(fs.psim)
+
+	// Calculate and estimate of the uper bound on the  size to avoid multiple allocations.
+	sz := 2 + // Magic and Version
+		(binary.MaxVarintLen64 * 6) + // FS data
+		binary.MaxVarintLen64 + fs.tsl + // NumSubjects + total subject length
+		numSubjects*(binary.MaxVarintLen64*4) + // psi record
+		len(fs.blks)*((binary.MaxVarintLen64*6)+512) + // msg blocks, 512 is est for dmap
+		binary.MaxVarintLen64 + 8 // last index + checksum
 
+	buf := make([]byte, hdrLen, sz)
+	buf[0], buf[1] = fullStateMagic, fullStateVersion
 	buf = binary.AppendUvarint(buf, fs.state.Msgs)
 	buf = binary.AppendUvarint(buf, fs.state.Bytes)
 	buf = binary.AppendUvarint(buf, fs.state.FirstSeq)
@@ -6961,9 +6975,7 @@ func (fs *fileStore) writeFullState() error {
 	buf = binary.AppendVarint(buf, timestampNormalized(fs.state.LastTime))
 
 	// Do per subject information map if applicable.
-	numSubjects := len(fs.psim)
 	buf = binary.AppendUvarint(buf, uint64(numSubjects))
-
 	if numSubjects > 0 {
 		for subj, psi := range fs.psim {
 			buf = binary.AppendUvarint(buf, uint64(len(subj)))
@@ -6985,6 +6997,7 @@ func (fs *fileStore) writeFullState() error {
 
 	// Use basetime to save some space.
 	baseTime := timestampNormalized(fs.state.FirstTime)
+	var scratch [8 * 1024]byte
 
 	for _, mb := range fs.blks {
 		mb.mu.RLock()
@@ -6998,7 +7011,6 @@ func (fs *fileStore) writeFullState() error {
 		numDeleted := mb.dmap.Size()
 		buf = binary.AppendUvarint(buf, uint64(numDeleted))
 		if numDeleted > 0 {
-			var scratch [8 * 1024]byte
 			dmap, _ := mb.dmap.Encode(scratch[:0])
 			buf = append(buf, dmap...)
 		}
@@ -7038,6 +7050,10 @@ func (fs *fileStore) writeFullState() error {
 	// Release lock.
 	fs.mu.Unlock()
 
+	if cap(buf) > sz {
+		fs.warn("WriteFullState reallocated from %d to %d", sz, cap(buf))
+	}
+
 	// Write to a tmp file and rename.
 	const tmpPre = streamStreamStateFile + tsep
 	f, err := os.CreateTemp(filepath.Join(fs.fcfg.StoreDir, msgDir), tmpPre)
diff --git a/server/route.go b/server/route.go
index 02b4bc1f..493f03e4 100644
--- a/server/route.go
+++ b/server/route.go
@@ -1860,17 +1860,6 @@ const (
 func (s *Server) addRoute(c *client, didSolicit bool, info *Info, accName string) bool {
 	id := info.ID
 
-	var acc *Account
-	if accName != _EMPTY_ {
-		var err error
-		acc, err = s.LookupAccount(accName)
-		if err != nil {
-			c.sendErrAndErr(fmt.Sprintf("Unable to lookup account %q: %v", accName, err))
-			c.closeConnection(MissingAccount)
-			return false
-		}
-	}
-
 	s.mu.Lock()
 	if !s.isRunning() || s.routesReject {
 		s.mu.Unlock()
@@ -1946,9 +1935,6 @@ func (s *Server) addRoute(c *client, didSolicit bool, info *Info, accName string
 			if c.last.IsZero() {
 				c.last = time.Now()
 			}
-			if acc != nil {
-				c.acc = acc
-			}
 			c.mu.Unlock()
 
 			// Store this route with key being the route id hash + account name


Test patch:
diff --git a/server/filestore_test.go b/server/filestore_test.go
index 611a2e9b..cc25c761 100644
--- a/server/filestore_test.go
+++ b/server/filestore_test.go
@@ -29,6 +29,7 @@ import (
 	"os"
 	"path/filepath"
 	"reflect"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -6334,6 +6335,112 @@ func TestFileStoreCompactAndPSIMWhenDeletingBlocks(t *testing.T) {
 	require_Equal(t, psi.fblk, psi.lblk)
 }
 
+func TestFileStoreTrackSubjLenForPSIM(t *testing.T) {
+	sd := t.TempDir()
+	fs, err := newFileStore(
+		FileStoreConfig{StoreDir: sd},
+		StreamConfig{Name: "zzz", Subjects: []string{">"}, Storage: FileStorage})
+	require_NoError(t, err)
+	defer fs.Stop()
+
+	// Place 1000 msgs with varying subjects.
+	// Make sure we track the subject length properly.
+	smap := make(map[string]int, 1000)
+	buf := make([]byte, 10)
+	for i := 0; i < 1000; i++ {
+		var b strings.Builder
+		// 1-6 tokens.
+		numTokens := rand.Intn(6) + 1
+		for i := 0; i < numTokens; i++ {
+			tlen := rand.Intn(4) + 2
+			tok := buf[:tlen]
+			crand.Read(tok)
+			b.WriteString(hex.EncodeToString(tok))
+			if i != numTokens-1 {
+				b.WriteString(".")
+			}
+		}
+		subj := b.String()
+		smap[subj] = len(subj)
+		fs.StoreMsg(subj, nil, nil)
+	}
+
+	check := func() {
+		t.Helper()
+
+		var total int
+		for _, slen := range smap {
+			total += slen
+		}
+		fs.mu.RLock()
+		tsl := fs.tsl
+		fs.mu.RUnlock()
+		require_Equal(t, tsl, total)
+	}
+
+	check()
+
+	// Delete ~half
+	var smv StoreMsg
+	for i := 0; i < 500; i++ {
+		seq := uint64(rand.Intn(1000) + 1)
+		sm, err := fs.LoadMsg(seq, &smv)
+		if err != nil {
+			continue
+		}
+		fs.RemoveMsg(seq)
+		delete(smap, sm.subj)
+	}
+
+	check()
+
+	// Make sure we can recover same after restart.
+	fs.Stop()
+	fs, err = newFileStore(
+		FileStoreConfig{StoreDir: sd},
+		StreamConfig{Name: "zzz", Subjects: []string{">"}, Storage: FileStorage})
+	require_NoError(t, err)
+	defer fs.Stop()
+
+	check()
+
+	// Drain the rest through purge.
+	fs.Purge()
+	smap = nil
+	check()
+}
+
+// This was used to make sure our estimate was correct, but not needed normally.
+func TestFileStoreLargeFullStatePSIM(t *testing.T) {
+	t.Skip()
+
+	sd := t.TempDir()
+	fs, err := newFileStore(
+		FileStoreConfig{StoreDir: sd},
+		StreamConfig{Name: "zzz", Subjects: []string{">"}, Storage: FileStorage})
+	require_NoError(t, err)
+	defer fs.Stop()
+
+	buf := make([]byte, 20)
+	for i := 0; i < 100_000; i++ {
+		var b strings.Builder
+		// 1-6 tokens.
+		numTokens := rand.Intn(6) + 1
+		for i := 0; i < numTokens; i++ {
+			tlen := rand.Intn(8) + 2
+			tok := buf[:tlen]
+			crand.Read(tok)
+			b.WriteString(hex.EncodeToString(tok))
+			if i != numTokens-1 {
+				b.WriteString(".")
+			}
+		}
+		subj := b.String()
+		fs.StoreMsg(subj, nil, nil)
+	}
+	fs.Stop()
+}
+
 ///////////////////////////////////////////////////////////////////////////
 // Benchmarks
 ///////////////////////////////////////////////////////////////////////////
diff --git a/server/norace_test.go b/server/norace_test.go
index bdf6a9ed..9427c4f5 100644
--- a/server/norace_test.go
+++ b/server/norace_test.go
@@ -8023,7 +8023,7 @@ func TestNoRaceRoutePool(t *testing.T) {
 	t.Logf("Gain      : %.2fx", perf2/perf1)
 }
 
-func testNoRaceRoutePerAccount(t *testing.T, useWildCard bool) {
+func TestNoRaceRoutePerAccount(t *testing.T) {
 	var dur1 time.Duration
 	var dur2 time.Duration
 
@@ -8047,7 +8047,6 @@ func testNoRaceRoutePerAccount(t *testing.T, useWildCard bool) {
 	} {
 		t.Run(test.name, func(t *testing.T) {
 			tmpl := `
-			server_name: "%s"
 			port: -1
 			accounts {
 				%s { users: [{user: "0", password: "0"}] }
@@ -8069,13 +8068,13 @@ func testNoRaceRoutePerAccount(t *testing.T, useWildCard bool) {
 			} else {
 				racc = _EMPTY_
 			}
-			conf1 := createConfFile(t, []byte(fmt.Sprintf(tmpl, "A",
+			conf1 := createConfFile(t, []byte(fmt.Sprintf(tmpl,
 				accounts[0], accounts[1], accounts[2], accounts[3],
 				accounts[4], _EMPTY_, racc)))
 			s1, o1 := RunServerWithConfig(conf1)
 			defer s1.Shutdown()
 
-			conf2 := createConfFile(t, []byte(fmt.Sprintf(tmpl, "B",
+			conf2 := createConfFile(t, []byte(fmt.Sprintf(tmpl,
 				accounts[0], accounts[1], accounts[2], accounts[3], accounts[4],
 				fmt.Sprintf("routes: [\"nats://127.0.0.1:%d\"]", o1.Cluster.Port),
 				racc)))
@@ -8092,14 +8091,7 @@ func testNoRaceRoutePerAccount(t *testing.T, useWildCard bool) {
 
 				s2nc := natsConnect(t, s2.ClientURL(), nats.UserInfo(user, user))
 				count := 0
-				var subj string
-				var checkSubj string
-				if useWildCard {
-					subj, checkSubj = "foo.*", "foo.0"
-				} else {
-					subj, checkSubj = "foo", "foo"
-				}
-				natsSub(t, s2nc, subj, func(_ *nats.Msg) {
+				natsSub(t, s2nc, "foo", func(_ *nats.Msg) {
 					if count++; count == total {
 						wg.Done()
 					}
@@ -8108,7 +8100,7 @@ func testNoRaceRoutePerAccount(t *testing.T, useWildCard bool) {
 
 				s1nc := natsConnect(t, s1.ClientURL(), nats.UserInfo(user, user))
 
-				checkSubInterest(t, s1, acc, checkSubj, time.Second)
+				checkSubInterest(t, s1, acc, "foo", time.Second)
 				return s2nc, s1nc
 			}
 
@@ -8127,13 +8119,7 @@ func testNoRaceRoutePerAccount(t *testing.T, useWildCard bool) {
 			for i := 0; i < 5; i++ {
 				go func(idx int) {
 					for i := 0; i < total; i++ {
-						var subj string
-						if useWildCard {
-							subj = fmt.Sprintf("foo.%d", i)
-						} else {
-							subj = "foo"
-						}
-						snd[idx].Publish(subj, payload)
+						snd[idx].Publish("foo", payload)
 					}
 				}(i)
 			}
@@ -8154,14 +8140,6 @@ func testNoRaceRoutePerAccount(t *testing.T, useWildCard bool) {
 	t.Logf("Gain                  : %.2fx", perf2/perf1)
 }
 
-func TestNoRaceRoutePerAccount(t *testing.T) {
-	testNoRaceRoutePerAccount(t, false)
-}
-
-func TestNoRaceRoutePerAccountSubWithWildcard(t *testing.T) {
-	testNoRaceRoutePerAccount(t, true)
-}
-
 // This test, which checks that messages are not duplicated when pooling or
 // per-account routes are reloaded, would cause a DATA RACE that is not
 // specific to the changes for pooling/per_account. For this reason, this
diff --git a/server/parser_test.go b/server/parser_test.go
index 458d76ce..14a5f94d 100644
--- a/server/parser_test.go
+++ b/server/parser_test.go
@@ -831,10 +831,7 @@ func TestMaxControlLine(t *testing.T) {
 				c.setNoReconnect()
 				c.flags.set(connectReceived)
 				c.kind = test.kind
-				switch test.kind {
-				case ROUTER:
-					c.route = &route{}
-				case GATEWAY:
+				if test.kind == GATEWAY {
 					c.gw = &gateway{outbound: false, connected: true, insim: make(map[string]*insie)}
 				}
 				c.mcl = 8
diff --git a/server/routes_test.go b/server/routes_test.go
index 2cb62583..2798b091 100644
--- a/server/routes_test.go
+++ b/server/routes_test.go
@@ -1858,10 +1858,6 @@ func TestRoutePoolAndPerAccountErrors(t *testing.T) {
 
 	conf1 := createConfFile(t, []byte(`
 		port: -1
-		accounts {
-			abc { users: [{user:abc, password: pwd}] }
-			def { users: [{user:def, password: pwd}] }
-		}
 		cluster {
 			port: -1
 			name: "local"
@@ -1876,10 +1872,6 @@ func TestRoutePoolAndPerAccountErrors(t *testing.T) {
 
 	conf2 := createConfFile(t, []byte(fmt.Sprintf(`
 		port: -1
-		accounts {
-			abc { users: [{user:abc, password: pwd}] }
-			def { users: [{user:def, password: pwd}] }
-		}
 		cluster {
 			port: -1
 			name: "local"