Repository: nats-io/nats-server
Pull request #2296 (2021-06-18 21:42:12)

Allow for more advanced purge operations.
This will allow purge by subject and the ability to specify the sequence or number of messages to keep.

Signed-off-by: Derek Collison <derek@nats.io>




Patch:
diff --git a/server/filestore.go b/server/filestore.go
index 8390ec99..91d17ffc 100644
--- a/server/filestore.go
+++ b/server/filestore.go
@@ -936,13 +936,7 @@ func (fs *fileStore) FilteredState(sseq uint64, subj string) SimpleState {
 		}
 	} else {
 		// Fallback to linear scan.
-		var eq func(string, string) bool
-		if wc {
-			eq = subjectIsSubsetMatch
-		} else {
-			eq = func(a, b string) bool { return a == b }
-		}
-
+		eq := compareFn(subj)
 		for seq := sseq; seq <= lseq; seq++ {
 			if sm, _ := fs.msgForSeq(seq); sm != nil && eq(sm.subj, subj) {
 				ss.Msgs++
@@ -3109,6 +3103,57 @@ func (fs *fileStore) dmapEntries() int {
 	return total
 }
 
+// Fixed helper for iterating.
+func subjectsEqual(a, b string) bool {
+	return a == b
+}
+
+func subjectsAll(a, b string) bool {
+	return true
+}
+
+func compareFn(subject string) func(string, string) bool {
+	if subject == _EMPTY_ || subject == ">" {
+		return subjectsAll
+	}
+	if subjectHasWildcard(subject) {
+		return subjectIsSubsetMatch
+	}
+	return subjectsEqual
+}
+
+// PurgeEx will remove messages based on subject filters, sequence and number of messages to keep.
+// Will return the number of purged messages.
+func (fs *fileStore) PurgeEx(subject string, sequence, keep uint64) (purged uint64, err error) {
+	if subject == _EMPTY_ || subject == ">" && keep == 0 {
+		return fs.Purge()
+	}
+	eq := compareFn(subject)
+	if ss := fs.FilteredState(1, subject); ss.Msgs > 0 {
+		if keep > 0 {
+			if keep > ss.Msgs {
+				return 0, nil
+			}
+			ss.Msgs -= keep
+		}
+		last := ss.Last
+		if sequence > 0 {
+			last = sequence + 1
+		}
+		for seq := ss.First; seq <= last; seq++ {
+			if sm, _ := fs.msgForSeq(seq); sm != nil && eq(sm.subj, subject) {
+				if ok, _ := fs.removeMsg(sm.seq, false, true); ok {
+					purged++
+					if purged >= ss.Msgs {
+						break
+					}
+				}
+			}
+		}
+	}
+	return purged, nil
+}
+
 // Purge will remove all messages from this store.
 // Will return the number of purged messages.
 func (fs *fileStore) Purge() (uint64, error) {
diff --git a/server/jetstream_api.go b/server/jetstream_api.go
index 6aaeea4d..94a19293 100644
--- a/server/jetstream_api.go
+++ b/server/jetstream_api.go
@@ -352,6 +352,20 @@ type JSApiStreamListResponse struct {
 
 const JSApiStreamListResponseType = "io.nats.jetstream.api.v1.stream_list_response"
 
+// JSApiStreamPurgeRequest is optional request information to the purge API.
+// Subject will filter the purge request to only messages that match the subject, which can have wildcards.
+// Sequence will purge up to but not including this sequence and can be combined with subject filtering.
+// Keep will specify how many messages to keep. This can also be combined with subject filtering.
+// Note that Sequence and Keep are mutually exclusive, so both can not be set at the same time.
+type JSApiStreamPurgeRequest struct {
+	// Purge up to but not including sequence.
+	Sequence uint64 `json:"seq,omitempty"`
+	// Subject to match against messages for the purge command.
+	Subject string `json:"filter,omitempty"`
+	// Number of messages to keep.
+	Keep uint64 `json:"keep,omitempty"`
+}
+
 type JSApiStreamPurgeResponse struct {
 	ApiResponse
 	Success bool   `json:"success,omitempty"`
@@ -2500,11 +2514,22 @@ func (s *Server) jsStreamPurgeRequest(sub *subscription, c *client, subject, rep
 		return
 	}
 
+	var purgeRequest *JSApiStreamPurgeRequest
 	if !isEmptyRequest(msg) {
-		resp.Error = ApiErrors[JSNotEmptyRequestErr]
-		s.sendAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp))
-		return
+		var req JSApiStreamPurgeRequest
+		if err := json.Unmarshal(msg, &req); err != nil {
+			resp.Error = ApiErrors[JSInvalidJSONErr]
+			s.sendAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp))
+			return
+		}
+		if req.Sequence > 0 && req.Keep > 0 {
+			resp.Error = ApiErrors[JSBadRequestErr]
+			s.sendAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp))
+			return
+		}
+		purgeRequest = &req
 	}
