accumulate h264 frames

2 files changed, 43 insertions, 30 deletions

diff --git a/kvmd/apps/vnc/server.py b/kvmd/apps/vnc/server.py
index fe340e65..59cd0f4e 100644
--- a/kvmd/apps/vnc/server.py
+++ b/kvmd/apps/vnc/server.py
@@ -123,6 +123,7 @@ class _Client(RfbClient):  # pylint: disable=too-many-instance-attributes
 
         self.__fb_requested = False
         self.__fb_stub: Optional[Tuple[int, str]] = None
+        self.__fb_h264_data = b""
 
         # Эти состояния шарить не обязательно - бекенд исключает дублирующиеся события.
         # Все это нужно только чтобы не посылать лишние жсоны в сокет KVMD
@@ -195,12 +196,12 @@ class _Client(RfbClient):  # pylint: disable=too-many-instance-attributes
         while True:
             try:
                 streaming = False
-                async for (online, width, height, data) in streamer.read_stream():
+                async for frame in streamer.read_stream():
                     if not streaming:
                         logger.info("[streamer] %s: Streaming ...", self._remote)
                         streaming = True
-                    if online:
-                        await self.__send_fb_real(width, height, data, streamer.get_format())
+                    if frame["online"]:
+                        await self.__send_fb_real(frame, streamer.get_format())
                     else:
                         await self.__send_fb_stub("No signal")
             except StreamerError as err:
@@ -229,29 +230,47 @@ class _Client(RfbClient):  # pylint: disable=too-many-instance-attributes
         get_logger(0).info("[streamer] %s: Using default %s", self._remote, streamer)
         return streamer
 
-    async def __send_fb_real(self, width: int, height: int, data: bytes, fmt: int) -> None:
+    async def __send_fb_real(self, frame: Dict, fmt: int) -> None:
         async with self.__lock:
             if self.__fb_requested:
-                if self._width != width or self._height != height:
-                    self.__shared_params.width = width
-                    self.__shared_params.height = height
-                    if not self._encodings.has_resize:
-                        msg = f"Resoultion changed: {self._width}x{self._height} -> {width}x{height}\nPlease reconnect"
-                        await self.__send_fb_stub(msg, no_lock=True)
-                        return
-                    await self._send_resize(width, height)
+                if not (await self.__resize_fb_unsafe(frame)):
+                    return
 
                 if fmt == StreamFormats.JPEG:
-                    await self._send_fb_jpeg(data)
+                    await self._send_fb_jpeg(frame["data"])
                 elif fmt == StreamFormats.H264:
                     if not self._encodings.has_h264:
+                        self.__fb_h264_data = b""
                         raise StreamerPermError("The client doesn't want to accept H264 anymore")
-                    await self._send_fb_h264(data)
+                    await self._send_fb_h264(self.__fb_h264_data)
                 else:
                     raise RuntimeError(f"Unknown format: {fmt}")
 
                 self.__fb_stub = None
                 self.__fb_requested = False
+                self.__fb_h264_data = b""
+
+            elif self._encodings.has_h264 and fmt == StreamFormats.H264:
+                if frame["key"]:
+                    self.__fb_h264_data = frame["data"]
+                elif len(self.__fb_h264_data) + len(frame["data"]) > 4194304:  # 4Mb
+                    get_logger(0).info("Accumulated H264 buffer is too big; resetting ...")
+                    self.__fb_h264_data = frame["data"]
+                else:
+                    self.__fb_h264_data += frame["data"]
+
+    async def __resize_fb_unsafe(self, frame: Dict) -> bool:
+        width = frame["width"]
+        height = frame["height"]
+        if self._width != width or self._height != height:
+            self.__shared_params.width = width
+            self.__shared_params.height = height
+            if not self._encodings.has_resize:
+                msg = f"Resoultion changed: {self._width}x{self._height} -> {width}x{height}\nPlease reconnect"
+                await self.__send_fb_stub(msg, no_lock=True)
+                return False
+            await self._send_resize(width, height)
+        return True
 
     async def __send_fb_stub(self, text: str, no_lock: bool=False) -> None:
         if not no_lock:
diff --git a/kvmd/clients/streamer.py b/kvmd/clients/streamer.py
index a0a07bf5..a42dcb20 100644
--- a/kvmd/clients/streamer.py
+++ b/kvmd/clients/streamer.py
@@ -22,7 +22,6 @@
 
 import types
 
-from typing import Tuple
 from typing import Dict
 from typing import AsyncGenerator
 
@@ -62,7 +61,7 @@ class BaseStreamerClient:
     def get_format(self) -> int:
         raise NotImplementedError()
 
-    async def read_stream(self) -> AsyncGenerator[Tuple[bool, int, int, bytes], None]:
+    async def read_stream(self) -> AsyncGenerator[Dict, None]:
         if self is not None:  # XXX: Vulture and pylint hack
             raise NotImplementedError()
         yield
@@ -90,7 +89,7 @@ class HttpStreamerClient(BaseStreamerClient):
     def get_format(self) -> int:
         return StreamFormats.JPEG
 
-    async def read_stream(self) -> AsyncGenerator[Tuple[bool, int, int, bytes], None]:
+    async def read_stream(self) -> AsyncGenerator[Dict, None]:
         try:
             async with self.__make_http_session() as session:
                 async with session.get(
@@ -110,12 +109,12 @@ class HttpStreamerClient(BaseStreamerClient):
                         if not data:
                             break
 
-                        yield (
-                            (frame.headers["X-UStreamer-Online"] == "true"),
-                            int(frame.headers["X-UStreamer-Width"]),
-                            int(frame.headers["X-UStreamer-Height"]),
-                            data,
-                        )
+                        yield {
+                            "online": (frame.headers["X-UStreamer-Online"] == "true"),
+                            "width": int(frame.headers["X-UStreamer-Width"]),
+                            "height": int(frame.headers["X-UStreamer-Height"]),
+                            "data": data,
+                        }
         except Exception as err:  # Тут бывают и ассерты, и KeyError, и прочая херня
             if isinstance(err, StreamerTempError):
                 raise
@@ -178,7 +177,7 @@ class MemsinkStreamerClient(BaseStreamerClient):
     def get_format(self) -> int:
         return self.__fmt
 
-    async def read_stream(self) -> AsyncGenerator[Tuple[bool, int, int, bytes], None]:
+    async def read_stream(self) -> AsyncGenerator[Dict, None]:
         if ustreamer is None:
             raise StreamerPermError("Missing ustreamer library")
         try:
@@ -187,12 +186,7 @@ class MemsinkStreamerClient(BaseStreamerClient):
                     frame = await aiotools.run_async(sink.wait_frame)
                     if frame is not None:
                         self.__check_format(frame["format"])
-                        yield (
-                            frame["online"],
-                            frame["width"],
-                            frame["height"],
-                            frame["data"],
-                        )
+                        yield frame
         except StreamerPermError:
             raise
         except FileNotFoundError as err: