Repository: nats-io/nats-server
Pull request #2878 (2022-02-17 16:53:32)

Cleanup key files when removing message blocks.
Signed-off-by: Derek Collison <derek@nats.io>

/cc @nats-io/core



Patch:
diff --git a/server/consumer.go b/server/consumer.go
index 00515fc1..f64314ef 100644
--- a/server/consumer.go
+++ b/server/consumer.go
@@ -2085,16 +2085,6 @@ func (o *consumer) needAck(sseq uint64) bool {
 	var asflr, osseq uint64
 	var pending map[uint64]*Pending
 	o.mu.RLock()
-
-	// Check first if we are filtered, and if so check if this is even applicable to us.
-	if o.isFiltered() && o.mset != nil {
-		subj, _, _, _, err := o.mset.store.LoadMsg(sseq)
-		if err != nil || !o.isFilteredMatch(subj) {
-			o.mu.RUnlock()
-			return false
-		}
-	}
-
 	if o.isLeader() {
 		asflr, osseq = o.asflr, o.sseq
 		pending = o.pending
diff --git a/server/filestore.go b/server/filestore.go
index 15e18a17..b37b183b 100644
--- a/server/filestore.go
+++ b/server/filestore.go
@@ -122,6 +122,7 @@ type msgBlock struct {
 	msgs    uint64 // User visible message count.
 	fss     map[string]*SimpleState
 	sfn     string
+	kfn     string
 	lwits   int64
 	lwts    int64
 	llts    int64
@@ -192,6 +193,8 @@ const (
 	fssScanAll = "*.fss"
 	// used to store our block encryption key.
 	keyScan = "%d.key"
+	// to look for orphans
+	keyScanAll = "*.key"
 	// This is where we keep state on consumers.
 	consumerDir = "obs"
 	// Index file for a consumer.
@@ -918,6 +921,23 @@ func (fs *fileStore) recoverMsgs() error {
 			os.Remove(fn)
 		}
 	}
+	// Same bug for keyfiles but for these we just need to identify orphans.
+	if kms, err := filepath.Glob(path.Join(mdir, keyScanAll)); err == nil && len(kms) > 0 {
+		valid := make(map[uint64]bool)
+		for _, mb := range fs.blks {
+			valid[mb.index] = true
+		}
+		for _, fn := range kms {
+			var index uint64
+			shouldRemove := true
+			if n, err := fmt.Sscanf(path.Base(fn), keyScan, &index); err == nil && n == 1 && valid[index] {
+				shouldRemove = false
+			}
+			if shouldRemove {
+				os.Remove(fn)
+			}
+		}
+	}
 
 	// Limits checks and enforcement.
 	fs.enforceMsgLimit()
@@ -1564,6 +1584,7 @@ func (fs *fileStore) genEncryptionKeysForBlock(mb *msgBlock) error {
 	if err := ioutil.WriteFile(keyFile, encrypted, defaultFilePerms); err != nil {
 		return err
 	}
+	mb.kfn = keyFile
 	return nil
 }
 
@@ -1794,9 +1815,11 @@ func (fs *fileStore) enforceBytesLimit() {
 	}
 }
 
-// Lock should be held.
+// Lock should be held on entry but will be released during actual remove.
 func (fs *fileStore) deleteFirstMsg() (bool, error) {
-	return fs.removeMsg(fs.state.FirstSeq, false, false)
+	fs.mu.Unlock()
+	defer fs.mu.Lock()
+	return fs.removeMsg(fs.state.FirstSeq, false, true)
 }
 
 // RemoveMsg will remove the message from this store.
@@ -1921,8 +1944,7 @@ func (fs *fileStore) removeMsg(seq uint64, secure, needFSLock bool) (bool, error
 	}
 
 	// Optimize for FIFO case.
-	fifo := seq == mb.first.seq
-	if fifo {
+	if seq == mb.first.seq {
 		mb.selectNextFirst()
 		if mb.isEmpty() {
 			fs.removeMsgBlock(mb)
@@ -1964,13 +1986,6 @@ func (fs *fileStore) removeMsg(seq uint64, secure, needFSLock bool) (bool, error
 			fs.rebuildStateLocked(ld)
 		}
 	}
-	// Check if we need to write the index file and we are flush in place (fip).
-	if shouldWriteIndex && fs.fip {
-		// Check if this is the first message, common during expirations etc.
-		if !fifo || time.Now().UnixNano()-mb.lwits > int64(2*time.Second) {
-			mb.writeIndexInfoLocked()
-		}
-	}
 	mb.mu.Unlock()
 
 	// Kick outside of lock.
@@ -1983,6 +1998,8 @@ func (fs *fileStore) removeMsg(seq uint64, secure, needFSLock bool) (bool, error
 			case fch <- struct{}{}:
 			default:
 			}
+		} else {
+			mb.writeIndexInfo()
 		}
 	}
 
@@ -4344,6 +4361,9 @@ func (mb *msgBlock) dirtyCloseWithRemove(remove bool) {
 			os.Remove(mb.sfn)
 			mb.sfn = _EMPTY_
 		}
+		if mb.kfn != _EMPTY_ {
+			os.Remove(mb.kfn)
+		}
 	}
 }
 


