Repository: nats-io/nats-server
Pull request #2381 (2021-07-28 19:50:27)

Allow delivery policy of DeliverLastPerSubject, which is helpful for K/V on JetStream

Signed-off-by: Derek Collison <derek@nats.io>

/cc @nats-io/core



Patch:
diff --git a/server/consumer.go b/server/consumer.go
index 4ce0004d..d63105bd 100644
--- a/server/consumer.go
+++ b/server/consumer.go
@@ -86,8 +86,10 @@ const (
 	DeliverNew
 	// DeliverByStartSequence will look for a defined starting sequence to start.
 	DeliverByStartSequence
-	// DeliverByStartTime will select the first messsage with a timestamp >= to StartTime
+	// DeliverByStartTime will select the first messsage with a timestamp >= to StartTime.
 	DeliverByStartTime
+	// DeliverLastPerSubject will start the consumer with the last message for all subjects received.
+	DeliverLastPerSubject
 )
 
 func (dp DeliverPolicy) String() string {
@@ -102,6 +104,8 @@ func (dp DeliverPolicy) String() string {
 		return "by_start_sequence"
 	case DeliverByStartTime:
 		return "by_start_time"
+	case DeliverLastPerSubject:
+		return "last_per_subject"
 	default:
 		return "undefined"
 	}
@@ -186,6 +190,7 @@ type consumer struct {
 	asflr             uint64
 	sgap              uint64
 	dsubj             string
+	lss               *lastSeqSkipList
 	rlimit            *rate.Limiter
 	reqSub            *subscription
 	ackSub            *subscription
@@ -362,38 +367,68 @@ func (mset *stream) addConsumerWithAssignment(config *ConsumerConfig, oname stri
 	switch config.DeliverPolicy {
 	case DeliverAll:
 		if config.OptStartSeq > 0 {
-			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}", "consumer delivery policy is deliver all, but optional start sequence is also set")
+			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}",
+				"consumer delivery policy is deliver all, but optional start sequence is also set")
 		}
 		if config.OptStartTime != nil {
-			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}", "consumer delivery policy is deliver all, but optional start time is also set")
+			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}",
+				"consumer delivery policy is deliver all, but optional start time is also set")
 		}
 	case DeliverLast:
 		if config.OptStartSeq > 0 {
-			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}", "consumer delivery policy is deliver last, but optional start sequence is also set")
+			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}",
+				"consumer delivery policy is deliver last, but optional start sequence is also set")
 		}
 		if config.OptStartTime != nil {
-			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}", "consumer delivery policy is deliver last, but optional start time is also set")
+			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}",
+				"consumer delivery policy is deliver last, but optional start time is also set")
+		}
+	case DeliverLastPerSubject:
+		if config.OptStartSeq > 0 {
+			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}",
+				"consumer delivery policy is deliver last per subject, but optional start sequence is also set")
+		}
+		if config.OptStartTime != nil {
+			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}",
+				"consumer delivery policy is deliver last per subject, but optional start time is also set")
+		}
+		badConfig := config.FilterSubject == _EMPTY_
+		if !badConfig {
+			subjects, ext := mset.allSubjects()
+			if len(subjects) == 1 && !ext && subjects[0] == config.FilterSubject && subjectIsLiteral(subjects[0]) {
+				badConfig = true
+			}
+		}
+		if badConfig {
+			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}",
+				"consumer delivery policy is deliver last per subject, but filter subject is not set")
 		}
 	case DeliverNew:
 		if config.OptStartSeq > 0 {
-			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}", "consumer delivery policy is deliver new, but optional start sequence is also set")
+			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}",
+				"consumer delivery policy is deliver new, but optional start sequence is also set")
 		}
 		if config.OptStartTime != nil {
-			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}", "consumer delivery policy is deliver new, but optional start time is also set")
+			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}",
+				"consumer delivery policy is deliver new, but optional start time is also set")
 		}
 	case DeliverByStartSequence:
 		if config.OptStartSeq == 0 {
-			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}", "consumer delivery policy is deliver by start sequence, but optional start sequence is not set")
+			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}",
+				"consumer delivery policy is deliver by start sequence, but optional start sequence is not set")
 		}
 		if config.OptStartTime != nil {
-			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}", "consumer delivery policy is deliver by start sequence, but optional start time is also set")
+			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}",
+				"consumer delivery policy is deliver by start sequence, but optional start time is also set")
 		}
 	case DeliverByStartTime:
 		if config.OptStartTime == nil {
-			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}", "consumer delivery policy is deliver by start time, but optional start time is not set")
+			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}",
+				"consumer delivery policy is deliver by start time, but optional start time is not set")
 		}
 		if config.OptStartSeq != 0 {
-			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}", "consumer delivery policy is deliver by start time, but optional start sequence is also set")
+			return nil, ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}",
+				"consumer delivery policy is deliver by start time, but optional start sequence is also set")
 		}
 	}
 
