Repository: nats-io/nats-server

Issue:
Batch publish - support isolated reads
### Observed behavior

We currently don't support isolated reads for batch publish, meaning once you batch publish 10 messages any readers through message GET requests or consumers will observe these messages as they are applied sequentially and not as one atomic action. Due to that a reader could observe parts of the batch's effects, even if the full batch was committed.

Likely this is already an issue when using `Nats-Rollup`, since the rollup happens separately after persisting the message that contains the rollup. Which would mean you could observe the full stream including the rollup message in a read, even if the rollup should have purged data and not allowed you to see the previous messages within the same read.

### Expected behavior

Support isolated reads in the server, guaranteeing that readers are also "aware" of the batch semantics and either observe none of a batch's effects, or observe it in its entirety after commit.

### Server and client version

`main`

### Host environment

_No response_

### Steps to reproduce

_No response_

Discussion:

Solution patch:
diff --git a/server/consumer.go b/server/consumer.go
index 64d00a31..65309830 100644
--- a/server/consumer.go
+++ b/server/consumer.go
@@ -4438,9 +4438,6 @@ func (o *consumer) getNextMsg() (*jsPubMsg, uint64, error) {
 		return pmsg, 1, err
 	}
 
-	// Hold onto this since we release the lock.
-	store := o.mset.store
-
 	var sseq uint64
 	var err error
 	var sm *StoreMsg
@@ -4450,13 +4447,13 @@ func (o *consumer) getNextMsg() (*jsPubMsg, uint64, error) {
 	filters, subjf, fseq := o.filters, o.subjf, o.sseq
 	// Check if we are multi-filtered or not.
 	if filters != nil {
-		sm, sseq, err = store.LoadNextMsgMulti(filters, fseq, &pmsg.StoreMsg)
+		sm, sseq, err = o.mset.store.LoadNextMsgMulti(filters, fseq, &pmsg.StoreMsg)
 	} else if len(subjf) > 0 { // Means single filtered subject since o.filters means > 1.
 		filter, wc := subjf[0].subject, subjf[0].hasWildcard
-		sm, sseq, err = store.LoadNextMsg(filter, wc, fseq, &pmsg.StoreMsg)
+		sm, sseq, err = o.mset.store.LoadNextMsg(filter, wc, fseq, &pmsg.StoreMsg)
 	} else {
 		// No filter here.
-		sm, sseq, err = store.LoadNextMsg(_EMPTY_, false, fseq, &pmsg.StoreMsg)
+		sm, sseq, err = o.mset.store.LoadNextMsg(_EMPTY_, false, fseq, &pmsg.StoreMsg)
 	}
 	if sm == nil {
 		pmsg.returnToPool()
diff --git a/server/jetstream.go b/server/jetstream.go
index 21458ab1..35ec097a 100644
--- a/server/jetstream.go
+++ b/server/jetstream.go
@@ -1499,7 +1499,7 @@ func (a *Account) EnableJetStream(limits map[string]JetStreamAccountLimits) erro
 						mset.accName(), mset.name(), seq, err)
 					break
 				}
-				mset.processJetStreamMsg(sm.subj, _EMPTY_, sm.hdr, sm.msg, 0, 0, nil, false)
+				mset.processJetStreamMsg(sm.subj, _EMPTY_, sm.hdr, sm.msg, 0, 0, nil, false, true)
 			}
 			store.Delete(true)
 		SKIP:
