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()