@@ -1930,7 +1965,16 @@ func (o *consumer) getNextMsg() (subj string, hdr, msg []byte, seq uint64, dc ui
 	}
 	for {
 		seq, dc := o.sseq, uint64(1)
-		if o.hasRedeliveries() {
+		if o.hasSkipListPending() {
+			seq = o.lss.seqs[0]
+			if len(o.lss.seqs) == 1 {
+				o.sseq = o.lss.resume
+				o.lss = nil
+				o.updateSkipped()
+			} else {
+				o.lss.seqs = o.lss.seqs[1:]
+			}
+		} else if o.hasRedeliveries() {
 			seq = o.getNextToRedeliver()
 			dc = o.incDeliveryCount(seq)
 			if o.maxdc > 0 && dc > o.maxdc {
@@ -2568,12 +2612,34 @@ func ackReplyInfo(subject string) (sseq, dseq, dc uint64) {
 
 // NextSeq returns the next delivered sequence number for this consumer.
 func (o *consumer) nextSeq() uint64 {
-	o.mu.Lock()
+	o.mu.RLock()
 	dseq := o.dseq
-	o.mu.Unlock()
+	o.mu.RUnlock()
 	return dseq
 }
 
+// Used to hold skip list when deliver policy is last per subject.
+type lastSeqSkipList struct {
+	resume uint64
+	seqs   []uint64
+}
+
+// Will create a skip list for us from a store's subjects state.
+func createLastSeqSkipList(mss map[string]SimpleState) []uint64 {
+	seqs := make([]uint64, 0, len(mss))
+	for _, ss := range mss {
+		seqs = append(seqs, ss.Last)
+	}
+	sort.Slice(seqs, func(i, j int) bool { return seqs[i] < seqs[j] })
+	return seqs
+}
+
+// Let's us know we have a skip list, which is for deliver last per subject and we are just starting.
+// Lock should be held.
+func (o *consumer) hasSkipListPending() bool {
+	return o.lss != nil && len(o.lss.seqs) > 0
+}
+
 // Will select the starting sequence.
 func (o *consumer) selectStartingSeqNo() {
 	if o.mset == nil || o.mset.store == nil {
@@ -2590,6 +2656,17 @@ func (o *consumer) selectStartingSeqNo() {
 					ss := o.mset.store.FilteredState(1, o.cfg.FilterSubject)
 					o.sseq = ss.Last
 				}
+			} else if o.cfg.DeliverPolicy == DeliverLastPerSubject {
+				if mss := o.mset.store.SubjectsState(o.cfg.FilterSubject); len(mss) > 0 {
+					o.lss = &lastSeqSkipList{
+						resume: stats.LastSeq + 1,
+						seqs:   createLastSeqSkipList(mss),
+					}
+					o.sseq = o.lss.seqs[0]
+				} else {
+					// If no mapping info just set to last.
+					o.sseq = stats.LastSeq
+				}
 			} else if o.cfg.OptStartTime != nil {
 				// If we are here we are time based.
 				// TODO(dlc) - Once clustered can't rely on this.
@@ -2933,6 +3010,7 @@ func (o *consumer) setInitialPendingAndStart() {
 	if mset == nil || mset.store == nil {
 		return
 	}
+
 	// notFiltered means we want all messages.
 	notFiltered := o.cfg.FilterSubject == _EMPTY_
 	if !notFiltered {
@@ -2951,11 +3029,17 @@ func (o *consumer) setInitialPendingAndStart() {
 		}
 	} else {
 		// Here we are filtered.
-		ss := mset.store.FilteredState(o.sseq, o.cfg.FilterSubject)
-		if ss.Msgs > 0 {
+		dp := o.cfg.DeliverPolicy
+		if dp == DeliverLastPerSubject && o.hasSkipListPending() && o.sseq < o.lss.resume {
+			if o.lss != nil {
+				ss := mset.store.FilteredState(o.lss.resume, o.cfg.FilterSubject)
+				o.sseq = o.lss.seqs[0]
+				o.sgap = ss.Msgs + uint64(len(o.lss.seqs))
+			}
+		} else if ss := mset.store.FilteredState(o.sseq, o.cfg.FilterSubject); ss.Msgs > 0 {
 			o.sgap = ss.Msgs
 			// See if we should update our starting sequence.
-			if dp := o.cfg.DeliverPolicy; dp == DeliverLast {
+			if dp == DeliverLast || dp == DeliverLastPerSubject {
 				o.sseq = ss.Last
 			} else if dp == DeliverNew {
 				o.sseq = ss.Last + 1
@@ -2963,6 +3047,10 @@ func (o *consumer) setInitialPendingAndStart() {
 				// DeliverAll, DeliverByStartSequence, DeliverByStartTime
 				o.sseq = ss.First
 			}
+			// Cleanup lss when we take over in clustered mode.
+			if dp == DeliverLastPerSubject && o.hasSkipListPending() && o.sseq >= o.lss.resume {
+				o.lss = nil
+			}
 		}
 		o.updateSkipped()
 	}
diff --git a/server/filestore.go b/server/filestore.go
index a8f6de54..d0578ee0 100644
--- a/server/filestore.go
+++ b/server/filestore.go
@@ -1125,6 +1125,33 @@ func (fs *fileStore) perSubjectState(subj string) (total, first, last uint64) {
 	return total, first, last
 }
 
+// SubjectsState returns a map of SimpleState for all matching subjects.
+func (fs *fileStore) SubjectsState(subject string) map[string]SimpleState {
+	fs.mu.RLock()
+	defer fs.mu.RUnlock()
+
+	if !fs.tms || fs.state.Msgs == 0 {
+		return nil
+	}
+
+	fss := make(map[string]SimpleState)
+	for _, mb := range fs.blks {
+		mb.mu.RLock()
+		for subj, ss := range mb.fss {
+			oss := fss[subj]
+			if oss.First == 0 { // New
+				fss[subj] = *ss
+			} else {
+				// Merge here.
+				oss.Last, oss.Msgs = ss.Last, oss.Msgs+ss.Msgs
+				fss[subj] = oss
+			}
+		}
+		mb.mu.RUnlock()
+	}
+	return fss
+}
+
 // RegisterStorageUpdates registers a callback for updates to storage changes.
 // It will present number of messages and bytes as a signed integer and an
 // optional sequence number of the message if a single.
diff --git a/server/jetstream_api.go b/server/jetstream_api.go
index d0a55c84..56c64b23 100644
--- a/server/jetstream_api.go
+++ b/server/jetstream_api.go
@@ -3018,7 +3018,6 @@ func (s *Server) jsConsumerCreateRequest(sub *subscription, c *client, acc *Acco
 }
 
 func (s *Server) jsConsumerCreate(sub *subscription, c *client, a *Account, subject, reply string, rmsg []byte, expectDurable bool) {
-
 	if c == nil || !s.JetStreamEnabled() {
 		return
 	}
diff --git a/server/jetstream_cluster.go b/server/jetstream_cluster.go
index 21c06150..9d371076 100644
--- a/server/jetstream_cluster.go
+++ b/server/jetstream_cluster.go
@@ -3990,6 +3990,23 @@ func (s *Server) jsClusteredConsumerRequest(ci *ClientInfo, acc *Account, subjec
 		return
 	}
 
+	// Also short circuit if DeliverLastPerSubject is set with no FilterSubject.
+	if cfg.DeliverPolicy == DeliverLastPerSubject {
+		badConfig := cfg.FilterSubject == _EMPTY_
+		if !badConfig {
+			subjects := sa.Config.Subjects
+			if len(subjects) == 1 && subjects[0] == cfg.FilterSubject && subjectIsLiteral(subjects[0]) {
+				badConfig = true
+			}
+		}
+		if badConfig {
+			resp.Error = ApiErrors[JSConsumerInvalidPolicyErrF].NewT("{err}",
+				"consumer delivery policy is deliver last per subject, but FilterSubject is not set")
+			s.sendAPIErrResponse(ci, acc, subject, reply, string(rmsg), s.jsonResponse(&resp))
+			return
+		}
+	}
+
 	// Setup proper default for ack wait if we are in explicit ack mode.
 	if cfg.AckWait == 0 && (cfg.AckPolicy == AckExplicit || cfg.AckPolicy == AckAll) {
 		cfg.AckWait = JsAckWaitDefault
diff --git a/server/memstore.go b/server/memstore.go
index a84b5651..d075d2ec 100644
--- a/server/memstore.go
+++ b/server/memstore.go
@@ -317,6 +317,24 @@ func (ms *memStore) filteredStateLocked(sseq uint64, subj string) SimpleState {
 	return ss
 }
 
+// SubjectsState returns a map of SimpleState for all matching subjects.
+func (ms *memStore) SubjectsState(subject string) map[string]SimpleState {
+	ms.mu.RLock()
+	defer ms.mu.RUnlock()
+
+	if len(ms.fss) == 0 {
+		return nil
+	}
+
+	fss := make(map[string]SimpleState)
+	for subj, ss := range ms.fss {
+		if subject == _EMPTY_ || subject == fwcs || subjectIsSubsetMatch(subj, subject) {
+			fss[subj] = *ss
+		}
+	}
+	return fss
+}
+
 // Will check the msg limit for this tracked subject.
 // Lock should be held.
 func (ms *memStore) enforcePerSubjectLimit(ss *SimpleState) {
diff --git a/server/store.go b/server/store.go
index 2139d1a6..5da8274a 100644
--- a/server/store.go
+++ b/server/store.go
@@ -79,7 +79,8 @@ type StreamStore interface {
 	Compact(seq uint64) (uint64, error)
 	Truncate(seq uint64) error
 	GetSeqFromTime(t time.Time) uint64
-	FilteredState(sseq uint64, subject string) SimpleState
+	FilteredState(seq uint64, subject string) SimpleState
+	SubjectsState(filterSubject string) map[string]SimpleState
 	State() StreamState
 	FastState(*StreamState)
 	Type() StorageType
@@ -111,7 +112,7 @@ type DiscardPolicy int
 const (
 	// DiscardOld will remove older messages to return to the limits.
 	DiscardOld = iota
-	//DiscardNew will error on a StoreMsg call
+	// DiscardNew will error on a StoreMsg call
 	DiscardNew
 )
 
@@ -389,6 +390,7 @@ const (
 	deliverNewPolicyString       = "new"
 	deliverByStartSequenceString = "by_start_sequence"
 	deliverByStartTimeString     = "by_start_time"
+	deliverLastPerPolicyString   = "last_per_subject"
 	deliverUndefinedString       = "undefined"
 )
 
@@ -398,6 +400,8 @@ func (p *DeliverPolicy) UnmarshalJSON(data []byte) error {
 		*p = DeliverAll
 	case jsonString(deliverLastPolicyString):
 		*p = DeliverLast
+	case jsonString(deliverLastPerPolicyString):
+		*p = DeliverLastPerSubject
 	case jsonString(deliverNewPolicyString):
 		*p = DeliverNew
 	case jsonString(deliverByStartSequenceString):
@@ -417,6 +421,8 @@ func (p DeliverPolicy) MarshalJSON() ([]byte, error) {
 		return json.Marshal(deliverAllPolicyString)
 	case DeliverLast:
 		return json.Marshal(deliverLastPolicyString)
+	case DeliverLastPerSubject:
+		return json.Marshal(deliverLastPerPolicyString)
 	case DeliverNew:
 		return json.Marshal(deliverNewPolicyString)
 	case DeliverByStartSequence:


Test patch:
diff --git a/server/jetstream_test.go b/server/jetstream_test.go
index 73fe9dca..f3c14343 100644
--- a/server/jetstream_test.go
+++ b/server/jetstream_test.go
@@ -12169,6 +12169,148 @@ func TestJetStreamConsumerBadNumPending(t *testing.T) {
 	checkForNoPending(mon)
 }
 
+func TestJetStreamDeliverLastPerSubject(t *testing.T) {
+	for _, st := range []StorageType{FileStorage, MemoryStorage} {
+		t.Run(st.String(), func(t *testing.T) {
+			s := RunBasicJetStreamServer()
+			defer s.Shutdown()
+
+			if config := s.JetStreamConfig(); config != nil {
+				defer removeDir(t, config.StoreDir)
+			}
+
+			// Client for API requests.
+			nc, js := jsClientConnect(t, s)
+			defer nc.Close()
+
+			cfg := StreamConfig{
+				Name:       "KV",
+				Subjects:   []string{"kv.>"},
+				Storage:    st,
+				MaxMsgsPer: 5,
+			}
+
+			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)
+			}
+
+			// Interleave them on purpose.
+			for i := 1; i <= 11; i++ {
+				msg := []byte(fmt.Sprintf("%d", i))
+				js.PublishAsync("kv.b1.foo", msg)
+				js.PublishAsync("kv.b2.foo", msg)
+
+				js.PublishAsync("kv.b1.bar", msg)
+				js.PublishAsync("kv.b2.bar", msg)
+
+				js.PublishAsync("kv.b1.baz", msg)
+				js.PublishAsync("kv.b2.baz", msg)
+			}
+
+			select {
+			case <-js.PublishAsyncComplete():
+			case <-time.After(2 * time.Second):
+				t.Fatalf("Did not receive completion signal")
+			}
+
+			// Do quick check that config needs FilteredSubjects otherwise bad config.
+			badReq := CreateConsumerRequest{
+				Stream: "KV",
+				Config: ConsumerConfig{
+					DeliverSubject: "b",
+					DeliverPolicy:  DeliverLastPerSubject,
+				},
+			}
+			req, err = json.Marshal(badReq)
+			if err != nil {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+			resp, err := nc.Request(fmt.Sprintf(JSApiConsumerCreateT, "KV"), req, time.Second)
+			if err != nil {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+			var ccResp JSApiConsumerCreateResponse
+			if err = json.Unmarshal(resp.Data, &ccResp); err != nil {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+			if ccResp.Error == nil || !strings.Contains(ccResp.Error.Description, "filter subject is not set") {
+				t.Fatalf("Expected an error, got none")
+			}
+
+			// Now let's consume these via last per subject.
+			obsReq := CreateConsumerRequest{
+				Stream: "KV",
+				Config: ConsumerConfig{
+					DeliverSubject: "d",
+					DeliverPolicy:  DeliverLastPerSubject,
+					FilterSubject:  "kv.b1.*",
+				},
+			}
+			req, err = json.Marshal(obsReq)
+			if err != nil {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+			resp, err = nc.Request(fmt.Sprintf(JSApiConsumerCreateT, "KV"), req, time.Second)
+			if err != nil {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+			ccResp.Error = nil
+			if err = json.Unmarshal(resp.Data, &ccResp); err != nil {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+
+			sub, _ := nc.SubscribeSync("d")
+			defer sub.Unsubscribe()
+
+			// Helper to check messages are correct.
+			checkNext := func(subject string, sseq uint64, v string) {
+				t.Helper()
+				m, err := sub.NextMsg(time.Second)
+				if err != nil {
+					t.Fatalf("Error receiving message: %v", err)
+				}
+				if m.Subject != subject {
+					t.Fatalf("Expected subject %q but got %q", subject, m.Subject)
+				}
+				meta, err := m.Metadata()
+				if err != nil {
+					t.Fatalf("didn't get metadata: %s", err)
+				}
+				if meta.Sequence.Stream != sseq {
+					t.Fatalf("Expected stream seq %d but got %d", sseq, meta.Sequence.Stream)
+				}
+				if string(m.Data) != v {
+					t.Fatalf("Expected data of %q but got %q", v, m.Data)
+				}
+			}
+
+			checkSubsPending(t, sub, 3)
+
+			// Now make sure they are what we expect.
+			checkNext("kv.b1.foo", 61, "11")
+			checkNext("kv.b1.bar", 63, "11")
+			checkNext("kv.b1.baz", 65, "11")
+
+			msg := []byte(fmt.Sprintf("%d", 22))
+			js.Publish("kv.b2.foo", msg) // Not filtered through..
+			js.Publish("kv.b1.bar", msg)
+
+			checkSubsPending(t, sub, 1)
+			checkNext("kv.b1.bar", 68, "22")
+		})
+	}
+}
+
 // We had a report of a consumer delete crashing the server when in interest retention mode.
 // This I believe is only really possible in clustered mode, but we will force the issue here.
 func TestJetStreamConsumerCleanupWithRetentionPolicy(t *testing.T) {