Repository: nats-io/nats-server
Pull request #6461 (2025-02-05 22:48:53)

[IMPROVED] Optimized calculating source sequences for sourced stream with lots of deletes.
Added LoadPrevMsg to optimize when we walk backwards looking for source sequences and we have lots of interior deletes.

Signed-off-by: Derek Collison <derek@nats.io>



Patch:
diff --git a/server/filestore.go b/server/filestore.go
index 14f23e63..bfb8cdc8 100644
--- a/server/filestore.go
+++ b/server/filestore.go
@@ -7306,6 +7306,57 @@ func (fs *fileStore) LoadNextMsg(filter string, wc bool, start uint64, sm *Store
 	return nil, fs.state.LastSeq, ErrStoreEOF
 }
 
+// Will load the next non-deleted msg starting at the start sequence and walking backwards.
+func (fs *fileStore) LoadPrevMsg(start uint64, smp *StoreMsg) (sm *StoreMsg, err error) {
+	fs.mu.RLock()
+	defer fs.mu.RUnlock()
+
+	if fs.closed {
+		return nil, ErrStoreClosed
+	}
+	if fs.state.Msgs == 0 || start < fs.state.FirstSeq {
+		return nil, ErrStoreEOF
+	}
+
+	if start > fs.state.LastSeq {
+		start = fs.state.LastSeq
+	}
+	if smp == nil {
+		smp = new(StoreMsg)
+	}
+
+	if bi, _ := fs.selectMsgBlockWithIndex(start); bi >= 0 {
+		for i := bi; i >= 0; i-- {
+			mb := fs.blks[i]
+			mb.mu.Lock()
+			// Need messages loaded from here on out.
+			if mb.cacheNotLoaded() {
+				if err := mb.loadMsgsWithLock(); err != nil {
+					mb.mu.Unlock()
+					return nil, err
+				}
+			}
+
+			lseq, fseq := atomic.LoadUint64(&mb.last.seq), atomic.LoadUint64(&mb.first.seq)
+			if start > lseq {
+				start = lseq
+			}
+			for seq := start; seq >= fseq; seq-- {
+				if mb.dmap.Exists(seq) {
+					continue
+				}
+				if sm, err := mb.cacheLookup(seq, smp); err == nil {
+					mb.mu.Unlock()
+					return sm, nil
+				}
+			}
+			mb.mu.Unlock()
+		}
+	}
+
+	return nil, ErrStoreEOF
+}
+
 // Type returns the type of the underlying store.
 func (fs *fileStore) Type() StorageType {
 	return FileStorage
diff --git a/server/memstore.go b/server/memstore.go
index efccd8ca..82d2ae5c 100644
--- a/server/memstore.go
+++ b/server/memstore.go
@@ -1509,6 +1509,33 @@ func (ms *memStore) LoadNextMsg(filter string, wc bool, start uint64, smp *Store
 	return nil, ms.state.LastSeq, ErrStoreEOF
 }
 
+// Will load the next non-deleted msg starting at the start sequence and walking backwards.
+func (ms *memStore) LoadPrevMsg(start uint64, smp *StoreMsg) (sm *StoreMsg, err error) {
+	ms.mu.RLock()
+	defer ms.mu.RUnlock()
+
+	if ms.msgs == nil {
+		return nil, ErrStoreClosed
+	}
+	if ms.state.Msgs == 0 || start < ms.state.FirstSeq {
+		return nil, ErrStoreEOF
+	}
+	if start > ms.state.LastSeq {
+		start = ms.state.LastSeq
+	}
+
+	for seq := start; seq >= ms.state.FirstSeq; seq-- {
+		if sm, ok := ms.msgs[seq]; ok {
+			if smp == nil {
+				smp = new(StoreMsg)
+			}
+			sm.copy(smp)
+			return smp, nil
+		}
+	}
+	return nil, ErrStoreEOF
+}
+
 // RemoveMsg will remove the message from this store.
 // Will return the number of bytes removed.
 func (ms *memStore) RemoveMsg(seq uint64) (bool, error) {
diff --git a/server/store.go b/server/store.go
index 56252011..15b66122 100644
--- a/server/store.go
+++ b/server/store.go
@@ -96,6 +96,7 @@ type StreamStore interface {
 	LoadNextMsg(filter string, wc bool, start uint64, smp *StoreMsg) (sm *StoreMsg, skip uint64, err error)
 	LoadNextMsgMulti(sl *Sublist, start uint64, smp *StoreMsg) (sm *StoreMsg, skip uint64, err error)
 	LoadLastMsg(subject string, sm *StoreMsg) (*StoreMsg, error)
+	LoadPrevMsg(start uint64, smp *StoreMsg) (sm *StoreMsg, err error)
 	RemoveMsg(seq uint64) (bool, error)
 	EraseMsg(seq uint64) (bool, error)
 	Purge() (uint64, error)
diff --git a/server/stream.go b/server/stream.go
index 77dbcb1b..7e40dc89 100644
--- a/server/stream.go
+++ b/server/stream.go
@@ -3663,16 +3663,21 @@ func (mset *stream) setStartingSequenceForSources(iNames map[string]struct{}) {
 	}
 
 	var smv StoreMsg
-	for seq := state.LastSeq; seq >= state.FirstSeq; seq-- {
-		sm, err := mset.store.LoadMsg(seq, &smv)
-		if err != nil || len(sm.hdr) == 0 {
+	for seq := state.LastSeq; seq >= state.FirstSeq; {
+		sm, err := mset.store.LoadPrevMsg(seq, &smv)
+		if err == ErrStoreEOF || err != nil {
+			break
+		}
+		seq = sm.seq - 1
+		if len(sm.hdr) == 0 {
 			continue
 		}
+
 		ss := getHeader(JSStreamSource, sm.hdr)
 		if len(ss) == 0 {
 			continue
 		}
-		streamName, indexName, sseq := streamAndSeq(string(ss))
+		streamName, indexName, sseq := streamAndSeq(bytesToString(ss))
 
 		if _, ok := iNames[indexName]; ok {
 			si := mset.sources[indexName]
@@ -3778,9 +3783,13 @@ func (mset *stream) startingSequenceForSources() {
 	}
 
 	var smv StoreMsg
-	for seq := state.LastSeq; seq >= state.FirstSeq; seq-- {
-		sm, err := mset.store.LoadMsg(seq, &smv)
-		if err != nil || sm == nil || len(sm.hdr) == 0 {
+	for seq := state.LastSeq; ; {
+		sm, err := mset.store.LoadPrevMsg(seq, &smv)
+		if err == ErrStoreEOF || err != nil {
+			break
+		}
+		seq = sm.seq - 1
+		if len(sm.hdr) == 0 {
 			continue
 		}
 		ss := getHeader(JSStreamSource, sm.hdr)
@@ -3788,7 +3797,7 @@ func (mset *stream) startingSequenceForSources() {
 			continue
 		}
 
-		streamName, iName, sseq := streamAndSeq(string(ss))
+		streamName, iName, sseq := streamAndSeq(bytesToString(ss))
 		if iName == _EMPTY_ { // Pre-2.10 message header means it's a match for any source using that stream name
 			for _, ssi := range mset.cfg.Sources {
 				if streamName == ssi.Name || (ssi.External != nil && streamName == ssi.Name+":"+getHash(ssi.External.ApiPrefix)) {


Test patch:
diff --git a/server/norace_test.go b/server/norace_test.go
index 855016cc..76215a14 100644
--- a/server/norace_test.go
+++ b/server/norace_test.go
@@ -9014,6 +9014,7 @@ func TestNoRaceStoreStreamEncoderDecoder(t *testing.T) {
 	}
 	ms, err := newMemStore(cfg)
 	require_NoError(t, err)
+	defer ms.Stop()
 
 	fs, err := newFileStore(
 		FileStoreConfig{StoreDir: t.TempDir()},
@@ -11298,3 +11299,70 @@ func TestNoRaceJetStreamClusterLargeMetaSnapshotTiming(t *testing.T) {
 	require_NoError(t, n.InstallSnapshot(snap))
 	t.Logf("Took %v to snap meta with size of %v\n", time.Since(start), friendlyBytes(len(snap)))
 }
+
+func TestNoRaceStoreReverseWalkWithDeletesPerf(t *testing.T) {
+	cfg := StreamConfig{Name: "zzz", Subjects: []string{"foo.*"}, Storage: FileStorage}
+
+	fs, err := newFileStore(FileStoreConfig{StoreDir: t.TempDir()}, cfg)
+	require_NoError(t, err)
+	defer fs.Stop()
+
+	cfg.Storage = MemoryStorage
+	ms, err := newMemStore(&cfg)
+	require_NoError(t, err)
+	defer ms.Stop()
+
+	msg := []byte("Hello")
+
+	for _, store := range []StreamStore{fs, ms} {
+		store.StoreMsg("foo.A", nil, msg, 0)
+		for i := 0; i < 1_000_000; i++ {
+			store.StoreMsg("foo.B", nil, msg, 0)
+		}
+		store.StoreMsg("foo.C", nil, msg, 0)
+
+		var ss StreamState
+		store.FastState(&ss)
+		require_Equal(t, ss.Msgs, 1_000_002)
+
+		// Create a bunch of interior deletes.
+		p, err := store.PurgeEx("foo.B", 1, 0)
+		require_NoError(t, err)
+		require_Equal(t, p, 1_000_000)
+
+		// Now simulate a walk backwards as we currently do when searching for starting sequence numbers in sourced streams.
+		start := time.Now()
+		var smv StoreMsg
+		for seq := ss.LastSeq; seq > 0; seq-- {
+			_, err := store.LoadMsg(seq, &smv)
+			if err == errDeletedMsg || err == ErrStoreMsgNotFound {
+				continue
+			}
+			require_NoError(t, err)
+		}
+		elapsed := time.Since(start)
+
+		// Now use the optimized load prev.
+		seq, seen := ss.LastSeq, 0
+		start = time.Now()
+		for {
+			sm, err := store.LoadPrevMsg(seq, &smv)
+			if err == ErrStoreEOF {
+				break
+			}
+			require_NoError(t, err)
+			seq = sm.seq - 1
+			seen++
+		}
+		elapsedNew := time.Since(start)
+		require_Equal(t, seen, 2)
+
+		switch store.(type) {
+		case *memStore:
+			require_True(t, elapsedNew < elapsed)
+		case *fileStore:
+			// Bigger gains for filestore, 10x
+			require_True(t, elapsedNew*10 < elapsed)
+		}
+	}
+}