Repository: nats-io/nats-server

Issue:
Batch publish - `Nats-Expected-Last-Sequence` and `Nats-Rollup` header support
### Observed behavior

Headers like `Nats-Expected-Last-Sequence` and `Nats-Rollup` should be properly supported. This currently doesn't work because only replicas check these headers. The leader should check, and be able to fail the batch on a mismatch.

### Expected behavior

Support these headers.

### Server and client version

main

### Host environment

_No response_

### Steps to reproduce

_No response_

Hints from issue comments:

Solution patch:
diff --git a/server/jetstream_batching.go b/server/jetstream_batching.go
index 2240bf37..6b7fc01a 100644
--- a/server/jetstream_batching.go
+++ b/server/jetstream_batching.go
@@ -82,8 +82,7 @@ func (b *batchGroup) cleanupLocked(batchId string, batches *batching) {
 type batchStagedDiff struct {
 	msgIds             map[string]struct{}
 	counter            map[string]*msgCounterRunningTotal
-	inflight           map[uint64]uint64
-	subjectsInBatch    map[string]struct{}
+	inflight           map[string]*inflightSubjectRunningTotal
 	expectedPerSubject map[string]*batchExpectedPerSubject
 }
 
@@ -116,10 +115,15 @@ func (diff *batchStagedDiff) commit(mset *stream) {
 	// Track inflight.
 	if len(diff.inflight) > 0 {
 		if mset.inflight == nil {
-			mset.inflight = make(map[uint64]uint64, len(diff.inflight))
+			mset.inflight = make(map[string]*inflightSubjectRunningTotal, len(diff.inflight))
 		}
-		for clseq, size := range diff.inflight {
-			mset.inflight[clseq] = size
+		for subj, i := range diff.inflight {
+			if c, ok := mset.inflight[subj]; ok {
+				c.bytes += i.bytes
+				c.ops += i.ops
+			} else {
+				mset.inflight[subj] = i
+			}
 		}
 	}
 
@@ -143,8 +147,8 @@ func (diff *batchStagedDiff) commit(mset *stream) {
 // mset.clMu lock must be held.
 func checkMsgHeadersPreClusteredProposal(
 	diff *batchStagedDiff, mset *stream, subject string, hdr []byte, msg []byte, sourced bool, name string,
-	jsa *jsAccount, allowTTL bool, allowMsgCounter bool, stype StorageType, store StreamStore,
-	interestPolicy bool, discard DiscardPolicy, maxMsgSize int, maxMsgs int64, maxBytes int64,
+	jsa *jsAccount, allowRollup, denyPurge, allowTTL, allowMsgCounter bool,
+	discard DiscardPolicy, discardNewPer bool, maxMsgSize int, maxMsgs int64, maxMsgsPer int64, maxBytes int64,
 ) ([]byte, []byte, uint64, *ApiError, error) {
 	var incr *big.Int
 
@@ -235,7 +239,7 @@ func checkMsgHeadersPreClusteredProposal(
 		apiErr := NewJSMessageIncrMissingError()
 		return hdr, msg, 0, apiErr, apiErr
 	}
-	if incr != nil && allowMsgCounter && store != nil {
+	if incr != nil && allowMsgCounter {
 		var initial big.Int
 		var sources CounterSources
 
@@ -254,7 +258,7 @@ func checkMsgHeadersPreClusteredProposal(
 		} else {
 			// Load last message, and store as inflight running total.
 			var smv StoreMsg
-			sm, err := store.LoadLastMsg(subject, &smv)
+			sm, err := mset.store.LoadLastMsg(subject, &smv)
 			if err == nil && sm != nil {
 				var val CounterValue
 				// Return an error if the counter is broken somehow.
@@ -343,48 +347,16 @@ func checkMsgHeadersPreClusteredProposal(
 		}
 	}
 
-	// Check if we have an interest policy and discard new with max msgs or bytes.
-	// We need to deny here otherwise it could succeed on some peers and not others
-	// depending on consumer ack state. So we deny here, if we allow that means we know
-	// it would succeed on every peer.
-	if interestPolicy && discard == DiscardNew && (maxMsgs > 0 || maxBytes > 0) {
-		// Track inflight.
-		if diff.inflight == nil {
-			diff.inflight = make(map[uint64]uint64)
-		}
-		if stype == FileStorage {
-			diff.inflight[mset.clseq] = fileStoreMsgSizeRaw(len(subject), len(hdr), len(msg))
-		} else {
-			diff.inflight[mset.clseq] = memStoreMsgSizeRaw(len(subject), len(hdr), len(msg))
-		}
-
-		var state StreamState
-		mset.store.FastState(&state)
-
-		var err error
-		if maxMsgs > 0 && state.Msgs+uint64(len(mset.inflight)+len(diff.inflight)) > uint64(maxMsgs) {
-			err = ErrMaxMsgs
-		} else if maxBytes > 0 {
-			// TODO(dlc) - Could track this rollup independently.
-			var bytesPending uint64
-			for _, nb := range mset.inflight {
-				bytesPending += nb
-			}
-			for _, nb := range diff.inflight {
-				bytesPending += nb
-			}
-			if state.Bytes+bytesPending > uint64(maxBytes) {
-				err = ErrMaxBytes
-			}
-		}
-		if err != nil {
-			return hdr, msg, 0, NewJSStreamStoreFailedError(err, Unless(err)), err
+	if len(hdr) > 0 {
+		// Expected last sequence.
+		if seq, exists := getExpectedLastSeq(hdr); exists && seq != mset.clseq-mset.clfs {
+			mlseq := mset.clseq - mset.clfs
+			err := fmt.Errorf("last sequence mismatch: %d vs %d", seq, mlseq)
+			return hdr, msg, 0, NewJSStreamWrongLastSequenceError(mlseq), err
 		}
-	}
 
-	if len(hdr) > 0 {
 		// Expected last sequence per subject.
-		if seq, exists := getExpectedLastSeqPerSubject(hdr); exists && store != nil {
+		if seq, exists := getExpectedLastSeqPerSubject(hdr); exists {
 			// Allow override of the subject used for the check.
 			seqSubj := subject
 			if optSubj := getExpectedLastSeqPerSubjectForSubject(hdr); optSubj != _EMPTY_ {
@@ -393,7 +365,7 @@ func checkMsgHeadersPreClusteredProposal(
 
 			// The subject is already written to in this batch, we can't allow
 			// expected checks since they would be incorrect.
-			if _, ok := diff.subjectsInBatch[seqSubj]; ok {
+			if _, ok := diff.inflight[seqSubj]; ok {
 				err := errors.New("last sequence by subject mismatch")
 				return hdr, msg, 0, NewJSStreamWrongLastSequenceConstantError(), err
 			}
@@ -405,6 +377,13 @@ func checkMsgHeadersPreClusteredProposal(
 				return hdr, msg, 0, NewJSStreamWrongLastSequenceConstantError(), err
 			}
 
+			// If the subject is already in process but without expected headers, block as we would have
+			// multiple messages inflight with the same subject.
+			if _, ok := mset.inflight[seqSubj]; ok {
+				err := errors.New("last sequence by subject mismatch")
+				return hdr, msg, 0, NewJSStreamWrongLastSequenceConstantError(), err
+			}
+
 			// If we've already done an expected-check on this subject, use the cached result.
 			if e, ok := diff.expectedPerSubject[seqSubj]; ok {
 				if e.sseq != seq {
@@ -415,7 +394,7 @@ func checkMsgHeadersPreClusteredProposal(
 			} else {
 				var smv StoreMsg
 				var fseq uint64
-				sm, err := store.LoadLastMsg(seqSubj, &smv)
+				sm, err := mset.store.LoadLastMsg(seqSubj, &smv)
 				if sm != nil {
 					fseq = sm.seq
 				}
@@ -435,14 +414,98 @@ func checkMsgHeadersPreClusteredProposal(
 				}
 			}
 		}
+
+		// Check for any rollups.
+		if rollup := getRollup(hdr); rollup != _EMPTY_ {
+			if !allowRollup || denyPurge {
+				err := errors.New("rollup not permitted")
+				return hdr, msg, 0, NewJSStreamRollupFailedError(err), err
+			}
+			switch rollup {
+			case JSMsgRollupSubject:
+				// Rolling up the subject is only allowed if the first occurrence of this subject in the batch.
+				if _, ok := diff.inflight[subject]; ok {
+					err := errors.New("batch rollup sub invalid")
+					return hdr, msg, 0, NewJSStreamRollupFailedError(err), err
+				}
+			case JSMsgRollupAll:
+				// Rolling up the whole stream is only allowed if this is the first message of the batch.
+				if len(diff.inflight) > 0 {
+					err := errors.New("batch rollup all invalid")
+					return hdr, msg, 0, NewJSStreamRollupFailedError(err), err
+				}
+			default:
+				err := fmt.Errorf("rollup value invalid: %q", rollup)
+				return hdr, msg, 0, NewJSStreamRollupFailedError(err), err
+			}
+		}
 	}
 
+	// Track inflight.
 	// Store the subject to ensure other messages in this batch using
-	// an expected check on the same subject fail.
-	if diff.subjectsInBatch == nil {
-		diff.subjectsInBatch = map[string]struct{}{subject: {}}
+	// an expected check or rollup on the same subject fail.
+	if diff.inflight == nil {
+		diff.inflight = make(map[string]*inflightSubjectRunningTotal, 1)
+	}
+	var sz uint64
+	if mset.store.Type() == FileStorage {
+		sz = fileStoreMsgSizeRaw(len(subject), len(hdr), len(msg))
 	} else {
-		diff.subjectsInBatch[subject] = struct{}{}
+		sz = memStoreMsgSizeRaw(len(subject), len(hdr), len(msg))
+	}
+	var (
+		i   *inflightSubjectRunningTotal
+		ok  bool
+		err error
+	)
+	if i, ok = diff.inflight[subject]; ok {
+		i.bytes += sz
+		i.ops++
+	} else {
+		i = &inflightSubjectRunningTotal{bytes: sz, ops: 1}
+		diff.inflight[subject] = i
+	}
+
+	// Check if we have discard new with max msgs or bytes.
+	// We need to deny here otherwise we'd need to bump CLFS, and it could succeed on some
+	// peers and not others depending on consumer ack state (if interest policy).
+	// So we deny here, if we allow that means we know it would succeed on every peer.
+	if discard == DiscardNew && (maxMsgs > 0 || maxBytes > 0) {
+		// Error if over DiscardNew per subject threshold.
+		if discardNewPer {
+			totalMsgsForSubject := i.ops
+			if i, ok = mset.inflight[subject]; ok {
+				totalMsgsForSubject += i.ops
+			}
+			if maxMsgsPer > 0 && totalMsgsForSubject > uint64(maxMsgsPer) {
+				err = ErrMaxMsgsPerSubject
+				return hdr, msg, 0, NewJSStreamStoreFailedError(err, Unless(err)), err
+			}
+		}
+
+		// Track usual max msgs/bytes thresholds for DiscardNew.
+		var state StreamState
+		mset.store.FastState(&state)
+
+		totalMsgs := state.Msgs
+		totalBytes := state.Bytes
+		for _, i = range mset.inflight {
+			totalMsgs += i.ops
+			totalBytes += i.bytes
+		}
+		for _, i = range diff.inflight {
+			totalMsgs += i.ops
+			totalBytes += i.bytes
+		}
+
+		if maxMsgs > 0 && totalMsgs > uint64(maxMsgs) {
+			err = ErrMaxMsgs
+		} else if maxBytes > 0 && totalBytes > uint64(maxBytes) {
+			err = ErrMaxBytes
+		}
+		if err != nil {
+			return hdr, msg, 0, NewJSStreamStoreFailedError(err, Unless(err)), err
+		}
 	}
 
 	return hdr, msg, 0, nil, nil
diff --git a/server/jetstream_cluster.go b/server/jetstream_cluster.go
index 5cac62be..970f07a9 100644
--- a/server/jetstream_cluster.go
+++ b/server/jetstream_cluster.go
@@ -3026,7 +3026,26 @@ func (js *jetStream) applyStreamEntries(mset *stream, ce *CommittedEntry, isReco
 				// But do not want to acquire lock since tracking this will be rare.
 				if mset.inflight != nil {
 					mset.clMu.Lock()
-					delete(mset.inflight, lseq)
+					if i, found := mset.inflight[subject]; found {
+						// Decrement from pending operations. Once it reaches zero, it can be deleted.
+						if i.ops > 0 {
+							var sz uint64
+							if mset.store.Type() == FileStorage {
+								sz = fileStoreMsgSizeRaw(len(subject), len(hdr), len(msg))
+							} else {
+								sz = memStoreMsgSizeRaw(len(subject), len(hdr), len(msg))
+							}
+							if i.bytes >= sz {
+								i.bytes -= sz
+							} else {
+								i.bytes = 0
+							}
+							i.ops--
+						}
+						if i.ops == 0 {
+							delete(mset.inflight, subject)
+						}
+					}
 					mset.clMu.Unlock()
 				}
 
@@ -8032,11 +8051,11 @@ func (mset *stream) processClusteredInboundMsg(subject, reply string, hdr, msg [
 
 	mset.mu.RLock()
 	canRespond := !mset.cfg.NoAck && len(reply) > 0
-	name, stype, store := mset.cfg.Name, mset.cfg.Storage, mset.store
-	interestPolicy, discard, maxMsgs, maxBytes := mset.cfg.Retention != LimitsPolicy, mset.cfg.Discard, mset.cfg.MaxMsgs, mset.cfg.MaxBytes
+	name, stype := mset.cfg.Name, mset.cfg.Storage
+	discard, discardNewPer, maxMsgs, maxMsgsPer, maxBytes := mset.cfg.Discard, mset.cfg.DiscardNewPer, mset.cfg.MaxMsgs, mset.cfg.MaxMsgsPer, mset.cfg.MaxBytes
 	s, js, jsa, st, r, tierName, outq, node := mset.srv, mset.js, mset.jsa, mset.cfg.Storage, mset.cfg.Replicas, mset.tier, mset.outq, mset.node
 	maxMsgSize, lseq := int(mset.cfg.MaxMsgSize), mset.lseq
-	isLeader, isSealed, allowTTL, allowMsgCounter, allowAtomicPublish := mset.isLeader(), mset.cfg.Sealed, mset.cfg.AllowMsgTTL, mset.cfg.AllowMsgCounter, mset.cfg.AllowAtomicPublish
+	isLeader, isSealed, allowRollup, denyPurge, allowTTL, allowMsgCounter, allowAtomicPublish := mset.isLeader(), mset.cfg.Sealed, mset.cfg.AllowRollup, mset.cfg.DenyPurge, mset.cfg.AllowMsgTTL, mset.cfg.AllowMsgCounter, mset.cfg.AllowAtomicPublish
 	mset.mu.RUnlock()
 
 	// This should not happen but possible now that we allow scale up, and scale down where this could trigger.
@@ -8304,14 +8323,11 @@ func (mset *stream) processClusteredInboundMsg(subject, reply string, hdr, msg [
 			if msgId := getMsgId(bhdr); msgId != _EMPTY_ {
 				return errorOnUnsupported(seq, JSMsgId)
 			}
-			if _, ok = getExpectedLastSeq(hdr); ok {
-				return errorOnUnsupported(seq, JSExpectedLastSeq)
-			}
 			if getExpectedLastMsgId(hdr) != _EMPTY_ {
 				return errorOnUnsupported(seq, JSExpectedLastMsgId)
 			}
 
-			if bhdr, bmsg, _, apiErr, err = checkMsgHeadersPreClusteredProposal(diff, mset, subject, bhdr, bmsg, sourced, name, jsa, allowTTL, allowMsgCounter, stype, store, interestPolicy, discard, maxMsgSize, maxMsgs, maxBytes); err != nil {
+			if bhdr, bmsg, _, apiErr, err = checkMsgHeadersPreClusteredProposal(diff, mset, subject, bhdr, bmsg, sourced, name, jsa, allowRollup, denyPurge, allowTTL, allowMsgCounter, discard, discardNewPer, maxMsgSize, maxMsgs, maxMsgsPer, maxBytes); err != nil {
 				// TODO(mvv): reset in-memory expected header maps
 				mset.clseq -= seq - 1
 				mset.clMu.Unlock()
@@ -8373,7 +8389,7 @@ func (mset *stream) processClusteredInboundMsg(subject, reply string, hdr, msg [
 		err    error
 	)
 	diff := &batchStagedDiff{}
-	if hdr, msg, dseq, apiErr, err = checkMsgHeadersPreClusteredProposal(diff, mset, subject, hdr, msg, sourced, name, jsa, allowTTL, allowMsgCounter, stype, store, interestPolicy, discard, maxMsgSize, maxMsgs, maxBytes); err != nil {
+	if hdr, msg, dseq, apiErr, err = checkMsgHeadersPreClusteredProposal(diff, mset, subject, hdr, msg, sourced, name, jsa, allowRollup, denyPurge, allowTTL, allowMsgCounter, discard, discardNewPer, maxMsgSize, maxMsgs, maxMsgsPer, maxBytes); err != nil {
 		// TODO(mvv): reset in-memory expected header maps
 		mset.clMu.Unlock()
 		if err == errMsgIdDuplicate && dseq > 0 {
diff --git a/server/stream.go b/server/stream.go
index f2e9dc50..ecdcae14 100644
--- a/server/stream.go
+++ b/server/stream.go
@@ -401,14 +401,14 @@ type stream struct {
 	clMu      sync.Mutex        // The mutex for clseq and clfs.
 	clseq     uint64            // The current last seq being proposed to the NRG layer.
 	clfs      uint64            // The count (offset) of the number of failed NRG sequences used to compute clseq.
-	inflight  map[uint64]uint64 // Inflight message sizes per clseq.
 	lqsent    time.Time         // The time at which the last lost quorum advisory was sent. Used to rate limit.
 	uch       chan struct{}     // The channel to signal updates to the monitor routine.
 	inMonitor bool              // True if the monitor routine has been started.
 
-	clusteredCounterTotal       map[string]*msgCounterRunningTotal // Inflight counter totals.
-	expectedPerSubjectSequence  map[uint64]string                  // Inflight 'expected per subject' subjects per clseq.
-	expectedPerSubjectInProcess map[string]struct{}                // Current 'expected per subject' subjects in process.
+	inflight                    map[string]*inflightSubjectRunningTotal // Inflight message sizes per subject.
+	clusteredCounterTotal       map[string]*msgCounterRunningTotal      // Inflight counter totals.
+	expectedPerSubjectSequence  map[uint64]string                       // Inflight 'expected per subject' subjects per clseq.
+	expectedPerSubjectInProcess map[string]struct{}                     // Current 'expected per subject' subjects in process.
 
 	// Direct get subscription.
 	directLeaderSub *subscription
@@ -422,6 +422,12 @@ type stream struct {
 	batches *batching // Inflight batches prior to committing them.
 }
 
+// inflightSubjectRunningTotal stores a running total of inflight messages for a specific subject.
+type inflightSubjectRunningTotal struct {
+	bytes uint64 // Running total of inflight bytes for inflight messages.
+	ops   uint64 // Inflight operations, i.e. inflight messages for this subject. If this reaches zero, we can remove the running total.
+}
+
 // msgCounterRunningTotal stores a running total and a number of inflight
 // but not yet applied clustered proposals/operations for this counter.
 type msgCounterRunningTotal struct {
@@ -5118,15 +5124,16 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 	allowMsgCounter := mset.cfg.AllowMsgCounter
 	// Snapshot if we are the leader and if we can respond.
 	isLeader, isSealed := mset.isLeaderNodeState(), mset.cfg.Sealed
+	isClustered := mset.isClustered()
+	canConsistencyCheck := !isClustered || traceOnly
 	canRespond := doAck && len(reply) > 0 && isLeader
 	outq := mset.outq
 
 	var resp = &JSPubAckResponse{}
 
 	// Bail here if sealed.
-	if isSealed {
+	if canConsistencyCheck && isSealed {
 		mset.mu.Unlock()
-		bumpCLFS()
 		if canRespond && outq != nil {
 			resp.PubAck = &PubAck{Stream: name}
 			resp.Error = ApiErrors[JSStreamSealedErr]
@@ -5188,7 +5195,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 	var batchId string
 	var batchSeq uint64
 	var rollupSub, rollupAll bool
-	isClustered := mset.isClustered()
 
 	if len(hdr) > 0 {
 
@@ -5200,13 +5206,12 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 		// Certain checks have already been performed if in clustered mode, so only check if not.
 		// Note, for cluster mode but with message tracing (without message delivery), we need
 		// to do this check here since it was not done in processClusteredInboundMsg().
-		if !isClustered || traceOnly {
+		if canConsistencyCheck {
 			// Counter increments.
 			// 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()
-				bumpCLFS()
 				apiErr := NewJSMessageIncrInvalidError()
 				if canRespond {
 					resp.PubAck = &PubAck{Stream: name}
@@ -5219,7 +5224,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 				// Only do checks if the message isn't sourced. Otherwise, we need to store verbatim.
 				if !allowMsgCounter {
 					mset.mu.Unlock()
-					bumpCLFS()
 					apiErr := NewJSMessageIncrDisabledError()
 					if canRespond {
 						resp.PubAck = &PubAck{Stream: name}
@@ -5230,7 +5234,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 					return apiErr
 				} else if len(msg) > 0 {
 					mset.mu.Unlock()
-					bumpCLFS()
 					apiErr := NewJSMessageIncrPayloadError()
 					if canRespond {
 						resp.PubAck = &PubAck{Stream: name}
@@ -5255,7 +5258,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 
 					if doErr {
 						mset.mu.Unlock()
-						bumpCLFS()
 						apiErr := NewJSMessageIncrInvalidError()
 						if canRespond {
 							resp.PubAck = &PubAck{Stream: name}
@@ -5271,7 +5273,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 			// Expected stream.
 			if sname := getExpectedStream(hdr); sname != _EMPTY_ && sname != name {
 				mset.mu.Unlock()
-				bumpCLFS()
 				if canRespond {
 					resp.PubAck = &PubAck{Stream: name}
 					resp.Error = NewJSStreamNotMatchError()
@@ -5280,28 +5281,70 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 				}
 				return errStreamMismatch
 			}
-		}
 
-		// TTL'd messages are rejected entirely if TTLs are not enabled on the stream.
-		// 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()
-			bumpCLFS()
-			if canRespond {
-				resp.PubAck = &PubAck{Stream: name}
-				resp.Error = NewJSMessageTTLDisabledError()
-				b, _ := json.Marshal(resp)
-				outq.sendMsg(reply, b)
+			// TTL'd messages are rejected entirely if TTLs are not enabled on the stream.
+			// 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()
+					b, _ := json.Marshal(resp)
+					outq.sendMsg(reply, b)
+				}
+				return errMsgTTLDisabled
+			}
+
+			// Expected last sequence per subject.
+			if seq, exists := getExpectedLastSeqPerSubject(hdr); exists {
+				// Allow override of the subject used for the check.
+				seqSubj := subject
+				if optSubj := getExpectedLastSeqPerSubjectForSubject(hdr); optSubj != _EMPTY_ {
+					seqSubj = optSubj
+				}
+
+				// TODO(dlc) - We could make a new store func that does this all in one.
+				var smv StoreMsg
+				var fseq uint64
+				sm, err := store.LoadLastMsg(seqSubj, &smv)
+				if sm != nil {
+					fseq = sm.seq
+				}
+				if err == ErrStoreMsgNotFound && seq == 0 {
+					fseq, err = 0, nil
+				}
+				if err != nil || fseq != seq {
+					mset.mu.Unlock()
+					if canRespond {
+						resp.PubAck = &PubAck{Stream: name}
+						resp.Error = NewJSStreamWrongLastSequenceError(fseq)
+						b, _ := json.Marshal(resp)
+						outq.sendMsg(reply, b)
+					}
+					return fmt.Errorf("last sequence by subject mismatch: %d vs %d", seq, fseq)
+				}
+			}
+
+			// 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)
+					b, _ := json.Marshal(resp)
+					outq.sendMsg(reply, b)
+				}
+				return fmt.Errorf("last sequence mismatch: %d vs %d", seq, mlseq)
 			}
-			return errMsgTTLDisabled
 		}
 
 		// Dedupe detection. This is done at the cluster level for dedupe detection above the
 		// lower layers. But we still need to pull out the msgId.
 		if msgId = getMsgId(hdr); msgId != _EMPTY_ {
 			// Do real check only if not clustered or traceOnly flag is set.
-			if !isClustered || traceOnly {
+			if canConsistencyCheck {
 				var seq uint64
 				mset.ddMu.Lock()
 				dde := mset.checkMsgId(msgId)
@@ -5311,7 +5354,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 				mset.ddMu.Unlock()
 				if seq > 0 {
 					mset.mu.Unlock()
-					bumpCLFS()
 					if canRespond {
 						response := append(pubAck, strconv.FormatUint(dde.seq, 10)...)
 						response = append(response, ",\"duplicate\": true}"...)
@@ -5322,59 +5364,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 			}
 		}
 
-		// Expected last sequence per subject.
-		if seq, exists := getExpectedLastSeqPerSubject(hdr); exists {
-			// Allow override of the subject used for the check.
-			seqSubj := subject
-			if optSubj := getExpectedLastSeqPerSubjectForSubject(hdr); optSubj != _EMPTY_ {
-				seqSubj = optSubj
-			}
-
-			// TODO(dlc) - We could make a new store func that does this all in one.
-			var smv StoreMsg
-			var fseq uint64
-			sm, err := store.LoadLastMsg(seqSubj, &smv)
-			if sm != nil {
-				fseq = sm.seq
-			}
-			if err == ErrStoreMsgNotFound {
-				if seq == 0 {
-					fseq, err = 0, nil
-				} else if mset.isClustered() {
-					// Do not bump clfs in case message was not found and could have been deleted.
-					var ss StreamState
-					store.FastState(&ss)
-					if seq <= ss.LastSeq {
-						fseq, err = seq, nil
-					}
-				}
-			}
-			if err != nil || fseq != seq {
-				mset.mu.Unlock()
-				bumpCLFS()
-				if canRespond {
-					resp.PubAck = &PubAck{Stream: name}
-					resp.Error = NewJSStreamWrongLastSequenceError(fseq)
-					b, _ := json.Marshal(resp)
-					outq.sendMsg(reply, b)
-				}
-				return fmt.Errorf("last sequence by subject mismatch: %d vs %d", seq, fseq)
-			}
-		}
-
-		// Expected last sequence.
-		if seq, exists := getExpectedLastSeq(hdr); exists && seq != mset.lseq {
-			mlseq := mset.lseq
-			mset.mu.Unlock()
-			bumpCLFS()
-			if canRespond {
-				resp.PubAck = &PubAck{Stream: name}
-				resp.Error = NewJSStreamWrongLastSequenceError(mlseq)
-				b, _ := json.Marshal(resp)
-				outq.sendMsg(reply, b)
-			}
-			return fmt.Errorf("last sequence mismatch: %d vs %d", seq, mlseq)
-		}
 		// Expected last msgId.
 		if lmsgId := getExpectedLastMsgId(hdr); lmsgId != _EMPTY_ {
 			if lmsgId != mset.lmsgId {
@@ -5392,9 +5381,8 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 		}
 		// Check for any rollups.
 		if rollup := getRollup(hdr); rollup != _EMPTY_ {
-			if !mset.cfg.AllowRollup || mset.cfg.DenyPurge {
+			if canConsistencyCheck && (!mset.cfg.AllowRollup || mset.cfg.DenyPurge) {
 				mset.mu.Unlock()
-				bumpCLFS()
 				if canRespond {
 					resp.PubAck = &PubAck{Stream: name}
 					resp.Error = NewJSStreamRollupFailedError(errors.New("rollup not permitted"))
@@ -5410,23 +5398,22 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 				rollupAll = true
 			default:
 				mset.mu.Unlock()
-				bumpCLFS()
-				err := fmt.Errorf("rollup value invalid: %q", rollup)
-				if canRespond {
-					resp.PubAck = &PubAck{Stream: name}
-					resp.Error = NewJSStreamRollupFailedError(err)
-					b, _ := json.Marshal(resp)
-					outq.sendMsg(reply, b)
+				if canConsistencyCheck {
+					err := fmt.Errorf("rollup value invalid: %q", rollup)
+					if canRespond {
+						resp.PubAck = &PubAck{Stream: name}
+						resp.Error = NewJSStreamRollupFailedError(err)
+						b, _ := json.Marshal(resp)
+						outq.sendMsg(reply, b)
+					}
+					return err
 				}
-				return err
 			}
 		}
 	}
 
-	if !isClustered && incr == nil && allowMsgCounter {
+	if canConsistencyCheck && incr == nil && allowMsgCounter {
 		mset.mu.Unlock()
-		// FIXME(mvv): we're not clustered, does bumping CLFS result in issues?
-		bumpCLFS()
 		apiErr := NewJSMessageIncrMissingError()
 		if canRespond {
 			resp.PubAck = &PubAck{Stream: name}
@@ -5446,7 +5433,7 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 
 	// Apply increment for counter.
 	// But only if it's allowed for this stream. This can happen when we store verbatim for a sourced stream.
-	if !isClustered && incr != nil && allowMsgCounter {
+	if canConsistencyCheck && incr != nil && allowMsgCounter {
 		var initial big.Int
 		var sources CounterSources
 		var smv StoreMsg
@@ -5456,7 +5443,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 			// Return an error if the counter is broken somehow.
 			if json.Unmarshal(sm.msg, &val) != nil {
 				mset.mu.Unlock()
-				bumpCLFS()
 				apiErr := NewJSMessageCounterBrokenError()
 				if canRespond {
 					resp.PubAck = &PubAck{Stream: name}
@@ -5469,7 +5455,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()
-					bumpCLFS()
 					apiErr := NewJSMessageCounterBrokenError()
 					if canRespond {
 						resp.PubAck = &PubAck{Stream: name}
@@ -5495,7 +5480,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 			var val CounterValue
 			if json.Unmarshal(msg, &val) != nil {
 				mset.mu.Unlock()
-				bumpCLFS()
 				apiErr := NewJSMessageCounterBrokenError()
 				if canRespond {
 					resp.PubAck = &PubAck{Stream: name}
@@ -5533,7 +5517,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 			nhdr, err := json.Marshal(sources)
 			if err != nil {
 				mset.mu.Unlock()
-				bumpCLFS()
 				if canRespond {
 					resp.PubAck = &PubAck{Stream: name}
 					resp.Error = NewJSMessageCounterBrokenError()
@@ -5551,7 +5534,6 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 		hdrLen, msgLen := int64(len(hdr)), int64(len(msg))
 		if hdrLen > maxPayload || msgLen > maxPayload-hdrLen {
 			mset.mu.Unlock()
-			bumpCLFS()
 			if canRespond {
 				resp.PubAck = &PubAck{Stream: name}
 				resp.Error = NewJSStreamMessageExceedsMaximumError()
@@ -5564,9 +5546,8 @@ 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 maxMsgSize >= 0 && (len(hdr) > maxMsgSize || len(msg) > maxMsgSize-len(hdr)) {
+	if canConsistencyCheck && maxMsgSize >= 0 && (len(hdr) > maxMsgSize || len(msg) > maxMsgSize-len(hdr)) {
 		mset.mu.Unlock()
-		bumpCLFS()
 		if canRespond {
 			resp.PubAck = &PubAck{Stream: name}
 			resp.Error = NewJSStreamMessageExceedsMaximumError()
@@ -5576,9 +5557,8 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 		return ErrMaxPayload
 	}
 
-	if len(hdr) > math.MaxUint16 {
+	if canConsistencyCheck && len(hdr) > math.MaxUint16 {
 		mset.mu.Unlock()
-		bumpCLFS()
 		if canRespond {
 			resp.PubAck = &PubAck{Stream: name}
 			resp.Error = NewJSStreamHeaderExceedsMaximumError()
@@ -5589,10 +5569,10 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 	}
 
 	// Check to see if we have exceeded our limits.
-	if js.limitsExceeded(stype) {
+	// Don't error and log/stepdown if we're tracing when clustered.
+	if !isClustered && js.limitsExceeded(stype) {
 		s.resourcesExceededError(stype)
 		mset.mu.Unlock()
-		bumpCLFS()
 		if canRespond {
 			resp.PubAck = &PubAck{Stream: name}
 			resp.Error = NewJSInsufficientResourcesError()
@@ -5675,6 +5655,7 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 	}
 
 	// If clustered this was already checked and we do not want to check here and possibly introduce skew.
+	// Don't error and log if we're tracing when clustered.
 	if !isClustered {
 		if exceeded, err := jsa.wouldExceedLimits(stype, tierName, mset.cfg.Replicas, subject, hdr, msg); exceeded {
 			if err == nil {



Test patch:
diff --git a/server/jetstream_batching_test.go b/server/jetstream_batching_test.go
index eaf95d3e..92c9c39f 100644
--- a/server/jetstream_batching_test.go
+++ b/server/jetstream_batching_test.go
@@ -16,6 +16,7 @@
 package server
 
 import (
+	"bytes"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -625,9 +626,8 @@ func TestJetStreamAtomicBatchPublishDenyHeaders(t *testing.T) {
 
 		// We might support these headers later on, but for now error.
 		for key, value := range map[string]string{
-			"Nats-Msg-Id":                 "msgId",
-			"Nats-Expected-Last-Sequence": "0",
-			"Nats-Expected-Last-Msg-Id":   "msgId",
+			"Nats-Msg-Id":               "msgId",
+			"Nats-Expected-Last-Msg-Id": "msgId",
 		} {
 			t.Run(key, func(t *testing.T) {
 				m := nats.NewMsg("foo")
@@ -648,17 +648,21 @@ func TestJetStreamAtomicBatchPublishStageAndCommit(t *testing.T) {
 	type BatchItem struct {
 		subject string
 		header  nats.Header
+		msg     []byte
 		err     error
 	}
 
 	type BatchTest struct {
-		title              string
-		allowTTL           bool
-		allowMsgCounter    bool
-		interestDiscardNew bool
-		init               func(mset *stream)
-		batch              []BatchItem
-		validate           func(mset *stream, commit bool)
+		title             string
+		allowRollup       bool
+		denyPurge         bool
+		allowTTL          bool
+		allowMsgCounter   bool
+		discardNew        bool
+		discardNewPerSubj bool
+		init              func(mset *stream)
+		batch             []BatchItem
+		validate          func(mset *stream, commit bool)
 	}
 
 	tests := []BatchTest{
@@ -754,23 +758,93 @@ func TestJetStreamAtomicBatchPublishStageAndCommit(t *testing.T) {
 			},
 		},
 		{
-			title:              "interest-discard-new",
-			interestDiscardNew: true,
+			title:      "discard-new",
+			discardNew: true,
 			batch: []BatchItem{
+				{subject: "foo1"},
+				{subject: "foo2"},
+			},
+			validate: func(mset *stream, commit bool) {
+				if !commit {
+					require_Len(t, len(mset.inflight), 0)
+				} else {
+					require_Len(t, len(mset.inflight), 2)
+					require_Equal(t, *mset.inflight["foo1"], inflightSubjectRunningTotal{bytes: 20, ops: 1})
+					require_Equal(t, *mset.inflight["foo2"], inflightSubjectRunningTotal{bytes: 20, ops: 1})
+				}
+			},
+		},
+		{
+			title:      "discard-new-max-msgs",
+			discardNew: true,
+			batch: []BatchItem{
+				{subject: "foo"},
 				{subject: "foo"},
-				{subject: "foo_2"},
+				{subject: "foo", err: ErrMaxMsgs},
+			},
+			validate: func(mset *stream, commit bool) {
+				if !commit {
+					require_Len(t, len(mset.inflight), 0)
+				} else {
+					require_Len(t, len(mset.inflight), 1)
+					require_Equal(t, *mset.inflight["foo"], inflightSubjectRunningTotal{bytes: 19 * 3, ops: 3})
+				}
+			},
+		},
+		{
+			title:      "discard-new-max-bytes",
+			discardNew: true,
+			batch: []BatchItem{
+				{subject: "foo1", msg: bytes.Repeat([]byte("A"), 10)},
+				{subject: "foo2", msg: bytes.Repeat([]byte("A"), 11), err: ErrMaxBytes},
 			},
 			validate: func(mset *stream, commit bool) {
 				if !commit {
 					require_Len(t, len(mset.inflight), 0)
 				} else {
 					require_Len(t, len(mset.inflight), 2)
-					require_Equal(t, mset.inflight[0], 19)
-					require_Equal(t, mset.inflight[1], 21)
+					require_Equal(t, *mset.inflight["foo1"], inflightSubjectRunningTotal{bytes: 30, ops: 1})
+					require_Equal(t, *mset.inflight["foo2"], inflightSubjectRunningTotal{bytes: 31, ops: 1})
+				}
+			},
+		},
+		{
+			title:             "discard-new-max-msgs-per-subj",
+			discardNew:        true,
+			discardNewPerSubj: true,
+			batch: []BatchItem{
+				{subject: "foo"},
+				{subject: "foo", err: ErrMaxMsgsPerSubject},
+			},
+			validate: func(mset *stream, commit bool) {
+				if !commit {
+					require_Len(t, len(mset.inflight), 0)
+				} else {
+					require_Len(t, len(mset.inflight), 1)
+					require_Equal(t, *mset.inflight["foo"], inflightSubjectRunningTotal{bytes: 19 * 2, ops: 2})
 				}
 			},
 		},
-		// TODO(mvv): expected-last-sequence header
+		{
+			title: "expect-last-seq",
+			batch: []BatchItem{
+				{subject: "foo", header: nats.Header{JSExpectedLastSeq: {"0"}}},
+				{subject: "bar", header: nats.Header{JSExpectedLastSeq: {"1"}}},
+			},
+		},
+		{
+			title: "expect-last-seq-invalid-first",
+			batch: []BatchItem{
+				{subject: "foo", header: nats.Header{JSExpectedLastSeq: {"1"}}, err: errors.New("last sequence mismatch: 1 vs 0")},
+			},
+		},
+		{
+			title: "expect-last-seq-invalid",
+			batch: []BatchItem{
+				{subject: "foo", header: nats.Header{JSExpectedLastSeq: {"0"}}},
+				{subject: "bar", header: nats.Header{JSExpectedLastSeq: {"0"}}, err: errors.New("last sequence mismatch: 0 vs 1")},
+			},
+		},
 		{
 			title: "expect-per-subj-simple",
 			batch: []BatchItem{
@@ -868,6 +942,71 @@ func TestJetStreamAtomicBatchPublishStageAndCommit(t *testing.T) {
 				require_Len(t, len(mset.expectedPerSubjectInProcess), 1)
 			},
 		},
+		{
+			title: "expect-per-subj-inflight",
+			init: func(mset *stream) {
+				mset.inflight = map[string]*inflightSubjectRunningTotal{"bar": {bytes: 33, ops: 1}}
+			},
+			batch: []BatchItem{
+				{subject: "foo", header: nats.Header{JSExpectedLastSubjSeq: {"10"}, JSExpectedLastSubjSeqSubj: {"bar"}}, err: errors.New("last sequence by subject mismatch")},
+			},
+			validate: func(mset *stream, commit bool) {
+				require_Len(t, len(mset.expectedPerSubjectSequence), 0)
+				require_Len(t, len(mset.expectedPerSubjectInProcess), 0)
+				require_Len(t, len(mset.inflight), 1)
+			},
+		},
+		{
+			title:       "rollup-deny-purge",
+			allowRollup: true,
+			denyPurge:   true,
+			batch: []BatchItem{
+				{subject: "foo", header: nats.Header{JSMsgRollup: {JSMsgRollupSubject}}, err: errors.New("rollup not permitted")},
+			},
+		},
+		{
+			title:       "rollup-invalid",
+			allowRollup: true,
+			denyPurge:   false,
+			batch: []BatchItem{
+				{subject: "foo", header: nats.Header{JSMsgRollup: {"invalid"}}, err: fmt.Errorf("rollup value invalid: %q", "invalid")},
+			},
+		},
+		{
+			title:       "rollup-all-first",
+			allowRollup: true,
+			denyPurge:   false,
+			batch: []BatchItem{
+				{subject: "foo", header: nats.Header{JSMsgRollup: {JSMsgRollupAll}}},
+			},
+		},
+		{
+			title:       "rollup-all-not-first",
+			allowRollup: true,
+			denyPurge:   false,
+			batch: []BatchItem{
+				{subject: "foo"},
+				{subject: "bar", header: nats.Header{JSMsgRollup: {JSMsgRollupAll}}, err: errors.New("batch rollup all invalid")},
+			},
+		},
+		{
+			title:       "rollup-sub-unique",
+			allowRollup: true,
+			denyPurge:   false,
+			batch: []BatchItem{
+				{subject: "foo", header: nats.Header{JSMsgRollup: {JSMsgRollupSubject}}},
+				{subject: "bar", header: nats.Header{JSMsgRollup: {JSMsgRollupSubject}}},
+			},
+		},
+		{
+			title:       "rollup-sub-overlap",
+			allowRollup: true,
+			denyPurge:   false,
+			batch: []BatchItem{
+				{subject: "foo", header: nats.Header{JSMsgRollup: {JSMsgRollupSubject}}},
+				{subject: "foo", header: nats.Header{JSMsgRollup: {JSMsgRollupSubject}}, err: errors.New("batch rollup sub invalid")},
+			},
+		},
 	}
 
 	for _, test := range tests {
@@ -887,22 +1026,24 @@ func TestJetStreamAtomicBatchPublishStageAndCommit(t *testing.T) {
 
 			mset, err := s.globalAccount().lookupStream("TEST")
 			require_NoError(t, err)
-			store := mset.store
-			require_NotNil(t, store)
 
 			if test.init != nil {
 				test.init(mset)
 			}
 
 			var (
-				interestPolicy bool
-				discard        DiscardPolicy
-				maxMsgs        int64
-				maxBytes       int64
+				discard       DiscardPolicy
+				discardNewPer bool
+				maxMsgs       int64 = -1
+				maxMsgsPer    int64 = -1
+				maxBytes      int64 = -1
 			)
-			if test.interestDiscardNew {
-				interestPolicy, discard = true, DiscardNew
-				maxMsgs, maxBytes = 10, 1024
+			if test.discardNew {
+				discard, maxMsgs, maxBytes = DiscardNew, 2, 60
+			}
+			if test.discardNewPerSubj {
+				require_True(t, test.discardNew)
+				discardNewPer, maxMsgs, maxMsgsPer = true, -1, 1
 			}
 
 			diff := &batchStagedDiff{}
@@ -915,7 +1056,7 @@ func TestJetStreamAtomicBatchPublishStageAndCommit(t *testing.T) {
 						hdr = genHeader(hdr, key, value)
 					}
 				}
-				_, _, _, _, err = checkMsgHeadersPreClusteredProposal(diff, mset, m.subject, hdr, nil, false, "TEST", nil, test.allowTTL, test.allowMsgCounter, MemoryStorage, store, interestPolicy, discard, -1, maxMsgs, maxBytes)
+				_, _, _, _, err = checkMsgHeadersPreClusteredProposal(diff, mset, m.subject, hdr, m.msg, false, "TEST", nil, test.allowRollup, test.denyPurge, test.allowTTL, test.allowMsgCounter, discard, discardNewPer, -1, maxMsgs, maxMsgsPer, maxBytes)
 				if m.err != nil {
 					require_Error(t, err, m.err)
 				} else {
diff --git a/server/jetstream_cluster_4_test.go b/server/jetstream_cluster_4_test.go
index cd2a861c..c5bb983e 100644
--- a/server/jetstream_cluster_4_test.go
+++ b/server/jetstream_cluster_4_test.go
@@ -4767,10 +4767,20 @@ func TestJetStreamClusterExpectedPerSubjectConsistency(t *testing.T) {
 	_, err = js.Publish("foo", nil, nats.ExpectLastSequencePerSubject(0))
 	require_Error(t, err, NewJSStreamWrongLastSequenceConstantError())
 
+	// Block updates when subject already inflight without expected headers.
+	mset.clMu.Lock()
+	mset.expectedPerSubjectSequence = nil
+	mset.expectedPerSubjectInProcess = nil
+	mset.inflight = map[string]*inflightSubjectRunningTotal{"foo": {bytes: 33, ops: 1}}
+	mset.clMu.Unlock()
+	_, err = js.Publish("foo", nil, nats.ExpectLastSequencePerSubject(0))
+	require_Error(t, err, NewJSStreamWrongLastSequenceConstantError())
+
 	// Allow updates when ready and subject not already in process.
 	mset.clMu.Lock()
 	mset.expectedPerSubjectSequence = nil
 	mset.expectedPerSubjectInProcess = nil
+	mset.inflight = nil
 	mset.clMu.Unlock()
 	pa, err := js.Publish("foo", nil, nats.ExpectLastSequencePerSubject(0))
 	require_NoError(t, err)