diff --git a/server/jetstream_api.go b/server/jetstream_api.go
index dbf0e6b1..4235802b 100644
--- a/server/jetstream_api.go
+++ b/server/jetstream_api.go
@@ -3633,6 +3633,10 @@ func (s *Server) jsMsgGetRequest(sub *subscription, c *client, _ *Account, subje
 	var svp StoreMsg
 	var sm *StoreMsg
 
+	// Ensure this read request is isolated and doesn't interleave with writes.
+	mset.mu.RLock()
+	defer mset.mu.RUnlock()
+
 	// If AsOfTime is set, perform this first to get the sequence.
 	var seq uint64
 	if req.StartTime != nil {
diff --git a/server/jetstream_cluster.go b/server/jetstream_cluster.go
index 52bac861..bbd9c928 100644
--- a/server/jetstream_cluster.go
+++ b/server/jetstream_cluster.go
@@ -3251,12 +3251,15 @@ func (js *jetStream) applyStreamEntries(mset *stream, ce *CommittedEntry, isReco
 				if err != nil {
 					panic(err.Error())
 				}
+
+				// Ensure the whole batch is fully isolated, and reads
+				// can only happen after the full batch is committed.
+				mset.mu.Lock()
+
 				// Initialize if unset.
 				if batch == nil {
 					batch = &batchApply{}
-					mset.mu.Lock()
 					mset.batchApply = batch
-					mset.mu.Unlock()
 				}
 
 				batch.mu.Lock()
@@ -3271,6 +3274,7 @@ func (js *jetStream) applyStreamEntries(mset *stream, ce *CommittedEntry, isReco
 				if batchSeq != batch.count {
 					batch.rejectBatchStateLocked(mset)
 					batch.mu.Unlock()
+					mset.mu.Unlock()
 					continue
 				}
 
@@ -3287,10 +3291,12 @@ func (js *jetStream) applyStreamEntries(mset *stream, ce *CommittedEntry, isReco
 						_, _, op, buf, err = decodeBatchMsg(entry.Data[1:])
 						if err != nil {
 							batch.mu.Unlock()
+							mset.mu.Unlock()
 							panic(err.Error())
 						}
-						if err = js.applyStreamMsgOp(mset, op, buf, isRecovering); err != nil {
+						if err = js.applyStreamMsgOp(mset, op, buf, isRecovering, false); err != nil {
 							batch.mu.Unlock()
+							mset.mu.Unlock()
 							// Make sure to return remaining entries to the pool on an error.
 							for _, nce := range batch.entries[j:] {
 								nce.ReturnToPool()
@@ -3313,16 +3319,19 @@ func (js *jetStream) applyStreamEntries(mset *stream, ce *CommittedEntry, isReco
 					_, _, op, buf, err = decodeBatchMsg(entry.Data[1:])
 					if err != nil {
 						batch.mu.Unlock()
+						mset.mu.Unlock()
 						panic(err.Error())
 					}
-					if err = js.applyStreamMsgOp(mset, op, buf, isRecovering); err != nil {
+					if err = js.applyStreamMsgOp(mset, op, buf, isRecovering, false); err != nil {
 						batch.mu.Unlock()
+						mset.mu.Unlock()
 						return 0, err
 					}
 				}
 				// Clear state, batch was successful.
 				batch.clearBatchStateLocked()
 				batch.mu.Unlock()
+				mset.mu.Unlock()
 				continue
 			} else if batch != nil && batch.id != _EMPTY_ {
 				// If a batch is abandoned without a commit, reject it.
@@ -3332,7 +3341,7 @@ func (js *jetStream) applyStreamEntries(mset *stream, ce *CommittedEntry, isReco
 			switch op {
 			case streamMsgOp, compressedStreamMsgOp:
 				mbuf := buf[1:]
-				if err := js.applyStreamMsgOp(mset, op, mbuf, isRecovering); err != nil {
+				if err := js.applyStreamMsgOp(mset, op, mbuf, isRecovering, true); err != nil {
 					return 0, err
 				}
 
@@ -3515,7 +3524,7 @@ func (js *jetStream) applyStreamEntries(mset *stream, ce *CommittedEntry, isReco
 	return 0, nil
 }
 
-func (js *jetStream) applyStreamMsgOp(mset *stream, op entryOp, mbuf []byte, isRecovering bool) error {
+func (js *jetStream) applyStreamMsgOp(mset *stream, op entryOp, mbuf []byte, isRecovering bool, needLock bool) error {
 	s := js.srv
 
 	if op == compressedStreamMsgOp {
@@ -3528,6 +3537,11 @@ func (js *jetStream) applyStreamMsgOp(mset *stream, op entryOp, mbuf []byte, isR
 
 	subject, reply, hdr, msg, lseq, ts, sourced, err := decodeStreamMsg(mbuf)
 	if err != nil {
+		// We're going to panic below, but if we're already holding the stream lock, we should let go now.
+		// Otherwise we'll deadlock when trying to get the raft node.
+		if !needLock {
+			mset.mu.Unlock()
+		}
 		if node := mset.raftNode(); node != nil {
 			s.Errorf("JetStream cluster could not decode stream msg for '%s > %s' [%s]",
 				mset.account(), mset.name(), node.Group())
@@ -3538,22 +3552,38 @@ func (js *jetStream) applyStreamMsgOp(mset *stream, op entryOp, mbuf []byte, isR
 	// Check for flowcontrol here.
 	if len(msg) == 0 && len(hdr) > 0 && reply != _EMPTY_ && isControlHdr(hdr) {
 		if !isRecovering {
+			if needLock {
+				mset.mu.RLock()
+			}
 			mset.sendFlowControlReply(reply)
+			if needLock {
+				mset.mu.RUnlock()
+			}
 		}
 		return nil
 	}
 
+	if needLock {
+		mset.mu.RLock()
+	}
 	// Grab last sequence and CLFS.
 	last, clfs := mset.lastSeqAndCLFS()
+	if needLock {
+		mset.mu.RUnlock()
+	}
 
 	// We can skip if we know this is less than what we already have.
 	if lseq-clfs < last {
 		s.Debugf("Apply stream entries for '%s > %s' skipping message with sequence %d with last of %d",
-			mset.account(), mset.name(), lseq+1-clfs, last)
-		mset.mu.Lock()
+			mset.accountLocked(needLock), mset.nameLocked(needLock), lseq+1-clfs, last)
+		if needLock {
+			mset.mu.Lock()
+		}
 		// Check for any preAcks in case we are interest based.
 		mset.clearAllPreAcks(lseq + 1 - clfs)
-		mset.mu.Unlock()
+		if needLock {
+			mset.mu.Unlock()
+		}
 		return nil
 	}
 
@@ -3568,10 +3598,14 @@ func (js *jetStream) applyStreamMsgOp(mset *stream, op entryOp, mbuf []byte, isR
 	if subject == _EMPTY_ && ts == 0 && len(msg) == 0 && len(hdr) == 0 {
 		// Skip and update our lseq.
 		last := mset.store.SkipMsg()
-		mset.mu.Lock()
+		if needLock {
+			mset.mu.Lock()
+		}
 		mset.setLastSeq(last)
 		mset.clearAllPreAcks(last)
-		mset.mu.Unlock()
+		if needLock {
+			mset.mu.Unlock()
+		}
 		return nil
 	}
 
@@ -3583,7 +3617,7 @@ func (js *jetStream) applyStreamMsgOp(mset *stream, op entryOp, mbuf []byte, isR
 		mt = mset.getAndDeleteMsgTrace(lseq)
 	}
 	// Process the actual message here.
-	err = mset.processJetStreamMsg(subject, reply, hdr, msg, lseq, ts, mt, sourced)
+	err = mset.processJetStreamMsg(subject, reply, hdr, msg, lseq, ts, mt, sourced, needLock)
 
 	// If we have inflight make sure to clear after processing.
 	// TODO(dlc) - technically check on inflight != nil could cause datarace.
@@ -3650,7 +3684,7 @@ func (js *jetStream) applyStreamMsgOp(mset *stream, op entryOp, mbuf []byte, isR
 			if state.Msgs == 0 {
 				mset.store.Compact(lseq + 1)
 				// Retry
-				err = mset.processJetStreamMsg(subject, reply, hdr, msg, lseq, ts, mt, sourced)
+				err = mset.processJetStreamMsg(subject, reply, hdr, msg, lseq, ts, mt, sourced, needLock)
 			}
 			// FIXME(dlc) - We could just run a catchup with a request defining the span between what we expected
 			// and what we got.
@@ -3661,7 +3695,7 @@ func (js *jetStream) applyStreamMsgOp(mset *stream, op entryOp, mbuf []byte, isR
 			return err
 		}
 		s.Debugf("Apply stream entries for '%s > %s' got error processing message: %v",
-			mset.account(), mset.name(), err)
+			mset.accountLocked(needLock), mset.nameLocked(needLock), err)
 	}
 	return nil
 }
@@ -8707,7 +8741,7 @@ func (mset *stream) processClusteredInboundMsg(subject, reply string, hdr, msg [
 	// We also invoke this in clustering mode for message tracing when not
 	// performing message delivery.
 	if node == nil || mt.traceOnly() {
-		return mset.processJetStreamMsg(subject, reply, hdr, msg, 0, 0, mt, sourced)
+		return mset.processJetStreamMsg(subject, reply, hdr, msg, 0, 0, mt, sourced, true)
 	}
 
 	// If message tracing (with message delivery), we will need to send the
diff --git a/server/stream.go b/server/stream.go
index 86c6f96f..d6b4dd3e 100644
--- a/server/stream.go
+++ b/server/stream.go
@@ -1118,9 +1118,17 @@ func (mset *stream) stopClusterSubs() {
 
 // account gets the account for this stream.
 func (mset *stream) account() *Account {
-	mset.mu.RLock()
+	return mset.accountLocked(true)
+}
+
+func (mset *stream) accountLocked(needLock bool) *Account {
+	if needLock {
+		mset.mu.RLock()
+	}
 	jsa := mset.jsa
-	mset.mu.RUnlock()
+	if needLock {
+		mset.mu.RUnlock()
+	}
 	if jsa == nil {
 		return nil
 	}
@@ -1229,8 +1237,9 @@ func (mset *stream) rebuildDedupe() {
 	}
 }
 
+// Lock should be held.
 func (mset *stream) lastSeqAndCLFS() (uint64, uint64) {
-	return mset.lastSeq(), mset.getCLFS()
+	return mset.lseq, mset.getCLFS()
 }
 
 func (mset *stream) getCLFS() uint64 {
@@ -2412,17 +2421,21 @@ func (mset *stream) getCfgName() string {
 
 // Purge will remove all messages from the stream and underlying store based on the request.
 func (mset *stream) purge(preq *JSApiStreamPurgeRequest) (purged uint64, err error) {
-	mset.mu.RLock()
+	return mset.purgeLocked(preq, true)
+}
+
+func (mset *stream) purgeLocked(preq *JSApiStreamPurgeRequest, needLock bool) (purged uint64, err error) {
+	if needLock {
+		mset.mu.Lock()
+		defer mset.mu.Unlock()
+	}
 	if mset.closed.Load() {
-		mset.mu.RUnlock()
 		return 0, errStreamClosed
 	}
 	if mset.cfg.Sealed {
-		mset.mu.RUnlock()
 		return 0, errors.New("sealed stream")
 	}
 	store, mlseq := mset.store, mset.lseq
-	mset.mu.RUnlock()
 
 	if preq != nil {
 		purged, err = mset.store.PurgeEx(preq.Subject, preq.Sequence, preq.Keep)
@@ -2438,7 +2451,6 @@ func (mset *stream) purge(preq *JSApiStreamPurgeRequest) (purged uint64, err err
 	store.FastState(&state)
 	fseq, lseq := state.FirstSeq, state.LastSeq
 
-	mset.mu.Lock()
 	// Check if our last has moved past what our original last sequence was, if so reset.
 	if lseq > mlseq {
 		mset.setLastSeq(lseq)
@@ -2446,7 +2458,6 @@ func (mset *stream) purge(preq *JSApiStreamPurgeRequest) (purged uint64, err err
 
 	// Clear any pending acks below first seq.
 	mset.clearAllPreAcksBelowFloor(fseq)
-	mset.mu.Unlock()
 
 	// Purge consumers.
 	// Check for filtered purge.
@@ -2822,7 +2833,7 @@ func (mset *stream) processInboundMirrorMsg(m *inMsg) bool {
 			err = node.Propose(encodeStreamMsg(m.subj, _EMPTY_, m.hdr, m.msg, sseq-1, ts, true))
 		}
 	} else {
-		err = mset.processJetStreamMsg(m.subj, _EMPTY_, m.hdr, m.msg, sseq-1, ts, nil, true)
+		err = mset.processJetStreamMsg(m.subj, _EMPTY_, m.hdr, m.msg, sseq-1, ts, nil, true, true)
 	}
 	if err != nil {
 		if strings.Contains(err.Error(), "no space left") {
@@ -3670,12 +3681,11 @@ func (m *inMsg) isControlMsg() bool {
 }
 
 // Sends a reply to a flow control request.
+// Lock should be held.
 func (mset *stream) sendFlowControlReply(reply string) {
-	mset.mu.RLock()
 	if mset.isLeader() && mset.outq != nil {
 		mset.outq.sendMsg(reply, nil)
 	}
-	mset.mu.RUnlock()
 }
 
 // handleFlowControl will properly handle flow control messages for both R==1 and R>1.
@@ -3793,7 +3803,7 @@ func (mset *stream) processInboundSourceMsg(si *sourceInfo, m *inMsg) bool {
 	if node != nil {
 		err = mset.processClusteredInboundMsg(m.subj, _EMPTY_, hdr, msg, nil, true)
 	} else {
-		err = mset.processJetStreamMsg(m.subj, _EMPTY_, hdr, msg, 0, 0, nil, true)
+		err = mset.processJetStreamMsg(m.subj, _EMPTY_, hdr, msg, 0, 0, nil, true, true)
 	}
 
 	if err != nil {
@@ -4444,7 +4454,7 @@ func (mset *stream) setupStore(fsCfg *FileStoreConfig) error {
 				mset.processClusteredInboundMsg(im.subj, im.rply, im.hdr, im.msg, im.mt, false)
 			}
 		} else {
-			mset.processJetStreamMsg(im.subj, im.rply, im.hdr, im.msg, 0, 0, im.mt, false)
+			mset.processJetStreamMsg(im.subj, im.rply, im.hdr, im.msg, 0, 0, im.mt, false, true)
 		}
 	})
 	mset.mu.Unlock()
@@ -4956,9 +4966,10 @@ func (mset *stream) getDirectMulti(req *JSApiMsgGetRequest, reply string) {
 	// TODO(dlc) - Make configurable?
 	const maxAllowedResponses = 1024
 
-	// We hold the lock here to try to avoid changes out from underneath of us.
+	// Ensure this read request is isolated and doesn't interleave with writes.
 	mset.mu.RLock()
 	defer mset.mu.RUnlock()
+
 	// Grab store and name.
 	store, name, s := mset.store, mset.cfg.Name, mset.srv
 
@@ -5008,7 +5019,7 @@ func (mset *stream) getDirectMulti(req *JSApiMsgGetRequest, reply string) {
 		return
 	}
 
-	np, lseq, sentBytes, sent := uint64(len(seqs)-1), uint64(0), 0, 0
+	np, lseq, sentBytes, sent := uint64(len(seqs)), uint64(0), 0, 0
 	for _, seq := range seqs {
 		if seq < req.Seq {
 			if np > 0 {
@@ -5027,6 +5038,10 @@ func (mset *stream) getDirectMulti(req *JSApiMsgGetRequest, reply string) {
 		hdr := sm.hdr
 		ts := time.Unix(0, sm.ts).UTC()
 
+		// Decrement num pending. This is an optimization, and we do not continue to look it up for these operations.
+		if np > 0 {
+			np--
+		}
 		if len(hdr) == 0 {
 			hdr = fmt.Appendf(nil, dgb, name, sm.subj, sm.seq, ts.Format(time.RFC3339Nano), np, lseq)
 		} else {
@@ -5038,10 +5053,6 @@ func (mset *stream) getDirectMulti(req *JSApiMsgGetRequest, reply string) {
 			hdr = genHeader(hdr, JSNumPending, strconv.FormatUint(np, 10))
 			hdr = genHeader(hdr, JSLastSequence, strconv.FormatUint(lseq, 10))
 		}
-		// Decrement num pending. This is optimization and we do not continue to look it up for these operations.
-		if np > 0 {
-			np--
-		}
 		// Track our lseq
 		lseq = sm.seq
 		// Send out our message.
@@ -5071,9 +5082,11 @@ func (mset *stream) getDirectRequest(req *JSApiMsgGetRequest, reply string) {
 		return
 	}
 
+	// Ensure this read request is isolated and doesn't interleave with writes.
 	mset.mu.RLock()
+	defer mset.mu.RUnlock()
+
 	store, name, s := mset.store, mset.cfg.Name, mset.srv
-	mset.mu.RUnlock()
 
 	var seq uint64
 	// Lookup start seq if AsOfTime is set.
@@ -5144,6 +5157,10 @@ func (mset *stream) getDirectRequest(req *JSApiMsgGetRequest, reply string) {
 		if !req.NoHeaders {
 			hdr = sm.hdr
 			if isBatchRequest {
+				// Decrement num pending. This is an optimization, and we do not continue to look it up for these operations.
+				if np > 0 {
+					np--
+				}
 				if len(hdr) == 0 {
 					hdr = fmt.Appendf(nil, dgb, name, sm.subj, sm.seq, ts.Format(time.RFC3339Nano), np, lseq)
 				} else {
@@ -5155,8 +5172,6 @@ func (mset *stream) getDirectRequest(req *JSApiMsgGetRequest, reply string) {
 					hdr = genHeader(hdr, JSNumPending, strconv.FormatUint(np, 10))
 					hdr = genHeader(hdr, JSLastSequence, strconv.FormatUint(lseq, 10))
 				}
-				// Decrement num pending. This is optimization and we do not continue to look it up for these operations.
-				np--
 			} else {
 				if len(hdr) == 0 {
 					hdr = fmt.Appendf(nil, dg, name, sm.subj, sm.seq, ts.Format(time.RFC3339Nano))
@@ -5220,7 +5235,7 @@ var (
 )
 
 // processJetStreamMsg is where we try to actually process the stream msg.
-func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte, lseq uint64, ts int64, mt *msgTrace, sourced bool) (retErr error) {
+func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte, lseq uint64, ts int64, mt *msgTrace, sourced bool, needLock bool) (retErr error) {
 	if mt != nil {
 		// Only the leader/standalone will have mt!=nil. On exit, send the
 		// message trace event.
@@ -5229,11 +5244,17 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 		}()
 	}
 
-	if mset.closed.Load() {
+	if mset == nil || mset.closed.Load() {
 		return errStreamClosed
 	}
 
-	mset.mu.Lock()
+	// Hold lock while storing the message, and potentially purging as part of a rollup.
+	// If we're writing an atomic batch of multiple messages, the lock is already held.
+	if needLock {
+		mset.mu.Lock()
+		defer mset.mu.Unlock()
+	}
+
 	s, store := mset.srv, mset.store
 
 	traceOnly := mt.traceOnly()
@@ -5290,7 +5311,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 
 	// Bail here if sealed.
 	if canConsistencyCheck && isSealed {
-		mset.mu.Unlock()
 		if canRespond && outq != nil {
 			resp.PubAck = &PubAck{Stream: name}
 			resp.Error = ApiErrors[JSStreamSealedErr]
@@ -5305,7 +5325,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 
 	// If this is a non-clustered msg and we are not considered active, meaning no active subscription, do not process.
 	if lseq == 0 && ts == 0 && !mset.active {
-		mset.mu.Unlock()
 		return nil
 	}
 
@@ -5328,7 +5347,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 			}
 			// Really is a mismatch.
 			if isMisMatch {
-				mset.mu.Unlock()
 				if canRespond && outq != nil {
 					resp.PubAck = &PubAck{Stream: name}
 					resp.Error = ApiErrors[JSStreamSequenceNotMatchErr]
@@ -5360,7 +5378,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 			// Only supported on counter streams, and payload must be empty (if not coming from a source).
 			var ok bool
 			if incr, ok = getMessageIncr(hdr); !ok {
-				mset.mu.Unlock()
 				apiErr := NewJSMessageIncrInvalidError()
 				if canRespond {
 					resp.PubAck = &PubAck{Stream: name}
@@ -5372,7 +5389,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 			} else if incr != nil && !sourced {
 				// Only do checks if the message isn't sourced. Otherwise, we need to store verbatim.
 				if !allowMsgCounter {
-					mset.mu.Unlock()
 					apiErr := NewJSMessageIncrDisabledError()
 					if canRespond {
 						resp.PubAck = &PubAck{Stream: name}
@@ -5382,7 +5398,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 					}
 					return apiErr
 				} else if len(msg) > 0 {
-					mset.mu.Unlock()
 					apiErr := NewJSMessageIncrPayloadError()
 					if canRespond {
 						resp.PubAck = &PubAck{Stream: name}
@@ -5406,7 +5421,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 					}
 
 					if doErr {
-						mset.mu.Unlock()
 						apiErr := NewJSMessageIncrInvalidError()
 						if canRespond {
 							resp.PubAck = &PubAck{Stream: name}
@@ -5421,7 +5435,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 
 			// Expected stream.
 			if sname := getExpectedStream(hdr); sname != _EMPTY_ && sname != name {
-				mset.mu.Unlock()
 				if canRespond {
 					resp.PubAck = &PubAck{Stream: name}
 					resp.Error = NewJSStreamNotMatchError()
@@ -5435,7 +5448,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 			// Shouldn't happen in clustered mode since we should have already caught this
 			// in processClusteredInboundMsg, but needed here for non-clustered etc.
 			if ttl, _ := getMessageTTL(hdr); !sourced && ttl != 0 && !mset.cfg.AllowMsgTTL {
-				mset.mu.Unlock()
 				if canRespond {
 					resp.PubAck = &PubAck{Stream: name}
 					resp.Error = NewJSMessageTTLDisabledError()
@@ -5464,7 +5476,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 					fseq, err = 0, nil
 				}
 				if err != nil || fseq != seq {
-					mset.mu.Unlock()
 					if canRespond {
 						resp.PubAck = &PubAck{Stream: name}
 						resp.Error = NewJSStreamWrongLastSequenceError(fseq)
@@ -5474,7 +5485,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 					return fmt.Errorf("last sequence by subject mismatch: %d vs %d", seq, fseq)
 				}
 			} else if getExpectedLastSeqPerSubjectForSubject(hdr) != _EMPTY_ {
-				mset.mu.Unlock()
 				apiErr := NewJSStreamExpectedLastSeqPerSubjectInvalidError()
 				if canRespond {
 					resp.PubAck = &PubAck{Stream: name}
@@ -5488,7 +5498,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 			// Expected last sequence.
 			if seq, exists := getExpectedLastSeq(hdr); exists && seq != mset.lseq {
 				mlseq := mset.lseq
-				mset.mu.Unlock()
 				if canRespond {
 					resp.PubAck = &PubAck{Stream: name}
 					resp.Error = NewJSStreamWrongLastSequenceError(mlseq)
@@ -5500,7 +5509,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 
 			// Message scheduling.
 			if schedule, ok := getMessageSchedule(hdr); !ok {
-				mset.mu.Unlock()
 				apiErr := NewJSMessageSchedulesPatternInvalidError()
 				if !allowMsgSchedules {
 					apiErr = NewJSMessageSchedulesDisabledError()
@@ -5514,7 +5522,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 				return apiErr
 			} else if !schedule.IsZero() {
 				if !allowMsgSchedules {
-					mset.mu.Unlock()
 					apiErr := NewJSMessageSchedulesDisabledError()
 					if canRespond {
 						resp.PubAck = &PubAck{Stream: name}
@@ -5524,7 +5531,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 					}
 					return apiErr
 				} else if scheduleTtl, ok := getMessageScheduleTTL(hdr); !ok {
-					mset.mu.Unlock()
 					apiErr := NewJSMessageSchedulesTTLInvalidError()
 					if canRespond {
 						resp.PubAck = &PubAck{Stream: name}
@@ -5534,7 +5540,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 					}
 					return apiErr
 				} else if scheduleTtl != _EMPTY_ && !mset.cfg.AllowMsgTTL {
-					mset.mu.Unlock()
 					if canRespond {
 						resp.PubAck = &PubAck{Stream: name}
 						resp.Error = NewJSMessageTTLDisabledError()
@@ -5544,7 +5549,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 					return errMsgTTLDisabled
 				} else if scheduleTarget := getMessageScheduleTarget(hdr); scheduleTarget == _EMPTY_ ||
 					!IsValidPublishSubject(scheduleTarget) || SubjectsCollide(scheduleTarget, subject) {
-					mset.mu.Unlock()
 					apiErr := NewJSMessageSchedulesTargetInvalidError()
 					if canRespond {
 						resp.PubAck = &PubAck{Stream: name}
@@ -5558,7 +5562,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 						return SubjectsCollide(subj, scheduleTarget)
 					})
 					if !match {
-						mset.mu.Unlock()
 						apiErr := NewJSMessageSchedulesTargetInvalidError()
 						if canRespond {
 							resp.PubAck = &PubAck{Stream: name}
@@ -5574,7 +5577,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 					if rollup := getRollup(hdr); rollup == _EMPTY_ {
 						hdr = genHeader(hdr, JSMsgRollup, JSMsgRollupSubject)
 					} else if rollup != JSMsgRollupSubject {
-						mset.mu.Unlock()
 						apiErr := NewJSMessageSchedulesRollupInvalidError()
 						if canRespond {
 							resp.PubAck = &PubAck{Stream: name}
@@ -5601,7 +5603,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 				}
 				mset.ddMu.Unlock()
 				if seq > 0 {
-					mset.mu.Unlock()
 					if canRespond {
 						response := append(pubAck, strconv.FormatUint(dde.seq, 10)...)
 						response = append(response, ",\"duplicate\": true}"...)
@@ -5616,7 +5617,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 		if lmsgId := getExpectedLastMsgId(hdr); lmsgId != _EMPTY_ {
 			if lmsgId != mset.lmsgId {
 				last := mset.lmsgId
-				mset.mu.Unlock()
 				bumpCLFS()
 				if canRespond {
 					resp.PubAck = &PubAck{Stream: name}
@@ -5630,7 +5630,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 		// Check for any rollups.
 		if rollup := getRollup(hdr); rollup != _EMPTY_ {
 			if canConsistencyCheck && (!mset.cfg.AllowRollup || mset.cfg.DenyPurge) {
-				mset.mu.Unlock()
 				err := errors.New("rollup not permitted")
 				if canRespond {
 					resp.PubAck = &PubAck{Stream: name}
@@ -5647,7 +5646,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 				rollupAll = true
 			default:
 				if canConsistencyCheck {
-					mset.mu.Unlock()
 					err := fmt.Errorf("rollup value invalid: %q", rollup)
 					if canRespond {
 						resp.PubAck = &PubAck{Stream: name}
@@ -5662,7 +5660,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 	}
 
 	if canConsistencyCheck && incr == nil && allowMsgCounter {
-		mset.mu.Unlock()
 		apiErr := NewJSMessageIncrMissingError()
 		if canRespond {
 			resp.PubAck = &PubAck{Stream: name}
@@ -5691,7 +5688,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 			var val CounterValue
 			// Return an error if the counter is broken somehow.
 			if json.Unmarshal(sm.msg, &val) != nil {
-				mset.mu.Unlock()
 				apiErr := NewJSMessageCounterBrokenError()
 				if canRespond {
 					resp.PubAck = &PubAck{Stream: name}
@@ -5703,7 +5699,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 			}
 			if ncs := sliceHeader(JSMessageCounterSources, sm.hdr); len(ncs) > 0 {
 				if err := json.Unmarshal(ncs, &sources); err != nil {
-					mset.mu.Unlock()
 					apiErr := NewJSMessageCounterBrokenError()
 					if canRespond {
 						resp.PubAck = &PubAck{Stream: name}
@@ -5728,7 +5723,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 			}
 			var val CounterValue
 			if json.Unmarshal(msg, &val) != nil {
-				mset.mu.Unlock()
 				apiErr := NewJSMessageCounterBrokenError()
 				if canRespond {
 					resp.PubAck = &PubAck{Stream: name}
@@ -5765,7 +5759,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 		if len(sources) > 0 {
 			nhdr, err := json.Marshal(sources)
 			if err != nil {
-				mset.mu.Unlock()
 				if canRespond {
 					resp.PubAck = &PubAck{Stream: name}
 					resp.Error = NewJSMessageCounterBrokenError()
@@ -5782,7 +5775,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 		maxPayload := int64(mset.srv.getOpts().MaxPayload)
 		hdrLen, msgLen := int64(len(hdr)), int64(len(msg))
 		if hdrLen > maxPayload || msgLen > maxPayload-hdrLen {
-			mset.mu.Unlock()
 			if canRespond {
 				resp.PubAck = &PubAck{Stream: name}
 				resp.Error = NewJSStreamMessageExceedsMaximumError()
@@ -5796,7 +5788,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 	// Check to see if we are over the max msg size.
 	// Subtract to prevent against overflows.
 	if canConsistencyCheck && maxMsgSize >= 0 && (len(hdr) > maxMsgSize || len(msg) > maxMsgSize-len(hdr)) {
-		mset.mu.Unlock()
 		if canRespond {
 			resp.PubAck = &PubAck{Stream: name}
 			resp.Error = NewJSStreamMessageExceedsMaximumError()
@@ -5807,7 +5798,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 	}
 
 	if canConsistencyCheck && len(hdr) > math.MaxUint16 {
-		mset.mu.Unlock()
 		if canRespond {
 			resp.PubAck = &PubAck{Stream: name}
 			resp.Error = NewJSStreamHeaderExceedsMaximumError()
@@ -5821,7 +5811,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 	// Don't error and log/stepdown if we're tracing when clustered.
 	if !isClustered && js.limitsExceeded(stype) {
 		s.resourcesExceededError(stype)
-		mset.mu.Unlock()
 		if canRespond {
 			resp.PubAck = &PubAck{Stream: name}
 			resp.Error = NewJSInsufficientResourcesError()
@@ -5829,7 +5818,7 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 			outq.sendMsg(reply, response)
 		}
 		// Stepdown regardless.
-		if node := mset.raftNode(); node != nil {
+		if node := mset.node; node != nil {
 			node.StepDown()
 		}
 		return NewJSInsufficientResourcesError()
@@ -5851,7 +5840,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 
 	mt.updateJetStreamEvent(subject, noInterest)
 	if traceOnly {
-		mset.mu.Unlock()
 		return nil
 	}
 
@@ -5872,7 +5860,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 			}
 			outq.sendMsg(reply, response)
 		}
-		mset.mu.Unlock()
 		return nil
 	}
 
@@ -5917,7 +5904,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 				response, _ = json.Marshal(resp)
 				outq.send(newJSPubMsg(reply, _EMPTY_, _EMPTY_, nil, response, nil, 0))
 			}
-			mset.mu.Unlock()
 			return err
 		}
 	}
@@ -5931,7 +5917,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 			response, _ = json.Marshal(resp)
 			outq.send(newJSPubMsg(reply, _EMPTY_, _EMPTY_, nil, response, nil, 0))
 		}
-		mset.mu.Unlock()
 		return err
 	}
 
@@ -5961,10 +5946,9 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 
 	if err != nil {
 		if isPermissionError(err) {
-			mset.mu.Unlock()
 			// messages in block cache could be lost in the worst case.
 			// In the clustered mode it is very highly unlikely as a result of replication.
-			mset.srv.DisableJetStream()
+			go mset.srv.DisableJetStream()
 			mset.srv.Warnf("Filesystem permission denied while writing msg, disabling JetStream: %v", err)
 			return err
 		}
@@ -5973,7 +5957,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 		mset.store.FastState(&state)
 		mset.lseq = state.LastSeq
 		mset.lmsgId = olmsgId
-		mset.mu.Unlock()
 		bumpCLFS()
 
 		switch err {
@@ -5994,7 +5977,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 	}
 
 	// If here we succeeded in storing the message.
-	mset.mu.Unlock()
 
 	// If we have a msgId make sure to save.
 	// This will replace our estimate from the cluster layer if we are clustered.
@@ -6015,14 +5997,14 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 
 	// No errors, this is the normal path.
 	if rollupSub {
-		mset.purge(&JSApiStreamPurgeRequest{Subject: subject, Keep: 1})
+		mset.purgeLocked(&JSApiStreamPurgeRequest{Subject: subject, Keep: 1}, false)
 	} else if rollupAll {
-		mset.purge(&JSApiStreamPurgeRequest{Keep: 1})
+		mset.purgeLocked(&JSApiStreamPurgeRequest{Keep: 1}, false)
 	} else if scheduleNext := sliceHeader(JSScheduleNext, hdr); len(scheduleNext) > 0 && bytesToString(scheduleNext) == JSScheduleNextPurge {
 		// Purge the message schedule.
 		scheduler := getMessageScheduler(hdr)
 		if scheduler != _EMPTY_ {
-			mset.purge(&JSApiStreamPurgeRequest{Subject: scheduler})
+			mset.purgeLocked(&JSApiStreamPurgeRequest{Subject: scheduler}, false)
 		}
 	}
 
@@ -6406,6 +6388,11 @@ func (mset *stream) processJetStreamBatchMsg(batchId, subject, reply string, hdr
 
 	// Commit batch.
 	if !isClustered {
+		mset.clMu.Unlock()
+
+		// Ensure the whole batch is fully isolated, and reads
+		// can only happen after the full batch is committed.
+		mset.mu.Lock()
 		for seq := uint64(1); seq <= batchSeq; seq++ {
 			if seq == batchSeq && b.store.Type() != FileStorage {
 				bsubj, bhdr, bmsg = subject, hdr, msg
@@ -6421,8 +6408,9 @@ func (mset *stream) processJetStreamBatchMsg(batchId, subject, reply string, hdr
 			if seq == batchSeq {
 				_reply = reply
 			}
-			mset.processJetStreamMsg(bsubj, _reply, bhdr, bmsg, 0, 0, mt, false)
+			mset.processJetStreamMsg(bsubj, _reply, bhdr, bmsg, 0, 0, mt, false, false)
 		}
+		mset.mu.Unlock()
 	} else {
 		// Do a single multi proposal. This ensures we get to push all entries to the proposal queue in-order
 		// and not interleaved with other proposals.
@@ -6440,8 +6428,8 @@ func (mset *stream) processJetStreamBatchMsg(batchId, subject, reply string, hdr
 			lerr := fmt.Errorf("JetStream stream '%s > %s' has high message lag", jsa.acc().Name, name)
 			s.RateLimitWarnf("%s", lerr.Error())
 		}
+		mset.clMu.Unlock()
 	}
-	mset.clMu.Unlock()
 	b.cleanupLocked(batchId, batches)
 	batches.mu.Unlock()
 	return nil
@@ -6638,11 +6626,17 @@ func (mset *stream) accName() string {
 
 // Name returns the stream name.
 func (mset *stream) name() string {
+	return mset.nameLocked(true)
+}
+
+func (mset *stream) nameLocked(needLock bool) string {
 	if mset == nil {
 		return _EMPTY_
 	}
-	mset.mu.RLock()
-	defer mset.mu.RUnlock()
+	if needLock {
+		mset.mu.RLock()
+		defer mset.mu.RUnlock()
+	}
 	return mset.cfg.Name
 }
 
@@ -6749,7 +6743,7 @@ func (mset *stream) internalLoop() {
 				} else if isClustered {
 					mset.processClusteredInboundMsg(im.subj, im.rply, im.hdr, im.msg, im.mt, false)
 				} else {
-					mset.processJetStreamMsg(im.subj, im.rply, im.hdr, im.msg, 0, 0, im.mt, false)
+					mset.processJetStreamMsg(im.subj, im.rply, im.hdr, im.msg, 0, 0, im.mt, false, true)
 				}
 				im.returnToPool()
 			}



Test patch:
diff --git a/server/jetstream_batching_test.go b/server/jetstream_batching_test.go
index 230629ba..7e17dc9c 100644
--- a/server/jetstream_batching_test.go
+++ b/server/jetstream_batching_test.go
@@ -23,6 +23,7 @@ import (
 	"math/big"
 	"strconv"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -1408,7 +1409,7 @@ func TestJetStreamAtomicBatchPublishSingleServerRecovery(t *testing.T) {
 	require_True(t, commitReady)
 
 	// Simulate the first message of the batch is committed.
-	err = mset.processJetStreamMsg("foo", _EMPTY_, hdr1, nil, 0, 0, nil, false)
+	err = mset.processJetStreamMsg("foo", _EMPTY_, hdr1, nil, 0, 0, nil, false, true)
 	require_NoError(t, err)
 
 	// Simulate a hard kill, upon recovery the rest of the batch should be applied.
@@ -1855,3 +1856,134 @@ func TestJetStreamAtomicBatchPublishContinuousBatchesStillMoveAppliedUp(t *testi
 		return checkState(t, c, globalAccountName, "TEST")
 	})
 }
+
+func TestJetStreamRollupIsolatedRead(t *testing.T) {
+	const (
+		DirectGet = iota
+		DirectBatchGet
+		DirectMultiGet
+		Consumer
+	)
+
+	test := func(t *testing.T, mode int) {
+		c := createJetStreamClusterExplicit(t, "R3S", 3)
+		defer c.shutdown()
+
+		nc, js := jsClientConnect(t, c.randomServer())
+		defer nc.Close()
+
+		_, err := js.AddStream(&nats.StreamConfig{
+			Name:        "TEST",
+			Subjects:    []string{"foo.*"},
+			AllowRollup: true,
+			Replicas:    3,
+		})
+		require_NoError(t, err)
+
+		rs := c.randomNonStreamLeader(globalAccountName, "TEST")
+		mset, err := rs.globalAccount().lookupStream("TEST")
+		require_NoError(t, err)
+
+		// Reconnect to the selected server.
+		nc.Close()
+		nc, js = jsClientConnect(t, rs)
+		defer nc.Close()
+
+		m := nats.NewMsg("foo.0")
+		pubAck, err := js.PublishMsg(m)
+		require_NoError(t, err)
+		require_Equal(t, pubAck.Sequence, 1)
+
+		// mset.processJetStreamMsg will first store the new message, then update dedupe state, and then do the rollup.
+		// We block the replica such that it can store the new message but can't do the rollup yet.
+		// A read should wait for the rollup to complete.
+		mset.ddMu.Lock()
+		m.Subject = "foo.1"
+		m.Header.Set(JSMsgRollup, JSMsgRollupAll)
+		m.Header.Set(JSMsgId, "msgId")
+		_, _ = js.PublishMsg(m)
+		err = checkForErr(2*time.Second, 200*time.Millisecond, func() error {
+			var state StreamState
+			mset.store.FastState(&state)
+			if state.LastSeq != 2 {
+				return fmt.Errorf("expected last seq 2, got: %d", state.LastSeq)
+			}
+			return nil
+		})
+		if err != nil {
+			mset.ddMu.Unlock()
+			require_NoError(t, err)
+		}
+
+		// We're now subscribing and going to do a request, while the write is still halfway.
+		sub, err := nc.SubscribeSync("reply")
+		if err != nil {
+			mset.ddMu.Unlock()
+			require_NoError(t, err)
+		}
+		defer sub.Drain()
+		if err = nc.Flush(); err != nil {
+			mset.ddMu.Unlock()
+			require_NoError(t, err)
+		}
+
+		// Run two goroutines, one will unblock the first write after a short sleep,
+		// the second will do the read/consumer request.
+		var (
+			ready   sync.WaitGroup
+			run     sync.WaitGroup
+			cleanup sync.WaitGroup
+		)
+		ready.Add(2)
+		run.Add(1)
+		cleanup.Add(1)
+		go func() {
+			ready.Done()
+			defer cleanup.Done()
+			run.Wait()
+			time.Sleep(250 * time.Millisecond)
+			mset.ddMu.Unlock()
+		}()
+		go func() {
+			ready.Done()
+			run.Wait()
+			switch mode {
+			case DirectGet:
+				mset.getDirectRequest(&JSApiMsgGetRequest{LastFor: "foo.0"}, "reply")
+			case DirectBatchGet:
+				mset.getDirectRequest(&JSApiMsgGetRequest{NextFor: "foo.*", Batch: 2}, "reply")
+			case DirectMultiGet:
+				mset.getDirectRequest(&JSApiMsgGetRequest{MultiLastFor: []string{"foo.*"}, Batch: 2}, "reply")
+			case Consumer:
+				mset.addConsumer(&ConsumerConfig{DeliverSubject: "reply", Direct: true})
+			}
+		}()
+
+		// Wait for both to be ready, then run both of them.
+		ready.Wait()
+		run.Done()
+
+		msg, err := sub.NextMsg(time.Second)
+		// Make sure cleanup has happened before validating.
+		cleanup.Wait()
+		require_NoError(t, err)
+		if mode == DirectGet {
+			require_Equal(t, msg.Header.Get("Status"), "404")
+		} else if mode != Consumer {
+			require_Equal(t, msg.Header.Get(JSNumPending), "0")
+			require_Equal(t, msg.Header.Get(JSSequence), "2")
+		} else {
+			// $JS.ACK.<stream>.<consumer>.<delivered>.<sseq>.<cseq>.<tm>.<pending>
+			require_True(t, strings.HasPrefix(msg.Reply, jsAckPre))
+			tokens := strings.Split(msg.Reply, ".")
+			require_Len(t, len(tokens), 9)
+			require_Equal(t, tokens[8], "0") // pending
+			require_Equal(t, tokens[5], "2") // sseq
+		}
+	}
+
+	t.Run("DirectGet", func(t *testing.T) { test(t, DirectGet) })
+	t.Run("DirectBatchGet", func(t *testing.T) { test(t, DirectBatchGet) })
+	t.Run("DirectMultiGet", func(t *testing.T) { test(t, DirectMultiGet) })
+	t.Run("Consumer", func(t *testing.T) { test(t, Consumer) })
+}
diff --git a/server/jetstream_cluster_4_test.go b/server/jetstream_cluster_4_test.go
index 7c5191d6..9a2db7f8 100644
--- a/server/jetstream_cluster_4_test.go
+++ b/server/jetstream_cluster_4_test.go
@@ -4308,7 +4308,7 @@ func TestJetStreamClusterDontInstallSnapshotWhenStoppingStream(t *testing.T) {
 	validateStreamState(snap)
 
 	// Simulate a message being stored, but not calling Applied yet.
-	err = mset.processJetStreamMsg("foo", _EMPTY_, nil, nil, 1, time.Now().UnixNano(), nil, false)
+	err = mset.processJetStreamMsg("foo", _EMPTY_, nil, nil, 1, time.Now().UnixNano(), nil, false, true)
 	require_NoError(t, err)
 
 	// Simulate the stream being stopped before we're able to call Applied.
diff --git a/server/jetstream_test.go b/server/jetstream_test.go
index 2d802e06..ca169df1 100644
--- a/server/jetstream_test.go
+++ b/server/jetstream_test.go
@@ -16868,9 +16868,11 @@ func TestJetStreamDirectGetBatch(t *testing.T) {
 				require_Equal(t, expected[i], msg.Header.Get(JSSubject))
 				// Should have Data field non-zero
 				require_True(t, len(msg.Data) > 0)
-				// Check we have NumPending and its correct.
+				// Check we have NumPending and it's correct.
+				if np > 0 {
+					np--
+				}
 				require_Equal(t, strconv.Itoa(np), msg.Header.Get(JSNumPending))
-				np--
 
 			} else {
 				// Check for properly formatted EOB marker.
@@ -16880,7 +16882,7 @@ func TestJetStreamDirectGetBatch(t *testing.T) {
 				require_Equal(t, msg.Header.Get("Status"), "204")
 				// Check description is EOB
 				require_Equal(t, msg.Header.Get("Description"), "EOB")
-				// Check we have NumPending and its correct.
+				// Check we have NumPending and it's correct.
 				require_Equal(t, strconv.Itoa(np), msg.Header.Get(JSNumPending))
 			}
 		}
@@ -17258,11 +17260,11 @@ func TestJetStreamDirectGetMulti(t *testing.T) {
 						require_Equal(t, strconv.Itoa(expected[i].seq), msg.Header.Get(JSSequence))
 						// Should have Data field non-zero
 						require_True(t, len(msg.Data) > 0)
-						// Check we have NumPending and its correct.
-						require_Equal(t, strconv.Itoa(np), msg.Header.Get(JSNumPending))
+						// Check we have NumPending and it's correct.
 						if np > 0 {
 							np--
 						}
+						require_Equal(t, strconv.Itoa(np), msg.Header.Get(JSNumPending))
 					} else {
 						// Check for properly formatted EOB marker.
 						// Should have no body.
@@ -17271,17 +17273,17 @@ func TestJetStreamDirectGetMulti(t *testing.T) {
 						require_Equal(t, msg.Header.Get("Status"), "204")
 						// Check description is EOB
 						require_Equal(t, msg.Header.Get("Description"), "EOB")
-						// Check we have NumPending and its correct.
+						// Check we have NumPending and it's correct.
 						require_Equal(t, strconv.Itoa(np), msg.Header.Get(JSNumPending))
 					}
 				}
 			}
 
 			sub := sendRequest(&JSApiMsgGetRequest{MultiLastFor: []string{"foo.*"}})
-			checkResponses(sub, 2, p{"foo.foo", 97}, p{"foo.bar", 98}, p{"foo.baz", 99}, eob)
+			checkResponses(sub, 3, p{"foo.foo", 97}, p{"foo.bar", 98}, p{"foo.baz", 99}, eob)
 			// Check with UpToSeq
 			sub = sendRequest(&JSApiMsgGetRequest{MultiLastFor: []string{"foo.*"}, UpToSeq: 3})
-			checkResponses(sub, 2, p{"foo.foo", 1}, p{"foo.bar", 2}, p{"foo.baz", 3}, eob)
+			checkResponses(sub, 3, p{"foo.foo", 1}, p{"foo.bar", 2}, p{"foo.baz", 3}, eob)
 
 			// Test No Results.
 			sub = sendRequest(&JSApiMsgGetRequest{MultiLastFor: []string{"bar.*"}})
@@ -17467,7 +17469,7 @@ func TestJetStreamDirectGetMultiPaging(t *testing.T) {
 	}
 
 	// Setup variables that control procesPartial
-	start, seq, np, b, bsz := 1, 1, sent-1, 0, 128
+	start, seq, np, b, bsz := 1, 1, sent, 0, 128
 
 	processPartial := func(expected int) {
 		t.Helper()
@@ -17480,11 +17482,11 @@ func TestJetStreamDirectGetMultiPaging(t *testing.T) {
 			require_NoError(t, err)
 			// Make sure sequence is correct.
 			require_Equal(t, strconv.Itoa(int(seq)), msg.Header.Get(JSSequence))
-			// Check we have NumPending and its correct.
-			require_Equal(t, strconv.Itoa(int(np)), msg.Header.Get(JSNumPending))
+			// Check we have NumPending and it's correct.
 			if np > 0 {
 				np--
 			}
+			require_Equal(t, strconv.Itoa(np), msg.Header.Get(JSNumPending))
 		}
 		// Now check EOB
 		msg, err := sub.NextMsg(10 * time.Millisecond)
@@ -17493,11 +17495,11 @@ func TestJetStreamDirectGetMultiPaging(t *testing.T) {
 		require_Equal(t, msg.Header.Get("Status"), "204")
 		// Check description is EOB
 		require_Equal(t, msg.Header.Get("Description"), "EOB")
-		// Check we have NumPending and its correct.
+		// Check we have NumPending and it's correct.
 		require_Equal(t, strconv.Itoa(np), msg.Header.Get(JSNumPending))
-		// Check we have LastSequence and its correct.
+		// Check we have LastSequence and it's correct.
 		require_Equal(t, strconv.Itoa(seq-1), msg.Header.Get(JSLastSequence))
-		// Check we have UpToSequence and its correct.
+		// Check we have UpToSequence and it's correct.
 		require_Equal(t, strconv.Itoa(sent), msg.Header.Get(JSUpToSequence))
 		// Update start
 		start = seq
@@ -17510,7 +17512,7 @@ func TestJetStreamDirectGetMultiPaging(t *testing.T) {
 	processPartial(116 + 1)
 
 	// Now reset and test that batch is honored as well.
-	start, seq, np, b = 1, 1, sent-1, 100
+	start, seq, np, b = 1, 1, sent, 100
 	for i := 0; i < 5; i++ {
 		processPartial(b + 1) // 100 + EOB
 	}