Repository: nats-io/nats-server
Pull request #4835 (2023-11-30 19:12:27)

[IMPROVED] MQTT: concurrently fetch retained messages for a new SUB
Speed up sending out retained messages to a sub by loading them from JetStream concurrently.

- added `jsa.newRequestExMulti` that fans out multiple requests (on the same subject, at least for now) and gathers all responses
- further changed the processing of retained messages in `as.processSubs`: now we 
  - a) find all messages that we need and that are not in `rmsCache` first
  - (b) load them to `rms` with one call
  - (c) serialize to subs/filters and fail the filter if it misses messages from (a) in `rms`.

kozlovic: @levb I just pushed a change that adds a test for the retained messages cache cleanup, that would have caught the panic during the interface conversion. Going to finish reviewing your recent changes now.

Patch:
diff --git a/server/mqtt.go b/server/mqtt.go
index 730d1404..e82ddb29 100644
--- a/server/mqtt.go
+++ b/server/mqtt.go
@@ -186,17 +186,18 @@ const (
 
 	mqttInitialPubHeader        = 16 // An overkill, should need 7 bytes max
 	mqttProcessSubTooLong       = 100 * time.Millisecond
-	mqttRetainedCacheTTL        = 2 * time.Minute
+	mqttDefaultRetainedCacheTTL = 2 * time.Minute
 	mqttRetainedTransferTimeout = 10 * time.Second
 )
 
 var (
-	mqttPingResponse  = []byte{mqttPacketPingResp, 0x0}
-	mqttProtoName     = []byte("MQTT")
-	mqttOldProtoName  = []byte("MQIsdp")
-	mqttSessJailDur   = mqttSessFlappingJailDur
-	mqttFlapCleanItvl = mqttSessFlappingCleanupInterval
-	mqttJSAPITimeout  = 4 * time.Second
+	mqttPingResponse     = []byte{mqttPacketPingResp, 0x0}
+	mqttProtoName        = []byte("MQTT")
+	mqttOldProtoName     = []byte("MQIsdp")
+	mqttSessJailDur      = mqttSessFlappingJailDur
+	mqttFlapCleanItvl    = mqttSessFlappingCleanupInterval
+	mqttJSAPITimeout     = 4 * time.Second
+	mqttRetainedCacheTTL = mqttDefaultRetainedCacheTTL
 )
 
 var (
@@ -242,20 +243,25 @@ type mqttAccountSessionManager struct {
 	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
-	rmsCache   sync.Map                       // map[string(subject)]mqttRetainedMsg
+	rmsCache   sync.Map                       // map[subject]*mqttRetainedMsg
 	jsa        mqttJSA
 	rrmLastSeq uint64        // Restore retained messages expected last sequence
 	rrmDoneCh  chan struct{} // To notify the caller that all retained messages have been loaded
 	domainTk   string        // Domain (with trailing "."), or possibly empty. This is added to session subject.
 }
 
+type mqttJSAResponse struct {
+	reply string // will be used to map to the original request in jsa.NewRequestExMulti
+	value any
+}
+
 type mqttJSA struct {
 	mu        sync.Mutex
 	id        string
 	c         *client
 	sendq     *ipQueue[*mqttJSPubMsg]
 	rplyr     string
-	replies   sync.Map
+	replies   sync.Map // [string]chan *mqttJSAResponse
 	nuid      *nuid.NUID
 	quitCh    chan struct{}
 	domain    string // Domain or possibly empty. This is added to session subject.
@@ -1232,7 +1238,7 @@ func (s *Server) mqttCreateAccountSessionManager(acc *Account, quitCh chan struc
 	// Start the go routine that will clean up cached retained messages that expired.
 	s.startGoRoutine(func() {
 		defer s.grWG.Done()
-		as.cleaupRetainedMessageCache(s, closeCh)
+		as.cleanupRetainedMessageCache(s, closeCh)
 	})
 
 	lookupStream := func(stream, txt string) (*StreamInfo, error) {
@@ -1523,51 +1529,106 @@ func (jsa *mqttJSA) prefixDomain(subject string) string {
 	return subject
 }
 
-func (jsa *mqttJSA) newRequestEx(kind, subject, cidHash string, hdr int, msg []byte, timeout time.Duration) (interface{}, error) {
-	var sb strings.Builder
-	// Either we use nuid.Next() which uses a global lock, or our own nuid object, but
-	// then it needs to be "write" protected. This approach will reduce across account
-	// contention since we won't use the global nuid's lock.
-	jsa.mu.Lock()
-	uid := jsa.nuid.Next()
-	sb.WriteString(jsa.rplyr)
-	jsa.mu.Unlock()
-
-	sb.WriteString(kind)
-	sb.WriteByte(btsep)
-	if cidHash != _EMPTY_ {
-		sb.WriteString(cidHash)
-		sb.WriteByte(btsep)
+func (jsa *mqttJSA) newRequestEx(kind, subject, cidHash string, hdr int, msg []byte, timeout time.Duration) (any, error) {
+	responses, err := jsa.newRequestExMulti(kind, subject, cidHash, []int{hdr}, [][]byte{msg}, timeout)
+	if err != nil {
+		return nil, err
+	}
+	if len(responses) != 1 {
+		return nil, fmt.Errorf("unreachable: invalid number of responses (%d)", len(responses))
 	}
-	sb.WriteString(uid)
-	reply := sb.String()
+	return responses[0].value, nil
+}
 
-	ch := make(chan interface{}, 1)
-	jsa.replies.Store(reply, ch)
+// newRequestExMulti sends multiple messages on the same subject and waits for
+// all responses. It returns the same number of responses in the same order as
+// msgs parameter. In case of a timeout it returns an error as well as all
+// responses received as a sparsely populated array, matching msgs, with nils
+// for the values that have not yet been received.
+//
+// Note that each response may represent an error and should be inspected as
+// such by the caller.
+func (jsa *mqttJSA) newRequestExMulti(kind, subject, cidHash string, hdrs []int, msgs [][]byte, timeout time.Duration) ([]*mqttJSAResponse, error) {
+	if len(hdrs) != len(msgs) {
+		return nil, fmt.Errorf("unreachable: invalid number of messages (%d) or header offsets (%d)", len(msgs), len(hdrs))
+	}
+	responseCh := make(chan *mqttJSAResponse, len(msgs))
+
+	// Generate and queue all outgoing requests, have all results reported to
+	// responseCh, and store a map of reply subjects to the original subjects'
+	// indices.
+	r2i := map[string]int{}
+	for i, msg := range msgs {
+		hdr := hdrs[i]
+		var sb strings.Builder
+		// Either we use nuid.Next() which uses a global lock, or our own nuid object, but
+		// then it needs to be "write" protected. This approach will reduce across account
+		// contention since we won't use the global nuid's lock.
+		jsa.mu.Lock()
+		uid := jsa.nuid.Next()
+		sb.WriteString(jsa.rplyr)
+		jsa.mu.Unlock()
+
+		sb.WriteString(kind)
+		sb.WriteByte(btsep)
+		if cidHash != _EMPTY_ {
+			sb.WriteString(cidHash)
+			sb.WriteByte(btsep)
+		}
+		sb.WriteString(uid)
+		reply := sb.String()
 
-	subject = jsa.prefixDomain(subject)
-	jsa.sendq.push(&mqttJSPubMsg{
-		subj:  subject,
-		reply: reply,
-		hdr:   hdr,
-		msg:   msg,
-	})
+		// Add responseCh to the reply channel map. It will be cleaned out on
+		// timeout (see below), or in processJSAPIReplies upon receiving the
+		// response.
+		jsa.replies.Store(reply, responseCh)
+
+		subject = jsa.prefixDomain(subject)
+		jsa.sendq.push(&mqttJSPubMsg{
+			subj:  subject,
+			reply: reply,
+			hdr:   hdr,
+			msg:   msg,
+		})
+		r2i[reply] = i
+	}
 
-	var i interface{}
-	// We don't want to use time.After() which causes memory growth because the timer
-	// can't be stopped and will need to expire to then be garbage collected.
+	// Wait for all responses to come back, or for the timeout to expire. We
+	// don't want to use time.After() which causes memory growth because the
+	// timer can't be stopped and will need to expire to then be garbage
+	// collected.
+	c := 0
+	responses := make([]*mqttJSAResponse, len(msgs))
+	start := time.Now()
 	t := time.NewTimer(timeout)
-	select {
-	case i = <-ch:
-		// Ensure we stop the timer so it can be quickly garbage collected.
-		t.Stop()
-	case <-jsa.quitCh:
-		return nil, ErrServerNotRunning
-	case <-t.C:
-		jsa.replies.Delete(reply)
-		return nil, fmt.Errorf("timeout for request type %q on %q (reply=%q)", kind, subject, reply)
+	defer t.Stop()
+	for {
+		select {
+		case r := <-responseCh:
+			i := r2i[r.reply]
+			responses[i] = r
+			c++
+			if c == len(msgs) {
+				return responses, nil
+			}
+
+		case <-jsa.quitCh:
+			return nil, ErrServerNotRunning
+
+		case <-t.C:
+			var reply string
+			now := time.Now()
+			for reply = range r2i { // preserve the last value for Errorf
+				jsa.replies.Delete(reply)
+			}
+
+			if len(msgs) == 1 {
+				return responses, fmt.Errorf("timeout after %v: request type %q on %q (reply=%q)", now.Sub(start), kind, subject, reply)
+			} else {
+				return responses, fmt.Errorf("timeout after %v: request type %q on %q: got %d out of %d", now.Sub(start), kind, subject, c, len(msgs))
+			}
+		}
 	}
-	return i, nil
 }
 
 func (jsa *mqttJSA) sendAck(ackSubject string) {
@@ -1667,6 +1728,30 @@ func (jsa *mqttJSA) loadLastMsgFor(streamName string, subject string) (*StoredMs
 	return lmr.Message, lmr.ToError()
 }
 
+func (jsa *mqttJSA) loadLastMsgForMulti(streamName string, subjects []string) ([]*JSApiMsgGetResponse, error) {
+	marshaled := make([][]byte, 0, len(subjects))
+	headerBytes := make([]int, 0, len(subjects))
+	for _, subject := range subjects {
+		mreq := &JSApiMsgGetRequest{LastFor: subject}
+		bb, err := json.Marshal(mreq)
+		if err != nil {
+			return nil, err
+		}
+		marshaled = append(marshaled, bb)
+		headerBytes = append(headerBytes, 0)
+	}
+
+	all, err := jsa.newRequestExMulti(mqttJSAMsgLoad, fmt.Sprintf(JSApiMsgGetT, streamName), _EMPTY_, headerBytes, marshaled, mqttJSAPITimeout)
+	// all has the same order as subjects, preserve it as we unmarshal
+	responses := make([]*JSApiMsgGetResponse, len(all))
+	for i, v := range all {
+		if v != nil {
+			responses[i] = v.value.(*JSApiMsgGetResponse)
+		}
+	}
+	return responses, err
+}
+
 func (jsa *mqttJSA) loadNextMsgFor(streamName string, subject string) (*StoredMsg, error) {
 	mreq := &JSApiMsgGetRequest{NextFor: subject}
 	req, err := json.Marshal(mreq)
@@ -1771,68 +1856,71 @@ func (as *mqttAccountSessionManager) processJSAPIReplies(_ *subscription, pc *cl
 		return
 	}
 	jsa.replies.Delete(subject)
-	ch := chi.(chan interface{})
+	ch := chi.(chan *mqttJSAResponse)
+	out := func(value any) {
+		ch <- &mqttJSAResponse{reply: subject, value: value}
+	}
 	switch token {
 	case mqttJSAStreamCreate:
 		var resp = &JSApiStreamCreateResponse{}
 		if err := json.Unmarshal(msg, resp); err != nil {
 			resp.Error = NewJSInvalidJSONError()
 		}
-		ch <- resp
+		out(resp)
 	case mqttJSAStreamUpdate:
 		var resp = &JSApiStreamUpdateResponse{}
 		if err := json.Unmarshal(msg, resp); err != nil {
 			resp.Error = NewJSInvalidJSONError()
 		}
-		ch <- resp
+		out(resp)
 	case mqttJSAStreamLookup:
 		var resp = &JSApiStreamInfoResponse{}
 		if err := json.Unmarshal(msg, &resp); err != nil {
 			resp.Error = NewJSInvalidJSONError()
 		}
-		ch <- resp
+		out(resp)
 	case mqttJSAStreamDel:
 		var resp = &JSApiStreamDeleteResponse{}
 		if err := json.Unmarshal(msg, &resp); err != nil {
 			resp.Error = NewJSInvalidJSONError()
 		}
-		ch <- resp
+		out(resp)
 	case mqttJSAConsumerCreate:
 		var resp = &JSApiConsumerCreateResponse{}
 		if err := json.Unmarshal(msg, resp); err != nil {
 			resp.Error = NewJSInvalidJSONError()
 		}
-		ch <- resp
+		out(resp)
 	case mqttJSAConsumerDel:
 		var resp = &JSApiConsumerDeleteResponse{}
 		if err := json.Unmarshal(msg, resp); err != nil {
 			resp.Error = NewJSInvalidJSONError()
 		}
-		ch <- resp
+		out(resp)
 	case mqttJSAMsgStore, mqttJSASessPersist:
 		var resp = &JSPubAckResponse{}
 		if err := json.Unmarshal(msg, resp); err != nil {
 			resp.Error = NewJSInvalidJSONError()
 		}
-		ch <- resp
+		out(resp)
 	case mqttJSAMsgLoad:
 		var resp = &JSApiMsgGetResponse{}
-		if err := json.Unmarshal(msg, resp); err != nil {
+		if err := json.Unmarshal(msg, &resp); err != nil {
 			resp.Error = NewJSInvalidJSONError()
 		}
-		ch <- resp
+		out(resp)
 	case mqttJSAStreamNames:
 		var resp = &JSApiStreamNamesResponse{}
 		if err := json.Unmarshal(msg, resp); err != nil {
 			resp.Error = NewJSInvalidJSONError()
 		}
-		ch <- resp
+		out(resp)
 	case mqttJSAMsgDelete:
 		var resp = &JSApiMsgDeleteResponse{}
 		if err := json.Unmarshal(msg, resp); err != nil {
 			resp.Error = NewJSInvalidJSONError()
 		}
-		ch <- resp
+		out(resp)
 	default:
 		pc.Warnf("Unknown reply code %q", token)
 	}
@@ -1856,8 +1944,8 @@ func (as *mqttAccountSessionManager) processRetainedMsg(_ *subscription, c *clie
 	// At this point we either recover from our own server, or process a remote retained message.
 	seq, _, _ := ackReplyInfo(reply)
 
-	// Handle this retained message
-	as.handleRetainedMsg(rm.Subject, &mqttRetainedMsgRef{sseq: seq}, rm)
+	// Handle this retained message, no need to copy the bytes.
+	as.handleRetainedMsg(rm.Subject, &mqttRetainedMsgRef{sseq: seq}, rm, false)
 
 	// If we were recovering (lastSeq > 0), then check if we are done.
 	if as.rrmLastSeq > 0 && seq >= as.rrmLastSeq {
@@ -2001,7 +2089,7 @@ func (as *mqttAccountSessionManager) createSubscription(subject string, cb msgHa
 // only used when the server shutdown.
 //
 // No lock held on entry.
-func (as *mqttAccountSessionManager) cleaupRetainedMessageCache(s *Server, closeCh chan struct{}) {
+func (as *mqttAccountSessionManager) cleanupRetainedMessageCache(s *Server, closeCh chan struct{}) {
 	tt := time.NewTicker(mqttRetainedCacheTTL)
 	defer tt.Stop()
 	for {
@@ -2013,7 +2101,7 @@ func (as *mqttAccountSessionManager) cleaupRetainedMessageCache(s *Server, close
 			i, maxScan := 0, 10*1000
 			now := time.Now()
 			as.rmsCache.Range(func(key, value interface{}) bool {
-				rm := value.(mqttRetainedMsg)
+				rm := value.(*mqttRetainedMsg)
 				if now.After(rm.expiresFromCache) {
 					as.rmsCache.Delete(key)
 				}
@@ -2123,7 +2211,7 @@ func (as *mqttAccountSessionManager) sendJSAPIrequests(s *Server, c *client, acc
 // If a message for this topic already existed, the existing record is updated
 // with the provided information.
 // Lock not held on entry.
-func (as *mqttAccountSessionManager) handleRetainedMsg(key string, rf *mqttRetainedMsgRef, rm *mqttRetainedMsg) {
+func (as *mqttAccountSessionManager) handleRetainedMsg(key string, rf *mqttRetainedMsgRef, rm *mqttRetainedMsg, copyBytesToCache bool) {
 	as.mu.Lock()
 	defer as.mu.Unlock()
 	if as.retmsgs == nil {
@@ -2151,11 +2239,7 @@ func (as *mqttAccountSessionManager) handleRetainedMsg(key string, rf *mqttRetai
 			// Update the in-memory retained message cache but only for messages
 			// that are already in the cache, i.e. have been (recently) used.
 			if rm != nil {
-				if _, ok := as.rmsCache.Load(key); ok {
-					toStore := *rm
-					toStore.expiresFromCache = time.Now().Add(mqttRetainedCacheTTL)
-					as.rmsCache.Store(key, toStore)
-				}
+				as.setCachedRetainedMsg(key, rm, true, copyBytesToCache)
 			}
 			return
 		}
@@ -2274,19 +2358,21 @@ func (as *mqttAccountSessionManager) removeSession(sess *mqttSession, lock bool)
 // Session lock held on entry. Acquires the subs lock and holds it for
 // the duration. Non-MQTT messages coming into mqttDeliverMsgCbQoS0 will be
 // waiting.
+func (sess *mqttSession) processQOS12Sub(
+	c *client, // subscribing client.
+	subject, sid []byte, isReserved bool, qos byte, jsDurName string, h msgHandler, // subscription parameters.
+) (*subscription, error) {
+	return sess.processSub(c, subject, sid, isReserved, qos, jsDurName, h, false, false, nil, false, nil)
+}
+
 func (sess *mqttSession) processSub(
-	// subscribing client.
-	c *client,
-	// subscription parameters.
-	subject, sid []byte, isReserved bool, qos byte, jsDurName string, h msgHandler,
-	// do we need to scan for shadow subscriptions? (we don't do it for QOS1+)
-	initShadow bool,
-	// len(rms) > 0 means to deliver retained messages for the subscription.
-	rms map[string]*mqttRetainedMsg,
-	// trace serialized retained messages in the log.
-	trace bool,
-	// the retained messages are kept in the account session manager.
-	as *mqttAccountSessionManager,
+	c *client, // subscribing client.
+	subject, sid []byte, isReserved bool, qos byte, jsDurName string, h msgHandler, // subscription parameters.
+	initShadow bool, // do we need to scan for shadow subscriptions? (not for QOS1+)
+	serializeRMS bool, // do we need to serialize RMS?
+	rms map[string]*mqttRetainedMsg, // preloaded rms (can be empty, or missing items if errors)
+	trace bool, // trace serialized retained messages in the log?
+	as *mqttAccountSessionManager, // needed only for rms serialization.
 ) (*subscription, error) {
 	start := time.Now()
 	defer func() {
@@ -2350,6 +2436,10 @@ func (sess *mqttSession) processSub(
 func (as *mqttAccountSessionManager) processSubs(sess *mqttSession, c *client,
 	filters []*mqttFilter, fromSubProto, trace bool) ([]*subscription, error) {
 
+	c.mu.Lock()
+	acc := c.acc
+	c.mu.Unlock()
+
 	// Helper to determine if we need to create a separate top-level
 	// subscription for a wildcard.
 	fwc := func(subject string) (bool, string, string) {
@@ -2364,29 +2454,7 @@ func (as *mqttAccountSessionManager) processSubs(sess *mqttSession, c *client,
 		return true, fwcsubject, fwcsid
 	}
 
-	// Cache and a helper to load retained messages for a given subject.
-	rms := make(map[string]*mqttRetainedMsg)
-	loadRMS := func(subject []byte) error {
-		sub := &subscription{
-			client:  c,
-			subject: subject,
-			sid:     subject,
-		}
-		c.mu.Lock()
-		acc := c.acc
-		c.mu.Unlock()
-		if err := c.addShadowSubscriptions(acc, sub, false); err != nil {
-			return err
-		}
-		// Best-effort loading the messages, logs on errors (to c.srv), loads
-		// once for subject.
-		as.loadRetainedMessagesForSubject(rms, subject, c.srv)
-		for _, ss := range sub.shadow {
-			as.loadRetainedMessagesForSubject(rms, ss.subject, c.srv)
-		}
-		return nil
-	}
-
+	rmSubjects := map[string]struct{}{}
 	// Preload retained messages for all requested subscriptions.  Also, since
 	// it's the first iteration over the filter list, do some cleanup.
 	for _, f := range filters {
@@ -2416,14 +2484,33 @@ func (as *mqttAccountSessionManager) processSubs(sess *mqttSession, c *client,
 			}
 		}
 
-		// Load retained messages.
+		// Find retained messages.
 		if fromSubProto {
-			if err := loadRMS([]byte(f.filter)); err != nil {
+			addRMSubjects := func(subject string) error {
+				sub := &subscription{
+					client:  c,
+					subject: []byte(subject),
+					sid:     []byte(subject),
+				}
+				if err := c.addShadowSubscriptions(acc, sub, false); err != nil {
+					return err
+				}
+
+				for _, sub := range append([]*subscription{sub}, sub.shadow...) {
+					as.addRetainedSubjectsForSubject(rmSubjects, bytesToString(sub.subject))
+					for _, ss := range sub.shadow {
+						as.addRetainedSubjectsForSubject(rmSubjects, bytesToString(ss.subject))
+					}
+				}
+				return nil
+			}
+
+			if err := addRMSubjects(f.filter); err != nil {
 				f.qos = mqttSubAckFailure
 				continue
 			}
 			if need, subject, _ := fwc(f.filter); need {
-				if err := loadRMS([]byte(subject)); err != nil {
+				if err := addRMSubjects(subject); err != nil {
 					f.qos = mqttSubAckFailure
 					continue
 				}
@@ -2431,6 +2518,14 @@ func (as *mqttAccountSessionManager) processSubs(sess *mqttSession, c *client,
 		}
 	}
 
+	serializeRMS := len(rmSubjects) > 0
+	var rms map[string]*mqttRetainedMsg
+	if serializeRMS {
+		// Make the best effort to load retained messages. We will identify
+		// errors in the next pass.
+		rms = as.loadRetainedMessages(rmSubjects, c)
+	}
+
 	// Small helper to add the consumer config to the session.
 	addJSConsToSess := func(sid string, cc *ConsumerConfig) {
 		if cc == nil {
@@ -2445,7 +2540,6 @@ func (as *mqttAccountSessionManager) processSubs(sess *mqttSession, c *client,
 	var err error
 	subs := make([]*subscription, 0, len(filters))
 	for _, f := range filters {
-
 		// Skip what's already been identified as a failure.
 		if f.qos == mqttSubAckFailure {
 			continue
@@ -2472,7 +2566,7 @@ func (as *mqttAccountSessionManager) processSubs(sess *mqttSession, c *client,
 			bsubject, bsid, isReserved, f.qos, // main subject
 			_EMPTY_, mqttDeliverMsgCbQoS0, // no jsDur for QOS0
 			processShadowSubs,
-			rms, trace, as) // rms is empty if not fromSubProto
+			serializeRMS, rms, trace, as)
 		sess.mu.Unlock()
 		as.mu.Unlock()
 
@@ -2506,7 +2600,7 @@ func (as *mqttAccountSessionManager) processSubs(sess *mqttSession, c *client,
 				[]byte(fwcsubject), []byte(fwcsid), isReserved, f.qos, // FWC (top-level wildcard) subject
 				_EMPTY_, mqttDeliverMsgCbQoS0, // no jsDur for QOS0
 				processShadowSubs,
-				rms, trace, as) // rms is empty if not fromSubProto
+				serializeRMS, rms, trace, as)
 			sess.mu.Unlock()
 			as.mu.Unlock()
 			if err != nil {
@@ -2532,6 +2626,7 @@ func (as *mqttAccountSessionManager) processSubs(sess *mqttSession, c *client,
 		subs = append(subs, sub)
 		addJSConsToSess(sid, jscons)
 	}
+
 	if fromSubProto {
 		err = sess.update(filters, true)
 	}
@@ -2548,20 +2643,21 @@ func (as *mqttAccountSessionManager) processSubs(sess *mqttSession, c *client,
 // Runs from the client's readLoop.
 // Account session manager lock held on entry.
 // Session lock held on entry.
-func (as *mqttAccountSessionManager) serializeRetainedMsgsForSub(rms map[string]*mqttRetainedMsg, sess *mqttSession, c *client, sub *subscription, trace bool) {
+func (as *mqttAccountSessionManager) serializeRetainedMsgsForSub(rms map[string]*mqttRetainedMsg, sess *mqttSession, c *client, sub *subscription, trace bool) error {
 	if len(as.retmsgs) == 0 || len(rms) == 0 {
-		return
+		return nil
 	}
 	result := as.sl.ReverseMatch(string(sub.subject))
 	if len(result.psubs) == 0 {
-		return
+		return nil
 	}
+	toTrace := []mqttPublish{}
 	for _, psub := range result.psubs {
 
 		rm := rms[string(psub.subject)]
 		if rm == nil {
-			// This should not happen since we pre-load messages into the cache
-			// before calling serialize.
+			// This should not happen since we pre-load messages into rms before
+			// calling serialize.
 			continue
 		}
 		var pi uint16
@@ -2591,65 +2687,90 @@ func (as *mqttAccountSessionManager) serializeRetainedMsgsForSub(rms map[string]
 		sub.mqtt.prm = append(sub.mqtt.prm, headerBytes, rm.Msg)
 		c.mu.Unlock()
 		if trace {
-			pp := mqttPublish{
+			toTrace = append(toTrace, mqttPublish{
 				topic: []byte(rm.Topic),
 				flags: flags,
 				pi:    pi,
 				sz:    len(rm.Msg),
-			}
-			c.traceOutOp("PUBLISH", []byte(mqttPubTrace(&pp)))
+			})
 		}
 	}
+	for _, pp := range toTrace {
+		c.traceOutOp("PUBLISH", []byte(mqttPubTrace(&pp)))
+	}
+	return nil
 }
 
-// Returns in the provided slice all publish retained message records that
+// Appends the stored message subjects for all retained message records that
 // match the given subscription's `subject` (which could have wildcards).
 //
 // Account session manager NOT lock held on entry.
-func (as *mqttAccountSessionManager) loadRetainedMessagesForSubject(rms map[string]*mqttRetainedMsg, topSubject []byte, log Logger) {
+func (as *mqttAccountSessionManager) addRetainedSubjectsForSubject(list map[string]struct{}, topSubject string) bool {
 	as.mu.RLock()
 	if len(as.retmsgs) == 0 {
 		as.mu.RUnlock()
-		return
+		return false
 	}
-	result := as.sl.ReverseMatch(string(topSubject))
+	result := as.sl.ReverseMatch(topSubject)
 	as.mu.RUnlock()
 
-	if len(result.psubs) == 0 {
-		return
-	}
+	added := false
 	for _, sub := range result.psubs {
 		subject := string(sub.subject)
-		if rms[subject] != nil {
-			continue // already loaded
+		if _, ok := list[subject]; ok {
+			continue
 		}
+		list[subject] = struct{}{}
+		added = true
+	}
 
-		// See if we have the retained message in the cache.
-		if rmv, _ := as.rmsCache.Load(subject); rmv != nil {
-			rm := rmv.(mqttRetainedMsg)
-			rms[subject] = &rm
-			continue
+	return added
+}
+
+type warner interface {
+	Warnf(format string, v ...any)
+}
+
+// Loads a list of retained messages given a list of stored message subjects.
+func (as *mqttAccountSessionManager) loadRetainedMessages(subjects map[string]struct{}, w warner) map[string]*mqttRetainedMsg {
+	rms := make(map[string]*mqttRetainedMsg, len(subjects))
+	ss := []string{}
+	for s := range subjects {
+		if rm := as.getCachedRetainedMsg(s); rm != nil {
+			rms[s] = rm
+		} else {
+			ss = append(ss, mqttRetainedMsgsStreamSubject+s)
 		}
+	}
 
-		// Load the retained message from the stream, and cache it for reuse in
-		// the near future.
-		loadSubject := mqttRetainedMsgsStreamSubject + subject
-		jsm, err := as.jsa.loadLastMsgFor(mqttRetainedMsgsStreamName, loadSubject)
-		if err != nil || jsm == nil {
-			log.Warnf("failed to load retained message for subject %q: %v", loadSubject, err)
+	if len(ss) == 0 {
+		return rms
+	}
+
+	results, err := as.jsa.loadLastMsgForMulti(mqttRetainedMsgsStreamName, ss)
+	// If an error occurred, warn, but then proceed with what we got.
+	if err != nil {
+		w.Warnf("error loading retained messages: %v", err)
+	}
+	for i, result := range results {
+		if result == nil {
+			continue // skip requests that timed out
+		}
+		if result.ToError() != nil {
+			w.Warnf("failed to load retained message for subject %q: %v", ss[i], err)
 			continue
 		}
 		var rm mqttRetainedMsg
-		if err := json.Unmarshal(jsm.Data, &rm); err != nil {
-			log.Warnf("failed to decode retained message for subject %q: %v", loadSubject, err)
+		if err := json.Unmarshal(result.Message.Data, &rm); err != nil {
+			w.Warnf("failed to decode retained message for subject %q: %v", ss[i], err)
 			continue
 		}
-
-		// Add the loaded retained message to the cache.
-		rm.expiresFromCache = time.Now().Add(mqttRetainedCacheTTL)
-		as.rmsCache.Store(subject, rm)
-		rms[subject] = &rm
+		// Add the loaded retained message to the cache, and to the results map.
+		key := ss[i][len(mqttRetainedMsgsStreamSubject):]
+		as.setCachedRetainedMsg(key, &rm, false, false)
+		rms[key] = &rm
 	}
+	return rms
 }
 
 // Creates the session stream (limit msgs of 1) for this client ID if it does
@@ -2833,6 +2954,32 @@ func (as *mqttAccountSessionManager) transferRetainedToPerKeySubjectStream(log *
 	return nil
 }
 
+func (as *mqttAccountSessionManager) getCachedRetainedMsg(subject string) *mqttRetainedMsg {
+	v, ok := as.rmsCache.Load(subject)
+	if !ok {
+		return nil
+	}
+	rm := v.(*mqttRetainedMsg)
+	if rm.expiresFromCache.Before(time.Now()) {
+		as.rmsCache.Delete(subject)
+		return nil
+	}
+	return rm
+}
+
+func (as *mqttAccountSessionManager) setCachedRetainedMsg(subject string, rm *mqttRetainedMsg, onlyReplace bool, copyBytesToCache bool) {
+	rm.expiresFromCache = time.Now().Add(mqttRetainedCacheTTL)
+	if onlyReplace {
+		if _, ok := as.rmsCache.Load(subject); !ok {
+			return
+		}
+	}
+	if copyBytesToCache {
+		rm.Msg = copyBytes(rm.Msg)
+	}
+	as.rmsCache.Store(subject, rm)
+}
+
 //////////////////////////////////////////////////////////////////////////////
 //
 // MQTT session related functions
@@ -4026,7 +4173,7 @@ func (c *client) mqttHandlePubRetain() {
 				sseq: smr.Sequence,
 			}
 			// Add/update the map
-			asm.handleRetainedMsg(key, rf, rm)
+			asm.handleRetainedMsg(key, rf, rm, true) // will copy the payload bytes if needs to update rmsCache
 		} else {
 			c.mu.Lock()
 			acc := c.acc
@@ -4893,9 +5040,8 @@ func (sess *mqttSession) processJSConsumer(c *client, subject, sid string,
 	// for the JS durable's deliver subject.
 	sess.mu.Lock()
 	sess.tmaxack = tmaxack
-	sub, err := sess.processSub(c, []byte(inbox), []byte(inbox),
-		isMQTTReservedSubscription(subject), qos, cc.Durable, mqttDeliverMsgCbQoS12,
-		false, nil, false, nil) // no shadow subs, no retained message delivery
+	sub, err := sess.processQOS12Sub(c, []byte(inbox), []byte(inbox),
+		isMQTTReservedSubscription(subject), qos, cc.Durable, mqttDeliverMsgCbQoS12)
 	sess.mu.Unlock()
 
 	if err != nil {


Test patch:
diff --git a/server/mqtt_test.go b/server/mqtt_test.go
index ae6ba887..7824e690 100644
--- a/server/mqtt_test.go
+++ b/server/mqtt_test.go
@@ -1938,7 +1938,7 @@ func testMQTTSub(t testing.TB, pi uint16, c net.Conn, r *mqttReader, filters []*
 		if err != nil {
 			t.Fatal(err)
 		}
-		if qos != expected[i] {
+		if expected != nil && qos != expected[i] {
 			t.Fatalf("For topic filter %q expected qos of %v, got %v",
 				filters[i].filter, expected[i], qos)
 		}
@@ -2086,13 +2086,42 @@ func testMQTTGetPubMsgExEx(t testing.TB, _ net.Conn, r *mqttReader, b byte, pl i
 			msgLen, r.pos, len(r.buf))
 	}
 	ppayload := r.buf[r.pos : r.pos+msgLen]
-	if !bytes.Equal(payload, ppayload) {
+	if payload != nil && !bytes.Equal(payload, ppayload) {
 		t.Fatalf("Expected payload %q, got %q", payload, ppayload)
 	}
 	r.pos += msgLen
 	return pflags, pi, ptopic
 }
 
+func testMQTTReadPubPacket(t testing.TB, r *mqttReader) (flags byte, pi uint16, topic string, payload []byte) {
+	t.Helper()
+	b, pl := testMQTTReadPacket(t, r)
+	if pt := b & mqttPacketMask; pt != mqttPacketPub {
+		t.Fatalf("Expected PUBLISH packet %x, got %x", mqttPacketPub, pt)
+	}
+	flags = b & mqttPacketFlagMask
+	start := r.pos
+	topic, err := r.readString("topic name")
+	if err != nil {
+		t.Fatal(err)
+	}
+	qos := (flags & mqttPubFlagQoS) >> 1
+	if qos > 0 {
+		pi, err = r.readUint16("packet identifier")
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+	msgLen := pl - (r.pos - start)
+	if r.pos+msgLen > len(r.buf) {
+		t.Fatalf("computed message length goes beyond buffer: ml=%v pos=%v lenBuf=%v",
+			msgLen, r.pos, len(r.buf))
+	}
+	payload = r.buf[r.pos : r.pos+msgLen]
+	r.pos += msgLen
+	return flags, pi, topic, payload
+}
+
 func testMQTTSendPIPacket(packetType byte, t testing.TB, c net.Conn, pi uint16) {
 	t.Helper()
 	w := newMQTTWriter(0)
@@ -3215,7 +3244,7 @@ func TestMQTTRetainedMsgNetworkUpdates(t *testing.T) {
 			for _, a := range test.order {
 				if a.add {
 					rf := &mqttRetainedMsgRef{sseq: a.seq}
-					asm.handleRetainedMsg(test.subject, rf, nil)
+					asm.handleRetainedMsg(test.subject, rf, nil, false)
 				} else {
 					asm.handleRetainedMsgDel(test.subject, a.seq)
 				}
@@ -3228,7 +3257,7 @@ func TestMQTTRetainedMsgNetworkUpdates(t *testing.T) {
 		t.Run("clear_"+subject, func(t *testing.T) {
 			// Now add a new message, which should clear the floor.
 			rf := &mqttRetainedMsgRef{sseq: 3}
-			asm.handleRetainedMsg(subject, rf, nil)
+			asm.handleRetainedMsg(subject, rf, nil, false)
 			check(t, subject, true, 3, 0)
 			// Now do a non network delete and make sure it is gone.
 			asm.handleRetainedMsgDel(subject, 0)
@@ -3249,7 +3278,7 @@ func TestMQTTRetainedMsgDel(t *testing.T) {
 	var i uint64
 	for i = 0; i < 3; i++ {
 		rf := &mqttRetainedMsgRef{sseq: i}
-		asm.handleRetainedMsg("subject", rf, nil)
+		asm.handleRetainedMsg("subject", rf, nil, false)
 	}
 	asm.handleRetainedMsgDel("subject", 2)
 	if asm.sl.count > 0 {
@@ -3277,7 +3306,8 @@ func TestMQTTRetainedMsgMigration(t *testing.T) {
 	require_NoError(t, err)
 
 	// Publish some retained messages on the old "$MQTT.rmsgs" subject.
-	for i := 0; i < 100; i++ {
+	const N = 100
+	for i := 0; i < N; i++ {
 		msg := fmt.Sprintf(
 			`{"origin":"b5IQZNtG","subject":"test%d","topic":"test%d","msg":"YmFy","flags":1}`, i, i,
 		)
@@ -3293,8 +3323,8 @@ func TestMQTTRetainedMsgMigration(t *testing.T) {
 	if si.State.NumSubjects != 1 {
 		t.Fatalf("expected 1 subject, got %d", si.State.NumSubjects)
 	}
-	if n := si.State.Subjects[`$MQTT.rmsgs`]; n != 100 {
-		t.Fatalf("expected to find 100 messages on the original subject but found %d", n)
+	if n := si.State.Subjects[`$MQTT.rmsgs`]; n != N {
+		t.Fatalf("expected to find %d messages on the original subject but found %d", N, n)
 	}
 
 	// Create an MQTT client, this will cause a migration to take place.
@@ -3304,22 +3334,22 @@ func TestMQTTRetainedMsgMigration(t *testing.T) {
 
 	testMQTTSub(t, 1, mc, rc, []*mqttFilter{{filter: "+", qos: 0}}, []byte{0})
 	topics := map[string]struct{}{}
-	for i := 0; i < 100; i++ {
+	for i := 0; i < N; i++ {
 		_, _, topic := testMQTTGetPubMsgEx(t, mc, rc, _EMPTY_, []byte("bar"))
 		topics[topic] = struct{}{}
 	}
-	if len(topics) != 100 {
+	if len(topics) != N {
 		t.Fatalf("Unexpected topics: %v", topics)
 	}
 
-	// Now look at the stream, there should be 100 messages on the new
+	// Now look at the stream, there should be N messages on the new
 	// divided subjects and none on the old undivided subject.
 	si, err = js.StreamInfo(mqttRetainedMsgsStreamName, &nats.StreamInfoRequest{
 		SubjectsFilter: `$MQTT.>`,
 	})
 	require_NoError(t, err)
-	if si.State.NumSubjects != 100 {
-		t.Fatalf("expected 100 subjects, got %d", si.State.NumSubjects)
+	if si.State.NumSubjects != N {
+		t.Fatalf("expected %d subjects, got %d", N, si.State.NumSubjects)
 	}
 	if n := si.State.Subjects[`$MQTT.rmsgs`]; n > 0 {
 		t.Fatalf("expected to find no messages on the original subject but found %d", n)
@@ -3327,7 +3357,7 @@ func TestMQTTRetainedMsgMigration(t *testing.T) {
 
 	// Check that the message counts look right. There should be one
 	// retained message per key.
-	for i := 0; i < 100; i++ {
+	for i := 0; i < N; i++ {
 		expected := fmt.Sprintf(`$MQTT.rmsgs.test%d`, i)
 		n, ok := si.State.Subjects[expected]
 		if !ok {
@@ -4953,12 +4983,12 @@ func TestMQTTPersistRetainedMsg(t *testing.T) {
 	defer c.Close()
 	testMQTTCheckConnAck(t, r, mqttConnAckRCConnectionAccepted, false)
 
-	testMQTTPublish(t, c, r, 1, false, true, "foo", 1, []byte("foo1"))
-	testMQTTPublish(t, c, r, 1, false, true, "foo", 1, []byte("foo2"))
-	testMQTTPublish(t, c, r, 1, false, true, "bar", 1, []byte("bar1"))
-	testMQTTPublish(t, c, r, 0, false, true, "baz", 1, []byte("baz1"))
+	testMQTTPublish(t, c, r, 1, false, true, "moo/foo", 1, []byte("foo1"))
+	testMQTTPublish(t, c, r, 1, false, true, "moo/foo", 1, []byte("foo2"))
+	testMQTTPublish(t, c, r, 1, false, true, "moo/bar", 1, []byte("bar1"))
+	testMQTTPublish(t, c, r, 0, false, true, "moo/baz", 1, []byte("baz1"))
 	// Remove bar
-	testMQTTPublish(t, c, r, 1, false, true, "bar", 1, nil)
+	testMQTTPublish(t, c, r, 1, false, true, "moo/bar", 1, nil)
 	testMQTTFlush(t, c, nil, r)
 	testMQTTDisconnect(t, c, nil)
 	c.Close()
@@ -4969,26 +4999,108 @@ func TestMQTTPersistRetainedMsg(t *testing.T) {
 	o.MQTT.Port = -1
 	o.StoreDir = dir
 	s = testMQTTRunServer(t, o)
-	// There is already the defer for shutdown at top of function
 
 	cisub := &mqttConnInfo{clientID: "sub", cleanSess: false}
 	c, r = testMQTTConnect(t, cisub, o.MQTT.Host, o.MQTT.Port)
 	defer c.Close()
 	testMQTTCheckConnAck(t, r, mqttConnAckRCConnectionAccepted, false)
 
-	testMQTTSub(t, 1, c, r, []*mqttFilter{{filter: "foo", qos: 1}}, []byte{1})
-	testMQTTCheckPubMsg(t, c, r, "foo", mqttPubFlagRetain|mqttPubQos1, []byte("foo2"))
+	t.Run("many subs one topic", func(t *testing.T) {
+		testMQTTSub(t, 1, c, r, []*mqttFilter{{filter: "moo/foo", qos: 1}}, []byte{1})
+		testMQTTCheckPubMsg(t, c, r, "moo/foo", mqttPubFlagRetain|mqttPubQos1, []byte("foo2"))
 
-	testMQTTSub(t, 1, c, r, []*mqttFilter{{filter: "baz", qos: 1}}, []byte{1})
-	testMQTTCheckPubMsg(t, c, r, "baz", mqttPubFlagRetain, []byte("baz1"))
+		testMQTTSub(t, 1, c, r, []*mqttFilter{{filter: "moo/baz", qos: 1}}, []byte{1})
+		testMQTTCheckPubMsg(t, c, r, "moo/baz", mqttPubFlagRetain, []byte("baz1"))
 
-	testMQTTSub(t, 1, c, r, []*mqttFilter{{filter: "bar", qos: 1}}, []byte{1})
-	testMQTTExpectNothing(t, r)
+		testMQTTSub(t, 1, c, r, []*mqttFilter{{filter: "moo/bar", qos: 1}}, []byte{1})
+		testMQTTExpectNothing(t, r)
+	})
+
+	runSingleSubscribe := func(filters []*mqttFilter) func(t *testing.T) {
+		return func(t *testing.T) {
+			testMQTTSub(t, 1, c, r, filters, nil)
+
+			for i := 0; i < 2; i++ {
+				flags, pi, topic, payload := testMQTTReadPubPacket(t, r)
+				if (flags & mqttPubFlagRetain) == 0 {
+					t.Fatalf("Expected flags to have retain set, got %v", flags)
+				}
+				if (flags & mqttPubFlagDup) != 0 {
+					t.Fatalf("Expected flags to not have Dup set, got %v", flags)
+				}
+				var expQOS byte
+				var expPayload []byte
+				switch topic {
+				case "moo/foo":
+					expQOS = mqttPubQos1
+					expPayload = []byte("foo2")
+					testMQTTSendPIPacket(mqttPacketPubAck, t, c, pi)
+				case "moo/baz":
+					expQOS = 0
+					expPayload = []byte("baz1")
+				default:
+					t.Fatalf("Unexpected topic: %v", topic)
+				}
+				if (flags & mqttPubFlagQoS) != expQOS {
+					t.Fatalf("Expected flags to have QOS %x set, got %x", expQOS, flags)
+				}
+				if string(payload) != string(expPayload) {
+					t.Fatalf("Expected payload to be %q, got %v", expPayload, string(payload))
+				}
+			}
+			testMQTTExpectNothing(t, r)
+		}
+	}
+
+	t.Run("one sub many topics", runSingleSubscribe([]*mqttFilter{
+		{filter: "moo/foo", qos: 1},
+		{filter: "moo/baz", qos: 1},
+		{filter: "moo/bar", qos: 1},
+	}))
+
+	t.Run("one sub wildcard plus", runSingleSubscribe([]*mqttFilter{
+		{filter: "moo/+", qos: 1},
+	}))
+
+	t.Run("one sub wildcard hash", runSingleSubscribe([]*mqttFilter{
+		{filter: "moo/#", qos: 1},
+	}))
 
 	testMQTTDisconnect(t, c, nil)
 	c.Close()
 }
 
+func TestMQTTRetainedMsgCleanup(t *testing.T) {
+	mqttRetainedCacheTTL = 250 * time.Millisecond
+	defer func() { mqttRetainedCacheTTL = mqttDefaultRetainedCacheTTL }()
+
+	o := testMQTTDefaultOptions()
+	s := testMQTTRunServer(t, o)
+	defer testMQTTShutdownServer(s)
+
+	ci := &mqttConnInfo{clientID: "cache", cleanSess: true}
+	c, r := testMQTTConnect(t, ci, o.MQTT.Host, o.MQTT.Port)
+	defer c.Close()
+	testMQTTCheckConnAck(t, r, mqttConnAckRCConnectionAccepted, false)
+
+	// Send a retained message.
+	testMQTTPublish(t, c, r, 1, false, true, "foo", 1, []byte("msg"))
+	testMQTTFlush(t, c, nil, r)
+
+	// Start a subscription.
+	testMQTTSub(t, 1, c, r, []*mqttFilter{{filter: "foo", qos: 1}}, []byte{1})
+	testMQTTCheckPubMsg(t, c, r, "foo", mqttPubQos1|mqttPubFlagRetain, []byte("msg"))
+
+	time.Sleep(2 * mqttRetainedCacheTTL)
+
+	// Make sure not in cache anymore
+	cli := testMQTTGetClient(t, s, "cache")
+	asm := cli.mqtt.asm
+	if v, ok := asm.rmsCache.Load("foo"); ok {
+		t.Fatalf("Should not be in cache, got %+v", v)
+	}
+}
+
 func TestMQTTConnAckFirstPacket(t *testing.T) {
 	o := testMQTTDefaultOptions()
 	o.NoLog, o.Debug, o.Trace = true, false, false