GithubFollow/One-KVM
1
0
Fork 0
mirror of https://github.com/mofeng-git/One-KVM.git synced 2026-01-29 00:51:53 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

refactoring

Browse Source
This commit is contained in:
Maxim Devaev
2024-10-22 00:51:03 +03:00
parent cda32a083f
commit 0e4a70e7b9
6 changed files with 234 additions and 232 deletions
Download Patch File Download Diff File Expand all files Collapse all files

122
kvmd/apps/kvmd/streamer.py
View File

@@ -20,12 +20,10 @@
# ========================================================================== #
import io
import signal
import asyncio
import asyncio.subprocess
import dataclasses
import functools
import copy
from typing import AsyncGenerator
@@ -33,10 +31,12 @@ from typing import Any
import aiohttp
from PIL import Image as PilImage
from ...logging import get_logger
from ...clients.streamer import StreamerSnapshot
from ...clients.streamer import HttpStreamerClient
from ...clients.streamer import HttpStreamerClientSession
from ... import tools
from ... import aiotools
from ... import aioproc
@@ -44,40 +44,6 @@ from ... import htclient
# =====
@dataclasses.dataclass(frozen=True)
class StreamerSnapshot:
online: bool
width: int
height: int
headers: tuple[tuple[str, str], ...]
data: bytes
async def make_preview(self, max_width: int, max_height: int, quality: int) -> bytes:
assert max_width >= 0
assert max_height >= 0
assert quality > 0
if max_width == 0 and max_height == 0:
max_width = self.width // 5
max_height = self.height // 5
else:
max_width = min((max_width or self.width), self.width)
max_height = min((max_height or self.height), self.height)
if (max_width, max_height) == (self.width, self.height):
return self.data
return (await aiotools.run_async(self.__inner_make_preview, max_width, max_height, quality))
@functools.lru_cache(maxsize=1)
def __inner_make_preview(self, max_width: int, max_height: int, quality: int) -> bytes:
with io.BytesIO(self.data) as snapshot_bio:
with io.BytesIO() as preview_bio:
with PilImage.open(snapshot_bio) as image:
image.thumbnail((max_width, max_height), PilImage.Resampling.LANCZOS)
image.save(preview_bio, format="jpeg", quality=quality)
return preview_bio.getvalue()
class _StreamerParams:
__DESIRED_FPS = "desired_fps"
@@ -204,7 +170,6 @@ class Streamer: # pylint: disable=too-many-instance-attributes
self.__state_poll = state_poll
self.__unix_path = unix_path
self.__timeout = timeout
self.__snapshot_timeout = snapshot_timeout
self.__process_name_prefix = process_name_prefix
@@ -221,7 +186,13 @@ class Streamer: # pylint: disable=too-many-instance-attributes
self.__streamer_task: (asyncio.Task | None) = None
self.__streamer_proc: (asyncio.subprocess.Process | None) = None # pylint: disable=no-member
self.__http_session: (aiohttp.ClientSession | None) = None
self.__client = HttpStreamerClient(
name="jpeg",
unix_path=self.__unix_path,
timeout=timeout,
user_agent=htclient.make_user_agent("KVMD"),
)
self.__client_session: (HttpStreamerClientSession | None) = None
self.__snapshot: (StreamerSnapshot | None) = None
@@ -300,11 +271,9 @@ class Streamer: # pylint: disable=too-many-instance-attributes
async def get_state(self) -> dict:
streamer_state = None
if self.__streamer_task:
session = self.__ensure_http_session()
session = self.__ensure_client_session()
try:
async with session.get(self.__make_url("state")) as response:
htclient.raise_not_200(response)
streamer_state = (await response.json())["result"]
streamer_state = await session.get_state()
except (aiohttp.ClientConnectionError, aiohttp.ServerConnectionError):
pass
except Exception:
@@ -350,39 +319,15 @@ class Streamer: # pylint: disable=too-many-instance-attributes
if load:
return self.__snapshot
logger = get_logger()
session = self.__ensure_http_session()
session = self.__ensure_client_session()
try:
async with session.get(
self.__make_url("snapshot"),
timeout=aiohttp.ClientTimeout(total=self.__snapshot_timeout),
) as response:
htclient.raise_not_200(response)
online = (response.headers["X-UStreamer-Online"] == "true")
if online or allow_offline:
snapshot = StreamerSnapshot(
online=online,
width=int(response.headers["X-UStreamer-Width"]),
height=int(response.headers["X-UStreamer-Height"]),
headers=tuple(
(key, value)
for (key, value) in tools.sorted_kvs(dict(response.headers))
if key.lower().startswith("x-ustreamer-") or key.lower() in [
"x-timestamp",
"access-control-allow-origin",
"cache-control",
"pragma",
"expires",
]
),
data=bytes(await response.read()),
)
if save:
self.__snapshot = snapshot
self.__notifier.notify()
return snapshot
logger.error("Stream is offline, no signal or so")
snapshot = await session.take_snapshot(self.__snapshot_timeout)
if snapshot.online or allow_offline:
if save:
self.__snapshot = snapshot
self.__notifier.notify()
return snapshot
logger.error("Stream is offline, no signal or so")
except (aiohttp.ClientConnectionError, aiohttp.ServerConnectionError) as ex:
logger.error("Can't connect to streamer: %s", tools.efmt(ex))
except Exception:
@@ -397,25 +342,14 @@ class Streamer: # pylint: disable=too-many-instance-attributes
@aiotools.atomic_fg
async def cleanup(self) -> None:
await self.ensure_stop(immediately=True)
if self.__http_session:
await self.__http_session.close()
self.__http_session = None
if self.__client_session:
await self.__client_session.close()
self.__client_session = None
# =====
def __ensure_http_session(self) -> aiohttp.ClientSession:
if not self.__http_session:
kwargs: dict = {
"headers": {"User-Agent": htclient.make_user_agent("KVMD")},
"connector": aiohttp.UnixConnector(path=self.__unix_path),
"timeout": aiohttp.ClientTimeout(total=self.__timeout),
}
self.__http_session = aiohttp.ClientSession(**kwargs)
return self.__http_session
def __make_url(self, handle: str) -> str:
assert not handle.startswith("/"), handle
return f"http://localhost:0/{handle}"
def __ensure_client_session(self) -> HttpStreamerClientSession:
if not self.__client_session:
self.__client_session = self.__client.make_session()
return self.__client_session
# =====