Test patch:
diff --git a/server/jetstream_cluster_test.go b/server/jetstream_cluster_test.go
index fe598dde..511709ba 100644
--- a/server/jetstream_cluster_test.go
+++ b/server/jetstream_cluster_test.go
@@ -10667,75 +10667,6 @@ func TestJetStreamClusterStreamTagPlacement(t *testing.T) {
 	placeErr("C1", []string{"cloud:DO"})
 }
 
-func TestJetStreamClusterInterestRetentionWithFilteredConsumersExtra(t *testing.T) {
-	c := createJetStreamClusterExplicit(t, "R3S", 3)
-	defer c.shutdown()
-
-	subjectNameZero := "foo.bar"
-	subjectNameOne := "foo.baz"
-
-	// Client based API
-	nc, js := jsClientConnect(t, c.randomServer())
-	defer nc.Close()
-
-	_, err := js.AddStream(&nats.StreamConfig{Name: "TEST", Subjects: []string{"foo.*"}, Retention: nats.InterestPolicy, Replicas: 3})
-	require_NoError(t, err)
-
-	checkState := func(expected uint64) {
-		t.Helper()
-		checkFor(t, 5*time.Second, 100*time.Millisecond, func() error {
-			si, err := js.StreamInfo("TEST")
-			require_NoError(t, err)
-			if si.State.Msgs != expected {
-				return fmt.Errorf("Expected %d msgs, got %d", expected, si.State.Msgs)
-			}
-			return nil
-		})
-	}
-
-	subZero, err := js.PullSubscribe(subjectNameZero, "dlc-0")
-	require_NoError(t, err)
-
-	subOne, err := js.PullSubscribe(subjectNameOne, "dlc-1")
-	require_NoError(t, err)
-
-	msg := []byte("FILTERED")
-	// Now send a bunch of messages
-	for i := 0; i < 1000; i++ {
-		_, err = js.PublishAsync(subjectNameZero, msg)
-		require_NoError(t, err)
-		_, err = js.PublishAsync(subjectNameOne, msg)
-		require_NoError(t, err)
-	}
-
-	// should be 2000 in total
-	checkState(2000)
-
-	// fetch and acknowledge, count records to ensure no errors acknowledging
-	getAndAckBatch := func(sub *nats.Subscription) {
-		t.Helper()
-		successCounter := 0
-		msgs, err := sub.Fetch(1000)
-		require_NoError(t, err)
-
-		for _, m := range msgs {
-			err = m.Ack()
-			require_NoError(t, err)
-			successCounter++
-		}
-		if successCounter != 1000 {
-			t.Fatalf("Unexpected number of acknowledges %d for subscription %v", successCounter, sub)
-		}
-	}
-
-	// fetch records subscription zero
-	getAndAckBatch(subZero)
-	// fetch records for subscription one
-	getAndAckBatch(subOne)
-	// Make sure stream is zero.
-	checkState(0)
-}
-
 // Support functions
 
 // Used to setup superclusters for tests.
diff --git a/server/jetstream_test.go b/server/jetstream_test.go
index 96b03da8..e85f7bc0 100644
--- a/server/jetstream_test.go
+++ b/server/jetstream_test.go
@@ -4594,45 +4594,6 @@ func TestJetStreamPubAckPerf(t *testing.T) {
 	fmt.Printf("%.0f msgs/sec\n", float64(toSend)/tt.Seconds())
 }
 
