Repository: nats-io/nats-server
Pull request #6174 (2024-11-25 17:48:35)

Refactor duplicate `Nats-Msg-Id` tracking
The current duplicate tracking has a couple of problems, namely that the `ddmap` capacity is potentially never reclaimed if it grows suddenly, and secondly that we have to constantly monitor `ddarr`'s capacity and copy it to stop the array from growing infinitely too.

This PR makes the following changes:

- `ddmap` is now using the stree, which means there is efficient in-place deduplication of message IDs and no infinitely growing backing capacity;
- `ddarr` and `ddindex` are now gone, replaced with linked-list references and `ddnext` and `ddlast` tracking, so we also don't need to watch the capacity and reallocate this either;
- `ddtmr` is now only scheduled when there is actually something to do.

Also added a unit test to prove the behaviour as well as the linked-list stitching.

Signed-off-by: Neil Twigg <neil@nats.io>


derekcollison: Want me to take a look?
neilalexander: Yes please. I did take a look at `container/heap` as an option but it still requires a backing array and has the same problem that the current `ddarr` has in terms of backing capacity. There's also `container/list` but it isn't generic (although we could copy it in from the stdlib and _make_ it generic if you think that's a better path).

Patch:
diff --git a/server/jetstream_cluster.go b/server/jetstream_cluster.go
index 38c0c7f0..0aeeb04e 100644
--- a/server/jetstream_cluster.go
+++ b/server/jetstream_cluster.go
@@ -8135,7 +8135,7 @@ func (mset *stream) processClusteredInboundMsg(subject, reply string, hdr, msg [
 			}
 			// FIXME(dlc) - locking conflict with accessing mset.clseq
 			// For now we stage with zero, and will update in processStreamMsg.
-			mset.storeMsgIdLocked(&ddentry{msgId, 0, time.Now().UnixNano()})
+			mset.storeMsgIdLocked(&ddentry{msgId, 0, time.Now().UnixNano(), nil, nil})
 			mset.mu.Unlock()
 		}
 	}
@@ -8763,7 +8763,7 @@ func (mset *stream) processCatchupMsg(msg []byte) (uint64, error) {
 				mset.rebuildDedupe()
 				mset.mu.Unlock()
 			}
-			mset.storeMsgId(&ddentry{msgId, seq, ts})
+			mset.storeMsgId(&ddentry{msgId, seq, ts, nil, nil})
 		}
 	}
 
diff --git a/server/stream.go b/server/stream.go
index a1694f95..4a588621 100644
--- a/server/stream.go
+++ b/server/stream.go
@@ -33,6 +33,7 @@ import (
 
 	"github.com/goccy/go-json"
 	"github.com/klauspost/compress/s2"
+	"github.com/nats-io/nats-server/v2/server/stree"
 	"github.com/nats-io/nuid"
 )
 