6
kvmd/apps/vnc/__init__.py
View File

@@ -21,7 +21,7 @@
from ...clients.kvmd import KvmdClient
from ...clients.streamer import StreamFormats
from ...clients.streamer import StreamerFormats
from ...clients.streamer import BaseStreamerClient
from ...clients.streamer import HttpStreamerClient
from ...clients.streamer import MemsinkStreamerClient
@@ -51,8 +51,8 @@ def main(argv: (list[str] | None)=None) -> None:
return None
streamers: list[BaseStreamerClient] = list(filter(None, [
make_memsink_streamer("h264", StreamFormats.H264),
make_memsink_streamer("jpeg", StreamFormats.JPEG),
make_memsink_streamer("h264", StreamerFormats.H264),
make_memsink_streamer("jpeg", StreamerFormats.JPEG),
HttpStreamerClient(name="JPEG", user_agent=user_agent, **config.streamer._unpack()),
]))

20
kvmd/apps/vnc/server.py
View File

@@ -42,7 +42,7 @@ from ...clients.kvmd import KvmdClient
from ...clients.streamer import StreamerError
from ...clients.streamer import StreamerPermError
from ...clients.streamer import StreamFormats
from ...clients.streamer import StreamerFormats
from ...clients.streamer import BaseStreamerClient
from ... import tools
@@ -222,8 +222,8 @@ class _Client(RfbClient): # pylint: disable=too-many-instance-attributes
def __get_preferred_streamer(self) -> BaseStreamerClient:
formats = {
StreamFormats.JPEG: "has_tight",
StreamFormats.H264: "has_h264",
StreamerFormats.JPEG: "has_tight",
StreamerFormats.H264: "has_h264",
}
streamer: (BaseStreamerClient | None) = None
for streamer in self.__streamers:
@@ -249,7 +249,7 @@ class _Client(RfbClient): # pylint: disable=too-many-instance-attributes
"data": (await make_text_jpeg(self._width, self._height, self._encodings.tight_jpeg_quality, text)),
"width": self._width,
"height": self._height,
"format": StreamFormats.JPEG,
"format": StreamerFormats.JPEG,
}
async def __fb_sender_task_loop(self) -> None: # pylint: disable=too-many-branches
@@ -259,21 +259,21 @@ class _Client(RfbClient): # pylint: disable=too-many-instance-attributes
frame = await self.__fb_queue.get()
if (
last is None # pylint: disable=too-many-boolean-expressions
or frame["format"] == StreamFormats.JPEG
or frame["format"] == StreamerFormats.JPEG
or last["format"] != frame["format"]
or (frame["format"] == StreamFormats.H264 and (
or (frame["format"] == StreamerFormats.H264 and (
frame["key"]
or last["width"] != frame["width"]
or last["height"] != frame["height"]
or len(last["data"]) + len(frame["data"]) > 4194304
))
):
self.__fb_has_key = (frame["format"] == StreamFormats.H264 and frame["key"])
self.__fb_has_key = (frame["format"] == StreamerFormats.H264 and frame["key"])
last = frame
if self.__fb_queue.qsize() == 0:
break
continue
assert frame["format"] == StreamFormats.H264
assert frame["format"] == StreamerFormats.H264
last["data"] += frame["data"]
if self.__fb_queue.qsize() == 0:
break
@@ -295,9 +295,9 @@ class _Client(RfbClient): # pylint: disable=too-many-instance-attributes
await self._send_fb_allow_again()
continue
if last["format"] == StreamFormats.JPEG:
if last["format"] == StreamerFormats.JPEG:
await self._send_fb_jpeg(last["data"])
elif last["format"] == StreamFormats.H264:
elif last["format"] == StreamerFormats.H264:
if not self._encodings.has_h264:
raise RfbError("The client doesn't want to accept H264 anymore")
if self.__fb_has_key: