xiaozhi-esp32-server/main/xiaozhi-server/core/servers/mqtt_server.py

import asyncio
import ipaddress
import json
import secrets
import socket
import time
import weakref
from typing import Dict, Any, Set
from config.logger import setup_logging
from core.protocols.mqtt_connection import MQTTConnection
from core.transport.mqtt_transport import MQTTTransport, UDPAudioHandler
from core.services.connection_service import ConnectionService
from core.services.native_mqtt_connection_registry import (
    NativeMqttConnectionRegistry,
)
from core.services.native_mqtt_call_manager import NativeMqttCallManager
from core.utils.mqtt_auth import normalize_signature_key, parse_mqtt_endpoint
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend

logger = setup_logging()


class MQTTServer:
    """
    原生MQTT服务器：直接处理MQTT协议连接
    集成到xiaozhi-server架构中
    """

    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.logger = setup_logging()

        # 服务器配置
        server_config = config.get('mqtt_server', {})
        self.mqtt_port = server_config.get('port', 1883)
        self.udp_port = server_config.get('udp_port', self.mqtt_port)
        self.host = server_config.get('host', '0.0.0.0')
        self.public_endpoint = self._resolve_public_host(server_config)
        self.udp_bind_host = self.host
        self._udp_bind_host_config = server_config.get('udp_bind_host')
        self.signature_key = self._resolve_signature_key(config)
        if not self.signature_key:
            raise ValueError(
                "启用原生MQTT必须配置mqtt_server.signature_key或"
                "server.mqtt_signature_key"
            )
        self.message_queue_size = int(server_config.get('message_queue_size', 128))
        self.business_ready_timeout = float(
            server_config.get('business_ready_timeout', 30) or 30
        )
        self.close_timeout = max(
            0.1, float(server_config.get('close_timeout', 2) or 2)
        )
        self.shutdown_timeout = max(
            self.close_timeout,
            float(server_config.get('shutdown_timeout', 10) or 10),
        )
        self.goodbye_timeout = max(
            0.1, float(server_config.get('goodbye_timeout', 1) or 1)
        )
        self.max_connections = int(server_config.get('max_connections', 1000))
        self.max_pending_connections = int(
            server_config.get('max_pending_connections', 128)
        )
        self.max_payload_size = int(server_config.get('max_payload_size', 8192))

        # 连接管理
        self.connections: Dict[int, MQTTConnection] = {}
        self.client_id_map: Dict[str, int] = {}
        self.udp_handlers: Dict[int, UDPAudioHandler] = {}
        self.connection_id_counter = 0
        self._client_id_lock = asyncio.Lock()
        self._client_locks = weakref.WeakValueDictionary()
        self._client_reservations: Dict[str, int] = {}
        self._pending_connection_ids: Set[int] = set()
        self._cleanup_tasks: Dict[int, asyncio.Task] = {}
        self._draining_tasks: Set[asyncio.Task] = set()
        self._connection_handler_tasks: Set[asyncio.Task] = set()

        # 服务器实例
        self.mqtt_server = None
        self.udp_server = None

        # 连接服务
        self.connection_service = ConnectionService(config)
        self.connection_service.server = self

        # 活跃连接管理
        self.active_transports: Set[MQTTTransport] = set()
        self.connection_registry = NativeMqttConnectionRegistry()
        self.call_manager = NativeMqttCallManager(
            self.connection_registry,
            timeout_seconds=server_config.get("call_timeout", 60),
            silence_frame=self._create_silence_frame(),
        )

        # 心跳检查
        self.heartbeat_task = None
        self.heartbeat_interval = int(server_config.get('heartbeat_interval', 30))
        self._stop_event = asyncio.Event()
        self._started_event = asyncio.Event()
        self._is_running = False
        self._stopping = False
        self._shutdown_lock = asyncio.Lock()
        self._shutdown_task = None

    async def start(self):
        """启动MQTT服务器"""
        if self._is_running:
            logger.warning("MQTT服务器已经在运行中")
            return

        self._stop_event.clear()
        self._started_event.clear()
        self._stopping = False
        try:
            # 启动MQTT TCP服务器
            await self._start_mqtt_server()

            # 启动UDP服务器
            await self._start_udp_server()

            # 启动心跳检查
            self.heartbeat_task = asyncio.create_task(self._heartbeat_check())
            self._is_running = True
            self._started_event.set()

            logger.info(f"MQTT服务器启动成功: {self.host}:{self.mqtt_port}")
            logger.info(f"UDP服务器启动成功: {self.udp_bind_host}:{self.udp_port}")

            # 与 WebSocket 服务器保持一致，由上层持有该长期任务。
            await self._stop_event.wait()

        except Exception as e:
            logger.error(f"启动MQTT服务器失败: {e}")
            await self._await_shutdown_owner()
            raise
        finally:
            self._is_running = False
            self._started_event.clear()

    async def _start_mqtt_server(self):
        """启动MQTT TCP服务器"""
        self.mqtt_server = await asyncio.start_server(
            self._accept_mqtt_connection,
            self.host,
            self.mqtt_port
        )
        if self.mqtt_port == 0 and self.mqtt_server.sockets:
            self.mqtt_port = self.mqtt_server.sockets[0].getsockname()[1]

    async def _start_udp_server(self):
        """启动UDP服务器"""
        loop = asyncio.get_event_loop()

        sock = None
        bind_candidates = self._udp_bind_candidates()
        for index, bind_host in enumerate(bind_candidates):
            candidate = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            try:
                candidate.bind((bind_host, self.udp_port))
            except OSError as error:
                candidate.close()
                if index + 1 >= len(bind_candidates):
                    raise
                logger.warning(
                    "UDP无法绑定到对外地址 {}，回退监听 {}: {}",
                    bind_host,
                    bind_candidates[index + 1],
                    error,
                )
                continue
            sock = candidate
            self.udp_bind_host = bind_host
            break

        if sock is None:
            raise RuntimeError("UDP socket bind failed")
        sock.setblocking(False)

        # 创建UDP协议处理器
        transport, protocol = await loop.create_datagram_endpoint(
            lambda: UDPProtocol(self),
            sock=sock
        )

        self.udp_server = (transport, protocol)
        if self.udp_port == 0:
            self.udp_port = transport.get_extra_info("sockname")[1]

    def _udp_bind_candidates(self):
        """Prefer a local advertised IPv4 address so UDP replies keep that source IP."""
        explicit_bind_host = str(self._udp_bind_host_config or "").strip()
        if explicit_bind_host:
            return [explicit_bind_host]

        wildcard_hosts = {"", "0.0.0.0"}
        if self.host not in wildcard_hosts:
            return [self.host]

        try:
            public_ip = ipaddress.ip_address(self.public_endpoint)
        except ValueError:
            return [self.host]

        if public_ip.version == 4 and not public_ip.is_unspecified:
            return [self.public_endpoint, self.host]
        return [self.host]

    def _accept_mqtt_connection(self, reader, writer) -> None:
        """Track every accepted callback before its coroutine can yield."""
        task = asyncio.create_task(
            self._handle_mqtt_connection(reader, writer),
            name="mqtt-connection-handler",
        )
        self._connection_handler_tasks.add(task)

        def consume(completed_task):
            self._connection_handler_tasks.discard(completed_task)
            MQTTConnection._consume_background_task(
                completed_task, "MQTT连接处理任务"
            )

        task.add_done_callback(consume)

    async def _handle_mqtt_connection(self, reader, writer):
        """处理新的MQTT连接"""
        if self._stopping:
            await self._close_stream_writer(writer, "MQTT停服期间连接")
            return

        if (
            self.max_pending_connections > 0
            and len(self._pending_connection_ids) >= self.max_pending_connections
        ):
            logger.warning(
                "MQTT待认证连接数达到上限: {}", self.max_pending_connections
            )
            await self._close_stream_writer(writer, "MQTT待认证连接")
            return

        connection_id = self._generate_connection_id()

        try:
            # 获取客户端地址
            client_addr = writer.get_extra_info('peername')
            logger.info(f"新MQTT连接: {client_addr}, connection_id: {connection_id}")

            # 创建UDP音频处理器
            udp_handler = UDPAudioHandler(
                connection_id,
                self,
                {},  # 加密配置将在hello消息中设置
                allowed_remote_ip=client_addr[0] if client_addr else None,
            )

            # 创建MQTT连接处理器
            mqtt_connection = MQTTConnection(
                writer.get_extra_info('socket'),
                connection_id,
                self,
                udp_handler,
                reader=reader,
                writer=writer
            )
            mqtt_connection.business_task = asyncio.current_task()
            udp_handler.set_audio_interceptor(
                lambda payload, timestamp: self.call_manager.route_audio(
                    mqtt_connection.device_id, payload, timestamp
                )
            )

            # 创建MQTT传输层
            transport = MQTTTransport(mqtt_connection, udp_handler)

            # 注册连接
            self.connections[connection_id] = mqtt_connection
            self.active_transports.add(transport)
            self._pending_connection_ids.add(connection_id)

            try:
                await asyncio.wait_for(
                    mqtt_connection.connect_processed_event.wait(), timeout=10
                )
            except asyncio.TimeoutError:
                logger.warning(f"MQTT CONNECT超时: connection_id={connection_id}")
                return
            finally:
                self._pending_connection_ids.discard(connection_id)

            if not mqtt_connection.connect_accepted:
                return

            # 提取连接头信息
            headers = {
                'x-real-ip': client_addr[0] if client_addr else 'unknown',
                'connection-type': 'mqtt'
            }

            try:
                # 使用ConnectionService处理连接
                await self.connection_service.handle_connection(transport, headers)

            except Exception as e:
                logger.error(f"ConnectionService处理MQTT连接失败: {e}")

        except Exception as e:
            logger.error(f"处理MQTT连接失败: {e}")
        finally:
            # 清理连接
            await self._cleanup_connection(connection_id)

    async def _cleanup_connection(self, connection_id: int):
        """Await the single cleanup owner without coupling it to caller cancellation."""
        if (
            connection_id not in self._cleanup_tasks
            and not self._connection_has_resources(connection_id)
        ):
            return
        cleanup_task = self._ensure_cleanup_task(connection_id)

        if cleanup_task is asyncio.current_task():
            return
        await asyncio.shield(cleanup_task)

    def _ensure_cleanup_task(self, connection_id: int) -> asyncio.Task:
        cleanup_task = self._cleanup_tasks.get(connection_id)
        if cleanup_task is None:
            cleanup_task = asyncio.create_task(
                self._cleanup_connection_impl(connection_id),
                name=f"mqtt-cleanup-{connection_id}",
            )
            self._cleanup_tasks[connection_id] = cleanup_task
        return cleanup_task

    def _connection_has_resources(self, connection_id: int) -> bool:
        if connection_id in self._pending_connection_ids:
            return True
        if connection_id in self.connections:
            return True
        if connection_id in self.client_id_map.values():
            return True
        if connection_id in self._client_reservations.values():
            return True
        return any(
            getattr(
                getattr(transport, "_mqtt_connection", None),
                "connection_id",
                None,
            )
            == connection_id
            for transport in self.active_transports
        )

    async def _cleanup_connection_impl(self, connection_id: int):
        """Physically release one connection before removing its registry entries."""
        connection = self.connections.get(connection_id)
        transports_to_remove = [
            transport
            for transport in list(self.active_transports)
            if (
                hasattr(transport, '_mqtt_connection')
                and transport._mqtt_connection.connection_id == connection_id
            )
        ]
        try:
            self._pending_connection_ids.discard(connection_id)
            if connection is not None:
                if (
                    connection.udp_config
                    and connection.is_connected()
                    and connection.reply_topic
                ):
                    try:
                        await asyncio.wait_for(
                            connection.notify_device_idle(),
                            timeout=self.goodbye_timeout,
                        )
                    except asyncio.TimeoutError:
                        logger.warning("MQTT连接关闭前发送goodbye超时")
                    except Exception as e:
                        logger.warning("MQTT连接关闭前发送goodbye失败: {}", e)
                try:
                    await asyncio.wait_for(
                        connection.close(), timeout=self.close_timeout
                    )
                except asyncio.TimeoutError:
                    logger.warning(
                        "MQTT连接关闭超时，强制中止socket: {}", connection_id
                    )
                    abort = getattr(connection.protocol, "abort", None)
                    if callable(abort):
                        abort()
                    close_task = getattr(connection, "_close_task", None)
                    if close_task is not None and not close_task.done():
                        self.track_draining_task(
                            close_task,
                            f"MQTT连接关闭任务({connection_id})",
                        )

            # UDP 使用每个 Hello 独立的随机 cookie，清理该物理连接的全部别名。
            udp_handler = connection.udp_handler if connection is not None else None
            if udp_handler is not None:
                for session_id, registered in list(self.udp_handlers.items()):
                    if registered is udp_handler:
                        self.udp_handlers.pop(session_id, None)
                try:
                    await asyncio.wait_for(
                        udp_handler.close(), timeout=self.close_timeout
                    )
                except asyncio.TimeoutError:
                    logger.warning(
                        "MQTT UDP处理器关闭超时: {}", connection_id
                    )

            for transport in transports_to_remove:
                try:
                    await asyncio.wait_for(
                        transport.close(), timeout=self.close_timeout
                    )
                except asyncio.TimeoutError:
                    logger.warning("MQTT transport关闭超时: {}", connection_id)
                    abort = getattr(
                        getattr(transport, "_mqtt_connection", None),
                        "protocol",
                        None,
                    )
                    abort = getattr(abort, "abort", None)
                    if callable(abort):
                        abort()
        except asyncio.CancelledError:
            raise
        except Exception as e:
            logger.error(f"清理MQTT连接失败: {e}")
        finally:
            # A cleanup owner may be cancelled during process shutdown. Always
            # force the physical socket closed before publishing registry removal.
            if connection is not None:
                abort = getattr(connection.protocol, "abort", None)
                if callable(abort):
                    abort()
                close_task = getattr(connection, "_close_task", None)
                if close_task is not None and not close_task.done():
                    self.track_draining_task(
                        close_task,
                        f"MQTT连接关闭任务({connection_id})",
                    )

            for transport in transports_to_remove:
                self.active_transports.discard(transport)

            async with self._client_id_lock:
                for client_id, mapped_id in list(self.client_id_map.items()):
                    if mapped_id == connection_id:
                        self.client_id_map.pop(client_id, None)
                for client_id, reserved_id in list(
                    self._client_reservations.items()
                ):
                    if reserved_id == connection_id:
                        self._client_reservations.pop(client_id, None)
                self.connections.pop(connection_id, None)

            current_task = asyncio.current_task()
            if self._cleanup_tasks.get(connection_id) is current_task:
                self._cleanup_tasks.pop(connection_id, None)
            logger.info(f"MQTT连接清理完成: {connection_id}")

    async def _close_stream_writer(self, writer, label: str) -> None:
        """Bound close for sockets rejected before MQTTConnection exists."""
        writer.close()
        try:
            await asyncio.wait_for(
                writer.wait_closed(), timeout=self.close_timeout
            )
        except asyncio.TimeoutError:
            logger.warning("{}关闭超时，强制中止transport", label)
            transport = getattr(writer, "transport", None)
            if transport is not None:
                transport.abort()
        except Exception:
            pass

    def track_draining_task(self, task: asyncio.Task, label: str) -> None:
        """Keep cancellation-resistant work visible until it actually exits."""
        if task is None or task.done():
            if task is not None:
                MQTTConnection._consume_background_task(task, label)
            return
        if task in self._draining_tasks:
            return
        self._draining_tasks.add(task)

        def consume(completed_task):
            self._draining_tasks.discard(completed_task)
            MQTTConnection._consume_background_task(completed_task, label)

        task.add_done_callback(consume)

    def _generate_connection_id(self) -> int:
        """生成连接ID"""
        self.connection_id_counter += 1
        return self.connection_id_counter

    async def _get_client_lock(self, client_id: str) -> asyncio.Lock:
        async with self._client_id_lock:
            lock = self._client_locks.get(client_id)
            if lock is None:
                lock = asyncio.Lock()
                self._client_locks[client_id] = lock
            return lock

    async def on_client_connected(
        self, mqtt_connection: MQTTConnection, accept_callback
    ) -> bool:
        """Atomically replace one clientId before acknowledging the new owner."""
        client_id = mqtt_connection.client_id
        if not client_id or self._stopping:
            return False

        client_lock = await self._get_client_lock(client_id)
        async with client_lock:
            if mqtt_connection._closed or self._stopping:
                return False

            reservation_active = False
            async with self._client_id_lock:
                old_id = self.client_id_map.get(client_id)
                identity_count = len(
                    set(self.client_id_map) | set(self._client_reservations)
                )
                if (
                    old_id is None
                    and client_id not in self._client_reservations
                    and self.max_connections > 0
                    and identity_count >= self.max_connections
                ):
                    logger.warning(
                        "MQTT活跃客户端数达到上限，拒绝clientId: {}", client_id
                    )
                    return False
                self._client_reservations[client_id] = (
                    mqtt_connection.connection_id
                )
                reservation_active = True
            try:
                if old_id and old_id != mqtt_connection.connection_id:
                    logger.info(f"检测到重复clientId，关闭旧连接: {client_id}")
                    cleanup_task = self._ensure_cleanup_task(old_id)
                    try:
                        await asyncio.shield(cleanup_task)
                    except asyncio.CancelledError as cancelled:
                        # Keep the per-client lock and capacity reservation until
                        # the predecessor's physical close barrier has completed.
                        try:
                            await asyncio.shield(cleanup_task)
                        except Exception as cleanup_error:
                            logger.error(
                                "取消接管时等待旧连接清理失败: {}", cleanup_error
                            )
                        raise cancelled

                if mqtt_connection._closed or self._stopping:
                    return False

                async with self._client_id_lock:
                    if (
                        self._client_reservations.get(client_id)
                        != mqtt_connection.connection_id
                    ):
                        return False
                    self.client_id_map[client_id] = mqtt_connection.connection_id
                    self._client_reservations.pop(client_id, None)
                    reservation_active = False

                await accept_callback()
            except BaseException:
                async with self._client_id_lock:
                    if (
                        self.client_id_map.get(client_id)
                        == mqtt_connection.connection_id
                    ):
                        self.client_id_map.pop(client_id, None)
                raise
            finally:
                if reservation_active:
                    async with self._client_id_lock:
                        if (
                            self._client_reservations.get(client_id)
                            == mqtt_connection.connection_id
                        ):
                            self._client_reservations.pop(client_id, None)

            logger.info(f"MQTT客户端已连接: {client_id}")
            return True

    async def on_client_disconnected(self, mqtt_connection: MQTTConnection):
        """客户端断开连接回调"""
        logger.info(f"MQTT客户端已断开: {mqtt_connection.client_id}")
        # Registry removal belongs to the physical cleanup owner. Scheduling
        # rather than awaiting avoids a close -> disconnect -> cleanup cycle.
        connection_id = mqtt_connection.connection_id
        if (
            connection_id in self.connections
            and connection_id not in self._cleanup_tasks
        ):
            self._ensure_cleanup_task(connection_id)

    async def register_connection_context(self, context, transport) -> bool:
        existing = self.connection_registry.resolve_device_now(
            getattr(context, "device_id", None)
        )
        registered = await self.connection_registry.register(
            context, transport
        )
        if (
            registered
            and existing is not None
            and existing.transport is not transport
        ):
            await self.call_manager.end_call(
                context.device_id,
                "设备连接被接管",
                notify_device=False,
                notify_peer=True,
                expected_session_id=getattr(
                    existing.transport, "session_id", None
                ),
            )
            try:
                await existing.transport.close()
            except Exception as error:
                logger.warning(
                    "关闭被接管的MQTT连接失败: device_id={}, error={}",
                    context.device_id,
                    error,
                )
        return registered

    async def unregister_connection_context(self, context, transport) -> bool:
        removed = await self.connection_registry.unregister(context, transport)
        if removed:
            await self.call_manager.end_call(
                context.device_id,
                "设备连接已断开",
                notify_device=False,
                notify_peer=True,
                expected_session_id=getattr(transport, "session_id", None),
            )
        return removed

    async def resolve_connection_context(self, client_id: str):
        return await self.connection_registry.resolve(client_id)

    async def get_connection_status(self, client_ids):
        return await self.connection_registry.status(client_ids)

    async def request_device_call(
        self, caller_mac: str, target_mac: str, caller_nickname: str = ""
    ):
        return await self.call_manager.request_call(
            caller_mac, target_mac, caller_nickname
        )

    async def accept_device_call(self, device_id: str):
        return await self.call_manager.accept_call(device_id)

    async def end_native_mqtt_call(
        self,
        device_id: str,
        reason: str = "",
        notify_device: bool = False,
        expected_session_id: str = None,
        expected_generation: int = None,
    ) -> bool:
        return await self.call_manager.end_call(
            device_id,
            reason,
            notify_device=notify_device,
            notify_peer=True,
            expected_session_id=expected_session_id,
            expected_generation=expected_generation,
        )

    async def handle_logical_hello(self, mqtt_connection) -> None:
        await self.call_manager.handle_logical_hello(
            mqtt_connection.device_id,
            mqtt_connection.session_id,
        )

    def bind_udp_session(
        self, mqtt_connection: MQTTConnection, udp_handler: UDPAudioHandler
    ) -> int:
        """Bind one unpredictable UDP cookie to the current Hello session."""
        for session_id, registered in list(self.udp_handlers.items()):
            if registered is udp_handler:
                self.udp_handlers.pop(session_id, None)

        while True:
            session_id = secrets.randbits(32)
            if session_id and session_id not in self.udp_handlers:
                break
        udp_handler.connection_id = session_id
        self.udp_handlers[session_id] = udp_handler
        return session_id

    async def send_udp_message(self, data: bytes, remote_addr: tuple):
        """发送UDP消息"""
        if not self.udp_server:
            raise RuntimeError("UDP server is not running")
        transport, protocol = self.udp_server
        transport.sendto(data, remote_addr)

    async def send_encrypted_audio(self, connection_id: int, audio_data: bytes,
                                   timestamp: int, sequence: int, remote_addr: tuple,
                                   encryption_config: Dict[str, Any]):
        """发送加密音频数据"""
        self.send_encrypted_audio_nowait(
            connection_id,
            audio_data,
            timestamp,
            sequence,
            remote_addr,
            encryption_config,
        )

    def send_encrypted_audio_nowait(
        self,
        connection_id: int,
        audio_data: bytes,
        timestamp: int,
        sequence: int,
        remote_addr: tuple,
        encryption_config: Dict[str, Any],
    ):
        header = self._generate_udp_header(connection_id, len(audio_data), timestamp, sequence)
        key = encryption_config.get('key') if encryption_config else None
        if key:
            cipher = Cipher(algorithms.AES(key), modes.CTR(header), backend=default_backend())
            encryptor = cipher.encryptor()
            encrypted = encryptor.update(audio_data) + encryptor.finalize()
            message = header + encrypted
        else:
            message = header + audio_data

        if not self.udp_server:
            raise RuntimeError("UDP server is not running")
        transport, _ = self.udp_server
        transport.sendto(message, remote_addr)

    def _generate_udp_header(self, connection_id: int, length: int, timestamp: int, sequence: int) -> bytes:
        """生成UDP消息头"""
        header = bytearray(16)
        header[0] = 1  # type
        header[2:4] = length.to_bytes(2, 'big')  # payload length
        header[4:8] = connection_id.to_bytes(4, 'big')  # connection id
        header[8:12] = timestamp.to_bytes(4, 'big')  # timestamp
        header[12:16] = sequence.to_bytes(4, 'big')  # sequence
        return bytes(header)

    async def _heartbeat_check(self):
        """心跳检查任务"""
        try:
            while True:
                await asyncio.sleep(self.heartbeat_interval)

                # 检查所有连接的状态
                dead_connections = []
                for connection_id, connection in self.connections.items():
                    if not connection.is_connected():
                        dead_connections.append(connection_id)

                # 清理死连接
                for connection_id in dead_connections:
                    logger.info(f"清理死连接: {connection_id}")
                    await self._cleanup_connection(connection_id)

                await self.call_manager.cleanup_timeouts()

                # 记录活跃连接数
                active_count = len(self.connections)
                if active_count > 0:
                    logger.info(f"MQTT活跃连接数: {active_count}")

        except asyncio.CancelledError:
            pass
        except Exception as e:
            logger.error(f"心跳检查任务出错: {e}")

    async def stop(self):
        """停止MQTT服务器"""
        async with self._shutdown_lock:
            self._prune_finished_lifecycle_tasks()
            shutdown_in_progress = (
                self._shutdown_task is not None
                and not self._shutdown_task.done()
            )
            if (
                not self._is_running
                and self._shutdown_resources_released()
                and not shutdown_in_progress
            ):
                self._stop_event.set()
                return True

            logger.info("正在停止MQTT服务器...")
            shutdown_task = self._ensure_shutdown_owner()

        shutdown_result = await asyncio.shield(shutdown_task)
        # Existing lifecycle test doubles and older protocol servers return
        # None on success. Only an explicit False means cleanup is partial.
        complete = shutdown_result is not False
        if complete:
            logger.info("MQTT服务器已停止")
        else:
            status = self.get_server_status()
            logger.error(
                "MQTT监听器已停止，但仍有未退出任务: handlers={}, draining={}",
                status["connection_handler_tasks"],
                status["draining_tasks"],
            )
        return complete

    def _ensure_shutdown_owner(self) -> asyncio.Task:
        """Create or reuse the single physical shutdown owner."""
        if self._shutdown_task is None or self._shutdown_task.done():
            self._shutdown_task = asyncio.create_task(
                self._run_shutdown_owner()
            )
        return self._shutdown_task

    async def _await_shutdown_owner(self):
        """Wait for shutdown without letting caller cancellation kill it."""
        return await asyncio.shield(self._ensure_shutdown_owner())

    async def _run_shutdown_owner(self):
        try:
            return await self._shutdown_resources()
        finally:
            # A cancelled stop waiter must not leave start() blocked forever.
            self._stop_event.set()
            self._is_running = False
            self._started_event.clear()

    def _prune_finished_lifecycle_tasks(self) -> None:
        for connection_id, task in list(self._cleanup_tasks.items()):
            if task.done() and self._cleanup_tasks.get(connection_id) is task:
                self._cleanup_tasks.pop(connection_id, None)
        for task in list(self._connection_handler_tasks):
            if task.done():
                self._connection_handler_tasks.discard(task)
        for task in list(self._draining_tasks):
            if task.done():
                self._draining_tasks.discard(task)

    def _shutdown_resources_released(self) -> bool:
        """Return true only when every MQTT-owned resource is gone."""
        heartbeat_alive = (
            self.heartbeat_task is not None
            and not self.heartbeat_task.done()
        )
        return not (
            self.mqtt_server
            or self.udp_server
            or heartbeat_alive
            or self.connections
            or self.client_id_map
            or self.udp_handlers
            or self.active_transports
            or self._client_reservations
            or self._pending_connection_ids
            or self._cleanup_tasks
            or self.connection_registry.count
            or self.call_manager.count
            or self.call_manager.background_task_count
            or any(
                not task.done()
                for task in self._connection_handler_tasks
            )
            or any(not task.done() for task in self._draining_tasks)
        )

    async def _shutdown_resources(self):
        """幂等释放 MQTT/UDP 监听器、连接和后台任务。"""
        self._stopping = True
        loop = asyncio.get_running_loop()
        deadline = loop.time() + self.shutdown_timeout

        # Stop accepting before taking any connection/task snapshot.
        if self.mqtt_server:
            self.mqtt_server.close()
            await self.mqtt_server.wait_closed()
            self.mqtt_server = None
        # Let callbacks already queued by the event loop enter the synchronous
        # tracking wrapper before handler/connection snapshots are taken.
        await asyncio.sleep(0)

        if self.heartbeat_task and not self.heartbeat_task.done():
            self.heartbeat_task.cancel()
            await asyncio.gather(self.heartbeat_task, return_exceptions=True)
        self.heartbeat_task = None

        # Start every cleanup owner first, then wait against one global deadline.
        connection_ids = (
            set(self.connections)
            | set(self._cleanup_tasks)
            | set(self.client_id_map.values())
            | set(self._client_reservations.values())
            | set(self._pending_connection_ids)
        )
        connection_ids.update(
            connection_id
            for connection_id in (
                getattr(
                    getattr(transport, "_mqtt_connection", None),
                    "connection_id",
                    None,
                )
                for transport in self.active_transports
            )
            if connection_id is not None
        )
        cleanup_tasks = [
            self._ensure_cleanup_task(connection_id)
            for connection_id in connection_ids
            if self._connection_has_resources(connection_id)
            or connection_id in self._cleanup_tasks
        ]
        pending_cleanup = await self._wait_until_deadline(
            cleanup_tasks, deadline
        )
        if pending_cleanup:
            logger.warning(
                "MQTT停服清理超过全局时限，强制取消{}个清理任务",
                len(pending_cleanup),
            )
            for task in pending_cleanup:
                task.cancel()
            await asyncio.gather(*pending_cleanup, return_exceptions=True)

        # Accepted callbacks are tracked synchronously by _accept_mqtt_connection.
        handler_tasks = [
            task
            for task in self._connection_handler_tasks
            if not task.done()
        ]
        for task in handler_tasks:
            task.cancel()
        pending_handlers = await self._wait_until_deadline(
            handler_tasks, deadline
        )
        for task in pending_handlers:
            self.track_draining_task(task, "MQTT连接处理任务")

        draining_tasks = [
            task for task in self._draining_tasks if not task.done()
        ]
        for task in draining_tasks:
            task.cancel()
        pending_draining = await self._wait_until_deadline(
            draining_tasks, deadline
        )
        if pending_draining:
            logger.warning(
                "MQTT停服后仍有{}个取消不响应任务", len(pending_draining)
            )

        # Defensive closure for any transport whose connection registration was
        # corrupted by an earlier failure.
        transport_tasks = [
            asyncio.create_task(self._close_transport(transport))
            for transport in list(self.active_transports)
        ]
        pending_transports = await self._wait_until_deadline(
            transport_tasks, deadline
        )
        for task in pending_transports:
            task.cancel()
        if pending_transports:
            await asyncio.gather(
                *pending_transports, return_exceptions=True
            )

        # Close any orphaned UDP session left behind by a corrupted registry.
        for udp_handler in set(self.udp_handlers.values()):
            await udp_handler.close()
        self.udp_handlers.clear()

        # 关闭UDP服务器
        if self.udp_server:
            transport, protocol = self.udp_server
            transport.close()
            self.udp_server = None

        await self.call_manager.clear()
        await self.connection_registry.clear()
        self._prune_finished_lifecycle_tasks()
        return self._shutdown_resources_released()

    @staticmethod
    def _create_silence_frame():
        try:
            import opuslib_next

            encoder = opuslib_next.Encoder(
                16000, 1, opuslib_next.APPLICATION_AUDIO
            )
            return encoder.encode(bytes(960 * 2), 960)
        except Exception as error:
            logger.warning("Native MQTT通话静音帧初始化失败: {}", error)
            return None

    async def _wait_until_deadline(self, tasks, deadline):
        pending = {task for task in tasks if task is not None and not task.done()}
        if not pending:
            return set()
        remaining = max(0.0, deadline - asyncio.get_running_loop().time())
        _, pending = await asyncio.wait(pending, timeout=remaining)
        return pending

    async def _close_transport(self, transport) -> None:
        try:
            await asyncio.wait_for(
                transport.close(), timeout=self.close_timeout
            )
        except asyncio.TimeoutError:
            connection = getattr(transport, "_mqtt_connection", None)
            abort = getattr(
                getattr(connection, "protocol", None), "abort", None
            )
            if callable(abort):
                abort()
        finally:
            self.active_transports.discard(transport)

    async def update_config(self, new_config: Dict[str, Any]) -> bool:
        """更新非监听配置；监听地址变化由 MultiProtocolServer 重建实例。"""
        try:
            self.config = new_config
            server_config = new_config.get('mqtt_server', {})
            self.public_endpoint = self._resolve_public_host(server_config)
            self._udp_bind_host_config = server_config.get('udp_bind_host')
            signature_key = self._resolve_signature_key(new_config)
            if not signature_key:
                raise ValueError(
                    "启用原生MQTT必须配置mqtt_server.signature_key或"
                    "server.mqtt_signature_key"
                )
            self.signature_key = signature_key
            self.message_queue_size = int(server_config.get('message_queue_size', 128))
            self.business_ready_timeout = float(
                server_config.get('business_ready_timeout', 30) or 30
            )
            self.close_timeout = max(
                0.1, float(server_config.get('close_timeout', 2) or 2)
            )
            self.shutdown_timeout = max(
                self.close_timeout,
                float(server_config.get('shutdown_timeout', 10) or 10),
            )
            self.goodbye_timeout = max(
                0.1, float(server_config.get('goodbye_timeout', 1) or 1)
            )
            self.max_connections = int(server_config.get('max_connections', 1000))
            self.max_pending_connections = int(
                server_config.get('max_pending_connections', 128)
            )
            self.max_payload_size = int(server_config.get('max_payload_size', 8192))
            self.heartbeat_interval = int(server_config.get('heartbeat_interval', 30))

            # 已建立连接保留会话快照，新连接使用新配置。
            self.connection_service = ConnectionService(new_config)
            self.connection_service.server = getattr(self, "management_owner", self)
            return True
        except Exception as e:
            logger.error(f"更新MQTT服务器配置失败: {e}")
            return False

    @staticmethod
    def _resolve_signature_key(config: Dict[str, Any]) -> str:
        server_config = config.get('mqtt_server', {})
        return (
            normalize_signature_key(server_config.get('signature_key'))
            or normalize_signature_key(
                config.get('server', {}).get('mqtt_signature_key')
            )
        )

    @staticmethod
    def _resolve_public_host(server_config: Dict[str, Any]) -> str:
        host, _ = parse_mqtt_endpoint(
            server_config.get('public_endpoint'),
            None,
        )
        return host or 'localhost'

    async def apply_config(self, new_config: Dict[str, Any]) -> bool:
        return await self.update_config(new_config)

    def get_server_status(self) -> Dict[str, Any]:
        """获取服务器状态"""
        return {
            'type': 'mqtt',
            'host': self.host,
            'mqtt_port': self.mqtt_port,
            'udp_port': self.udp_port,
            'active_connections': len(self.connections),
            'active_transports': len(self.active_transports),
            'pending_connections': len(self._pending_connection_ids),
            'cleanup_tasks': len(self._cleanup_tasks),
            'draining_tasks': sum(
                not task.done() for task in self._draining_tasks
            ),
            'connection_handler_tasks': sum(
                not task.done()
                for task in self._connection_handler_tasks
            ),
            'stopping': self._stopping,
        }


