Repository: nats-io/nats-server
Pull request #2960 (2022-03-27 19:46:17)

Additional improvements to memory pooling and management. 
During contention to the head write block, the system could perform worse memory wise compared to simple go runtime and GC. Also had some references for the subject of messages being read back in to the underlying block which was bloating memory.

Additionally a fix for firstMatching() that did unnecessary work when matching all. This would cause excessive CPU and excessive load on the garbage collector.

Signed-off-by: Derek Collison <derek@nats.io>

/cc @nats-io/core



Patch:
diff --git a/server/filestore.go b/server/filestore.go
index 57ce7f05..1f00d45e 100644
--- a/server/filestore.go
+++ b/server/filestore.go
@@ -32,6 +32,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"sort"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -216,25 +217,25 @@ const (
 	blkKeySize  = 72
 
 	// Default stream block size.
-	defaultStreamBlockSize = 16 * 1024 * 1024 // 16MB
+	defaultLargeBlockSize = 16 * 1024 * 1024 // 16MB
 	// Default for workqueue or interest based.
-	defaultOtherBlockSize = 8 * 1024 * 1024 // 8MB
+	defaultMediumBlockSize = 8 * 1024 * 1024 // 8MB
 	// Default for KV based
 	defaultKVBlockSize = 8 * 1024 * 1024 // 8MB
-	// Default for minimum recycle size.
-	defaultMinRecycleBlockSize = 4 * 1024 * 1024 // 4MB
+	// For smaller reuse buffers. Usually being generated during contention on the lead write buffer.
+	// E.g. mirrors/sources etc.
+	defaultSmallBlockSize = 2 * 1024 * 1024 // 2MB
 	// max block size for now.
-	maxBlockSize = defaultStreamBlockSize
+	maxBlockSize = defaultLargeBlockSize
 	// Compact minimum threshold.
 	compactMinimum = 2 * 1024 * 1024 // 2MB
 	// FileStoreMinBlkSize is minimum size we will do for a blk size.
 	FileStoreMinBlkSize = 32 * 1000 // 32kib
 	// FileStoreMaxBlkSize is maximum size we will do for a blk size.
 	FileStoreMaxBlkSize = maxBlockSize
+
 	// Check for bad record length value due to corrupt data.
 	rlBadThresh = 32 * 1024 * 1024
-	// Time threshold to write index info.
-	wiThresh = int64(2 * time.Second)
 )
 
 func newFileStore(fcfg FileStoreConfig, cfg StreamConfig) (*fileStore, error) {
@@ -396,10 +397,10 @@ func dynBlkSize(retention RetentionPolicy, maxBytes int64) uint64 {
 
 	if retention == LimitsPolicy {
 		// TODO(dlc) - Make the blocksize relative to this if set.
-		return defaultStreamBlockSize
+		return defaultLargeBlockSize
 	} else {
 		// TODO(dlc) - Make the blocksize relative to this if set.
-		return defaultOtherBlockSize
+		return defaultMediumBlockSize
 	}
 }
 
@@ -485,20 +486,30 @@ func (fs *fileStore) writeStreamMeta() error {
 }
 
 // Pools to recycle the blocks to help with memory pressure.
-var blkPoolBig sync.Pool
-var blkPoolSmall sync.Pool
+var blkPoolBig sync.Pool    // 16MB
+var blkPoolMedium sync.Pool // 8MB
+var blkPoolSmall sync.Pool  // 2MB
 
 // Get a new msg block based on sz estimate.
 func getMsgBlockBuf(sz int) (buf []byte) {
 	var pb interface{}
-	if sz <= defaultOtherBlockSize {
-		pb, sz = blkPoolSmall.Get(), defaultOtherBlockSize
+	if sz <= defaultSmallBlockSize {
+		pb = blkPoolSmall.Get()
+	} else if sz <= defaultMediumBlockSize {
+		pb = blkPoolMedium.Get()
 	} else {
-		pb, sz = blkPoolBig.Get(), maxBlockSize
+		pb = blkPoolBig.Get()
 	}
 	if pb != nil {
 		buf = *(pb.(*[]byte))
 	} else {
+		// Here we need to make a new blk.
+		// If small leave as is..
+		if sz > defaultSmallBlockSize && sz <= defaultMediumBlockSize {
+			sz = defaultMediumBlockSize
+		} else if sz > defaultMediumBlockSize {
+			sz = defaultLargeBlockSize
+		}
 		buf = make([]byte, sz)
 	}
 	return buf[:0]
@@ -506,18 +517,22 @@ func getMsgBlockBuf(sz int) (buf []byte) {
 
 // Recycle the msg block.
 func recycleMsgBlockBuf(buf []byte) {
-	if buf == nil {
-		return
-	}
-	if cap(buf) < defaultMinRecycleBlockSize {
+	if buf == nil || cap(buf) < defaultSmallBlockSize {
 		return
 	}
-
+	// Make sure to reset before placing back into pool.
 	buf = buf[:0]
-	if cap(buf) > defaultOtherBlockSize {
-		blkPoolBig.Put(&buf)
-	} else {
+
+	// We need to make sure the load code gets a block that can fit the maximum for a size block.
+	// E.g. 8, 16 etc. otherwise we thrash and actually make things worse by pulling it out, and putting
+	// it right back in and making a new []byte.
+	// From above we know its already >= defaultSmallBlockSize
+	if sz := cap(buf); sz < defaultMediumBlockSize {
 		blkPoolSmall.Put(&buf)
+	} else if sz < defaultLargeBlockSize {
+		blkPoolMedium.Put(&buf)
+	} else {
+		blkPoolBig.Put(&buf)
 	}
 }
 
@@ -865,14 +880,16 @@ func (mb *msgBlock) rebuildStateLocked() (*LostStreamData, error) {
 			mb.bytes += uint64(rl)
 
 			// Do per subject info.
-			if mb.fss != nil {
-				if subj := string(data[:slen]); len(subj) > 0 {
-					if ss := mb.fss[subj]; ss != nil {
-						ss.Msgs++
-						ss.Last = seq
-					} else {
-						mb.fss[subj] = &SimpleState{Msgs: 1, First: seq, Last: seq}
-					}
+			if slen > 0 && mb.fss != nil {
+				// For the lookup, we cast the byte slice and there won't be any copy
+				if ss := mb.fss[string(data[:slen])]; ss != nil {
+					ss.Msgs++
+					ss.Last = seq
+				} else {
+					// This will either use a subject from the config, or make a copy
+					// so we don't reference the underlying buffer.
+					subj := mb.subjString(data[:slen])
+					mb.fss[subj] = &SimpleState{Msgs: 1, First: seq, Last: seq}
 				}
 			}
 		}
@@ -1063,7 +1080,7 @@ func (fs *fileStore) expireMsgsOnRecover() {
 			purged++
 			// Update fss
 			fs.removePerSubject(sm.subj)
-			mb.removeSeqPerSubject(sm.subj, seq)
+			mb.removeSeqPerSubject(sm.subj, seq, &smv)
 		}
 
 		// Check if empty after processing, could happen if tail of messages are all deleted.
@@ -1157,28 +1174,32 @@ func (mb *msgBlock) firstMatching(filter string, wc bool, start uint64, sm *Stor
 	mb.mu.Lock()
 	defer mb.mu.Unlock()
 
-	isAll, subs := filter == _EMPTY_ || filter == fwcs, []string{filter}
-	// If we have a wildcard match against all tracked subjects we know about.
-	if wc || isAll {
-		subs = subs[:0]
-		for subj := range mb.fss {
-			if isAll || subjectIsSubsetMatch(subj, filter) {
-				subs = append(subs, subj)
+	fseq, isAll, subs := start, filter == _EMPTY_ || filter == fwcs, []string{filter}
+
+	if !isAll {
+		// If we have a wildcard match against all tracked subjects we know about.
+		if wc {
+			subs = subs[:0]
+			for subj := range mb.fss {
+				if isAll || subjectIsSubsetMatch(subj, filter) {
+					subs = append(subs, subj)
+				}
 			}
 		}
-	}
-	fseq := mb.last.seq + 1
-	for _, subj := range subs {
-		ss := mb.fss[subj]
-		if ss == nil || start > ss.Last || ss.First >= fseq {
-			continue
-		}
-		if ss.First < start {
-			fseq = start
-		} else {
-			fseq = ss.First
+		fseq = mb.last.seq + 1
+		for _, subj := range subs {
+			ss := mb.fss[subj]
+			if ss == nil || start > ss.Last || ss.First >= fseq {
+				continue
+			}
+			if ss.First < start {
+				fseq = start
+			} else {
+				fseq = ss.First
+			}
 		}
 	}
+
 	if fseq > mb.last.seq {
 		return nil, false, ErrStoreMsgNotFound
 	}
@@ -1200,7 +1221,7 @@ func (mb *msgBlock) firstMatching(filter string, wc bool, start uint64, sm *Stor
 			continue
 		}
 		expireOk := seq == mb.last.seq && mb.llseq == seq
-		if len(subs) == 1 && fsm.subj == subs[0] {
+		if isAll || len(subs) == 1 && fsm.subj == subs[0] {
 			return fsm, expireOk, nil
 		}
 		for _, subj := range subs {
@@ -1744,30 +1765,26 @@ func (mb *msgBlock) skipMsg(seq uint64, now time.Time) {
 	}
 	var needsRecord bool
 
-	nowts := now.UnixNano()
-
 	mb.mu.Lock()
 	// If we are empty can just do meta.
 	if mb.msgs == 0 {
 		mb.last.seq = seq
-		mb.last.ts = nowts
+		mb.last.ts = now.UnixNano()
 		mb.first.seq = seq + 1
-		mb.first.ts = nowts
-		// Take care of index if needed.
-		if nowts-mb.lwits > wiThresh {
-			mb.writeIndexInfoLocked()
-		}
+		mb.first.ts = now.UnixNano()
 	} else {
 		needsRecord = true
 		if mb.dmap == nil {
 			mb.dmap = make(map[uint64]struct{})
 		}
 		mb.dmap[seq] = struct{}{}
+		mb.msgs--
+		mb.bytes -= emptyRecordLen
 	}
 	mb.mu.Unlock()
 
 	if needsRecord {
-		mb.writeMsgRecord(emptyRecordLen, seq|ebit, _EMPTY_, nil, nil, nowts, true)
+		mb.writeMsgRecord(emptyRecordLen, seq|ebit, _EMPTY_, nil, nil, now.UnixNano(), true)
 	} else {
 		mb.kickFlusher()
 	}
@@ -1778,14 +1795,18 @@ func (fs *fileStore) SkipMsg() uint64 {
 	fs.mu.Lock()
 	defer fs.mu.Unlock()
 
-	// Grab time and last seq.
-	now, seq := time.Now().UTC(), fs.state.LastSeq+1
-	fs.state.LastSeq, fs.state.LastTime = seq, now
+	// Grab time.
+	now := time.Now().UTC()
+	seq := fs.state.LastSeq + 1
+	fs.state.LastSeq = seq
+	fs.state.LastTime = now
 	if fs.state.Msgs == 0 {
-		fs.state.FirstSeq, fs.state.FirstTime = seq, now
+		fs.state.FirstSeq = seq
+		fs.state.FirstTime = now
 	}
 	if seq == fs.state.FirstSeq {
-		fs.state.FirstSeq, fs.state.FirstTime = seq+1, now
+		fs.state.FirstSeq = seq + 1
+		fs.state.FirstTime = now
 	}
 	fs.lmb.skipMsg(seq, now)
 
@@ -1967,7 +1988,7 @@ func (fs *fileStore) removeMsg(seq uint64, secure, needFSLock bool) (bool, error
 
 	// If we are tracking multiple subjects here make sure we update that accounting.
 	fs.removePerSubject(sm.subj)
-	mb.removeSeqPerSubject(sm.subj, seq)
+	mb.removeSeqPerSubject(sm.subj, seq, &smv)
 
 	var shouldWriteIndex, firstSeqNeedsUpdate bool
 
@@ -2024,7 +2045,7 @@ func (fs *fileStore) removeMsg(seq uint64, secure, needFSLock bool) (bool, error
 	// Check if we need to write the index file and we are flush in place (fip).
 	if shouldWriteIndex && fs.fip {
 		// Check if this is the first message, common during expirations etc.
-		if !fifo || time.Now().UnixNano()-mb.lwits > wiThresh {
+		if !fifo || time.Now().UnixNano()-mb.lwits > int64(2*time.Second) {
 			mb.writeIndexInfoLocked()
 		}
 	}
@@ -2048,7 +2069,11 @@ func (fs *fileStore) removeMsg(seq uint64, secure, needFSLock bool) (bool, error
 	// we don't lose track of the first sequence.
 	if firstSeqNeedsUpdate {
 		fs.selectNextFirst()
-		fs.lmb.writeIndexInfo()
+		// Write out the new first message block if we have one.
+		if len(fs.blks) > 0 {
+			fmb := fs.blks[0]
+			fmb.writeIndexInfo()
+		}
 	}
 	fs.mu.Unlock()
 
@@ -2538,16 +2563,16 @@ func (mb *msgBlock) clearCache() {
 		return
 	}
 
+	buf := mb.cache.buf
 	if mb.cache.off == 0 {
-		recycleMsgBlockBuf(mb.cache.buf)
 		mb.cache = nil
 	} else {
 		// Clear msgs and index.
-		recycleMsgBlockBuf(mb.cache.buf)
 		mb.cache.buf = nil
 		mb.cache.idx = nil
 		mb.cache.wp = 0
 	}
+	recycleMsgBlockBuf(buf)
 }
 
 // Called to possibly expire a message block cache.
@@ -2825,10 +2850,10 @@ func (mb *msgBlock) writeMsgRecord(rl, seq uint64, subj string, mhdr, msg []byte
 	// Set cache timestamp for last store.
 	mb.lwts = ts
 	// Decide if we write index info if flushing in place.
-	writeIndex := ts-mb.lwits > wiThresh
+	writeIndex := ts-mb.lwits > int64(2*time.Second)
 
 	// Accounting
-	mb.updateAccounting(seq, ts, rl)
+	mb.updateAccounting(seq&^ebit, ts, rl)
 
 	// Check if we are tracking per subject for our simple state.
 	if len(subj) > 0 && mb.fss != nil {
@@ -2931,11 +2956,6 @@ func (mb *msgBlock) blkSize() uint64 {
 // Update accounting on a write msg.
 // Lock should be held.
 func (mb *msgBlock) updateAccounting(seq uint64, ts int64, rl uint64) {
-	isDeleted := seq&ebit != 0
-	if isDeleted {
-		seq = seq &^ ebit
-	}
-
 	if mb.first.seq == 0 || mb.first.ts == 0 {
 		mb.first.seq = seq
 		mb.first.ts = ts
@@ -2943,12 +2963,9 @@ func (mb *msgBlock) updateAccounting(seq uint64, ts int64, rl uint64) {
 	// Need atomics here for selectMsgBlock speed.
 	atomic.StoreUint64(&mb.last.seq, seq)
 	mb.last.ts = ts
+	mb.bytes += rl
 	mb.rbytes += rl
-	// Only update this accounting if message is not a deleted message.
-	if !isDeleted {
-		mb.bytes += rl
-		mb.msgs++
-	}
+	mb.msgs++
 }
 
 // Lock should be held.
@@ -3233,6 +3250,7 @@ func (mb *msgBlock) flushPendingMsgsLocked() (*LostStreamData, error) {
 		if cap(mb.cache.buf) <= maxBufReuse {
 			buf = mb.cache.buf[:0]
 		} else {
+			recycleMsgBlockBuf(mb.cache.buf)
 			buf = nil
 		}
 		if moreBytes > 0 {
@@ -3496,7 +3514,7 @@ func (mb *msgBlock) cacheLookup(seq uint64, sm *StoreMsg) (*StoreMsg, error) {
 	}
 
 	// Parse from the raw buffer.
-	fsm, err := msgFromBuf(buf, sm, hh)
+	fsm, err := mb.msgFromBuf(buf, sm, hh)
 	if err != nil || fsm == nil {
 		return nil, err
 	}
@@ -3567,7 +3585,8 @@ func (fs *fileStore) msgForSeq(seq uint64, sm *StoreMsg) (*StoreMsg, error) {
 }
 
 // Internal function to return msg parts from a raw buffer.
-func msgFromBuf(buf []byte, sm *StoreMsg, hh hash.Hash64) (*StoreMsg, error) {
+// Lock should be held.
+func (mb *msgBlock) msgFromBuf(buf []byte, sm *StoreMsg, hh hash.Hash64) (*StoreMsg, error) {
 	if len(buf) < emptyRecordLen {
 		return nil, errBadMsg
 	}
@@ -3626,10 +3645,43 @@ func msgFromBuf(buf []byte, sm *StoreMsg, hh hash.Hash64) (*StoreMsg, error) {
 		sm.buf = append(sm.buf, data[slen:end]...)
 		sm.msg = sm.buf[0 : end-slen]
 	}
-	sm.subj, sm.seq, sm.ts = string(data[:slen]), seq, ts
+	sm.seq, sm.ts = seq, ts
+	// Treat subject a bit different to not reference underlying buf.
+	if slen > 0 {
+		sm.subj = mb.subjString(data[:slen])
+	}
+
 	return sm, nil
 }
 
+// Given the `key` byte slice, this function will return the subject
+// as a copy of `key` or a configured subject as to minimize memory allocations.
+// Lock should be held.
+func (mb *msgBlock) subjString(key []byte) string {
+	if len(key) == 0 {
+		return _EMPTY_
+	}
+
+	if lsubjs := len(mb.fs.cfg.Subjects); lsubjs > 0 {
+		if lsubjs == 1 {
+			// The cast for the comparison does not make a copy
+			if string(key) == mb.fs.cfg.Subjects[0] {
+				return mb.fs.cfg.Subjects[0]
+			}
+		} else {
+			for _, subj := range mb.fs.cfg.Subjects {
+				if string(key) == subj {
+					return subj
+				}
+			}
+		}
+	}
+	// Copy here to not reference underlying buffer.
+	var sb strings.Builder
+	sb.Write(key)
+	return sb.String()
+}
+
 // LoadMsg will lookup the message by sequence number and return it if found.
 func (fs *fileStore) LoadMsg(seq uint64, sm *StoreMsg) (*StoreMsg, error) {
 	return fs.msgForSeq(seq, sm)
@@ -4076,7 +4128,7 @@ func (fs *fileStore) PurgeEx(subject string, sequence, keep uint64) (purged uint
 				mb.bytes -= rl
 				// FSS updates.
 				fs.removePerSubject(sm.subj)
-				mb.removeSeqPerSubject(sm.subj, seq)
+				mb.removeSeqPerSubject(sm.subj, seq, &smv)
 				// Check for first message.
 				if seq == mb.first.seq {
 					mb.selectNextFirst()
@@ -4251,7 +4303,7 @@ func (fs *fileStore) Compact(seq uint64) (uint64, error) {
 			purged++
 			// Update fss
 			fs.removePerSubject(sm.subj)
-			smb.removeSeqPerSubject(sm.subj, mseq)
+			smb.removeSeqPerSubject(sm.subj, mseq, &smv)
 		}
 	}
 
@@ -4467,7 +4519,7 @@ func (mb *msgBlock) dirtyCloseWithRemove(remove bool) {
 
 // Remove a seq from the fss and select new first.
 // Lock should be held.
-func (mb *msgBlock) removeSeqPerSubject(subj string, seq uint64) {
+func (mb *msgBlock) removeSeqPerSubject(subj string, seq uint64, smp *StoreMsg) {
 	ss := mb.fss[subj]
 	if ss == nil {
 		return
@@ -4491,8 +4543,11 @@ func (mb *msgBlock) removeSeqPerSubject(subj string, seq uint64) {
 
 	// TODO(dlc) - Might want to optimize this.
 	var smv StoreMsg
+	if smp == nil {
+		smp = &smv
+	}
 	for tseq := seq + 1; tseq <= ss.Last; tseq++ {
-		if sm, _ := mb.cacheLookup(tseq, &smv); sm != nil {
+		if sm, _ := mb.cacheLookup(tseq, smp); sm != nil {
 			if sm.subj == subj {
 				ss.First = tseq
 				return
@@ -4592,14 +4647,15 @@ func (mb *msgBlock) readPerSubjectInfo() error {
 		return num
 	}
 
+	mb.mu.Lock()
 	for i, numEntries := uint64(0), readU64(); i < numEntries; i++ {
 		lsubj := readU64()
-		subj := buf[bi : bi+int(lsubj)]
+		// Make a copy or use a configured subject (to avoid mem allocation)
+		subj := mb.subjString(buf[bi : bi+int(lsubj)])
 		bi += int(lsubj)
 		msgs, first, last := readU64(), readU64(), readU64()
-		fss[string(subj)] = &SimpleState{Msgs: msgs, First: first, Last: last}
+		fss[subj] = &SimpleState{Msgs: msgs, First: first, Last: last}
 	}
-	mb.mu.Lock()
 	mb.fss = fss
 	mb.mu.Unlock()
 	return nil


Test patch:
diff --git a/server/filestore_test.go b/server/filestore_test.go
index 3017f848..494980fb 100644
--- a/server/filestore_test.go
+++ b/server/filestore_test.go
@@ -1265,8 +1265,12 @@ func TestFileStoreEraseMsg(t *testing.T) {
 	defer fp.Close()
 
 	fp.ReadAt(buf, 0)
-	//	nsubj, _, nmsg, seq, ts, err := msgFromBuf(buf, nil)
-	sm, err = msgFromBuf(buf, nil, nil)
+	fs.mu.RLock()
+	mb := fs.blks[0]
+	fs.mu.RUnlock()
+	mb.mu.Lock()
+	sm, err = mb.msgFromBuf(buf, nil, nil)
+	mb.mu.Unlock()
 	if err != nil {
 		t.Fatalf("error reading message from block: %v", err)
 	}
@@ -2296,7 +2300,7 @@ func TestFileStoreReadBackMsgPerf(t *testing.T) {
 	storedMsgSize := fileStoreMsgSize(subj, nil, msg)
 
 	// Make sure we store 2 blocks.
-	toStore := defaultStreamBlockSize * 2 / storedMsgSize
+	toStore := defaultLargeBlockSize * 2 / storedMsgSize
 
 	fmt.Printf("storing %d msgs of %s each, totalling %s\n",
 		toStore,
@@ -2327,7 +2331,7 @@ func TestFileStoreReadBackMsgPerf(t *testing.T) {
 
 	// We should not have cached here with no reads.
 	// Pick something towards end of the block.
-	index := defaultStreamBlockSize/storedMsgSize - 22
+	index := defaultLargeBlockSize/storedMsgSize - 22
 	start = time.Now()
 	fs.LoadMsg(index, nil)
 	fmt.Printf("Time to load first msg [%d] = %v\n", index, time.Since(start))
diff --git a/server/jetstream_cluster_test.go b/server/jetstream_cluster_test.go
index 8d0f8239..fc95a21d 100644
--- a/server/jetstream_cluster_test.go
+++ b/server/jetstream_cluster_test.go
@@ -12269,24 +12269,16 @@ func checkSubsPending(t *testing.T, sub *nats.Subscription, numExpected int) {
 	})
 }
 
-func fetchMsgs(t *testing.T, sub *nats.Subscription, numExpected int, totalWait time.Duration) []*nats.Msg {
+func fetchMsgs(t *testing.T, sub *nats.Subscription, numExpected int, wait time.Duration) []*nats.Msg {
 	t.Helper()
-	result := make([]*nats.Msg, 0, numExpected)
-	for start, count, wait := time.Now(), numExpected, totalWait; len(result) != numExpected; {
-		msgs, err := sub.Fetch(count, nats.MaxWait(wait))
-		if err != nil {
-			t.Fatal(err)
-		}
-		result = append(result, msgs...)
-		count -= len(msgs)
-		if wait = totalWait - time.Since(start); wait < 0 {
-			break
-		}
+	msgs, err := sub.Fetch(numExpected, nats.MaxWait(wait))
+	if err != nil {
+		t.Fatal(err)
 	}
-	if len(result) != numExpected {
-		t.Fatalf("Unexpected msg count, got %d, want %d", len(result), numExpected)
+	if len(msgs) != numExpected {
+		t.Fatalf("Unexpected msg count, got %d, want %d", len(msgs), numExpected)
 	}
-	return result
+	return msgs
 }
 
 func (c *cluster) restartServer(rs *Server) *Server {
diff --git a/server/jetstream_test.go b/server/jetstream_test.go
index 34898819..e52d7508 100644
--- a/server/jetstream_test.go
+++ b/server/jetstream_test.go
@@ -14921,7 +14921,7 @@ func TestJetStreamPullConsumersOneShotBehavior(t *testing.T) {
 	}
 	expectInstant := func(batch int, expires time.Duration, noWait bool, ne int, he *nats.Header) {
 		t.Helper()
-		expect(batch, expires, noWait, ne, he, 5*time.Millisecond, 0)
+		expect(batch, expires, noWait, ne, he, 2*time.Millisecond, 0)
 	}
 	expectOK := func(batch int, expires time.Duration, noWait bool, ne int) {
 		t.Helper()
@@ -16079,62 +16079,6 @@ func TestJetStreamConsumerAckSampling(t *testing.T) {
 	}
 }
 
-func TestJetStreamConsumerAckSamplingSpecifiedUsingUpdateConsumer(t *testing.T) {
-	t.Skip("sampling msg is not sent when sampling option is part of update flow")
-
-	s := RunBasicJetStreamServer()
-	if config := s.JetStreamConfig(); config != nil {
-		defer removeDir(t, config.StoreDir)
-	}
-	defer s.Shutdown()
-
-	nc, js := jsClientConnect(t, s)
-	defer nc.Close()
-
-	_, err := js.AddStream(&nats.StreamConfig{Name: "TEST", Subjects: []string{"foo"}})
-	require_NoError(t, err)
-
-	_, err = js.AddConsumer("TEST", &nats.ConsumerConfig{
-		Durable:       "dlc",
-		AckPolicy:     nats.AckExplicitPolicy,
-		FilterSubject: "foo",
-	})
-	require_NoError(t, err)
-
-	_, err = js.UpdateConsumer("TEST", &nats.ConsumerConfig{
-		Durable:         "dlc",
-		AckPolicy:       nats.AckExplicitPolicy,
-		FilterSubject:   "foo",
-		SampleFrequency: "100%",
-	})
-	require_NoError(t, err)
-
-	sub, err := js.PullSubscribe("foo", "dlc")
-	require_NoError(t, err)
-
-	_, err = js.Publish("foo", []byte("Hello"))
-	require_NoError(t, err)
-
-	msub, err := nc.SubscribeSync("$JS.EVENT.METRIC.>")
-	require_NoError(t, err)
-
-	for _, m := range fetchMsgs(t, sub, 1, time.Second) {
-		err = m.AckSync()
-		require_NoError(t, err)
-	}
-
-	m, err := msub.NextMsg(time.Second)
-	require_NoError(t, err)
-
-	var am JSConsumerAckMetric
-	err = json.Unmarshal(m.Data, &am)
-	require_NoError(t, err)
-
-	if am.Stream != "TEST" || am.Consumer != "dlc" || am.ConsumerSeq != 1 {
-		t.Fatalf("Not a proper ack metric: %+v", am)
-	}
-}
-
 func TestJetStreamRemoveExternalSource(t *testing.T) {
 	ho := DefaultTestOptions
 	ho.Port = 4000 //-1
@@ -16643,8 +16587,8 @@ func TestJetStreamConsumerStreamUpdate(t *testing.T) {
 			listen: 127.0.0.1:-1
 			server_name: %s
 			jetstream: {
-				max_mem_store: 2MB,
-				max_file_store: 8MB,
+				max_mem_store: 2MB, 
+				max_file_store: 8MB, 
 				store_dir: '%s',
 			}
 			cluster {
diff --git a/server/leafnode_test.go b/server/leafnode_test.go
index 1e47d93c..ec508faa 100644
--- a/server/leafnode_test.go
+++ b/server/leafnode_test.go
@@ -2860,7 +2860,7 @@ func TestLeafNodeWSBasic(t *testing.T) {
 
 			nc2 := natsConnect(t, ln.ClientURL())
 			defer nc2.Close()
-			msg1Payload := make([]byte, 2048)
+			msg1Payload := make([]byte, 512)
 			for i := 0; i < len(msg1Payload); i++ {
 				msg1Payload[i] = 'A'
 			}
@@ -2876,7 +2876,7 @@ func TestLeafNodeWSBasic(t *testing.T) {
 
 			checkSubInterest(t, s, globalAccountName, "bar", time.Second)
 
-			msg2Payload := make([]byte, 2048)
+			msg2Payload := make([]byte, 512)
 			for i := 0; i < len(msg2Payload); i++ {
 				msg2Payload[i] = 'B'
 			}
@@ -2893,10 +2893,10 @@ func TestLeafNodeWSBasic(t *testing.T) {
 				trackSizeConn.Unlock()
 
 				if test.acceptCompression && test.remoteCompression {
-					if size >= 1024 {
+					if size >= 100 {
 						t.Fatalf("Seems that there was no compression: size=%v", size)
 					}
-				} else if size < 2048 {
+				} else if size < 500 {
 					t.Fatalf("Seems compression was on while it should not: size=%v", size)
 				}
 			}
diff --git a/server/mqtt_test.go b/server/mqtt_test.go
index bc6e6db3..7cabba2b 100644
--- a/server/mqtt_test.go
+++ b/server/mqtt_test.go
@@ -4674,9 +4674,6 @@ func TestMQTTMaxAckPending(t *testing.T) {
 	s := testMQTTRunServer(t, o)
 	defer testMQTTShutdownRestartedServer(&s)
 
-	nc, js := jsClientConnect(t, s)
-	defer nc.Close()
-
 	dir := strings.TrimSuffix(s.JetStreamConfig().StoreDir, JetStreamStoreDir)
 
 	cisub := &mqttConnInfo{clientID: "sub", cleanSess: false}
@@ -4703,15 +4700,8 @@ func TestMQTTMaxAckPending(t *testing.T) {
 	testMQTTCheckPubMsg(t, c, r, "foo", mqttPubQos1, []byte("msg2"))
 	testMQTTDisconnect(t, c, nil)
 
-	// Give a chance to the server to "close" the consumer
-	checkFor(t, 2*time.Second, 15*time.Millisecond, func() error {
-		for ci := range js.ConsumersInfo("$MQTT_msgs") {
-			if ci.PushBound {
-				return fmt.Errorf("Consumer still connected")
-			}
-		}
-		return nil
-	})
+	// Give a chance to the server to register that this client is gone.
+	checkClientsCount(t, s, 1)
 
 	// Send 2 messages while sub is offline
 	testMQTTPublish(t, cp, rp, 1, false, false, "foo", 1, []byte("msg3"))
diff --git a/server/norace_test.go b/server/norace_test.go
index 600eb6d0..739dcca0 100644
--- a/server/norace_test.go
+++ b/server/norace_test.go
@@ -3449,7 +3449,9 @@ func TestNoRaceJetStreamClusterCorruptWAL(t *testing.T) {
 	}
 
 	// This will cause us to stepdown and truncate our WAL.
-	sub.Fetch(100)
+	if _, err = sub.Fetch(100); err != nil {
+		t.Fatal(err)
+	}
 
 	checkFor(t, 20*time.Second, 500*time.Millisecond, func() error {
 		// Make sure we changed leaders.
@@ -4522,3 +4524,56 @@ func TestNoRaceRebuildDeDupeAndMemoryPerf(t *testing.T) {
 	v, _ = s.Varz(nil)
 	fmt.Printf("Memory: %v\n", friendlyBytes(v.Mem))
 }
+
+func TestNoRaceMemoryUsageOnLimitedStreamWithMirror(t *testing.T) {
+	skip(t)
+
+	s := RunBasicJetStreamServer()
+	if config := s.JetStreamConfig(); config != nil {
+		defer removeDir(t, config.StoreDir)
+	}
+	defer s.Shutdown()
+
+	nc, js := jsClientConnect(t, s)
+	defer nc.Close()
+
+	_, err := js.AddStream(&nats.StreamConfig{Name: "DD", Subjects: []string{"ORDERS.*"}, MaxMsgs: 10_000})
+	require_NoError(t, err)
+
+	_, err = js.AddStream(&nats.StreamConfig{
+		Name:    "M",
+		Mirror:  &nats.StreamSource{Name: "DD"},
+		MaxMsgs: 10_000,
+	})
+	require_NoError(t, err)
+
+	m := nats.NewMsg("ORDERS.0")
+	m.Data = []byte(strings.Repeat("Z", 2048))
+
+	start := time.Now()
+
+	n := 1_000_000
+	for i := 0; i < n; i++ {
+		m.Subject = fmt.Sprintf("ORDERS.%d", i)
+		m.Header.Set(JSMsgId, strconv.Itoa(i))
+		_, err := js.PublishMsgAsync(m)
+		require_NoError(t, err)
+	}
+
+	select {
+	case <-js.PublishAsyncComplete():
+	case <-time.After(20 * time.Second):
+		t.Fatalf("Did not receive completion signal")
+	}
+
+	tt := time.Since(start)
+	si, err := js.StreamInfo("DD")
+	require_NoError(t, err)
+
+	fmt.Printf("Took %v to send %d msgs\n", tt, n)
+	fmt.Printf("%.0f msgs/s\n", float64(n)/tt.Seconds())
+	fmt.Printf("%.0f mb/s\n\n", float64(si.State.Bytes/(1024*1024))/tt.Seconds())
+
+	v, _ := s.Varz(nil)
+	fmt.Printf("Memory AFTER SEND: %v\n", friendlyBytes(v.Mem))
+}