-func TestJetStreamPubPerfWithFullStream(t *testing.T) {
-	// Comment out to run, holding place for now.
-	t.SkipNow()
-
-	s := RunBasicJetStreamServer()
-	defer s.Shutdown()
-
-	if config := s.JetStreamConfig(); config != nil {
-		defer removeDir(t, config.StoreDir)
-	}
-
-	nc, js := jsClientConnect(t, s)
-	defer nc.Close()
-
-	toSend, msg := 1_000_000, []byte("OK")
-
-	_, err := js.AddStream(&nats.StreamConfig{Name: "foo", MaxMsgs: int64(toSend)})
-	require_NoError(t, err)
-
-	start := time.Now()
-	for i := 0; i < toSend; i++ {
-		js.PublishAsync("foo", msg)
-	}
-	<-js.PublishAsyncComplete()
-	tt := time.Since(start)
-	fmt.Printf("time is %v\n", tt)
-	fmt.Printf("%.0f msgs/sec\n", float64(toSend)/tt.Seconds())
-
-	// Now do same amount but knowing we are at our limit.
-	start = time.Now()
-	for i := 0; i < toSend; i++ {
-		js.PublishAsync("foo", msg)
-	}
-	<-js.PublishAsyncComplete()
-	tt = time.Since(start)
-	fmt.Printf("\ntime is %v\n", tt)
-	fmt.Printf("%.0f msgs/sec\n", float64(toSend)/tt.Seconds())
-}
-
 func TestJetStreamSnapshotsAPIPerf(t *testing.T) {
 	// Comment out to run, holding place for now.
 	t.SkipNow()
diff --git a/server/norace_test.go b/server/norace_test.go
index 80164b76..117a9ee0 100644
--- a/server/norace_test.go
+++ b/server/norace_test.go
@@ -40,7 +40,9 @@ import (
 	"testing"
 	"time"
 
+	"crypto/hmac"
 	crand "crypto/rand"
+	"crypto/sha256"
 
 	"github.com/klauspost/compress/s2"
 	"github.com/nats-io/jwt/v2"
@@ -4324,3 +4326,50 @@ func TestNoRaceFileStoreSubjectInfoWithSnapshotCleanup(t *testing.T) {
 		t.Fatalf("Expected to find no fss files, found %d", len(fms))
 	}
 }
+
+func TestNoRaceFileStoreKeyFileCleanup(t *testing.T) {
+	storeDir := createDir(t, JetStreamStoreDir)
+	defer removeDir(t, storeDir)
+
+	prf := func(context []byte) ([]byte, error) {
+		h := hmac.New(sha256.New, []byte("dlc22"))
+		if _, err := h.Write(context); err != nil {
+			return nil, err
+		}
+		return h.Sum(nil), nil
+	}
+
+	fs, err := newFileStoreWithCreated(
+		FileStoreConfig{StoreDir: storeDir, BlockSize: 1024 * 1024},
+		StreamConfig{Name: "TEST", Storage: FileStorage},
+		time.Now(),
+		prf)
+	require_NoError(t, err)
+	defer fs.Stop()
+
+	n, msg := 10_000, []byte(strings.Repeat("Z", 1024))
+	for i := 0; i < n; i++ {
+		_, _, err := fs.StoreMsg(fmt.Sprintf("X.%d", i), nil, msg)
+		require_NoError(t, err)
+	}
+
+	var seqs []uint64
+	for i := 1; i <= n; i++ {
+		seqs = append(seqs, uint64(i))
+	}
+	// Randomly delete msgs, make sure we cleanup as we empty the message blocks.
+	rand.Shuffle(len(seqs), func(i, j int) { seqs[i], seqs[j] = seqs[j], seqs[i] })
+
+	for _, seq := range seqs {
+		_, err := fs.RemoveMsg(seq)
+		require_NoError(t, err)
+	}
+
+	// We will have cleanup the main .blk and .idx sans the lmb, but we should not have any *.fss files.
+	kms, err := filepath.Glob(path.Join(storeDir, msgDir, keyScanAll))
+	require_NoError(t, err)
+
+	if len(kms) > 1 {
+		t.Fatalf("Expected to find only 1 key file, found %d", len(kms))
+	}
+}