Repository: nats-io/nats-server
Pull request #4228 (2023-06-09 12:52:58)

Don't hold entire MQTT retained messages in memory
This PR separates out the small amount of necessary metadata for retained messages (stream sequence, floor) from the message itself, instead accessing the messages themselves with KV-like access patterns.

This should save quite a bit of memory where there are lots of retained messages since we only now need to hold a small amount of metadata instead of the entire messages.

Signed-off-by: Neil Twigg <neil@nats.io>


Patch:
diff --git a/server/mqtt.go b/server/mqtt.go
index 018d3577..75f3f8d4 100644
--- a/server/mqtt.go
+++ b/server/mqtt.go
@@ -23,6 +23,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"sort"
 	"strconv"
 	"strings"
 	"sync"
@@ -221,13 +222,13 @@ type mqttSessionManager struct {
 
 type mqttAccountSessionManager struct {
 	mu         sync.RWMutex
-	sessions   map[string]*mqttSession     // key is MQTT client ID
-	sessByHash map[string]*mqttSession     // key is MQTT client ID hash
-	sessLocked map[string]struct{}         // key is MQTT client ID and indicate that a session can not be taken by a new client at this time
-	flappers   map[string]int64            // When connection connects with client ID already in use
-	flapTimer  *time.Timer                 // Timer to perform some cleanup of the flappers map
-	sl         *Sublist                    // sublist allowing to find retained messages for given subscription
-	retmsgs    map[string]*mqttRetainedMsg // retained messages
+	sessions   map[string]*mqttSession        // key is MQTT client ID
+	sessByHash map[string]*mqttSession        // key is MQTT client ID hash
+	sessLocked map[string]struct{}            // key is MQTT client ID and indicate that a session can not be taken by a new client at this time
+	flappers   map[string]int64               // When connection connects with client ID already in use
+	flapTimer  *time.Timer                    // Timer to perform some cleanup of the flappers map
+	sl         *Sublist                       // sublist allowing to find retained messages for given subscription
+	retmsgs    map[string]*mqttRetainedMsgRef // retained messages
 	jsa        mqttJSA
 	rrmLastSeq uint64           // Restore retained messages expected last sequence
 	rrmDoneCh  chan struct{}    // To notify the caller that all retained messages have been loaded
@@ -293,8 +294,9 @@ type mqttRetainedMsg struct {
 	Msg     []byte `json:"msg,omitempty"`
 	Flags   byte   `json:"flags,omitempty"`
 	Source  string `json:"source,omitempty"`
+}
 
-	// non exported
+type mqttRetainedMsgRef struct {
 	sseq  uint64
 	floor uint64
 	sub   *subscription
@@ -1604,8 +1606,9 @@ func (as *mqttAccountSessionManager) processRetainedMsg(_ *subscription, c *clie
 	seq, _, _ := ackReplyInfo(reply)
 
 	// Handle this retained message
-	rm.sseq = seq
-	as.handleRetainedMsg(rm.Subject, rm)
+	rf := &mqttRetainedMsgRef{}
+	rf.sseq = seq
+	as.handleRetainedMsg(rm.Subject, rf)
 
 	// If we were recovering (lastSeq > 0), then check if we are done.
 	if as.rrmLastSeq > 0 && seq >= as.rrmLastSeq {
@@ -1873,11 +1876,11 @@ func (as *mqttAccountSessionManager) sendJSAPIrequests(s *Server, c *client, acc
 // or 0 if the record was added instead of updated.
 //
 // Lock not held on entry.
-func (as *mqttAccountSessionManager) handleRetainedMsg(key string, rm *mqttRetainedMsg) uint64 {
+func (as *mqttAccountSessionManager) handleRetainedMsg(key string, rm *mqttRetainedMsgRef) {
 	as.mu.Lock()
 	defer as.mu.Unlock()
 	if as.retmsgs == nil {
-		as.retmsgs = make(map[string]*mqttRetainedMsg)
+		as.retmsgs = make(map[string]*mqttRetainedMsgRef)
 		as.sl = NewSublistWithCache()
 	} else {
 		// Check if we already had one. If so, update the existing one.
@@ -1885,15 +1888,9 @@ func (as *mqttAccountSessionManager) handleRetainedMsg(key string, rm *mqttRetai
 			// If the new sequence is below the floor or the existing one,
 			// then ignore the new one.
 			if rm.sseq <= erm.sseq || rm.sseq <= erm.floor {
-				return 0
+				return
 			}
-			// Update the existing retained message record with the new rm record.
-			erm.Origin = rm.Origin
-			erm.Msg = rm.Msg
-			erm.Flags = rm.Flags
-			erm.Source = rm.Source
 			// Capture existing sequence number so we can return it as the old sequence.
-			oldSeq := erm.sseq
 			erm.sseq = rm.sseq
 			// Clear the floor
 			erm.floor = 0
@@ -1903,13 +1900,11 @@ func (as *mqttAccountSessionManager) handleRetainedMsg(key string, rm *mqttRetai
 				erm.sub = &subscription{subject: []byte(key)}
 				as.sl.Insert(erm.sub)
 			}
-			return oldSeq
 		}
 	}
 	rm.sub = &subscription{subject: []byte(key)}
 	as.retmsgs[key] = rm
 	as.sl.Insert(rm.sub)
-	return 0
 }
 
 // Removes the retained message for the given `subject` if present, and returns the
@@ -1922,7 +1917,7 @@ func (as *mqttAccountSessionManager) handleRetainedMsgDel(subject string, seq ui
 	var seqToRemove uint64
 	as.mu.Lock()
 	if as.retmsgs == nil {
-		as.retmsgs = make(map[string]*mqttRetainedMsg)
+		as.retmsgs = make(map[string]*mqttRetainedMsgRef)
 		as.sl = NewSublistWithCache()
 	}
 	if erm, ok := as.retmsgs[subject]; ok {
@@ -1941,8 +1936,8 @@ func (as *mqttAccountSessionManager) handleRetainedMsgDel(subject string, seq ui
 			seqToRemove = erm.sseq
 		}
 	} else if seq != 0 {
-		rm := &mqttRetainedMsg{Subject: subject, floor: seq}
-		as.retmsgs[subject] = rm
+		rf := &mqttRetainedMsgRef{floor: seq}
+		as.retmsgs[subject] = rf
 	}
 	as.mu.Unlock()
 	return seqToRemove
@@ -2193,11 +2188,16 @@ func (as *mqttAccountSessionManager) getRetainedPublishMsgs(subject string, rms
 		return
 	}
 	for _, sub := range result.psubs {
-		// Since this is a reverse match, the subscription objects here
-		// contain literals corresponding to the published subjects.
-		if rm, ok := as.retmsgs[string(sub.subject)]; ok {
-			*rms = append(*rms, rm)
+		subj := mqttRetainedMsgsStreamSubject + as.domainTk + string(sub.subject)
+		jsm, err := as.jsa.loadLastMsgFor(mqttRetainedMsgsStreamName, subj)
+		if err != nil || jsm == nil {
+			continue
 		}
+		var rm mqttRetainedMsg
+		if err := json.Unmarshal(jsm.Data, &rm); err != nil {
+			continue
+		}
+		*rms = append(*rms, &rm)
 	}
 }
 
@@ -3213,13 +3213,11 @@ func (c *client) mqttHandlePubRetain() {
 		smr, err := asm.jsa.storeMsg(mqttRetainedMsgsStreamSubject+asm.domainTk+key, -1, rmBytes)
 		if err == nil {
 			// Update the new sequence
-			rm.sseq = smr.Sequence
-			// Add/update the map
-			oldSeq := asm.handleRetainedMsg(key, rm)
-			// If this is a new message on the same subject, delete the old one.
-			if oldSeq != 0 {
-				asm.deleteRetainedMsg(oldSeq)
+			rf := &mqttRetainedMsgRef{
+				sseq: smr.Sequence,
 			}
+			// Add/update the map
+			asm.handleRetainedMsg(key, rf)
 		} else {
 			c.mu.Lock()
 			acc := c.acc
@@ -3256,14 +3254,50 @@ func (s *Server) mqttCheckPubRetainedPerms() {
 	}
 	s.mu.Unlock()
 
+	// First get a list of all of the sessions.
 	sm.mu.RLock()
-	defer sm.mu.RUnlock()
-
+	asms := make([]*mqttAccountSessionManager, 0, len(sm.sessions))
 	for _, asm := range sm.sessions {
+		asms = append(asms, asm)
+	}
+	sm.mu.RUnlock()
+
+	type retainedMsg struct {
+		subj string
+		rmsg *mqttRetainedMsgRef
+	}
+
+	// For each session we will obtain a list of retained messages.
+	var _rms [128]retainedMsg
+	rms := _rms[:0]
+	for _, asm := range asms {
+		// Get all of the retained messages. Then we will sort them so
+		// that they are in sequence order, which should help the file
+		// store to not have to load out-of-order blocks so often.
+		asm.mu.RLock()
+		rms = rms[:0] // reuse slice
+		for subj, rf := range asm.retmsgs {
+			rms = append(rms, retainedMsg{
+				subj: subj,
+				rmsg: rf,
+			})
+		}
+		asm.mu.RUnlock()
+		sort.Slice(rms, func(i, j int) bool {
+			return rms[i].rmsg.sseq < rms[j].rmsg.sseq
+		})
+
 		perms := map[string]*perm{}
 		deletes := map[string]uint64{}
-		asm.mu.Lock()
-		for subject, rm := range asm.retmsgs {
+		for _, rf := range rms {
+			jsm, err := asm.jsa.loadMsg(mqttRetainedMsgsStreamName, rf.rmsg.sseq)
+			if err != nil || jsm == nil {
+				continue
+			}
+			var rm mqttRetainedMsg
+			if err := json.Unmarshal(jsm.Data, &rm); err != nil {
+				continue
+			}
 			if rm.Source == _EMPTY_ {
 				continue
 			}
@@ -3277,7 +3311,7 @@ func (s *Server) mqttCheckPubRetainedPerms() {
 				}
 				// If there is permission and no longer allowed to publish in
 				// the subject, remove the publish retained message from the map.
-				if p != nil && !pubAllowed(p, subject) {
+				if p != nil && !pubAllowed(p, rf.subj) {
 					u = nil
 				}
 			}
@@ -3285,12 +3319,14 @@ func (s *Server) mqttCheckPubRetainedPerms() {
 			// Not present or permissions have changed such that the source can't
 			// publish on that subject anymore: remove it from the map.
 			if u == nil {
-				delete(asm.retmsgs, subject)
-				asm.sl.Remove(rm.sub)
-				deletes[subject] = rm.sseq
+				asm.mu.Lock()
+				delete(asm.retmsgs, rf.subj)
+				asm.sl.Remove(rf.rmsg.sub)
+				asm.mu.Unlock()
+				deletes[rf.subj] = rf.rmsg.sseq
 			}
 		}
-		asm.mu.Unlock()
+
 		for subject, seq := range deletes {
 			asm.deleteRetainedMsg(seq)
 			asm.notifyRetainedMsgDeleted(subject, seq)


Test patch:
diff --git a/server/mqtt_test.go b/server/mqtt_test.go
index 31f59da9..a4e793eb 100644
--- a/server/mqtt_test.go
+++ b/server/mqtt_test.go
@@ -2975,8 +2975,8 @@ func TestMQTTRetainedMsgNetworkUpdates(t *testing.T) {
 		t.Run(test.subject, func(t *testing.T) {
 			for _, a := range test.order {
 				if a.add {
-					rm := &mqttRetainedMsg{sseq: a.seq}
-					asm.handleRetainedMsg(test.subject, rm)
+					rf := &mqttRetainedMsgRef{sseq: a.seq}
+					asm.handleRetainedMsg(test.subject, rf)
 				} else {
 					asm.handleRetainedMsgDel(test.subject, a.seq)
 				}
@@ -2988,8 +2988,8 @@ func TestMQTTRetainedMsgNetworkUpdates(t *testing.T) {
 	for _, subject := range []string{"foo.5", "foo.6"} {
 		t.Run("clear_"+subject, func(t *testing.T) {
 			// Now add a new message, which should clear the floor.
-			rm := &mqttRetainedMsg{sseq: 3}
-			asm.handleRetainedMsg(subject, rm)
+			rf := &mqttRetainedMsgRef{sseq: 3}
+			asm.handleRetainedMsg(subject, rf)
 			check(t, subject, true, 3, 0)
 			// Now do a non network delete and make sure it is gone.
 			asm.handleRetainedMsgDel(subject, 0)