fix: 大帧名额按 PacketID 在 PUBACK 与断线时归还

internal/broker/hooks.go

@@ -25,7 +25,9 @@ func (h *nixHook) Provides(b byte) bool {
 		mqtt.OnPublishDropped,
 		mqtt.OnSessionEstablished,
 		mqtt.OnDisconnect,
+		mqtt.OnQosPublish,
 		mqtt.OnQosComplete,
+		mqtt.OnQosDropped,
 		mqtt.OnSubscribed,
 	}, []byte{b})
 }
@@ -169,13 +171,13 @@ func (h *nixHook) OnSubscribed(cl *mqtt.Client, pk packets.Packet, reasonCodes [
 
 func (h *nixHook) OnPublishDropped(cl *mqtt.Client, pk packets.Packet) {
 	h.b.log.Debug("publish dropped", "client", cl.ID, "topic", pk.TopicName, "size", len(pk.Payload))
-	if h.b.onDrop == nil {
-		return
-	}
 	h.b.connsMu.RLock()
 	st := h.b.byClient[cl]
 	h.b.connsMu.RUnlock()
-	if st == nil {
+	if st != nil {
+		h.b.reconcileLargeInflight(st)
+	}
+	if h.b.onDrop == nil || st == nil {
 		return
 	}
 	h.b.onDrop(context.Background(), st.endpointID, st.connID, append([]byte(nil), pk.Payload...))
@@ -263,7 +265,7 @@ func (h *nixHook) noteConnectionClose(st *connState) {
 	st.metricsCounted = false
 }
 
-func (h *nixHook) OnQosComplete(cl *mqtt.Client, pk packets.Packet) {
+func (h *nixHook) OnQosPublish(cl *mqtt.Client, pk packets.Packet, _ int64, _ int) {
 	if len(pk.Payload) <= largeFrameBytes {
 		return
 	}
@@ -273,5 +275,31 @@ func (h *nixHook) OnQosComplete(cl *mqtt.Client, pk packets.Packet) {
 	if st == nil {
 		return
 	}
-	h.b.releaseOneLarge(st)
+	st.mu.Lock()
+	if st.largePending > 0 {
+		st.largePending--
+	}
+	if st.largePIDs == nil {
+		st.largePIDs = make(map[uint16]struct{})
+	}
+	st.largePIDs[pk.PacketID] = struct{}{}
+	st.mu.Unlock()
+}
+
+func (h *nixHook) OnQosComplete(cl *mqtt.Client, pk packets.Packet) {
+	h.releaseLargeByPacketID(cl, pk.PacketID)
+}
+
+func (h *nixHook) OnQosDropped(cl *mqtt.Client, pk packets.Packet) {
+	h.releaseLargeByPacketID(cl, pk.PacketID)
+}
+
+func (h *nixHook) releaseLargeByPacketID(cl *mqtt.Client, id uint16) {
+	h.b.connsMu.RLock()
+	st := h.b.byClient[cl]
+	h.b.connsMu.RUnlock()
+	if st == nil {
+		return
+	}
+	h.b.releaseLargePID(st, id)
 }