+
 	mset, err := acc.lookupStream(stream)
 	if err != nil {
 		resp.Error = ApiErrors[JSStreamNotFoundErr].ErrOr(err)
@@ -2513,11 +2538,11 @@ func (s *Server) jsStreamPurgeRequest(sub *subscription, c *client, subject, rep
 	}
 
 	if s.JetStreamIsClustered() {
-		s.jsClusteredStreamPurgeRequest(ci, acc, mset, stream, subject, reply, rmsg)
+		s.jsClusteredStreamPurgeRequest(ci, acc, mset, stream, subject, reply, rmsg, purgeRequest)
 		return
 	}
 
-	purged, err := mset.purge()
+	purged, err := mset.purge(purgeRequest)
 	if err != nil {
 		resp.Error = ApiErrors[JSStreamGeneralErrorF].ErrOrNewT(err, "{err}", err)
 	} else {
diff --git a/server/jetstream_cluster.go b/server/jetstream_cluster.go
index 050fd81d..321a04bf 100644
--- a/server/jetstream_cluster.go
+++ b/server/jetstream_cluster.go
@@ -128,11 +128,12 @@ type consumerAssignment struct {
 
 // streamPurge is what the stream leader will replicate when purging a stream.
 type streamPurge struct {
-	Client  *ClientInfo `json:"client,omitempty"`
-	Stream  string      `json:"stream"`
-	LastSeq uint64      `json:"last_seq"`
-	Subject string      `json:"subject"`
-	Reply   string      `json:"reply"`
+	Client  *ClientInfo              `json:"client,omitempty"`
+	Stream  string                   `json:"stream"`
+	LastSeq uint64                   `json:"last_seq"`
+	Subject string                   `json:"subject"`
+	Reply   string                   `json:"reply"`
+	Request *JSApiStreamPurgeRequest `json:"request,omitempty"`
 }
 
 // streamMsgDelete is what the stream leader will replicate when deleting a message.
@@ -1655,7 +1656,7 @@ func (js *jetStream) applyStreamEntries(mset *stream, ce *CommittedEntry, isReco
 				}
 
 				s := js.server()
-				purged, err := mset.purge()
+				purged, err := mset.purge(sp.Request)
 				if err != nil {
 					s.Warnf("JetStream cluster failed to purge stream %q for account %q: %v", sp.Stream, sp.Client.serviceAccount(), err)
 				}
@@ -3482,7 +3483,16 @@ func (s *Server) jsClusteredStreamDeleteRequest(ci *ClientInfo, acc *Account, st
 	cc.meta.Propose(encodeDeleteStreamAssignment(sa))
 }
 
-func (s *Server) jsClusteredStreamPurgeRequest(ci *ClientInfo, acc *Account, mset *stream, stream, subject, reply string, rmsg []byte) {
+// Process a clustered purge request.
+func (s *Server) jsClusteredStreamPurgeRequest(
+	ci *ClientInfo,
+	acc *Account,
+	mset *stream,
+	stream, subject, reply string,
+	rmsg []byte,
+	preq *JSApiStreamPurgeRequest,
+) {
+
 	js, cc := s.getJetStreamCluster()
 	if js == nil || cc == nil {
 		return
@@ -3500,11 +3510,11 @@ func (s *Server) jsClusteredStreamPurgeRequest(ci *ClientInfo, acc *Account, mse
 	}
 
 	if n := sa.Group.node; n != nil {
-		sp := &streamPurge{Stream: stream, LastSeq: mset.state().LastSeq, Subject: subject, Reply: reply, Client: ci}
+		sp := &streamPurge{Stream: stream, LastSeq: mset.state().LastSeq, Subject: subject, Reply: reply, Client: ci, Request: preq}
 		n.Propose(encodeStreamPurge(sp))
 	} else if mset != nil {
 		var resp = JSApiStreamPurgeResponse{ApiResponse: ApiResponse{Type: JSApiStreamPurgeResponseType}}
-		purged, err := mset.purge()
+		purged, err := mset.purge(preq)
 		if err != nil {
 			resp.Error = ApiErrors[JSStreamGeneralErrorF].ErrOrNewT(err, "{err}", err)
 		} else {
diff --git a/server/memstore.go b/server/memstore.go
index 4d0063e8..7ef8e189 100644
--- a/server/memstore.go
+++ b/server/memstore.go
@@ -273,16 +273,36 @@ func (ms *memStore) FilteredState(sseq uint64, subj string) SimpleState {
 		return ss
 	}
 
-	// FIXME(dlc) - Optimize like filestore.
-	var eq func(string, string) bool
-	if subj == _EMPTY_ {
-		eq = func(a, b string) bool { return true }
-	} else if subjectHasWildcard(subj) {
-		eq = subjectIsSubsetMatch
-	} else {
-		eq = func(a, b string) bool { return a == b }
+	wc := subjectHasWildcard(subj)
+	subs := []string{subj}
+	if wc {
+		subs = subs[:0]
+		for fsubj := range ms.fss {
+			if subjectIsSubsetMatch(fsubj, subj) {
+				subs = append(subs, fsubj)
+			}
+		}
+	}
+	fseq, lseq := ms.state.LastSeq, uint64(0)
+	for _, subj := range subs {
+		ss := ms.fss[subj]
+		if ss == nil {
+			continue
+		}
+		if ss.First < fseq {
+			fseq = ss.First
+		}
+		if ss.Last > lseq {
+			lseq = ss.Last
+		}
+	}
+	if fseq < sseq {
+		fseq = sseq
 	}
-	for seq := sseq; seq <= ms.state.LastSeq; seq++ {
+
+	// FIXME(dlc) - Optimize better like filestore.
+	eq := compareFn(subj)
+	for seq := fseq; seq <= lseq; seq++ {
 		if sm, ok := ms.msgs[seq]; ok && eq(sm.subj, subj) {
 			ss.Msgs++
 			if ss.First == 0 {
@@ -366,6 +386,40 @@ func (ms *memStore) expireMsgs() {
 	}
 }
 
+// PurgeEx will remove messages based on subject filters, sequence and number of messages to keep.
+// Will return the number of purged messages.
+func (ms *memStore) PurgeEx(subject string, sequence, keep uint64) (purged uint64, err error) {
+	if subject == _EMPTY_ || subject == ">" && keep == 0 {
+		return ms.Purge()
+	}
+	eq := compareFn(subject)
+	if ss := ms.FilteredState(1, subject); ss.Msgs > 0 {
+		if keep > 0 {
+			if keep > ss.Msgs {
+				return 0, nil
+			}
+			ss.Msgs -= keep
+		}
+		last := ss.Last
+		if sequence > 0 {
+			last = sequence + 1
+		}
+		ms.mu.Lock()
+		for seq := ss.First; seq <= last; seq++ {
+			if sm, ok := ms.msgs[seq]; ok && eq(sm.subj, subject) {
+				if ok := ms.removeMsg(sm.seq, false); ok {
+					purged++
+					if purged >= ss.Msgs {
+						break
+					}
+				}
+			}
+		}
+		ms.mu.Unlock()
+	}
+	return purged, nil
+}
+
 // Purge will remove all messages from this store.
 // Will return the number of purged messages.
 func (ms *memStore) Purge() (uint64, error) {
diff --git a/server/store.go b/server/store.go
index 38b9a615..f39ff7c3 100644
--- a/server/store.go
+++ b/server/store.go
@@ -74,6 +74,7 @@ type StreamStore interface {
 	RemoveMsg(seq uint64) (bool, error)
 	EraseMsg(seq uint64) (bool, error)
 	Purge() (uint64, error)
+	PurgeEx(subject string, seq, keep uint64) (uint64, error)
 	Compact(seq uint64) (uint64, error)
 	Truncate(seq uint64) error
 	GetSeqFromTime(t time.Time) uint64
diff --git a/server/stream.go b/server/stream.go
index 13d94ad4..0156c339 100644
--- a/server/stream.go
+++ b/server/stream.go
@@ -954,7 +954,7 @@ func (mset *stream) update(config *StreamConfig) error {
 }
 
 // Purge will remove all messages from the stream and underlying store.
-func (mset *stream) purge() (uint64, error) {
+func (mset *stream) purge(preq *JSApiStreamPurgeRequest) (purged uint64, err error) {
 	mset.mu.Lock()
 	if mset.client == nil {
 		mset.mu.Unlock()
@@ -969,7 +969,11 @@ func (mset *stream) purge() (uint64, error) {
 	}
 	mset.mu.Unlock()
 
-	purged, err := mset.store.Purge()
+	if preq != nil {
+		purged, err = mset.store.PurgeEx(preq.Subject, preq.Sequence, preq.Keep)
+	} else {
+		purged, err = mset.store.Purge()
+	}
 	if err != nil {
 		return purged, err
 	}


Test patch:
diff --git a/server/jetstream_cluster_test.go b/server/jetstream_cluster_test.go
index f6a2870d..c346fca6 100644
--- a/server/jetstream_cluster_test.go
+++ b/server/jetstream_cluster_test.go
@@ -7344,6 +7344,121 @@ func TestJetStreamClusterSourceAndMirrorConsumersLeaderChange(t *testing.T) {
 		}
 		return nil
 	})
+}
+
+func TestJetStreamClusterExtendedStreamPurge(t *testing.T) {
+	for _, st := range []StorageType{FileStorage, MemoryStorage} {
+		t.Run(st.String(), func(t *testing.T) {
+			c := createJetStreamClusterExplicit(t, "JSC", 3)
+			defer c.shutdown()
+
+			nc, js := jsClientConnect(t, c.randomServer())
+			defer nc.Close()
+
+			cfg := StreamConfig{
+				Name:       "KV",
+				Subjects:   []string{"kv.>"},
+				Storage:    st,
+				Replicas:   2,
+				MaxMsgsPer: 100,
+			}
+			req, err := json.Marshal(cfg)
+			if err != nil {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+			// Do manually for now.
+			nc.Request(fmt.Sprintf(JSApiStreamCreateT, cfg.Name), req, time.Second)
+
+			si, err := js.StreamInfo("KV")
+			if err != nil {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+			if si == nil || si.Config.Name != "KV" {
+				t.Fatalf("StreamInfo is not correct %+v", si)
+			}
+
+			for i := 0; i < 1000; i++ {
+				js.PublishAsync("kv.foo", []byte("OK")) // 1 * i
+				js.PublishAsync("kv.bar", []byte("OK")) // 2 * i
+				js.PublishAsync("kv.baz", []byte("OK")) // 3 * i, so after first is 2700, last is 3000
+			}
+			for i := 0; i < 700; i++ {
+				js.PublishAsync(fmt.Sprintf("kv.%d", i+1), []byte("OK"))
+			}
+
+			checkFor(t, 5*time.Second, 50*time.Millisecond, func() error {
+				si, err := js.StreamInfo("KV")
+				if err != nil {
+					return err
+				}
+				if si.State.Msgs != 1000 {
+					return fmt.Errorf("Expected %d msgs, got %d", 300, si.State.Msgs)
+				}
+				return nil
+			})
+
+			shouldFail := func(preq *JSApiStreamPurgeRequest) {
+				req, _ := json.Marshal(preq)
+				resp, err := nc.Request(fmt.Sprintf(JSApiStreamPurgeT, "KV"), req, time.Second)
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				var pResp JSApiStreamPurgeResponse
+				if err = json.Unmarshal(resp.Data, &pResp); err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				if pResp.Success || pResp.Error == nil {
+					t.Fatalf("Expected an error response but got none")
+				}
+			}
+
+			// Sequence and Keep should be mutually exclusive.
+			shouldFail(&JSApiStreamPurgeRequest{Sequence: 10, Keep: 10})
+
+			purge := func(preq *JSApiStreamPurgeRequest, newTotal uint64) {
+				req, _ := json.Marshal(preq)
+				resp, err := nc.Request(fmt.Sprintf(JSApiStreamPurgeT, "KV"), req, time.Second)
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				var pResp JSApiStreamPurgeResponse
+				if err = json.Unmarshal(resp.Data, &pResp); err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				if !pResp.Success || pResp.Error != nil {
+					t.Fatalf("Got a bad response %+v", pResp)
+				}
+				si, err = js.StreamInfo("KV")
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				if si.State.Msgs != newTotal {
+					t.Fatalf("Expected total after purge to be %d but got %d", newTotal, si.State.Msgs)
+				}
+			}
+			expectLeft := func(subject string, expected uint64) {
+				ci, err := js.AddConsumer("KV", &nats.ConsumerConfig{Durable: "dlc", FilterSubject: subject, AckPolicy: nats.AckExplicitPolicy})
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				defer js.DeleteConsumer("KV", "dlc")
+				if ci.NumPending != expected {
+					t.Fatalf("Expected %d remaining but got %d", expected, ci.NumPending)
+				}
+			}
+
+			purge(&JSApiStreamPurgeRequest{Subject: "kv.foo"}, 900)
+			expectLeft("kv.foo", 0)
+
+			purge(&JSApiStreamPurgeRequest{Subject: "kv.bar", Keep: 1}, 801)
+			expectLeft("kv.bar", 1)
+
+			purge(&JSApiStreamPurgeRequest{Subject: "kv.baz", Sequence: 2850}, 751)
+			expectLeft("kv.baz", 50)
+
+			purge(&JSApiStreamPurgeRequest{Subject: "kv.*"}, 0)
+		})
+	}
 
 }
 
diff --git a/server/jetstream_test.go b/server/jetstream_test.go
index 5ec9deaa..3489548b 100644
--- a/server/jetstream_test.go
+++ b/server/jetstream_test.go
@@ -294,7 +294,7 @@ func TestJetStreamAddStreamDiscardNew(t *testing.T) {
 			}
 
 			// Now do bytes.
-			mset.purge()
+			mset.purge(nil)
 
 			big := make([]byte, 8192)
 			resp, _ = nc.Request(subj, big, 100*time.Millisecond)
@@ -1780,7 +1780,7 @@ func TestJetStreamWorkQueueRequest(t *testing.T) {
 				sub.NextMsg(time.Millisecond)
 			}
 			checkSubPending(0)
-			mset.purge()
+			mset.purge(nil)
 
 			// Now do expiration
 			req.Batch = 1
@@ -2213,7 +2213,7 @@ func TestJetStreamWorkQueueRequestBatch(t *testing.T) {
 			// Now queue up the request without messages and add them after.
 			sub, _ = nc.SubscribeSync(nats.NewInbox())
 			defer sub.Unsubscribe()
-			mset.purge()
+			mset.purge(nil)
 
 			nc.PublishRequest(o.requestNextMsgSubject(), sub.Subject, []byte(strconv.Itoa(batchSize)))
 			nc.Flush() // Make sure its registered.
@@ -2520,7 +2520,7 @@ func TestJetStreamAckReplyStreamPending(t *testing.T) {
 
 			expectPending(toSend*2 - 2)
 			// Purge and send a new one.
-			mset.purge()
+			mset.purge(nil)
 			nc.Flush()
 
 			sendStreamMsg(t, nc, "foo.1", "Hello World!")
@@ -2537,7 +2537,7 @@ func TestJetStreamAckReplyStreamPending(t *testing.T) {
 			expectPending(toSend - 4) // 203
 
 			// Test Expiration.
-			mset.purge()
+			mset.purge(nil)
 			for i := 0; i < toSend; i++ {
 				sendStreamMsg(t, nc, "foo.1", "Hello World!")
 			}
@@ -2584,7 +2584,7 @@ func TestJetStreamAckReplyStreamPending(t *testing.T) {
 			}
 			nc.Flush()
 			expectPending(100)
-			mset.purge()
+			mset.purge(nil)
 			sendStreamMsg(t, nc, "foo.22", "Hello World!")
 			expectPending(0)
 		})
@@ -3298,7 +3298,7 @@ func TestJetStreamPublishDeDupe(t *testing.T) {
 	if err := mset.update(&cfg); err != nil {
 		t.Fatalf("Unexpected error: %v", err)
 	}
-	mset.purge()
+	mset.purge(nil)
 
 	// Send 5 new messages.
 	sendMsg(10, "AAAA", "Hello DeDupe!")
@@ -3343,7 +3343,7 @@ func TestJetStreamPublishDeDupe(t *testing.T) {
 	}
 
 	// Purge should wipe the msgIds as well.
-	mset.purge()
+	mset.purge(nil)
 	nmids(0)
 }
 
@@ -5577,7 +5577,7 @@ func TestJetStreamStreamPurge(t *testing.T) {
 			if state := mset.state(); state.Msgs != 100 {
 				t.Fatalf("Expected %d messages, got %d", 100, state.Msgs)
 			}
-			mset.purge()
+			mset.purge(nil)
 			state := mset.state()
 			if state.Msgs != 0 {
 				t.Fatalf("Expected %d messages, got %d", 0, state.Msgs)
@@ -5668,7 +5668,7 @@ func TestJetStreamStreamPurgeWithConsumer(t *testing.T) {
 				t.Fatalf("Expected len(pending) to be 25, got %d", state.NumAckPending)
 			}
 			// Now do purge.
-			mset.purge()
+			mset.purge(nil)
 			if state := mset.state(); state.Msgs != 0 {
 				t.Fatalf("Expected %d messages, got %d", 0, state.Msgs)
 			}
@@ -5756,7 +5756,7 @@ func TestJetStreamStreamPurgeWithConsumerAndRedelivery(t *testing.T) {
 			// Now wait to make sure we are in a redelivered state.
 			time.Sleep(wcfg.AckWait * 2)
 			// Now do purge.
-			mset.purge()
+			mset.purge(nil)
 			if state := mset.state(); state.Msgs != 0 {
 				t.Fatalf("Expected %d messages, got %d", 0, state.Msgs)
 			}
@@ -8274,7 +8274,7 @@ func TestJetStreamDeleteMsg(t *testing.T) {
 			deleteAndCheck(3, 2)
 			deleteAndCheck(2, 4)
 
-			mset.purge()
+			mset.purge(nil)
 			// Put ten more one.
 			pubTen()
 			deleteAndCheck(11, 12)
@@ -10084,7 +10084,7 @@ func TestJetStreamConsumerMaxAckPending(t *testing.T) {
 			checkSubPending(maxAckPending)
 
 			o.stop()
-			mset.purge()
+			mset.purge(nil)
 
 			// Now test a consumer that is live while we publish messages to the stream.
 			o, err = mset.addConsumer(&ConsumerConfig{
@@ -10853,7 +10853,7 @@ func TestJetStreamMaxMsgsPerSubject(t *testing.T) {
 			pubAndCheck("baz.33", 5, 15)
 
 			// Now purge and make sure all is still good.
-			mset.purge()
+			mset.purge(nil)
 			pubAndCheck("foo", 1, 1)
 			pubAndCheck("foo", 4, 5)
 			pubAndCheck("baz.22", 5, 10)