NixMsg/internal/broker/downlink.go

package broker

import (
	"context"
	"time"

	"git.asio.asia/nixevol/NixMsg/internal/app/port"
)

type downItem struct {
	payload    []byte
	qos        byte
	disconnect port.DisconnectReason // 非空表示该帧写出后断开（B-04）
	sent       chan struct{}
}

func (st *connState) startDownLoop(b *Broker) {
	st.mu.Lock()
	if st.downCh != nil {
		st.mu.Unlock()
		return
	}
	st.downCh = make(chan downItem, downQueueMax)
	st.downStop = make(chan struct{})
	st.downDone = make(chan struct{})
	st.mu.Unlock()
	go st.downLoop(b)
}

func (st *connState) stopDownLoop() {
	st.mu.Lock()
	stop := st.downStop
	done := st.downDone
	ch := st.downCh
	st.mu.Unlock()
	if stop == nil {
		return
	}
	select {
	case <-stop:
	default:
		close(stop)
	}
	if done != nil {
		select {
		case <-done:
		case <-time.After(2 * time.Second):
		}
	}
	if ch != nil {
		for {
			select {
			case item := <-ch:
				st.downBytes.Add(-int64(len(item.payload)))
			default:
				return
			}
		}
	}
}

func (st *connState) enqueueDown(item downItem) error {
	st.mu.Lock()
	ch := st.downCh
	stop := st.downStop
	closing := st.closing
	st.mu.Unlock()
	if ch == nil || stop == nil {
		return ErrNoConnection
	}
	select {
	case <-stop:
		return ErrNoConnection
	default:
	}
	if closing && item.disconnect == "" {
		return ErrNoConnection
	}
	n := int64(len(item.payload))
	for {
		cur := st.downBytes.Load()
		if cur+n > downQueueBytes {
			return ErrBackpressure
		}
		if st.downBytes.CompareAndSwap(cur, cur+n) {
			break
		}
	}
	select {
	case ch <- item:
		return nil
	default:
		st.downBytes.Add(-n)
		return ErrBackpressure
	}
}

func (st *connState) downLoop(b *Broker) {
	defer close(st.downDone)
	for {
		select {
		case <-st.downStop:
			return
		case item, ok := <-st.downCh:
			if !ok {
				return
			}
			st.downBytes.Add(-int64(len(item.payload)))
			st.sendOne(b, item)
		}
	}
}

func (st *connState) sendOne(b *Broker, item downItem) {
	if b.closed.Load() {
		st.signalSent(item)
		return
	}
	large := len(item.payload) > largeFrameBytes
	if large {
		if err := b.acquireLarge(context.Background()); err != nil {
			if b.onDrop != nil {
				b.onDrop(context.Background(), st.endpointID, st.connID, item.payload)
			}
			st.signalSent(item)
			return
		}
		st.mu.Lock()
		st.largePending++
		st.mu.Unlock()
	}
	topic := downTopic(st.endpointID)
	before := st.sentPub.Load()
	var err error
	for {
		if b.closed.Load() {
			break
		}
		select {
		case <-st.downStop:
			err = ErrNoConnection
		default:
			err = b.server.Publish(topic, item.payload, false, item.qos)
			if err == nil {
				break
			}
			select {
			case <-st.downStop:
				err = ErrNoConnection
			case <-time.After(2 * time.Millisecond):
				continue
			}
		}
		break
	}
	if large {
		b.finishLargePublish(st)
	}
	if err != nil && b.onDrop != nil {
		b.onDrop(context.Background(), st.endpointID, st.connID, item.payload)
	}
	st.signalSent(item)
	if err == nil && item.disconnect != "" {
		st.waitPacketWritten(before)
		_ = b.Disconnect(context.Background(), st.endpointID, st.connID, item.disconnect)
	}
}

func (st *connState) waitPacketWritten(before int64) {
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		if st.sentPub.Load() > before {
			return
		}
		select {
		case <-st.downStop:
			return
		case <-time.After(2 * time.Millisecond):
		}
	}
}

func (st *connState) signalSent(item downItem) {
	if item.sent == nil {
		return
	}
	select {
	case <-item.sent:
	default:
		close(item.sent)
	}
}