fix: 积压时等本帧写出再断开并让 Shutdown 等待 0x8B

This commit is contained in:
Nixevol
2026-09-30 19:21:51 +08:00
parent 2d1dfd7cd8
commit ef46563f70
6 changed files with 375 additions and 22 deletions
+87 -5
View File
@@ -141,9 +141,14 @@ type connState struct {
downStop chan struct{}
downDone chan struct{}
downBytes atomic.Int64
sentPub atomic.Int64
wirePending atomic.Int64 // Publish 入 mochi outbound 后、OnPacketSent 前
mu sync.Mutex
// 带断开的下行帧:只等本帧 OnPacketSent,不用连接级计数。
writeWaitCh chan struct{}
writeWaitPayload []byte
writeWaitPID uint16 // 非 0 时优先按 packet id 匹配
handshakeTimer *time.Timer
}
@@ -232,28 +237,105 @@ func (b *Broker) Close() error {
}
// Shutdown 向所有连接发 MQTT 5 0x8B 后关闭。完整 HTTP 停机顺序见 L-03。
// ctx 未取消时先等下行队列与 wirePending 排空,再 DisconnectClient(此时 outbound 空,
// 0x8B 直写套接字),并在截止前等连接拆掉;ctx 已取消则发完即 Close,不等待。
func (b *Broker) Shutdown(ctx context.Context) error {
if b.closed.Load() {
return nil
}
if ctx == nil {
ctx = context.Background()
}
b.connsMu.RLock()
states := make([]*connState, 0, len(b.byClient))
clients := make([]*mqtt.Client, 0, len(b.byClient))
for cl := range b.byClient {
for cl, st := range b.byClient {
if cl != nil {
clients = append(clients, cl)
}
if st != nil {
states = append(states, st)
}
}
b.connsMu.RUnlock()
for _, st := range states {
st.mu.Lock()
st.closing = true
st.mu.Unlock()
}
alreadyCancelled := false
select {
case <-ctx.Done():
alreadyCancelled = true
default:
}
var waitErr error
if !alreadyCancelled {
if !b.waitConnsQuiet(ctx, states) {
waitErr = ctx.Err()
}
}
for _, cl := range clients {
_ = b.server.DisconnectClient(cl, packets.ErrServerShuttingDown)
}
if ctx != nil {
if !alreadyCancelled && waitErr == nil {
waitErr = b.waitConnsGone(ctx)
}
closeErr := b.Close()
if waitErr != nil {
return waitErr
}
return closeErr
}
func (b *Broker) waitConnsQuiet(ctx context.Context, states []*connState) bool {
for {
quiet := true
for _, st := range states {
if st.wirePending.Load() > 0 {
quiet = false
break
}
st.mu.Lock()
ch := st.downCh
st.mu.Unlock()
if ch != nil && len(ch) > 0 {
quiet = false
break
}
}
if quiet {
return true
}
select {
case <-ctx.Done():
default:
return false
case <-time.After(2 * time.Millisecond):
}
}
}
func (b *Broker) waitConnsGone(ctx context.Context) error {
for {
b.connsMu.RLock()
n := len(b.byClient)
b.connsMu.RUnlock()
if n == 0 {
return nil
}
select {
case <-ctx.Done():
return ctx.Err()
case <-time.After(2 * time.Millisecond):
}
}
return b.Close()
}
// AttachTCP 把裸 TCP/TLS 连接交给 mochi;阻塞到连接结束。