iTi-Flask-IOT/iti/applications/service/iot/rocketmq_mgr.py

from __future__ import annotations

import json
import logging
import os
from dataclasses import dataclass
from typing import Any, Callable, Optional, Union
from iti.applications.service.iot.alert import add_endpoint_alert_log, add_node_alert_log

try:
    from rocketmq.client import ConsumeStatus, Message, Producer, PushConsumer
except Exception:
    ConsumeStatus = None
    Message = None
    Producer = None
    PushConsumer = None

logger = logging.getLogger(__name__)

RocketMQBody = Union[str, bytes, bytearray, dict, list]
ConsumeCallback = Callable[[Any], Any]


@dataclass(frozen=True)
class RocketMQConfig:
    namesrv_addr: str
    producer_group: str
    consumer_group: str
    default_topic: str
    default_tags: str
    charset: str
    access_key: str
    secret_key: str
    security_token: str
    instance_name: str

    @staticmethod
    def from_env() -> "RocketMQConfig":
        return RocketMQConfig(
            namesrv_addr=os.getenv("ROCKETMQ_NAMESRV_ADDR", "127.0.0.1:9876").strip(),
            producer_group=os.getenv("ROCKETMQ_PRODUCER_GROUP", "iot-collect-group").strip(),
            consumer_group=os.getenv("ROCKETMQ_CONSUMER_GROUP", "iot-collect-group").strip(),
            default_topic=os.getenv("ROCKETMQ_TOPIC", "iot-collect-topic").strip(),
            default_tags=os.getenv("ROCKETMQ_TAGS", "*").strip() or "*",
            charset=os.getenv("ROCKETMQ_CHARSET", "utf-8").strip() or "utf-8",
            access_key=os.getenv("ROCKETMQ_ACCESS_KEY", "").strip(),
            secret_key=os.getenv("ROCKETMQ_SECRET_KEY", "").strip(),
            security_token=os.getenv("ROCKETMQ_SECURITY_TOKEN", "").strip(),
            instance_name=os.getenv("ROCKETMQ_INSTANCE_NAME", "").strip(),
        )


