refactoring

This commit is contained in:
Maxim Devaev 2022-11-03 17:00:27 +03:00
parent 08241e9255
commit 260f118820
2 changed files with 26 additions and 32 deletions

View File

@ -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,14 +160,16 @@ 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
assert self._encodings.tight_jpeg_quality > 0
@ -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:

View File

@ -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