@@ -275,28 +276,28 @@ type stream struct {
 	// Those subscriptions are for the subjects filters being listened to and captured by the stream.
 	sid atomic.Uint64
 
-	pubAck    []byte                  // The template (prefix) to generate the pubAck responses for this stream quickly.
-	outq      *jsOutQ                 // Queue of *jsPubMsg for sending messages.
-	msgs      *ipQueue[*inMsg]        // Intra-process queue for the ingress of messages.
-	gets      *ipQueue[*directGetReq] // Intra-process queue for the direct get requests.
-	store     StreamStore             // The storage for this stream.
-	ackq      *ipQueue[uint64]        // Intra-process queue for acks.
-	lseq      uint64                  // The sequence number of the last message stored in the stream.
-	lmsgId    string                  // The de-duplication message ID of the last message stored in the stream.
-	consumers map[string]*consumer    // The consumers for this stream.
-	numFilter int                     // The number of filtered consumers.
-	cfg       StreamConfig            // The stream's config.
-	cfgMu     sync.RWMutex            // Config mutex used to solve some races with consumer code
-	created   time.Time               // Time the stream was created.
-	stype     StorageType             // The storage type.
-	tier      string                  // The tier is the number of replicas for the stream (e.g. "R1" or "R3").
-	ddmap     map[string]*ddentry     // The dedupe map.
-	ddarr     []*ddentry              // The dedupe array.
-	ddindex   int                     // The dedupe index.
-	ddtmr     *time.Timer             // The dedupe timer.
-	qch       chan struct{}           // The quit channel.
-	mqch      chan struct{}           // The monitor's quit channel.
-	active    bool                    // Indicates that there are active internal subscriptions (for the subject filters)
+	pubAck    []byte                       // The template (prefix) to generate the pubAck responses for this stream quickly.
+	outq      *jsOutQ                      // Queue of *jsPubMsg for sending messages.
+	msgs      *ipQueue[*inMsg]             // Intra-process queue for the ingress of messages.
+	gets      *ipQueue[*directGetReq]      // Intra-process queue for the direct get requests.
+	store     StreamStore                  // The storage for this stream.
+	ackq      *ipQueue[uint64]             // Intra-process queue for acks.
+	lseq      uint64                       // The sequence number of the last message stored in the stream.
+	lmsgId    string                       // The de-duplication message ID of the last message stored in the stream.
+	consumers map[string]*consumer         // The consumers for this stream.
+	numFilter int                          // The number of filtered consumers.
+	cfg       StreamConfig                 // The stream's config.
+	cfgMu     sync.RWMutex                 // Config mutex used to solve some races with consumer code
+	created   time.Time                    // Time the stream was created.
+	stype     StorageType                  // The storage type.
+	tier      string                       // The tier is the number of replicas for the stream (e.g. "R1" or "R3").
+	ddmap     *stree.SubjectTree[*ddentry] // The dedupe map.
+	ddnext    *ddentry                     // The next ddentry due to expire.
+	ddlast    *ddentry                     // The last ddentry due to expire.
+	ddtmr     *time.Timer                  // The dedupe timer.
+	qch       chan struct{}                // The quit channel.
+	mqch      chan struct{}                // The monitor's quit channel.
+	active    bool                         // Indicates that there are active internal subscriptions (for the subject filters)
 	// and/or mirror/sources consumers are scheduled to be established or already started.
 	ddloaded bool        // set to true when the deduplication structures are been built.
 	closed   atomic.Bool // Set to true when stop() is called on the stream.
@@ -434,9 +435,11 @@ const (
 
 // Dedupe entry
 type ddentry struct {
-	id  string // The unique message ID provided by the client.
-	seq uint64 // The sequence number of the message.
-	ts  int64  // The timestamp of the message.
+	id   string   // The unique message ID provided by the client.
+	seq  uint64   // The sequence number of the message.
+	ts   int64    // The timestamp of the message.
+	prev *ddentry // Will be populated by storeMsgIdLocked; the previous entry due to expire before this one.
+	next *ddentry // Will be populated by storeMsgIdLocked; the next entry due to expire after this one.
 }
 
 // Replicas Range
@@ -1063,6 +1066,9 @@ func (mset *stream) rebuildDedupe() {
 	}
 
 	mset.ddloaded = true
+	mset.ddmap.Empty()
+	mset.ddnext = nil
+	mset.ddlast = nil
 
 	// We have some messages. Lookup starting sequence by duplicate time window.
 	sseq := mset.store.GetSeqFromTime(time.Now().Add(-mset.cfg.Duplicates))
@@ -1082,7 +1088,15 @@ func (mset *stream) rebuildDedupe() {
 		var msgId string
 		if len(sm.hdr) > 0 {
 			if msgId = getMsgId(sm.hdr); msgId != _EMPTY_ {
-				mset.storeMsgIdLocked(&ddentry{msgId, sm.seq, sm.ts})
+				new := &ddentry{msgId, sm.seq, sm.ts, mset.ddlast, nil}
+				if mset.ddlast != nil {
+					mset.ddlast.next = new
+				}
+				mset.storeMsgIdLocked(new)
+				if mset.ddnext == nil {
+					mset.ddnext = new
+				}
+				mset.ddlast = new
 			}
 		}
 		if seq == state.LastSeq {
@@ -4052,7 +4066,7 @@ func (mset *stream) numMsgIds() int {
 	if !mset.ddloaded {
 		mset.rebuildDedupe()
 	}
-	return len(mset.ddmap)
+	return mset.ddmap.Size()
 }
 
 // checkMsgId will process and check for duplicates.
@@ -4061,10 +4075,14 @@ func (mset *stream) checkMsgId(id string) *ddentry {
 	if !mset.ddloaded {
 		mset.rebuildDedupe()
 	}
-	if id == _EMPTY_ || len(mset.ddmap) == 0 {
+	if id == _EMPTY_ || mset.ddmap.Size() == 0 {
+		return nil
+	}
+	dde, found := mset.ddmap.Find(stringToBytes(id))
+	if !found || dde == nil {
 		return nil
 	}
-	return mset.ddmap[id]
+	return *dde
 }
 
 // Will purge the entries that are past the window.
@@ -4077,39 +4095,26 @@ func (mset *stream) purgeMsgIds() {
 	tmrNext := mset.cfg.Duplicates
 	window := int64(tmrNext)
 
-	for i, dde := range mset.ddarr[mset.ddindex:] {
-		if now-dde.ts >= window {
-			delete(mset.ddmap, dde.id)
-		} else {
-			mset.ddindex += i
-			// Check if we should garbage collect here if we are 1/3 total size.
-			if cap(mset.ddarr) > 3*(len(mset.ddarr)-mset.ddindex) {
-				mset.ddarr = append([]*ddentry(nil), mset.ddarr[mset.ddindex:]...)
-				mset.ddindex = 0
-			}
-			tmrNext = time.Duration(window - (now - dde.ts))
-			break
-		}
+	for ; mset.ddnext != nil && now-mset.ddnext.ts >= window; mset.ddnext = mset.ddnext.next {
+		mset.ddmap.Delete(stringToBytes(mset.ddnext.id))
 	}
-	if len(mset.ddmap) > 0 {
+
+	if next := mset.ddnext; next != nil {
 		// Make sure to not fire too quick
 		const minFire = 50 * time.Millisecond
-		if tmrNext < minFire {
-			tmrNext = minFire
-		}
-		if mset.ddtmr != nil {
-			mset.ddtmr.Reset(tmrNext)
-		} else {
-			mset.ddtmr = time.AfterFunc(tmrNext, mset.purgeMsgIds)
+		if mset.ddnext != nil {
+			expiresAt := time.Unix(0, mset.ddnext.ts).Add(mset.cfg.Duplicates)
+			until := time.Until(expiresAt)
+			if until < minFire {
+				until = minFire
+			}
+			mset.ddtmr = time.AfterFunc(until, mset.purgeMsgIds)
 		}
 	} else {
-		if mset.ddtmr != nil {
-			mset.ddtmr.Stop()
-			mset.ddtmr = nil
-		}
+		stopAndClearTimer(&mset.ddtmr)
 		mset.ddmap = nil
-		mset.ddarr = nil
-		mset.ddindex = 0
+		mset.ddnext = nil
+		mset.ddlast = nil
 	}
 }
 
@@ -4121,15 +4126,52 @@ func (mset *stream) storeMsgId(dde *ddentry) {
 }
 
 // storeMsgIdLocked will store the message id for duplicate detection.
-// Lock should he held.
+// Lock should be held.
 func (mset *stream) storeMsgIdLocked(dde *ddentry) {
 	if mset.ddmap == nil {
-		mset.ddmap = make(map[string]*ddentry)
+		mset.ddmap = stree.NewSubjectTree[*ddentry]()
 	}
-	mset.ddmap[dde.id] = dde
-	mset.ddarr = append(mset.ddarr, dde)
-	if mset.ddtmr == nil {
-		mset.ddtmr = time.AfterFunc(mset.cfg.Duplicates, mset.purgeMsgIds)
+
+	// Ordinarily a Nats-Msg-Id would have been checked by now for a new
+	// publish into a stream, but when the ddmap needs to be rebuilt, the
+	// rebuildDedupe() function will call this function for every stream
+	// message in the deduplication window. In that case, we will need to
+	// clean up the previous entry we just replaced and stitch the gap.
+	if o, updated := mset.ddmap.Insert(stringToBytes(dde.id), dde); updated && o != nil {
+		overwritten := *o
+		if overwritten.prev != nil {
+			overwritten.prev.next = overwritten.next
+		}
+		if overwritten.next != nil {
+			overwritten.next.prev = overwritten.prev
+		}
+		if mset.ddnext == overwritten {
+			mset.ddnext = overwritten.next
+			stopAndClearTimer(&mset.ddtmr)
+		}
+		if mset.ddlast == overwritten {
+			mset.ddlast.next = dde
+			mset.ddlast = dde
+		}
+	}
+
+	// Update the last and next references, as well as the entries
+	// that attach to them.
+	dde.prev = mset.ddlast
+	if last := mset.ddlast; last != nil {
+		last.next = dde
+	}
+	mset.ddlast = dde
+	if next := mset.ddnext; next == nil {
+		mset.ddnext = dde
+	}
+
+	// Work out how long the timer needs to be set for, if needed.
+	if mset.ddtmr == nil && mset.ddnext != nil {
+		expiresAt := time.Unix(0, mset.ddnext.ts).Add(mset.cfg.Duplicates)
+		if time.Now().Before(expiresAt) {
+			mset.ddtmr = time.AfterFunc(time.Until(expiresAt), mset.purgeMsgIds)
+		}
 	}
 }
 
@@ -4930,7 +4972,7 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 		mset.lmsgId = msgId
 		// If we have a msgId make sure to save.
 		if msgId != _EMPTY_ {
-			mset.storeMsgIdLocked(&ddentry{msgId, mset.lseq, ts})
+			mset.storeMsgIdLocked(&ddentry{msgId, mset.lseq, ts, nil, nil})
 		}
 		if canRespond {
 			response = append(pubAck, strconv.FormatUint(mset.lseq, 10)...)
@@ -5029,16 +5071,7 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 	// If we have a msgId make sure to save.
 	// This will replace our estimate from the cluster layer if we are clustered.
 	if msgId != _EMPTY_ {
-		if isClustered && isLeader && mset.ddmap != nil {
-			if dde := mset.ddmap[msgId]; dde != nil {
-				dde.seq, dde.ts = seq, ts
-			} else {
-				mset.storeMsgIdLocked(&ddentry{msgId, seq, ts})
-			}
-		} else {
-			// R1 or not leader..
-			mset.storeMsgIdLocked(&ddentry{msgId, seq, ts})
-		}
+		mset.storeMsgIdLocked(&ddentry{msgId, seq, ts, nil, nil})
 	}
 
 	// If here we succeeded in storing the message.
@@ -5570,8 +5603,8 @@ func (mset *stream) stop(deleteFlag, advisory bool) error {
 		mset.ddtmr.Stop()
 		mset.ddtmr = nil
 		mset.ddmap = nil
-		mset.ddarr = nil
-		mset.ddindex = 0
+		mset.ddnext = nil
+		mset.ddlast = nil
 	}
 
 	sysc := mset.sysc


Test patch:
diff --git a/server/jetstream_test.go b/server/jetstream_test.go
index e362be50..0aa31900 100644
--- a/server/jetstream_test.go
+++ b/server/jetstream_test.go
@@ -24458,7 +24458,7 @@ func TestJetStreamStreamConfigClone(t *testing.T) {
 		Name:             "name",
 		Placement:        &Placement{Cluster: "placement", Tags: []string{"tag"}},
 		Mirror:           &StreamSource{Name: "mirror"},
-		Sources:          []*StreamSource{&StreamSource{Name: "source"}},
+		Sources:          []*StreamSource{{Name: "source"}},
 		SubjectTransform: &SubjectTransformConfig{Source: "source", Destination: "dest"},
 		RePublish:        &RePublish{Source: "source", Destination: "dest", HeadersOnly: false},
 		Metadata:         make(map[string]string),
@@ -24721,3 +24721,166 @@ func TestJetStreamMemoryPurgeClearsSubjectsState(t *testing.T) {
 	require_NoError(t, err)
 	require_Len(t, len(si.State.Subjects), 0)
 }
+
+func TestJetStreamMsgIdDeduplication(t *testing.T) {
+	s := RunBasicJetStreamServer(t)
+	defer s.Shutdown()
+
+	nc, js := jsClientConnect(t, s)
+	defer nc.Close()
+
+	_, err := js.AddStream(&nats.StreamConfig{
+		Name:       "TEST",
+		Subjects:   []string{"test"},
+		Duplicates: time.Second / 2,
+	})
+	require_NoError(t, err)
+
+	require_Publish := func(id string) {
+		t.Helper()
+		msg := &nats.Msg{
+			Subject: "test",
+			Header:  nats.Header{},
+		}
+		msg.Header.Set("Nats-Msg-Id", id)
+		_, err := js.PublishMsg(msg)
+		require_NoError(t, err)
+	}
+
+	require_Messages := func(msgs uint64) {
+		t.Helper()
+		si, err := js.StreamInfo("TEST")
+		require_NoError(t, err)
+		require_Equal(t, si.State.Msgs, msgs)
+	}
+
+	mset, err := s.GlobalAccount().lookupStream("TEST")
+	require_NoError(t, err)
+
+	require_Stitched := func() {
+		t.Helper()
+		require_NotNil(t, mset.ddnext)
+		require_NotNil(t, mset.ddlast)
+		require_Equal(t, mset.ddnext.prev, nil)
+		require_Equal(t, mset.ddlast.next, nil)
+		for dde := mset.ddnext; dde != nil && dde != mset.ddlast; dde = dde.next {
+			if dde.next != nil {
+				require_Equal(t, dde.next.prev, dde)
+			}
+			if dde.prev != nil {
+				require_Equal(t, dde.prev.next, dde)
+			}
+		}
+	}
+
+	// At this point nothing has sent a message with a Nats-Msg-Id
+	// so there should be no state and no timer.
+	mset.mu.RLock()
+	require_Equal(t, mset.ddtmr, nil)
+	require_Equal(t, mset.ddnext, nil)
+	require_Equal(t, mset.ddlast, nil)
+	require_Equal(t, mset.ddtmr, nil)
+	mset.mu.RUnlock()
+
+	// Publish the first message with a Nats-Msg-Id. There should
+	// be exactly one message, which means it is both the next and
+	// the last.
+	require_Publish("one")
+	require_Messages(1)
+	mset.mu.RLock()
+	require_Equal(t, mset.ddmap.Size(), 1)
+	one, ok := mset.ddmap.Find(stringToBytes("one"))
+	require_True(t, ok)
+	require_Equal(t, mset.ddnext, *one)
+	require_Equal(t, mset.ddlast, *one)
+	require_Stitched()
+	require_NotNil(t, mset.ddtmr)
+	mset.mu.RUnlock()
+
+	// Now publish a duplicate. Nothing should change.
+	require_Publish("one")
+	require_Messages(1)
+	mset.mu.RLock()
+	require_Equal(t, mset.ddmap.Size(), 1)
+	require_Equal(t, mset.ddnext, *one)
+	require_Equal(t, mset.ddlast, *one)
+	require_Stitched()
+	require_NotNil(t, mset.ddtmr)
+	mset.mu.RUnlock()
+
+	// Publish a second message with a different Nats-Msg-Id.
+	// There will now be two messages and the ddentries for one
+	// and two should now be stitched together, with one being
+	// up front and two being last.
+	require_Publish("two")
+	require_Messages(2)
+	mset.mu.RLock()
+	require_Equal(t, mset.ddmap.Size(), 2)
+	two, ok := mset.ddmap.Find(stringToBytes("two"))
+	require_True(t, ok)
+	require_Equal(t, mset.ddnext, *one)
+	require_Equal(t, mset.ddlast, *two)
+	require_Stitched()
+	require_NotNil(t, mset.ddtmr)
+	mset.mu.RUnlock()
+
+	// Publish a third message with a different Nats-Msg-Id.
+	// There will now be three messages and the ddentries for
+	// all three should have been stitched together, with one
+	// still being up front and three now being last.
+	require_Publish("three")
+	require_Messages(3)
+	mset.mu.RLock()
+	require_Equal(t, mset.ddmap.Size(), 3)
+	three, ok := mset.ddmap.Find(stringToBytes("three"))
+	require_True(t, ok)
+	require_Equal(t, mset.ddnext, *one)
+	require_Equal(t, mset.ddlast, *three)
+	require_Stitched()
+	require_NotNil(t, mset.ddtmr)
+	mset.mu.RUnlock()
+
+	// Publish another duplicate for one, which is filtered
+	// out and ignored. Therefore all of the state should be
+	// unchanged from the previous publish.
+	require_Publish("one")
+	require_Messages(3)
+	mset.mu.RLock()
+	require_Equal(t, mset.ddmap.Size(), 3)
+	newone, ok := mset.ddmap.Find(stringToBytes("one"))
+	require_True(t, ok)
+	require_Equal(t, *one, *newone) // Entry not overwritten.
+	require_Equal(t, mset.ddnext, *one)
+	require_Equal(t, mset.ddlast, *three)
+	require_Stitched()
+	require_NotNil(t, mset.ddtmr)
+	mset.mu.RUnlock()
+
+	// Wait for the duplicate window to pass.
+	time.Sleep(time.Second)
+
+	// By this point the timer has fired and therefore the
+	// state for duplicates has been cleaned up by now.
+	mset.mu.RLock()
+	require_Equal(t, mset.ddmap.Size(), 0)
+	require_Equal(t, mset.ddtmr, nil)
+	require_Equal(t, mset.ddnext, nil)
+	require_Equal(t, mset.ddlast, nil)
+	mset.mu.RUnlock()
+
+	// This should now let us publish another message with
+	// a Nats-Msg-Id of one, which would have been a duplicate
+	// before but now isn't. Therefore there should now be 4
+	// messages in the stream.
+	require_Publish("one")
+	require_Messages(4)
+	mset.mu.RLock()
+	require_Equal(t, mset.ddmap.Size(), 1)
+	one, ok = mset.ddmap.Find(stringToBytes("one"))
+	require_True(t, ok)
+	require_Equal(t, mset.ddnext, *one)
+	require_Equal(t, mset.ddlast, *one)
+	require_Stitched()
+	require_NotNil(t, mset.ddtmr)
+	mset.mu.RUnlock()
+}