first commit

2026-07-31 16:59:09 +08:00
commit 48e540fefa
5 changed files with 231 additions and 0 deletions
@@ -0,0 +1,130 @@
+import asyncio
+from typing import Dict, Optional
+import logging
+
+logger = logging.getLogger(__name__)
+
+class StreamManager:
+    def __init__(self, 
+                 mediamtx_addr, 
+                 max_streams: int = 16,
+                 video_path: str = ""):
+        self.max_streams = max_streams
+        self._processes: Dict[str, asyncio.subprocess.Process] = {}
+        self._tasks: Dict[str, asyncio.Task] = {}
+        self._semaphore = asyncio.Semaphore(max_streams)
+        self._mediamtx_addr = mediamtx_addr
+        self._video_path = video_path
+
+    def _is_local_file(self, url: str) -> bool:
+        """判断是否为本地文件（绝对路径或相对路径）"""
+        # 如果以 / 开头，或者包含 .mp4/.avi/.mkv 等，认定为文件
+        return url.startswith('/') or url.startswith('./') or \
+               any(url.lower().endswith(ext) for ext in ['.mp4', '.avi', '.mkv', '.mov', '.flv'])
+
+    def _build_ffmpeg_cmd(self, source_url: str, output_url: str, loop: int = -1) -> list:
+        """根据源类型构建 ffmpeg 命令"""
+        cmd = ["ffmpeg"]
+
+        if self._is_local_file(source_url):
+            # ----- 本地文件模式（支持循环） -----
+            cmd.extend([
+                "-re",                     # 按原始帧率读取
+                "-stream_loop", str(loop), # -1 为无限循环
+                "-i", source_url
+            ])
+        else:
+            # ----- RTSP 网络流模式（支持自动重连） -----
+            cmd.extend([
+                "-rtsp_transport", "tcp",
+                "-stimeout", "3000000",   # 3 秒超时
+                "-reconnect", "1",
+                "-reconnect_streamed", "1",
+                "-reconnect_delay_max", "5",
+                "-i", source_url
+            ])
+
+        # 公共输出部分（流复制，推送到 MediaMTX）
+        cmd.extend([
+            "-c", "copy",
+            "-f", "rtsp",
+            output_url
+        ])
+        return cmd
+
+    async def start_stream(self, stream_id: str, source_url: str, duration: int = 0, loop: int = -1):
+        if len(self._processes) >= self.max_streams:
+            logger.warning(f"已达最大流上限 {self.max_streams}，拒绝新流 {stream_id}")
+            return False
+
+        if stream_id in self._processes:
+            logger.info(f"流 {stream_id} 已存在")
+            return True
+
+        async with self._semaphore:
+            output_url = f"{self._mediamtx_addr}{stream_id}"
+            cmd = self._build_ffmpeg_cmd(source_url, output_url, loop)
+
+            logging.info(f"[Stream {stream_id}] 启动命令: {' '.join(cmd)}")
+
+            process = await asyncio.create_subprocess_exec(
+                *cmd,
+                stdout=asyncio.subprocess.DEVNULL,
+                stderr=asyncio.subprocess.PIPE
+            )
+
+            self._processes[stream_id] = process
+
+            task = asyncio.create_task(self._monitor(stream_id, process, duration))
+            self._tasks[stream_id] = task
+            return True
+
+    async def _monitor(self, stream_id: str, process: asyncio.subprocess.Process, duration: int):
+        """监控单个 ffmpeg 进程的生命周期"""
+        try:
+            if duration > 0:
+                await asyncio.wait_for(process.wait(), timeout=duration)
+            else:
+                await process.wait()
+
+        except asyncio.TimeoutError:
+            logger.info(f"[Stream {stream_id}] 推流超时 ({duration}s), 自动停止")
+            process.terminate()
+            try:
+                await asyncio.wait_for(process.wait(), timeout=3)
+            except asyncio.TimeoutError:
+                process.kill()
+                await process.wait()
+
+        except Exception as e:
+            logger.error(f"[Stream {stream_id}] 监控异常: {e}")
+
+        finally:
+            await self._cleanup(stream_id)
+
+    async def stop_stream(self, stream_id: str):
+        """主动停止一路流"""
+        process = self._processes.get(stream_id)
+        if process:
+            logger.info(f"[Stream {stream_id}] 收到停止指令")
+            process.terminate()
+            try:
+                await asyncio.wait_for(process.wait(), timeout=5)
+            except asyncio.TimeoutError:
+                process.kill()
+                await process.wait()
+            await self._cleanup(stream_id)
+
+    async def _cleanup(self, stream_id: str):
+        if stream_id in self._processes:
+            del self._processes[stream_id]
+        if stream_id in self._tasks:
+            task = self._tasks.pop(stream_id)
+            if not task.done():
+                task.cancel()
+        logger.info(f"[Stream {stream_id}] 资源已释放")
+
+    async def stop_all(self):
+        """停止所有流 (用于应用关闭时) """
+        for sid in list(self._processes.keys()):
+            await self.stop_stream(sid)