refactoring

2 files changed, 25 insertions, 31 deletions

diff --git a/kvmd/apps/vnc/rfb/__init__.py b/kvmd/apps/vnc/rfb/__init__.py
index bbb3036c..54c9648f 100644
--- a/kvmd/apps/vnc/rfb/__init__.py
+++ b/kvmd/apps/vnc/rfb/__init__.py
@@ -25,6 +25,7 @@ import ssl
 
 from typing import Callable
 from typing import Coroutine
+from typing import AsyncGenerator
 
 from ....logging import get_logger
 
@@ -87,6 +88,9 @@ class RfbClient(RfbClientStream):  # pylint: disable=too-many-instance-attribute
 
         self.__reset_h264 = False
 
+        self.__fb_notifier = aiotools.AioNotifier()
+        self.__fb_cont_updates = False
+
     # =====
 
     async def _run(self, **coros: Coroutine) -> None:
@@ -156,13 +160,15 @@ class RfbClient(RfbClientStream):  # pylint: disable=too-many-instance-attribute
     async def _on_set_encodings(self) -> None:
         raise NotImplementedError
 
-    async def _on_fb_update_request(self) -> None:
-        raise NotImplementedError
+    # =====
 
-    async def _on_enable_cont_updates(self, enabled: bool) -> None:
-        raise NotImplementedError
+    async def _send_fb_allowed(self) -> AsyncGenerator[None, None]:
+        while True:
+            await self.__fb_notifier.wait()
+            yield
 
-    # =====
+    async def _send_fb_allow_again(self) -> None:
+        self.__fb_notifier.notify()
 
     async def _send_fb_jpeg(self, data: bytes) -> None:
         assert self._encodings.has_tight
@@ -178,6 +184,8 @@ class RfbClient(RfbClientStream):  # pylint: disable=too-many-instance-attribute
             length_bytes = bytes([0b10011111, length & 0x7F | 0x80, length >> 7 & 0x7F | 0x80, length >> 14 & 0xFF])
         await self._write_struct("JPEG length + data", "", length_bytes, data)
         self.__reset_h264 = True
+        if self.__fb_cont_updates:
+            self.__fb_notifier.notify()
 
     async def _send_fb_h264(self, data: bytes) -> None:
         assert self._encodings.has_h264
@@ -186,6 +194,8 @@ class RfbClient(RfbClientStream):  # pylint: disable=too-many-instance-attribute
         await self._write_struct("H264 length + flags", "LL", len(data), int(self.__reset_h264), drain=False)
         await self._write_struct("H264 data", "", data)
         self.__reset_h264 = False
+        if self.__fb_cont_updates:
+            self.__fb_notifier.notify()
 
     async def _send_resize(self, width: int, height: int) -> None:
         assert self._encodings.has_resize
@@ -443,7 +453,8 @@ class RfbClient(RfbClientStream):  # pylint: disable=too-many-instance-attribute
     async def __handle_fb_update_request(self) -> None:
         self.__check_encodings()
         await self._read_struct("FBUR", "? HH HH")  # Ignore any arguments, just perform the full update
-        await self._on_fb_update_request()
+        if not self.__fb_cont_updates:
+            self.__fb_notifier.notify()
 
     def __check_encodings(self) -> None:
         # JpegCompression may only be used when the client has advertized
@@ -482,8 +493,11 @@ class RfbClient(RfbClientStream):  # pylint: disable=too-many-instance-attribute
 
     async def __handle_enable_cont_updates(self) -> None:
         enabled = bool((await self._read_struct("enabled ContUpdates", "B HH HH"))[0])
-        await self._on_enable_cont_updates(enabled)
-        if not enabled:
+        get_logger(0).info("%s [main]: Applying ContUpdates=%s ...", self._remote, enabled)
+        self.__fb_cont_updates = enabled
+        if enabled:
+            self.__fb_notifier.notify()
+        else:
             await self._write_struct("disabled ContUpdates", "B", 150)
 
     async def __handle_qemu_event(self) -> None:
diff --git a/kvmd/apps/vnc/server.py b/kvmd/apps/vnc/server.py
index 0195cdd9..366a4a91 100644
--- a/kvmd/apps/vnc/server.py
+++ b/kvmd/apps/vnc/server.py
@@ -120,9 +120,7 @@ class _Client(RfbClient):  # pylint: disable=too-many-instance-attributes
         self.__kvmd_session: (KvmdClientSession | None) = None
         self.__kvmd_ws: (KvmdClientWs | None) = None
 
-        self.__fb_notifier = aiotools.AioNotifier()
         self.__fb_queue: "asyncio.Queue[dict]" = asyncio.Queue()
-        self.__fb_cont_updates = False
         self.__fb_key_required = False
 
         # Эти состояния шарить не обязательно - бекенд исключает дублирующиеся события.
@@ -252,9 +250,7 @@ class _Client(RfbClient):  # pylint: disable=too-many-instance-attributes
     async def __fb_sender_task_loop(self) -> None:  # pylint: disable=too-many-branches
         has_h264_key = False
         last: (dict | None) = None
-        while True:
-            await self.__fb_notifier.wait()
-
+        async for _ in self._send_fb_allowed():
             while True:
                 frame = await self.__fb_queue.get()
                 if (
@@ -288,31 +284,25 @@ class _Client(RfbClient):  # pylint: disable=too-many-instance-attributes
                             f" -> {last['width']}x{last['height']}\nPlease reconnect"
                         )
                         await self._send_fb_jpeg((await self.__make_text_frame(msg))["data"])
-                        if self.__fb_cont_updates:
-                            self.__fb_notifier.notify()
                         continue
                     await self._send_resize(last["width"], last["height"])
 
                 if len(last["data"]) == 0:
                     # Вдруг какой-то баг
-                    self.__fb_notifier.notify()
+                    await self._send_fb_allow_again()
                     continue
 
                 if last["format"] == StreamFormats.JPEG:
                     await self._send_fb_jpeg(last["data"])
-                    if self.__fb_cont_updates:
-                        self.__fb_notifier.notify()
                 elif last["format"] == StreamFormats.H264:
                     if not self._encodings.has_h264:
                         raise RfbError("The client doesn't want to accept H264 anymore")
                     if has_h264_key:
                         self.__fb_key_required = False
                         await self._send_fb_h264(last["data"])
-                        if self.__fb_cont_updates:
-                            self.__fb_notifier.notify()
                     else:
                         self.__fb_key_required = True
-                        self.__fb_notifier.notify()
+                        await self._send_fb_allow_again()
                 else:
                     raise RuntimeError(f"Unknown format: {last['format']}")
                 last["data"] = b""
@@ -424,16 +414,6 @@ class _Client(RfbClient):  # pylint: disable=too-many-instance-attributes
                            self._remote, quality, self.__desired_fps)
         await self.__kvmd_session.streamer.set_params(quality, self.__desired_fps)
 
-    async def _on_fb_update_request(self) -> None:
-        if not self.__fb_cont_updates:
-            self.__fb_notifier.notify()
-
-    async def _on_enable_cont_updates(self, enabled: bool) -> None:
-        get_logger(0).info("%s [main]: Applying ContUpdates=%s ...", self._remote, enabled)
-        self.__fb_cont_updates = enabled
-        if enabled:
-            self.__fb_notifier.notify()
-
 
 # =====
 class VncServer:  # pylint: disable=too-many-instance-attributes