(2.14) [IMPROVED] JetStream header indexing

Signed-off-by: Maurice van Veen <[email protected]>

Author: MauriceVanVeen

server/filestore.go

@@ -1556,7 +1556,7 @@ func (mb *msgBlock) rebuildStateLocked() (*LostStreamData, []uint64, error) {
 		hasHeaders := rl&hbit != 0
 		var ttl int64
 		if mb.fs.ttls != nil && len(hdr) > 0 {
-			ttl, _ = getMessageTTL(hdr)
+			ttl, _ = getMessageTTLNoIdx(hdr)
 		}
 		// Clear any headers bit that could be set.
 		rl &^= hbit
@@ -1649,7 +1649,7 @@ func (mb *msgBlock) rebuildStateLocked() (*LostStreamData, []uint64, error) {
 				mb.ttls++
 			}
 			if mb.fs.scheduling != nil {
-				if schedule, ok := getMessageSchedule(hdr); ok && !schedule.IsZero() {
+				if schedule, ok := getMessageScheduleNoIdx(hdr); ok && !schedule.IsZero() {
 					mb.fs.scheduling.add(seq, string(subj), schedule.UnixNano())
 					mb.schedules++
 				}
@@ -2092,7 +2092,7 @@ func (fs *fileStore) recoverTTLState() error {
 			if len(msg.hdr) == 0 {
 				continue
 			}
-			if ttl, _ := getMessageTTL(msg.hdr); ttl > 0 {
+			if ttl, _ := getMessageTTLNoIdx(msg.hdr); ttl > 0 {
 				expires := time.Duration(msg.ts) + (time.Second * time.Duration(ttl))
 				fs.ttls.Add(seq, int64(expires))
 			}
@@ -2173,7 +2173,7 @@ func (fs *fileStore) recoverMsgSchedulingState() error {
 			if len(msg.hdr) == 0 {
 				continue
 			}
-			if schedule, ok := getMessageSchedule(sm.hdr); ok && !schedule.IsZero() {
+			if schedule, ok := getMessageScheduleNoIdx(sm.hdr); ok && !schedule.IsZero() {
 				fs.scheduling.init(seq, sm.subj, schedule.UnixNano())
 			}
 		}
@@ -4455,7 +4455,7 @@ func (fs *fileStore) storeRawMsg(subj string, hdr, msg []byte, seq uint64, ts, t
 
 	// Message scheduling.
 	if fs.scheduling != nil {
-		if schedule, ok := getMessageSchedule(hdr); ok && !schedule.IsZero() {
+		if schedule, ok := getMessageScheduleNoIdx(hdr); ok && !schedule.IsZero() {
 			fs.scheduling.add(seq, subj, schedule.UnixNano())
 			fs.lmb.schedules++
 		} else {
@@ -5071,7 +5071,7 @@ func (fs *fileStore) removeMsg(seq uint64, secure, viaLimits, needFSLock bool) (
 	mb.removeSeqPerSubject(sm.subj, seq)
 	fs.removePerSubject(sm.subj)
 	if fs.ttls != nil {
-		if ttl, err := getMessageTTL(sm.hdr); err == nil {
+		if ttl, err := getMessageTTLNoIdx(sm.hdr); err == nil {
 			expires := time.Duration(sm.ts) + (time.Second * time.Duration(ttl))
 			fs.ttls.Remove(seq, int64(expires))
 		}
@@ -5997,7 +5997,7 @@ func (fs *fileStore) expireMsgs() {
 		var seq uint64
 		for sm, seq, _ = fs.LoadNextMsg(fwcs, true, 0, &smv); sm != nil && sm.ts <= minAge; sm, seq, _ = fs.LoadNextMsg(fwcs, true, seq+1, &smv) {
 			if len(sm.hdr) > 0 {
-				if ttl, err := getMessageTTL(sm.hdr); err == nil && ttl < 0 {
+				if ttl, err := getMessageTTLNoIdx(sm.hdr); err == nil && ttl < 0 {
 					// The message has a negative TTL, therefore it must "never expire".
 					minAge = ats.AccessTime() - maxAge
 					continue

server/jetstream.go

@@ -1473,8 +1473,8 @@ func (a *Account) EnableJetStream(limits map[string]JetStreamAccountLimits) erro
 			if err != nil || sm == nil {
 				goto SKIP
 			}
-			batchId = getBatchId(sm.hdr)
-			batchSeq, ok = getBatchSequence(sm.hdr)
+			batchId = getBatchIdNoIdx(sm.hdr)
+			batchSeq, ok = getBatchSequenceNoIdx(sm.hdr)
 			commit = len(sliceHeader(JSBatchCommit, sm.hdr)) != 0
 			if batchId == _EMPTY_ || !ok || commit {
 				goto SKIP
@@ -1520,7 +1520,10 @@ func (a *Account) EnableJetStream(limits map[string]JetStreamAccountLimits) erro
 				if commitEob && seq == state.LastSeq {
 					hdr = genHeader(hdr, JSBatchCommit, "1")
 				}
-				mset.processJetStreamMsg(sm.subj, _EMPTY_, hdr, sm.msg, 0, 0, nil, false, true)
+				var hdrIdx *jsHdrIndex
+				hdr, hdrIdx = indexJsHdr(hdr)
+				mset.processJetStreamMsg(sm.subj, _EMPTY_, hdr, hdrIdx, sm.msg, 0, 0, nil, false, true)
+				hdrIdx.returnToPool()
 			}
 			store.Delete(true)
 		SKIP:

server/jetstream_batching.go

@@ -238,7 +238,7 @@ func (batch *batchApply) rejectBatchState(mset *stream) {
 // mset.mu lock must NOT be held or used.
 // mset.clMu lock must be held.
 func checkMsgHeadersPreClusteredProposal(
-	diff *batchStagedDiff, mset *stream, subject string, hdr []byte, msg []byte, sourced bool, name string,
+	diff *batchStagedDiff, mset *stream, subject string, hdr []byte, hdrIdx *jsHdrIndex, msg []byte, sourced bool, name string,
 	jsa *jsAccount, allowRollup, denyPurge, allowTTL, allowMsgCounter, allowMsgSchedules bool,
 	discard DiscardPolicy, discardNewPer bool, maxMsgSize int, maxMsgs int64, maxMsgsPer int64, maxBytes int64,
 ) ([]byte, []byte, uint64, *ApiError, error) {
@@ -252,10 +252,13 @@ func checkMsgHeadersPreClusteredProposal(
 			err := fmt.Errorf("JetStream header size exceeds limits for '%s > %s'", jsa.acc().Name, mset.cfg.Name)
 			return hdr, msg, 0, NewJSStreamHeaderExceedsMaximumError(), err
 		}
+	}
+
+	if hdrIdx != nil {
 		// 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 {
+		if incr, ok = getMessageIncr(hdrIdx); !ok {
 			apiErr := NewJSMessageIncrInvalidError()
 			return hdr, msg, 0, apiErr, apiErr
 		} else if incr != nil && !sourced {
@@ -269,14 +272,14 @@ func checkMsgHeadersPreClusteredProposal(
 			} else {
 				// Check for incompatible headers.
 				var doErr bool
-				if getRollup(hdr) != _EMPTY_ ||
-					getExpectedStream(hdr) != _EMPTY_ ||
-					getExpectedLastMsgId(hdr) != _EMPTY_ ||
-					getExpectedLastSeqPerSubjectForSubject(hdr) != _EMPTY_ {
+				if getRollup(hdrIdx) != _EMPTY_ ||
+					getExpectedStream(hdrIdx) != _EMPTY_ ||
+					getExpectedLastMsgId(hdrIdx) != _EMPTY_ ||
+					getExpectedLastSeqPerSubjectForSubject(hdrIdx) != _EMPTY_ {
 					doErr = true
-				} else if _, ok = getExpectedLastSeq(hdr); ok {
+				} else if _, ok = getExpectedLastSeq(hdrIdx); ok {
 					doErr = true
-				} else if _, ok = getExpectedLastSeqPerSubject(hdr); ok {
+				} else if _, ok = getExpectedLastSeqPerSubject(hdrIdx); ok {
 					doErr = true
 				}
 
@@ -287,11 +290,11 @@ func checkMsgHeadersPreClusteredProposal(
 			}
 		}
 		// Expected stream name can also be pre-checked.
-		if sname := getExpectedStream(hdr); sname != _EMPTY_ && sname != name {
+		if sname := getExpectedStream(hdrIdx); sname != _EMPTY_ && sname != name {
 			return hdr, msg, 0, NewJSStreamNotMatchError(), errStreamMismatch
 		}
 		// TTL'd messages are rejected entirely if TTLs are not enabled on the stream, or if the TTL is invalid.
-		if ttl, err := getMessageTTL(hdr); !sourced && (ttl != 0 || err != nil) {
+		if ttl, err := getMessageTTL(hdrIdx); !sourced && (ttl != 0 || err != nil) {
 			if !allowTTL {
 				return hdr, msg, 0, NewJSMessageTTLDisabledError(), errMsgTTLDisabled
 			} else if err != nil {
@@ -300,7 +303,7 @@ func checkMsgHeadersPreClusteredProposal(
 		}
 		// Check for MsgIds here at the cluster level to avoid excessive CLFS accounting.
 		// Will help during restarts.
-		if msgId := getMsgId(hdr); msgId != _EMPTY_ {
+		if msgId := getMsgId(hdrIdx); msgId != _EMPTY_ {
 			// Dedupe if staged.
 			if _, ok = diff.msgIds[msgId]; ok {
 				return hdr, msg, 0, NewJSAtomicPublishContainsDuplicateMessageError(), errMsgIdDuplicate
@@ -439,9 +442,9 @@ func checkMsgHeadersPreClusteredProposal(
 		}
 	}
 
-	if len(hdr) > 0 {
+	if hdrIdx != nil {
 		// Expected last sequence.
-		if seq, exists := getExpectedLastSeq(hdr); exists && seq != mset.clseq-mset.clfs {
+		if seq, exists := getExpectedLastSeq(hdrIdx); 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
@@ -452,10 +455,10 @@ func checkMsgHeadersPreClusteredProposal(
 		}
 
 		// Expected last sequence per subject.
-		if seq, exists := getExpectedLastSeqPerSubject(hdr); exists {
+		if seq, exists := getExpectedLastSeqPerSubject(hdrIdx); exists {
 			// Allow override of the subject used for the check.
 			seqSubj := subject
-			if optSubj := getExpectedLastSeqPerSubjectForSubject(hdr); optSubj != _EMPTY_ {
+			if optSubj := getExpectedLastSeqPerSubjectForSubject(hdrIdx); optSubj != _EMPTY_ {
 				seqSubj = optSubj
 			}
 
@@ -509,13 +512,13 @@ func checkMsgHeadersPreClusteredProposal(
 					diff.expectedPerSubject[seqSubj] = e
 				}
 			}
-		} else if getExpectedLastSeqPerSubjectForSubject(hdr) != _EMPTY_ {
+		} else if getExpectedLastSeqPerSubjectForSubject(hdrIdx) != _EMPTY_ {
 			apiErr := NewJSStreamExpectedLastSeqPerSubjectInvalidError()
 			return hdr, msg, 0, apiErr, apiErr
 		}
 
 		// Message scheduling.
-		if schedule, ok := getMessageSchedule(hdr); !ok {
+		if schedule, ok := getMessageSchedule(hdrIdx); !ok {
 			apiErr := NewJSMessageSchedulesPatternInvalidError()
 			if !allowMsgSchedules {
 				apiErr = NewJSMessageSchedulesDisabledError()
@@ -525,12 +528,12 @@ func checkMsgHeadersPreClusteredProposal(
 			if !allowMsgSchedules {
 				apiErr := NewJSMessageSchedulesDisabledError()
 				return hdr, msg, 0, apiErr, apiErr
-			} else if scheduleTtl, ok := getMessageScheduleTTL(hdr); !ok {
+			} else if scheduleTtl, ok := getMessageScheduleTTL(hdrIdx); !ok {
 				apiErr := NewJSMessageSchedulesTTLInvalidError()
 				return hdr, msg, 0, apiErr, apiErr
 			} else if scheduleTtl != _EMPTY_ && !allowTTL {
 				return hdr, msg, 0, NewJSMessageTTLDisabledError(), errMsgTTLDisabled
-			} else if scheduleTarget := getMessageScheduleTarget(hdr); scheduleTarget == _EMPTY_ ||
+			} else if scheduleTarget := getMessageScheduleTarget(hdrIdx); scheduleTarget == _EMPTY_ ||
 				!IsValidPublishSubject(scheduleTarget) || SubjectsCollide(scheduleTarget, subject) {
 				apiErr := NewJSMessageSchedulesTargetInvalidError()
 				return hdr, msg, 0, apiErr, apiErr
@@ -547,7 +550,7 @@ func checkMsgHeadersPreClusteredProposal(
 
 				// Add a rollup sub header if it doesn't already exist.
 				// Otherwise, it must exist already as a rollup on the subject.
-				if rollup := getRollup(hdr); rollup == _EMPTY_ {
+				if rollup := getRollup(hdrIdx); rollup == _EMPTY_ {
 					hdr = genHeader(hdr, JSMsgRollup, JSMsgRollupSubject)
 				} else if rollup != JSMsgRollupSubject {
 					apiErr := NewJSMessageSchedulesRollupInvalidError()
@@ -557,7 +560,7 @@ func checkMsgHeadersPreClusteredProposal(
 		}
 
 		// Check for any rollups.
-		if rollup := getRollup(hdr); rollup != _EMPTY_ {
+		if rollup := getRollup(hdrIdx); rollup != _EMPTY_ {
 			if !allowRollup || denyPurge {
 				err := errors.New("rollup not permitted")
 				return hdr, msg, 0, NewJSStreamRollupFailedError(err), err

server/jetstream_batching_test.go

@@ -1388,7 +1388,9 @@ func TestJetStreamAtomicBatchPublishStageAndCommit(t *testing.T) {
 						hdr = genHeader(hdr, key, value)
 					}
 				}
-				_, _, _, _, err = checkMsgHeadersPreClusteredProposal(diff, mset, m.subject, hdr, m.msg, false, "TEST", nil, test.allowRollup, test.denyPurge, test.allowTTL, test.allowMsgCounter, test.allowMsgSchedules, discard, discardNewPer, -1, maxMsgs, maxMsgsPer, maxBytes)
+				_, hdrIdx := indexJsHdr(hdr)
+				_, _, _, _, err = checkMsgHeadersPreClusteredProposal(diff, mset, m.subject, hdr, hdrIdx, m.msg, false, "TEST", nil, test.allowRollup, test.denyPurge, test.allowTTL, test.allowMsgCounter, test.allowMsgSchedules, discard, discardNewPer, -1, maxMsgs, maxMsgsPer, maxBytes)
+				hdrIdx.returnToPool()
 				if m.err != nil {
 					require_Error(t, err, m.err)
 				} else if err != nil {
@@ -1582,7 +1584,9 @@ 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, true)
+	_, hdrIdx := indexJsHdr(hdr1)
+	err = mset.processJetStreamMsg("foo", _EMPTY_, hdr1, hdrIdx, nil, 0, 0, nil, false, true)
+	hdrIdx.returnToPool()
 	require_NoError(t, err)
 
 	// Simulate a hard kill, upon recovery the rest of the batch should be applied.
@@ -1672,7 +1676,9 @@ func TestJetStreamAtomicBatchPublishSingleServerRecoveryCommitEob(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, true)
+	_, hdrIdx := indexJsHdr(hdr1)
+	err = mset.processJetStreamMsg("foo", _EMPTY_, hdr1, hdrIdx, nil, 0, 0, nil, false, true)
+	hdrIdx.returnToPool()
 	require_NoError(t, err)
 
 	// Simulate a hard kill, upon recovery the rest of the batch should be applied.

server/jetstream_cluster.go

@@ -3665,7 +3665,10 @@ 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, needLock)
+	var hdrIdx *jsHdrIndex
+	hdr, hdrIdx = indexJsHdr(hdr)
+	err = mset.processJetStreamMsg(subject, reply, hdr, hdrIdx, msg, lseq, ts, mt, sourced, needLock)
+	hdrIdx.returnToPool()
 
 	// If we have inflight make sure to clear after processing.
 	// TODO(dlc) - technically check on inflight != nil could cause datarace.
@@ -3732,7 +3735,9 @@ 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, needLock)
+				hdr, hdrIdx = indexJsHdr(hdr)
+				err = mset.processJetStreamMsg(subject, reply, hdr, hdrIdx, msg, lseq, ts, mt, sourced, needLock)
+				hdrIdx.returnToPool()
 			}
 			// FIXME(dlc) - We could just run a catchup with a request defining the span between what we expected
 			// and what we got.
@@ -8776,7 +8781,7 @@ func (mset *stream) stateSnapshotLocked() []byte {
 const streamLagWarnThreshold = 10_000
 
 // processClusteredInboundMsg will propose the inbound message to the underlying raft group.
-func (mset *stream) processClusteredInboundMsg(subject, reply string, hdr, msg []byte, mt *msgTrace, sourced bool) (retErr error) {
+func (mset *stream) processClusteredInboundMsg(subject, reply string, hdr []byte, hdrIdx *jsHdrIndex, msg []byte, mt *msgTrace, sourced bool) (retErr error) {
 	// For possible error response.
 	var response []byte
 
@@ -8794,7 +8799,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, true)
+		return mset.processJetStreamMsg(subject, reply, hdr, hdrIdx, msg, 0, 0, mt, sourced, true)
 	}
 
 	// If message tracing (with message delivery), we will need to send the
@@ -8898,7 +8903,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, allowRollup, denyPurge, allowTTL, allowMsgCounter, allowMsgSchedules, discard, discardNewPer, maxMsgSize, maxMsgs, maxMsgsPer, maxBytes); err != nil {
+	if hdr, msg, dseq, apiErr, err = checkMsgHeadersPreClusteredProposal(diff, mset, subject, hdr, hdrIdx, msg, sourced, name, jsa, allowRollup, denyPurge, allowTTL, allowMsgCounter, allowMsgSchedules, discard, discardNewPer, maxMsgSize, maxMsgs, maxMsgsPer, maxBytes); err != nil {
 		mset.clMu.Unlock()
 		if err == errMsgIdDuplicate && dseq > 0 {
 			var buf [256]byte
@@ -9486,7 +9491,7 @@ func (mset *stream) processCatchupMsg(msg []byte) (uint64, error) {
 	// Find the message TTL if any.
 	// TODO(nat): If the TTL isn't valid by this stage then there isn't really a
 	// lot we can do about it, as we'd break the catchup if we reject the message.
-	ttl, _ := getMessageTTL(hdr)
+	ttl, _ := getMessageTTLNoIdx(hdr)
 
 	// Put into our store
 	// Messages to be skipped have no subject or timestamp.
@@ -9506,7 +9511,7 @@ func (mset *stream) processCatchupMsg(msg []byte) (uint64, error) {
 
 	// Check for MsgId and if we have one here make sure to update our internal map.
 	if len(hdr) > 0 {
-		if msgId := getMsgId(hdr); msgId != _EMPTY_ {
+		if msgId := getMsgIdNoIdx(hdr); msgId != _EMPTY_ {
 			mset.ddMu.Lock()
 			mset.storeMsgIdLocked(&ddentry{msgId, seq, ts})
 			mset.ddMu.Unlock()

server/jetstream_cluster_4_test.go

@@ -4379,7 +4379,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, true)
+	err = mset.processJetStreamMsg("foo", _EMPTY_, nil, 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.

server/jetstream_test.go

@@ -22254,3 +22254,33 @@ func TestJetStreamScheduledMessageNotDeactivated(t *testing.T) {
 		})
 	}
 }
+
+func TestJetStreamHdrIndexUpdateHdr(t *testing.T) {
+	updateKey := "Nats-Update-Header"
+	for _, test := range []struct {
+		title     string
+		updateHdr func(hdr []byte)
+	}{
+		{title: "SetHeader", updateHdr: func(hdr []byte) { setHeader(updateKey, "s", hdr) }},
+		{title: "GenHeader", updateHdr: func(hdr []byte) { genHeader(hdr, updateKey, "s") }},
+		{title: "RemoveHeaderIfPresent", updateHdr: func(hdr []byte) { removeHeaderIfPresent(hdr, updateKey) }},
+		{title: "RemoveHeaderIfPrefixPresent", updateHdr: func(hdr []byte) { removeHeaderIfPrefixPresent(hdr, updateKey) }},
+	} {
+		t.Run(test.title, func(t *testing.T) {
+			hdr := genHeader(nil, "Nats-Batch-Id", "uuid")
+			hdr = genHeader(hdr, updateKey, "long_value")
+			hdr = genHeader(hdr, "Nats-Batch-Sequence", "seq")
+
+			var idx *jsHdrIndex
+			hdr, idx = indexJsHdr(hdr)
+			defer idx.returnToPool()
+			require_NotNil(t, idx)
+			require_Equal(t, string(idx.batchId), "uuid")
+			require_Equal(t, string(idx.batchSeq), "seq")
+
+			test.updateHdr(hdr)
+			require_Equal(t, string(idx.batchId), "uuid")
+			require_Equal(t, string(idx.batchSeq), "seq")
+		})
+	}
+}