""" 截图处理模块 - 监听云端 Redis Stream 截图请求，截图后上传 COS 数据流： WVP → XADD edge_snap_request → Edge 消费 → capture frame → upload COS → SET snap:result:{id} """ import io import json import logging import threading import time import uuid from typing import Optional import cv2 import numpy as np from config.settings import get_settings, COSConfig logger = logging.getLogger(__name__) # Redis 常量 SNAP_REQUEST_STREAM = "edge_snap_request" SNAP_CONSUMER_GROUP = "edge_snap_group" SNAP_CONSUMER_NAME = "edge_snap_worker" SNAP_RESULT_KEY_PREFIX = "snap:result:" SNAP_CACHE_KEY_PREFIX = "snap:cache:" SNAP_RESULT_TTL = 60 # 结果 key 60s 过期 SNAP_CACHE_TTL = 300 # 缓存 key 5min 过期 class ScreenshotHandler: """截图处理器 从云端 Redis Stream 消费截图请求，通过 MultiStreamManager 获取最新帧， JPEG 编码后直传 COS，最后将结果 URL 写回 Redis。 """ def __init__(self, cloud_redis, stream_manager): """ Args: cloud_redis: 云端 Redis 客户端 (redis.Redis) stream_manager: MultiStreamManager 实例 """ self._cloud_redis = cloud_redis self._stream_manager = stream_manager self._settings = get_settings() self._cos_client = None self._cos_config: COSConfig = self._settings.cos self._thread: Optional[threading.Thread] = None self._stop_event = threading.Event() # ==================== 生命周期 ==================== def start(self): """启动截图监听线程""" if not self._cloud_redis: logger.warning("[截图] 云端 Redis 不可用，截图处理器未启动") return if not self._cos_config.secret_id or not self._cos_config.bucket: logger.warning("[截图] COS 配置不完整，截图处理器未启动") return self._init_cos_client() self._ensure_consumer_group() self._stop_event.clear() self._thread = threading.Thread( target=self._listen_loop, name="ScreenshotHandler", daemon=True, ) self._thread.start() logger.info("[截图] 截图处理器已启动") def stop(self): """停止截图监听线程""" self._stop_event.set() if self._thread and self._thread.is_alive(): self._thread.join(timeout=5) logger.info("[截图] 截图处理器已停止") # ==================== COS 初始化 ==================== def _init_cos_client(self): """初始化 COS 客户端""" try: from qcloud_cos import CosConfig, CosS3Client config = CosConfig( Region=self._cos_config.region, SecretId=self._cos_config.secret_id, SecretKey=self._cos_config.secret_key, ) self._cos_client = CosS3Client(config) logger.info("[截图] COS 客户端初始化成功: bucket=%s, region=%s", self._cos_config.bucket, self._cos_config.region) except Exception as e: logger.error("[截图] COS 客户端初始化失败: %s", e) # ==================== Consumer Group ==================== def _ensure_consumer_group(self): """确保 Redis Consumer Group 存在""" try: self._cloud_redis.xgroup_create( SNAP_REQUEST_STREAM, SNAP_CONSUMER_GROUP, id="0", mkstream=True ) logger.info("[截图] 创建 consumer group: %s", SNAP_CONSUMER_GROUP) except Exception as e: # BUSYGROUP = group already exists, safe to ignore if "BUSYGROUP" in str(e): logger.debug("[截图] consumer group 已存在") else: logger.error("[截图] 创建 consumer group 失败: %s", e) # ==================== 主循环 ==================== def _listen_loop(self): """XREADGROUP 循环消费截图请求""" backoff = 5 max_backoff = 60 while not self._stop_event.is_set(): try: results = self._cloud_redis.xreadgroup( SNAP_CONSUMER_GROUP, SNAP_CONSUMER_NAME, {SNAP_REQUEST_STREAM: ">"}, count=1, block=5000, ) if not results: continue backoff = 5 # 重置退避 for stream_name, messages in results: for msg_id, fields in messages: try: self._handle_request(fields) except Exception as e: logger.error("[截图] 处理请求失败: %s", e) finally: # ACK 消息 try: self._cloud_redis.xack( SNAP_REQUEST_STREAM, SNAP_CONSUMER_GROUP, msg_id ) except Exception: pass except Exception as e: if self._stop_event.is_set(): return logger.warning("[截图] 监听异常: %s, %ds 后重试", e, backoff) self._stop_event.wait(backoff) backoff = min(backoff * 2, max_backoff) # ==================== 请求处理 ==================== def _handle_request(self, fields: dict): """处理单个截图请求""" request_id = fields.get("request_id", "") camera_code = fields.get("camera_code", "") cos_path = fields.get("cos_path", "") if not request_id or not camera_code or not cos_path: logger.warning("[截图] 请求字段不完整: %s", fields) self._write_result(request_id, {"status": "error", "message": "请求字段不完整"}) return logger.info("[截图] 收到截图请求: request_id=%s, camera=%s", request_id, camera_code) # 1. 抓帧 jpeg_bytes = self._capture_frame(camera_code) if jpeg_bytes is None: self._write_result(request_id, { "status": "error", "message": f"摄像头流未连接: {camera_code}" }) return # 2. 上传 COS（失败重试 1 次） url = self._upload_to_cos(jpeg_bytes, cos_path) if url is None: # 重试一次 logger.info("[截图] COS 上传失败，重试一次...") url = self._upload_to_cos(jpeg_bytes, cos_path) if url is None: self._write_result(request_id, { "status": "error", "message": "COS 上传失败" }) return # 3. 写结果 result = {"status": "ok", "url": url} self._write_result(request_id, result) # 4. 写缓存 cache_key = SNAP_CACHE_KEY_PREFIX + camera_code cache_data = json.dumps({"url": url, "timestamp": int(time.time())}) try: self._cloud_redis.set(cache_key, cache_data, ex=SNAP_CACHE_TTL) except Exception as e: logger.warning("[截图] 写入缓存失败: %s", e) logger.info("[截图] 截图完成: camera=%s, url=%s", camera_code, url) # ==================== 抓帧 ==================== def _capture_frame(self, camera_code: str) -> Optional[bytes]: """从 MultiStreamManager 获取最新帧并编码为 JPEG""" stream = self._stream_manager.get_stream(camera_code) if stream is None: logger.warning("[截图] 未找到摄像头流: %s", camera_code) return None if not stream.is_connected: logger.warning("[截图] 摄像头未连接: %s", camera_code) return None frame = stream.get_latest_frame(timeout=2.0) if frame is None: # 回退：直接从缓冲区读一帧 frame = stream.read(timeout=2.0) if frame is None: logger.warning("[截图] 获取帧超时: %s", camera_code) return None # JPEG 编码 try: encode_params = [cv2.IMWRITE_JPEG_QUALITY, 85] success, buffer = cv2.imencode(".jpg", frame.image, encode_params) if not success: logger.error("[截图] JPEG 编码失败: %s", camera_code) return None return buffer.tobytes() except Exception as e: logger.error("[截图] 帧编码异常: %s", e) return None # ==================== COS 上传 ==================== def _upload_to_cos(self, jpeg_bytes: bytes, cos_path: str) -> Optional[str]: """上传 JPEG 到 COS，返回访问 URL""" if not self._cos_client: logger.error("[截图] COS 客户端未初始化") return None try: body = io.BytesIO(jpeg_bytes) self._cos_client.put_object( Bucket=self._cos_config.bucket, Body=body, Key=cos_path, ContentType="image/jpeg", ) url = "https://{}.cos.{}.myqcloud.com/{}".format( self._cos_config.bucket, self._cos_config.region, cos_path ) return url except Exception as e: logger.error("[截图] COS 上传失败: %s", e) return None # ==================== Redis 结果 ==================== def _write_result(self, request_id: str, data: dict): """将结果写入 Redis key（TTL 60s）""" if not request_id: return key = SNAP_RESULT_KEY_PREFIX + request_id try: self._cloud_redis.set(key, json.dumps(data), ex=SNAP_RESULT_TTL) except Exception as e: logger.error("[截图] 写入结果失败: %s", e)