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

2026-09-30 19:21:51 +08:00
parent 2d1dfd7cd8
commit ef46563f70
6 changed files with 375 additions and 22 deletions
@@ -87,6 +87,50 @@ func TestPublishThenDisconnectWritesThenCloses(t *testing.T) {
 	}
 }

+func TestPublishThenDisconnectAfterQueuedFrame(t *testing.T) {
+	b, w, done := startTCPClient(t, "ep-ptd-q")
+	defer func() { _ = b.Close() }()
+	defer func() {
+		_ = w.Close()
+		select {
+		case <-done:
+		case <-time.After(3 * time.Second):
+		}
+	}()
+	writeConnect(t, w, "ep-ptd-q", 30, 0)
+	readExactPacket(t, w, packets.Connack, 3*time.Second)
+	writeSubscribe(t, w, downTopic("ep-ptd-q"))
+	readExactPacket(t, w, packets.Suback, 3*time.Second)
+	waitSession(t, b, "ep-ptd-q")
+
+	first := []byte(`{"v":1,"type":"msg","id":"queued-ahead"}`)
+	fatal := []byte(`{"v":1,"type":"fatal","reason":"disabled"}`)
+	if err := b.PublishDown(context.Background(), "ep-ptd-q", "", first, port.PublishOpts{QoS: 1}); err != nil {
+		t.Fatal(err)
+	}
+	if err := b.PublishThenDisconnect(context.Background(), "ep-ptd-q", "", fatal, 1, port.DisconnectFatal); err != nil {
+		t.Fatal(err)
+	}
+
+	gotFirst := readDownPayload(t, w, 3*time.Second)
+	if !bytes.Equal(gotFirst, first) {
+		t.Fatalf("first got %s", gotFirst)
+	}
+	gotFatal := readDownPayload(t, w, 3*time.Second)
+	if !bytes.Equal(gotFatal, fatal) {
+		t.Fatalf("fatal got %s want %s (disconnected before fatal frame)", gotFatal, fatal)
+	}
+	_ = w.SetReadDeadline(time.Now().Add(3 * time.Second))
+	buf := make([]byte, 64)
+	n, err := io.ReadAtLeast(w, buf, 2)
+	if err != nil && n == 0 {
+		return
+	}
+	if n > 0 && buf[0]>>4 == packets.Disconnect {
+		return
+	}
+}
+
 func TestShutdownUsesServerShuttingDown(t *testing.T) {
 	b, w, done := startTCPClient(t, "ep-shut")
 	defer func() {
@@ -116,6 +160,120 @@ func TestShutdownUsesServerShuttingDown(t *testing.T) {
 	}
 }

+func TestShutdownWithBacklogDeliversServerShuttingDown(t *testing.T) {
+	b, w, done := startTCPClient(t, "ep-shut-bl")
+	defer func() {
+		_ = w.Close()
+		select {
+		case <-done:
+		case <-time.After(3 * time.Second):
+		}
+	}()
+	writeConnect(t, w, "ep-shut-bl", 30, 0)
+	readExactPacket(t, w, packets.Connack, 3*time.Second)
+	writeSubscribe(t, w, downTopic("ep-shut-bl"))
+	readExactPacket(t, w, packets.Suback, 3*time.Second)
+	waitSession(t, b, "ep-shut-bl")
+
+	payload := bytes.Repeat([]byte("b"), 1024)
+	for i := 0; i < 8; i++ {
+		if err := b.PublishDown(context.Background(), "ep-shut-bl", "", payload, port.PublishOpts{QoS: 1}); err != nil {
+			t.Fatalf("publish %d: %v", i, err)
+		}
+	}
+
+	saw8B := make(chan bool, 1)
+	go func() {
+		deadline := time.Now().Add(3 * time.Second)
+		for time.Now().Before(deadline) {
+			_ = w.SetReadDeadline(time.Now().Add(200 * time.Millisecond))
+			hdr := make([]byte, 1)
+			if _, err := io.ReadFull(w, hdr); err != nil {
+				continue
+			}
+			rem, err := readRemainingLengthConn(w)
+			if err != nil {
+				continue
+			}
+			body := make([]byte, rem)
+			if _, err := io.ReadFull(w, body); err != nil {
+				continue
+			}
+			switch hdr[0] >> 4 {
+			case packets.Publish:
+				qos := (hdr[0] >> 1) & 0x3
+				if qos > 0 {
+					pk := new(packets.Packet)
+					pk.ProtocolVersion = 5
+					pk.FixedHeader = packets.FixedHeader{Type: packets.Publish, Remaining: rem, Qos: qos}
+					if decErr := pk.PublishDecode(body); decErr == nil {
+						ack := packets.Packet{
+							FixedHeader:     packets.FixedHeader{Type: packets.Puback},
+							ProtocolVersion: 5,
+							PacketID:        pk.PacketID,
+						}
+						var ab bytes.Buffer
+						_ = ack.PubackEncode(&ab)
+						_, _ = w.Write(ab.Bytes())
+					}
+				}
+			case packets.Disconnect:
+				if rem >= 1 && body[0] == packets.ErrServerShuttingDown.Code {
+					saw8B <- true
+					return
+				}
+			}
+		}
+		saw8B <- false
+	}()
+	time.Sleep(20 * time.Millisecond)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+	err := b.Shutdown(ctx)
+
+	got := false
+	select {
+	case got = <-saw8B:
+	case <-time.After(4 * time.Second):
+		t.Fatal("reader hung")
+	}
+	if got {
+		if err != nil && !errors.Is(err, context.DeadlineExceeded) {
+			t.Fatalf("shutdown after 0x8B: %v", err)
+		}
+		return
+	}
+	if err == nil {
+		t.Fatal("expected 0x8B or shutdown deadline error, got neither")
+	}
+	if !errors.Is(err, context.DeadlineExceeded) && !errors.Is(err, context.Canceled) {
+		t.Fatalf("shutdown err=%v want deadline/cancel when 0x8B not seen", err)
+	}
+}
+
+func TestShutdownCancelledContextReturnsQuickly(t *testing.T) {
+	b, w, done := startTCPClient(t, "ep-shut-cancel")
+	defer func() {
+		_ = w.Close()
+		select {
+		case <-done:
+		case <-time.After(3 * time.Second):
+		}
+	}()
+	writeConnect(t, w, "ep-shut-cancel", 30, 0)
+	readExactPacket(t, w, packets.Connack, 3*time.Second)
+	waitSession(t, b, "ep-shut-cancel")
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	start := time.Now()
+	_ = b.Shutdown(ctx)
+	if time.Since(start) > 500*time.Millisecond {
+		t.Fatalf("cancelled shutdown took %s", time.Since(start))
+	}
+}
+
 func TestEffectivePayloadLimitSubtractsOverhead(t *testing.T) {
 	got := EffectivePayloadLimit(200, 0)
 	if got != 200-packetOverheadBudget {