class RocketMQMgr:
    _instance = None
    _initialized = False

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def init_app(self, app) -> None:
        if self._initialized:
            logger.warning("rocketmq连接已初始化，跳过重复初始化")
            return

        if Producer is None or PushConsumer is None or Message is None:
            logger.error(
                "rocketmq客户端库不可用，请安装 rocketmq-client-python（并确保本机可加载 librocketmq）"
            )
            self._initialized = False
            return

        cfg = RocketMQConfig.from_env()

        if not cfg.namesrv_addr:
            logger.error("rocketmq配置不完整，缺少: ROCKETMQ_NAMESRV_ADDR")
            self._initialized = False
            return

        self.cfg = cfg
        self._app = app
        self.producer: Optional[Any] = None
        self.consumer: Optional[Any] = None
        self._producer_started = False
        self._consumer_started = False

        self._initialized = True
        logger.info("rocketmq初始化成功")

    def _ensure_initialized(self) -> bool:
        if not self._initialized:
            logger.error("rocketmq未初始化，请先调用 iot_rocketmq.init_app(app)")
            return False
        return True

    def _apply_credentials_if_supported(self, client: Any) -> None:
        if not self.cfg.access_key or not self.cfg.secret_key:
            return
        if hasattr(client, "set_session_credentials"):
            try:
                client.set_session_credentials(
                    self.cfg.access_key, self.cfg.secret_key, self.cfg.security_token
                )
            except Exception as e:
                logger.warning(f"rocketmq设置ACL凭证失败: {e}")

    def start_producer(self) -> bool:
        if not self._ensure_initialized():
            return False
        if self._producer_started and self.producer:
            return True

        try:
            producer = Producer(self.cfg.producer_group)
            producer.set_name_server_address(self.cfg.namesrv_addr)
            if self.cfg.instance_name and hasattr(producer, "set_instance_name"):
                producer.set_instance_name(self.cfg.instance_name)
            self._apply_credentials_if_supported(producer)
            producer.start()
            self.producer = producer
            self._producer_started = True
            logger.info("rocketmq producer启动成功")
            return True
        except Exception as e:
            logger.error(f"rocketmq producer启动失败: {e}", exc_info=True)
            self.producer = None
            self._producer_started = False
            return False

    def shutdown_producer(self) -> None:
        if not self.producer:
            return
        try:
            self.producer.shutdown()
        except Exception as e:
            logger.warning(f"rocketmq producer关闭失败: {e}")
        finally:
            self.producer = None
            self._producer_started = False

    def _encode_body(self, body: RocketMQBody) -> bytes:
        if isinstance(body, bytes):
            return body
        if isinstance(body, bytearray):
            return bytes(body)
        if isinstance(body, str):
            return body.encode(self.cfg.charset, errors="replace")
        return json.dumps(body, ensure_ascii=False).encode(self.cfg.charset, errors="replace")

    def send_sync(
        self,
        body: RocketMQBody,
        topic: Optional[str] = None,
        tags: Optional[str] = None,
        keys: Optional[str] = None,
    ) -> Any:
        if not self._ensure_initialized():
            raise RuntimeError("rocketmq未初始化")
        if not self.start_producer() or not self.producer:
            raise RuntimeError("rocketmq producer未启动")

        real_topic = (topic or self.cfg.default_topic).strip()
        if not real_topic:
            logger.error("topic不能为空（可通过参数传入或设置ROCKETMQ_TOPIC）")
            raise ValueError("topic不能为空（可通过参数传入或设置ROCKETMQ_TOPIC）")

        msg = Message(real_topic)
        if tags:
            if hasattr(msg, "set_tags"):
                msg.set_tags(tags)
        if keys:
            if hasattr(msg, "set_keys"):
                msg.set_keys(keys)
        msg.set_body(self._encode_body(body))
        return self.producer.send_sync(msg)

    def start_consumer(
        self,
        topic: Optional[str] = None,
        tags: Optional[str] = None,
        callback: Optional[ConsumeCallback] = None,
    ) -> bool:
        if not self._ensure_initialized():
            return False
        if self._consumer_started and self.consumer:
            return True

        real_topic = (topic or self.cfg.default_topic).strip()
        if not real_topic:
            logger.error("consumer topic不能为空（可通过参数传入或设置ROCKETMQ_TOPIC）")
            return False

        selector_expression = (tags or self.cfg.default_tags).strip() or "*"
        consume_cb = callback or self._default_consume_callback

        try:
            consumer = PushConsumer(self.cfg.consumer_group)
            consumer.set_name_server_address(self.cfg.namesrv_addr)
            if self.cfg.instance_name and hasattr(consumer, "set_instance_name"):
                consumer.set_instance_name(self.cfg.instance_name)
            self._apply_credentials_if_supported(consumer)

            try:
                consumer.subscribe(real_topic, consume_cb, selector_expression)
            except TypeError:
                consumer.subscribe(real_topic, selector_expression, consume_cb)

            consumer.start()
            self.consumer = consumer
            self._consumer_started = True
            logger.info("rocketmq consumer启动成功")
            return True
        except Exception as e:
            logger.error(f"rocketmq consumer启动失败: {e}", exc_info=True)
            self.consumer = None
            self._consumer_started = False
            return False

    def shutdown_consumer(self) -> None:
        if not self.consumer:
            return
        try:
            self.consumer.shutdown()
        except Exception as e:
            logger.warning(f"rocketmq consumer关闭失败: {e}")
        finally:
            self.consumer = None
            self._consumer_started = False

    def _default_consume_callback(self, msg: Any) -> Any:
        topic = getattr(msg, "topic", "")
        tags = getattr(msg, "tags", "")
        keys = getattr(msg, "keys", "")
        body_text = getattr(msg, "body", b"").decode(self.cfg.charset, errors="replace")
        
        # logger.info(
        #     f"[RocketMQ] topic={topic} tags={tags} keys={keys} body={body_text}"
        # )
        try:
            body_dict = json.loads(body_text)
        except Exception:
            logger.error(f"[RocketMQ] 解析body失败: {body_text}")
            return False

        app = getattr(self, "_app", None)
        if app is None:
            logger.error("[RocketMQ] 未绑定Flask app，无法创建application context")
            return False

        with app.app_context():
            if tags == b"netData" or tags == "netData":
                add_endpoint_alert_log(body_dict["endpointId"])

            elif tags == b"collectData" or tags == "collectData":
                for node_data in body_dict["nodeDataList"]:
                    node_id = int(node_data["nodeId"])
                    alert_value = node_data["alertValue"]
                    add_node_alert_log(node_id, alert_value)

        if ConsumeStatus is not None and hasattr(ConsumeStatus, "CONSUME_SUCCESS"):
            return ConsumeStatus.CONSUME_SUCCESS
        return True

    def shutdown(self) -> None:
        self.shutdown_consumer()
        self.shutdown_producer()

    def __del__(self):
        try:
            self.shutdown()
        except Exception:
            pass


iot_rocketmq = RocketMQMgr()