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

internal/broker/broker.go

@@ -34,6 +34,11 @@ var ErrPayloadTooLarge = errors.New("broker: payload exceeds client limit")
 // ErrNoConnection 目标端没有当前连接。
 var ErrNoConnection = errors.New("broker: no active connection")
 
+// ErrLargeFrameTimeout 全局大帧名额在有界等待内拿不到。
+var ErrLargeFrameTimeout = errors.New("broker: large frame quota timeout")
+
+const largeAcquireWait = 5 * time.Second
+
 // AuthResult 是登录校验结论（N3 实现真实逻辑；N2 默认拒绝）。
 type AuthResult struct {
 	OK           bool
@@ -108,7 +113,8 @@ type connState struct {
 	sessionToken   string
 	handshook      bool
 	subscribedDown bool
-	largeHeld      int
+	largePIDs      map[uint16]struct{}
+	largePending   int
 	metricsCounted bool
 	mu             sync.Mutex
 
@@ -219,54 +225,98 @@ func (b *Broker) PublishDown(ctx context.Context, endpointID string, connID port
 	}
 	topic := downTopic(endpointID)
 	large := len(payload) > largeFrameBytes
-
 	if large {
-		select {
-		case b.largeSem <- struct{}{}:
-		case <-ctx.Done():
-			return ctx.Err()
+		if err := b.acquireLarge(ctx); err != nil {
+			return err
 		}
 		st.mu.Lock()
-		st.largeHeld++
+		st.largePending++
 		st.mu.Unlock()
 	}
 
 	if err := b.server.Publish(topic, payload, false, qos); err != nil {
 		if large {
-			b.releaseOneLarge(st)
+			b.finishLargePublish(st)
 		}
 		return err
 	}
-	if large && qos == 0 {
-		b.releaseOneLarge(st)
+	if large {
+		b.finishLargePublish(st)
 	}
 	return nil
 }
 
-func (b *Broker) releaseOneLarge(st *connState) {
+func (b *Broker) acquireLarge(ctx context.Context) error {
+	timer := time.NewTimer(largeAcquireWait)
+	defer timer.Stop()
+	select {
+	case b.largeSem <- struct{}{}:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-timer.C:
+		return ErrLargeFrameTimeout
+	}
+}
+
+func (b *Broker) releaseLargeSlot() {
+	select {
+	case <-b.largeSem:
+	default:
+	}
+}
+
+func (b *Broker) finishLargePublish(st *connState) {
+	b.reconcileLargeInflight(st)
 	st.mu.Lock()
-	if st.largeHeld > 0 {
-		st.largeHeld--
-		st.mu.Unlock()
-		select {
-		case <-b.largeSem:
-		default:
-		}
-		return
+	n := st.largePending
+	st.largePending = 0
+	st.mu.Unlock()
+	for i := 0; i < n; i++ {
+		b.releaseLargeSlot()
+	}
+}
+
+func (b *Broker) releaseLargePID(st *connState, id uint16) {
+	st.mu.Lock()
+	_, ok := st.largePIDs[id]
+	if ok {
+		delete(st.largePIDs, id)
 	}
 	st.mu.Unlock()
+	if ok {
+		b.releaseLargeSlot()
+	}
+}
+
+func (b *Broker) reconcileLargeInflight(st *connState) {
+	if st == nil {
+		return
+	}
+	st.mu.Lock()
+	ids := make([]uint16, 0, len(st.largePIDs))
+	for id := range st.largePIDs {
+		ids = append(ids, id)
+	}
+	st.mu.Unlock()
+	for _, id := range ids {
+		if st.client != nil {
+			if _, ok := st.client.State.Inflight.Get(id); ok {
+				continue
+			}
+		}
+		b.releaseLargePID(st, id)
+	}
 }
 
 func (b *Broker) releaseAllLarge(st *connState) {
 	st.mu.Lock()
-	n := st.largeHeld
-	st.largeHeld = 0
+	n := len(st.largePIDs) + st.largePending
+	st.largePIDs = nil
+	st.largePending = 0
 	st.mu.Unlock()
 	for i := 0; i < n; i++ {
-		select {
-		case <-b.largeSem:
-		default:
-		}
+		b.releaseLargeSlot()
 	}
 }