iAOP/core/edge-gateway/collector/engine.py

# -*- coding: utf-8 -*-
"""周期采集调度引擎 —— 只读采集 + 背压保护 + 健康度上报。

模板化封装要点：
- 采集配置全部外置（gateway.yaml），引擎不关心具体协议；
- 点位按驱动实例的 device 匹配规则分组，一次 tick 内按驱动批量读取；
- 严格只读：引擎只调用 Driver.read_points()，不存在任何控制指令路径；
- 背压保护：未确认(spool 待上行)记录超过阈值时丢弃新样本并计入丢失，
  防止 Kafka 故障时内存/磁盘无限增长；
- 每个 tick 结束记录轮次耗时与样本成败 → HealthMetrics（P99/丢失率/可用性）。
"""
from __future__ import annotations

import logging
import threading
import time
from typing import Dict, List, Optional, Tuple

from drivers.base import Driver, SampleValue
from point_dict.loader import Point, PointDict
from .metrics import HealthMetrics
from .spool import SpoolStore

logger = logging.getLogger("edge_gateway.engine")


class CollectorEngine:
    """只读采集调度引擎（单线程 tick，可替换为 asyncio 版本保持接口不变）。"""

    def __init__(
        self,
        point_dict: PointDict,
        driver_slots: List[Tuple[str, Driver, List[str]]],
        spool: SpoolStore,
        metrics: HealthMetrics,
        interval_ms: int = 1000,
        max_pending: int = 100_000,
    ):
        """
        Args:
            point_dict: 点位字典（已通过校验）；
            driver_slots: [(protocol, driver, device_prefixes)]，
                device_prefixes 为空列表 = 兜底驱动（接收未分配点位）；
            spool: 本地缓存（断点续传）；
            metrics: 健康度统计；
            interval_ms: 采集周期（模板配置，点位字典未覆盖时的默认值）；
            max_pending: 背压阈值（未确认 spool 记录数上限）。
        """
        self.point_dict = point_dict
        self.driver_slots = driver_slots
        self.spool = spool
        self.metrics = metrics
        self.interval_ms = max(50, int(interval_ms))
        self.max_pending = max(1, int(max_pending))
        self._stop = threading.Event()
        self._thread: Optional[threading.Thread] = None
        self._point_to_driver = self._build_routing()

    # ------------------------------------------------------------------
    def _build_routing(self) -> Dict[str, Driver]:
        """按设备前缀把点位路由到对应驱动实例（模板配置驱动）。"""
        routing: Dict[str, Driver] = {}
        for protocol, driver, prefixes in self.driver_slots:
            for p in self.point_dict.points:
                if p.point_id in routing:
                    continue
                if not prefixes or any(p.device_id.startswith(pre) for pre in prefixes):
                    routing[p.point_id] = driver
        return routing

    # ------------------------------------------------------------------
    def collect_once(self, sink=None) -> int:
        """执行一轮采集：读点位 → 写 spool → 上行。

        Args:
            sink: 可选 KafkaSink；为 None 时仅写入 spool（离线采集模式）。

        Returns:
            本轮成功读到的样本数。
        """
        started = time.monotonic()
        expected = len(self.point_dict.points)
        got = 0
        samples: List[dict] = []

        # 1) 按驱动分组读取（只读）
        by_driver: Dict[Driver, List[Point]] = {}
        for p in self.point_dict.points:
            drv = self._point_to_driver.get(p.point_id)
            if drv is None:
                continue
            by_driver.setdefault(drv, []).append(p)

        for drv, points in by_driver.items():
            try:
                values = drv.read_points(points)
            except Exception:  # 驱动级异常：本轮整体记失败，不中断网关
                logger.exception("驱动读取异常: %s", drv.protocol)
                self.metrics.record_round(time.monotonic() - started, expected, got, failed=True)
                return got
            for p in points:
                value = values.get(p.point_id)
                if value is None:
                    continue  # 未读到 → 计入丢失
                got += 1
                ts = time.time()
                samples.append(
                    {"device_id": p.device_id, "point_id": p.point_id,
                     "value": value, "ts": ts, "unit": p.unit}
                )

        # 2) 背压保护：待上行记录超阈值时丢弃新样本
        pending = self.spool.total_pending()
        if pending >= self.max_pending:
            dropped = len(samples)
            samples = []
            # 丢弃样本计入丢失率
            self.metrics.record_round(time.monotonic() - started, expected, got, failed=False)
            logger.warning("背压保护触发：spool 待上行 %d 条 ≥ 阈值 %d，丢弃本轮 %d 条样本",
                           pending, self.max_pending, dropped)
            return got

        # 3) 写 spool（断点续传落盘）
        for s in samples:
            self.spool.append(s["device_id"], s["point_id"], s["value"], s["ts"])

        # 4) 上行（Kafka）；失败由 sink 内部重试/保留 spool
        if sink is not None:
            try:
                sink.publish(samples)
            except Exception:
                logger.exception("上行异常，样本保留在 spool 等待重发")

        latency = time.monotonic() - started
        self.metrics.record_round(latency, expected, got)
        return got

    # ------------------------------------------------------------------
    def run_forever(self, sink=None) -> None:
        """tick 循环入口（供线程调用）。"""
        while not self._stop.is_set():
            try:
                self.collect_once(sink=sink)
            except Exception:
                logger.exception("采集轮次异常")
            # 下一轮 tick 对齐 interval_ms
            self._stop.wait(self.interval_ms / 1000.0)

    def start(self, sink=None) -> None:
        """后台启动采集线程。"""
        if self._thread is not None and self._thread.is_alive():
            return
        self._stop.clear()
        self._thread = threading.Thread(
            target=self.run_forever, args=(sink,), name="edge-gateway-collector", daemon=True
        )
        self._thread.start()

    def stop(self) -> None:
        """停止采集线程。"""
        self._stop.set()
        if self._thread is not None:
            self._thread.join(timeout=5)
            self._thread = None