class UDPProtocol(asyncio.DatagramProtocol):
    """UDP协议处理器"""

    def __init__(self, mqtt_server: MQTTServer):
        self.mqtt_server = mqtt_server
        self.transport = None

    def connection_made(self, transport):
        self.transport = transport

    def datagram_received(self, data: bytes, addr: tuple):
        """接收UDP数据报"""
        try:
            # 解析UDP消息头
            if len(data) < 16:
                return

            packet_type = data[0]
            if packet_type != 1:
                return

            payload_length = int.from_bytes(data[2:4], 'big')
            if len(data) != 16 + payload_length:
                logger.warning(
                    "UDP数据报长度不匹配: actual={}, expected={}",
                    len(data),
                    16 + payload_length,
                )
                return

            connection_id = int.from_bytes(data[4:8], 'big')
            timestamp = int.from_bytes(data[8:12], 'big')
            sequence = int.from_bytes(data[12:16], 'big')
            header = data[:16]
            encrypted_payload = data[16:16 + payload_length]

            # 找到对应的UDP处理器
            udp_handler = self.mqtt_server.udp_handlers.get(connection_id)
            if udp_handler:
                udp_handler.on_udp_message(header, encrypted_payload, payload_length, timestamp, sequence, addr)

        except Exception as e:
            logger.error(f"处理UDP数据报失败: {e}")

    def error_received(self, exc):
        logger.error(f"UDP协议错误: {exc}")