Repository: nats-io/nats-server
Pull request #3579 (2022-10-26 21:37:39)

[FIXED] MaxMsgPerSubject recovered on corrupt or missing index file.
On abnormal server exit, for streams or KV with max msgs per subject set we could recover more than N msgs per subject.
This fix allows for recovery of correct state on restart when index files are missing or not current.

Signed-off-by: Derek Collison <derek@nats.io>

/cc @nats-io/core



Patch:
diff --git a/server/filestore.go b/server/filestore.go
index d681caa7..1adbfbb3 100644
--- a/server/filestore.go
+++ b/server/filestore.go
@@ -1178,6 +1178,11 @@ func (fs *fileStore) recoverMsgs() error {
 		fs.startAgeChk()
 	}
 
+	// If we have max msgs per subject make sure the is also enforced.
+	if fs.cfg.MaxMsgsPer > 0 {
+		fs.enforceMsgPerSubjectLimit()
+	}
+
 	return nil
 }
 
@@ -2136,6 +2141,71 @@ func (fs *fileStore) enforceBytesLimit() {
 	}
 }
 
+// Will make sure we have limits honored for max msgs per subject on recovery.
+// We will make sure to go through all msg blocks etc. but in practice this
+// will most likely only be the last one, so can take a more conservative approach.
+// Lock should be held.
+func (fs *fileStore) enforceMsgPerSubjectLimit() {
+	maxMsgsPer := uint64(fs.cfg.MaxMsgsPer)
+
+	// We want to suppress callbacks from remove during this process
+	// since these should have already been deleted and accounted for.
+	cb := fs.scb
+	fs.scb = nil
+	defer func() { fs.scb = cb }()
+
+	// collect all that are not correct.
+	needAttention := make(map[string]*psi)
+	for subj, psi := range fs.psim {
+		if psi.total > maxMsgsPer {
+			needAttention[subj] = psi
+		}
+	}
+
+	// Collect all the msgBlks we alter.
+	blks := make(map[*msgBlock]struct{})
+
+	// For re-use below.
+	var sm StoreMsg
+
+	// Walk all subjects that need attention here.
+	for subj, info := range needAttention {
+		total, start, stop := info.total, info.fblk, info.lblk
+
+		for i := start; i <= stop; i++ {
+			mb := fs.bim[i]
+			if mb == nil {
+				continue
+			}
+			// Grab the ss entry for this subject in case sparse.
+			mb.mu.Lock()
+			mb.ensurePerSubjectInfoLoaded()
+			ss := mb.fss[subj]
+			mb.mu.Unlock()
+
+			for seq := ss.First; seq <= ss.Last && total > maxMsgsPer; {
+				m, _, err := mb.firstMatching(subj, false, seq, &sm)
+				if err == nil {
+					seq = m.seq + 1
+					if removed, _ := fs.removeMsg(m.seq, false, false); removed {
+						total--
+						blks[mb] = struct{}{}
+					}
+				} else {
+					// On error just do single increment.
+					seq++
+				}
+			}
+		}
+	}
+
+	// Now write updated index for all affected msgBlks.
+	for mb := range blks {
+		mb.writeIndexInfo()
+		mb.tryForceExpireCacheLocked()
+	}
+}
+
 // Lock should be held.
 func (fs *fileStore) deleteFirstMsg() (bool, error) {
 	return fs.removeMsg(fs.state.FirstSeq, false, false)


Test patch:
diff --git a/server/filestore_test.go b/server/filestore_test.go
index 0e897b9c..a3bb3bfe 100644
--- a/server/filestore_test.go
+++ b/server/filestore_test.go
@@ -4724,3 +4724,77 @@ func TestFileStoreOutOfSpaceRebuildState(t *testing.T) {
 		t.Fatalf("Subject state expected to be\n  %+v\nvs\n  %+v", ss, nss)
 	}
 }
+
+func TestFileStoreRebuildStateProperlyWithMaxMsgsPerSubject(t *testing.T) {
+	storeDir := createDir(t, JetStreamStoreDir)
+	defer removeDir(t, storeDir)
+
+	fs, err := newFileStore(
+		FileStoreConfig{StoreDir: storeDir, BlockSize: 4096},
+		StreamConfig{Name: "zzz", Subjects: []string{"foo", "bar", "baz"}, Storage: FileStorage, MaxMsgsPer: 1},
+	)
+	require_NoError(t, err)
+	defer fs.Stop()
+
+	// Send one to baz at beginning.
+	_, _, err = fs.StoreMsg("baz", nil, nil)
+	require_NoError(t, err)
+
+	ns := 1000
+	for i := 1; i <= ns; i++ {
+		_, _, err := fs.StoreMsg("foo", nil, nil)
+		require_NoError(t, err)
+		_, _, err = fs.StoreMsg("bar", nil, nil)
+		require_NoError(t, err)
+	}
+
+	checkState := func() {
+		var ss StreamState
+		fs.FastState(&ss)
+		if ss.NumSubjects != 3 {
+			t.Fatalf("Expected NumSubjects of 3, got %d", ss.NumSubjects)
+		}
+		if ss.Msgs != 3 {
+			t.Fatalf("Expected NumMsgs of 3, got %d", ss.Msgs)
+		}
+	}
+
+	checkState()
+
+	// Stop filestore but invalidate the idx files by removing them.
+	// This will simulate a server panic or kill -9 scenario.
+	fs.Stop()
+
+	fs.mu.RLock()
+	for _, mb := range fs.blks {
+		mb.removeIndexFile()
+	}
+	fs.mu.RUnlock()
+
+	fs, err = newFileStore(
+		FileStoreConfig{StoreDir: storeDir, BlockSize: 4096},
+		StreamConfig{Name: "zzz", Subjects: []string{"foo", "bar", "baz"}, Storage: FileStorage, MaxMsgsPer: 1},
+	)
+	require_NoError(t, err)
+	defer fs.Stop()
+
+	checkState()
+
+	// Make sure we wrote all index files from recovery.
+	fs.mu.RLock()
+	for _, mb := range fs.blks {
+		mb.mu.Lock()
+		if err := mb.readIndexInfo(); err != nil {
+			mb.mu.Unlock()
+			fs.mu.RUnlock()
+			t.Fatalf("Unexpected error reading index info: %v", err)
+		}
+		if mb.msgs == 0 {
+			mb.mu.Unlock()
+			fs.mu.RUnlock()
+			t.Fatalf("Expected msgs for all blks, got none for index %d", mb.index)
+		}
+		mb.mu.Unlock()
+	}
+	fs.mu.RUnlock()
+}