Optimise enforceMsgPerSubjectLimit (#6871)

When all of the state was lost and we need to rebuild and enforce a max
messages per subject limit, then the `enforceMsgPerSubjectLimit`
function could linearly scan multiple times and become very expensive if
there are lots of subjects. This refactors the function to use new
subject tree intersection logic to more quickly work out if there is any
overlap between the subjects needing attention and the subjects in a
given block, meaning we only need to scan the store once.

Signed-off-by: Neil Twigg <neil@nats.io>

Author: neilalexander

server/filestore.go

@@ -4492,11 +4492,18 @@ func (fs *fileStore) enforceMsgPerSubjectLimit(fireCallback bool) {
 	var numMsgs uint64
 
 	// collect all that are not correct.
-	needAttention := make(map[string]*psi)
+	needAttention := stree.NewSubjectTree[uint64]()
+	fblk, lblk := uint32(math.MaxUint32), uint32(0)
 	fs.psim.IterFast(func(subj []byte, psi *psi) bool {
 		numMsgs += psi.total
 		if psi.total > maxMsgsPer {
-			needAttention[string(subj)] = psi
+			needAttention.Insert(subj, psi.total)
+			if psi.fblk < fblk {
+				fblk = psi.fblk
+			}
+			if psi.lblk > lblk {
+				lblk = psi.lblk
+			}
 		}
 		return true
 	})
@@ -4517,55 +4524,71 @@ func (fs *fileStore) enforceMsgPerSubjectLimit(fireCallback bool) {
 		// Rebuild fs state too.
 		fs.rebuildStateLocked(nil)
 		// Need to redo blocks that need attention.
-		needAttention = make(map[string]*psi)
+		needAttention.Empty()
+		fblk, lblk = uint32(math.MaxUint32), uint32(0)
 		fs.psim.IterFast(func(subj []byte, psi *psi) bool {
 			if psi.total > maxMsgsPer {
-				needAttention[string(subj)] = psi
+				needAttention.Insert(subj, psi.total)
+				if psi.fblk < fblk {
+					fblk = psi.fblk
+				}
+				if psi.lblk > lblk {
+					lblk = psi.lblk
+				}
 			}
 			return true
 		})
 	}
 
+	// If nothing to do then stop.
+	if fblk == math.MaxUint32 {
+		return
+	}
+
 	// Collect all the msgBlks we alter.
 	blks := make(map[*msgBlock]struct{})
 
 	// For re-use below.
 	var sm StoreMsg
-
-	// Walk all subjects that need attention here.
-	for subj, info := range needAttention {
-		total, start, stop := info.total, info.fblk, info.lblk
-
-		for i := start; i <= stop; i++ {
-			mb := fs.bim[i]
-			if mb == nil {
-				continue
-			}
-			// Grab the ss entry for this subject in case sparse.
-			mb.mu.Lock()
-			mb.ensurePerSubjectInfoLoaded()
-			ss, ok := mb.fss.Find(stringToBytes(subj))
-			if ok && ss != nil && (ss.firstNeedsUpdate || ss.lastNeedsUpdate) {
-				mb.recalculateForSubj(subj, ss)
+	var fss *stree.SubjectTree[*SimpleState]
+	for i := fblk; i <= lblk; i++ {
+		mb := fs.bim[i]
+		if mb == nil {
+			continue
+		}
+		mb.mu.Lock()
+		mb.ensurePerSubjectInfoLoaded()
+		// It isn't safe to intersect mb.fss directly, because removeMsgViaLimits modifies it
+		// during the iteration, which can cause us to miss keys. We won't copy the entire
+		// SimpleState structs though but rather just take pointers for speed.
+		fss = fss.Empty()
+		mb.fss.IterFast(func(subject []byte, val *SimpleState) bool {
+			fss.Insert(subject, val)
+			return true
+		})
+		mb.mu.Unlock()
+		stree.LazyIntersect(needAttention, fss, func(subj []byte, total *uint64, ssptr **SimpleState) {
+			if ssptr == nil || total == nil {
+				return
 			}
-			mb.mu.Unlock()
-			if ss == nil {
-				continue
+			ss := *ssptr
+			if ss.firstNeedsUpdate || ss.lastNeedsUpdate {
+				mb.mu.Lock()
+				mb.recalculateForSubj(bytesToString(subj), ss)
+				mb.mu.Unlock()
 			}
-			for seq := ss.First; seq <= ss.Last && total > maxMsgsPer; {
-				m, _, err := mb.firstMatching(subj, false, seq, &sm)
-				if err == nil {
-					seq = m.seq + 1
-					if removed, _ := fs.removeMsgViaLimits(m.seq); removed {
-						total--
-						blks[mb] = struct{}{}
-					}
-				} else {
-					// On error just do single increment.
-					seq++
+			for first := ss.First; *total > maxMsgsPer && first <= ss.Last; {
+				m, _, err := mb.firstMatching(bytesToString(subj), false, first, &sm)
+				if err != nil {
+					break
+				}
+				first = m.seq + 1
+				if removed, _ := fs.removeMsgViaLimits(m.seq); removed {
+					blks[mb] = struct{}{}
+					*total--
 				}
 			}
-		}
+		})
 	}
 
 	// Expire the cache if we can.
@@ -8946,7 +8969,7 @@ func (mb *msgBlock) generatePerSubjectInfo() error {
 	}
 
 	// Create new one regardless.
-	mb.fss = stree.NewSubjectTree[SimpleState]()
+	mb.fss = mb.fss.Empty()
 
 	var smv StoreMsg
 	fseq, lseq := atomic.LoadUint64(&mb.first.seq), atomic.LoadUint64(&mb.last.seq)

server/jetstream_cluster_long_test.go

@@ -1143,3 +1143,50 @@ func TestLongClusterJetStreamRestartThenScaleStreamReplicas(t *testing.T) {
 	cancel()
 	wg.Wait()
 }
+
+func TestLongFileStoreEnforceMsgPerSubjectLimit(t *testing.T) {
+	td := t.TempDir()
+	fs, err := newFileStore(
+		FileStoreConfig{StoreDir: td, BlockSize: 1024},
+		StreamConfig{
+			Name: "zzz", Subjects: []string{"test.>"}, Storage: FileStorage,
+		},
+	)
+	require_NoError(t, err)
+	defer fs.Stop()
+
+	t.Logf("Starting publishes")
+	for i := 0; i < 100_000; i++ {
+		_, _, err := fs.StoreMsg(fmt.Sprintf("test.%06d", i), nil, nil, 0)
+		require_NoError(t, err)
+	}
+	// Now update some of them. Leave a bit of a mess with some big gaps.
+	for i := 0; i < 5_000_000; i++ {
+		n := rand.Int31n(100_000)
+		if n < 5000 {
+			continue
+		}
+		_, _, err := fs.StoreMsg(fmt.Sprintf("test.%06d", n), nil, nil, 0)
+		require_NoError(t, err)
+	}
+	t.Logf("Publish complete")
+
+	require_NoError(t, fs.Stop())
+	fs, err = newFileStore(
+		FileStoreConfig{StoreDir: td, BlockSize: 1024},
+		StreamConfig{
+			Name: "zzz", Subjects: []string{"test.>"}, Storage: FileStorage,
+			MaxMsgsPer: 1,
+		},
+	)
+	require_NoError(t, err)
+	defer fs.Stop()
+
+	// Mangle the filestore state and then see how long it takes to enforce
+	// the per-subject limit.
+	fs.state.Msgs++
+	start := time.Now()
+	fs.enforceMsgPerSubjectLimit(false)
+	require_LessThan(t, time.Since(start), time.Minute)
+	t.Logf("Took %s", time.Since(start))
+}

server/stree/stree.go

@@ -418,3 +418,29 @@ func (t *SubjectTree[T]) iter(n node, pre []byte, ordered bool, cb func(subject
 	}
 	return true
 }
+
+// LazyIntersect iterates the smaller of the two provided subject trees and
+// looks for matching entries in the other. It is lazy in that it does not
+// aggressively optimize against repeated walks, but is considerably faster
+// in most cases than intersecting against a potentially large sublist.
+func LazyIntersect[TL, TR any](tl *SubjectTree[TL], tr *SubjectTree[TR], cb func([]byte, *TL, *TR)) {
+	if tl.root == nil || tr.root == nil {
+		return
+	}
+	// Iterate over the smaller tree to reduce the number of rounds.
+	if tl.Size() <= tr.Size() {
+		tl.IterFast(func(key []byte, v1 *TL) bool {
+			if v2, ok := tr.Find(key); ok {
+				cb(key, v1, v2)
+			}
+			return true
+		})
+	} else {
+		tr.IterFast(func(key []byte, v2 *TR) bool {
+			if v1, ok := tl.Find(key); ok {
+				cb(key, v1, v2)
+			}
+			return true
+		})
+	}
+}

server/stree/stree_test.go

@@ -923,3 +923,36 @@ func TestSubjectTreeMatchHasFWCNoPanic(t *testing.T) {
 	st.Insert(subj, 1)
 	st.Match([]byte("."), func(subject []byte, val *int) {})
 }
+
+func TestSubjectTreeLazyIntersect(t *testing.T) {
+	st1 := NewSubjectTree[int]()
+	st2 := NewSubjectTree[int]()
+
+	// Should cause an intersection.
+	st1.Insert([]byte("foo.bar"), 1)
+	st2.Insert([]byte("foo.bar"), 1)
+
+	// Should cause an intersection.
+	st1.Insert([]byte("foo.bar.baz.qux"), 1)
+	st2.Insert([]byte("foo.bar.baz.qux"), 1)
+
+	// Should not cause any intersections.
+	st1.Insert([]byte("bar"), 1)
+	st2.Insert([]byte("baz"), 1)
+	st1.Insert([]byte("a.b.c"), 1)
+	st2.Insert([]byte("a.b.d"), 1)
+	st1.Insert([]byte("a.b.ee"), 1)
+	st2.Insert([]byte("a.b.e"), 1)
+	st1.Insert([]byte("bb.c.d"), 1)
+	st2.Insert([]byte("b.c.d"), 1)
+	st2.Insert([]byte("foo.bar.baz.qux.alice"), 1)
+	st2.Insert([]byte("foo.bar.baz.qux.bob"), 1)
+
+	intersected := map[string]int{}
+	LazyIntersect(st1, st2, func(key []byte, val1, val2 *int) {
+		intersected[string(key)]++
+	})
+	require_Equal(t, len(intersected), 2)
+	require_Equal(t, intersected["foo.bar"], 1)
+	require_Equal(t, intersected["foo.bar.baz.qux"], 1)
+}