[IMPROVED] Optimized calculating source sequences for sourced stream with lots of deletes. (#6461)

Added LoadPrevMsg to optimize when we walk backwards looking for source
sequences and we have lots of interior deletes.

Signed-off-by: Derek Collison <[email protected]>

Author: derekcollison

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

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) {

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)
+		}
+	}
+}

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)

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,17 +3783,21 @@ 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)
 		if len(ss) == 0 {
 			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)) {