GithubFollow/One-KVM
1
0
Fork 0
mirror of https://github.com/mofeng-git/One-KVM.git synced 2025-12-12 09:10:30 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

many fixes with asyncio

Browse Source
This commit is contained in:
Devaev Maxim 2019-06-19 04:15:43 +03:00
parent 376ab295bd
commit c7a2e445d0
7 changed files with 163 additions and 78 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

67
kvmd/aiotools.py
View File

@ -22,24 +22,50 @@
import asyncio import asyncio
import functools import functools
import contextlib
import typing import typing
from typing import List
from typing import Callable from typing import Callable
from typing import Coroutine from typing import Coroutine
from typing import Generator
from typing import AsyncGenerator
from typing import TypeVar from typing import TypeVar
from typing import Any from typing import Any
from . import aioregion
from .logging import get_logger
# ===== # =====
_AtomicF = TypeVar("_AtomicF", bound=Callable[..., Any]) _ATTR_SHORT_TASK = "_aiotools_short_task"
_MethodT = TypeVar("_MethodT", bound=Callable[..., Any])
_RetvalT = TypeVar("_RetvalT")
def atomic(method: _AtomicF) -> _AtomicF: # =====
def atomic(method: _MethodT) -> _MethodT:
@functools.wraps(method) @functools.wraps(method)
async def wrapper(*args: Any, **kwargs: Any) -> Any: async def wrapper(*args: Any, **kwargs: Any) -> Any:
return (await asyncio.shield(method(*args, **kwargs))) return (await asyncio.shield(method(*args, **kwargs)))
return typing.cast(_AtomicF, wrapper) return typing.cast(_MethodT, wrapper)
def muted(msg: str) -> Callable[[_MethodT], Callable[..., None]]:
def make_wrapper(method: _MethodT) -> Callable[..., None]:
@functools.wraps(method)
async def wrapper(*args: Any, **kwargs: Any) -> None:
try:
await method(*args, **kwargs)
except asyncio.CancelledError: # pylint: disable=try-except-raise
raise
except Exception:
get_logger(0).exception(msg)
return typing.cast(Callable[..., None], wrapper)
return make_wrapper
def tasked(method: Callable[..., Any]) -> Callable[..., asyncio.Task]: def tasked(method: Callable[..., Any]) -> Callable[..., asyncio.Task]:
@ -49,29 +75,46 @@ def tasked(method: Callable[..., Any]) -> Callable[..., asyncio.Task]:
return typing.cast(Callable[..., asyncio.Task], wrapper) return typing.cast(Callable[..., asyncio.Task], wrapper)
_ATTR_SHORT_TASK = "_aiotools_short_task" # =====
def create_short_task(coro: Coroutine) -> asyncio.Task: def create_short_task(coro: Coroutine) -> asyncio.Task:
task = asyncio.create_task(coro) task = asyncio.create_task(coro)
setattr(task, _ATTR_SHORT_TASK, True) setattr(task, _ATTR_SHORT_TASK, True)
return task return task
async def gather_short_tasks() -> None: def get_short_tasks() -> List[asyncio.Task]:
await asyncio.gather(*[ return [
task task
for task in asyncio.Task.all_tasks() for task in asyncio.Task.all_tasks()
if getattr(task, _ATTR_SHORT_TASK, False) if getattr(task, _ATTR_SHORT_TASK, False)
]) ]
_RetvalT = TypeVar("_RetvalT")
# =====
async def run_async(method: Callable[..., _RetvalT], *args: Any) -> _RetvalT: async def run_async(method: Callable[..., _RetvalT], *args: Any) -> _RetvalT:
return (await asyncio.get_running_loop().run_in_executor(None, method, *args)) return (await asyncio.get_running_loop().run_in_executor(None, method, *args))
def run_sync(coro: Coroutine[Any, Any, _RetvalT]) -> _RetvalT: def run_sync(coro: Coroutine[Any, Any, _RetvalT]) -> _RetvalT:
return asyncio.get_event_loop().run_until_complete(coro) return asyncio.get_event_loop().run_until_complete(coro)
# =====
@contextlib.contextmanager
def unregion_only_on_exception(region: aioregion.AioExclusiveRegion) -> Generator[None, None, None]:
region.enter()
try:
yield
except: # noqa: E722
region.exit()
raise
@contextlib.asynccontextmanager
async def unlock_only_on_exception(lock: asyncio.Lock) -> AsyncGenerator[None, None]:
await lock.acquire()
try:
yield
except: # noqa: E722
lock.release()
raise

2
kvmd/apps/kvmd/__init__.py
View File

@ -67,4 +67,4 @@ def main(argv: Optional[List[str]]=None) -> None:
streamer=Streamer(**config.streamer._unpack()), streamer=Streamer(**config.streamer._unpack()),
).run(**config.server._unpack()) ).run(**config.server._unpack())
get_logger().info("Bye-bye") get_logger(0).info("Bye-bye")

42
kvmd/apps/kvmd/atx.py
View File

@ -129,6 +129,16 @@ class Atx: # pylint: disable=too-many-instance-attributes
else: else:
await asyncio.sleep(60) await asyncio.sleep(60)
async def cleanup(self) -> None:
for (name, pin) in [
("power", self.__power_switch_pin),
("reset", self.__reset_switch_pin),
]:
try:
gpio.write(pin, False)
except Exception:
get_logger(0).exception("Can't cleanup %s pin %d", name, pin)
# ===== # =====
@_atx_working @_atx_working
@ -163,25 +173,33 @@ class Atx: # pylint: disable=too-many-instance-attributes
@_atx_working @_atx_working
async def click_power(self) -> None: async def click_power(self) -> None:
get_logger().info("Clicking power ...") await self.__click("power", self.__power_switch_pin, self.__click_delay)
await self.__click(self.__power_switch_pin, self.__click_delay)
@_atx_working @_atx_working
async def click_power_long(self) -> None: async def click_power_long(self) -> None:
get_logger().info("Clicking power (long press) ...") await self.__click("power_long", self.__power_switch_pin, self.__long_click_delay)
await self.__click(self.__power_switch_pin, self.__long_click_delay)
@_atx_working @_atx_working
async def click_reset(self) -> None: async def click_reset(self) -> None:
get_logger().info("Clicking reset") await self.__click("reset", self.__reset_switch_pin, self.__click_delay)
await self.__click(self.__reset_switch_pin, self.__click_delay)
# ===== # =====
@aiotools.tasked
@aiotools.atomic @aiotools.atomic
async def __click(self, pin: int, delay: float) -> None: async def __click(self, name: str, pin: int, delay: float) -> None:
with self.__region: with aiotools.unregion_only_on_exception(self.__region):
for flag in [True, False]: await self.__inner_click(name, pin, delay)
gpio.write(pin, flag)
await asyncio.sleep(delay) @aiotools.tasked
@aiotools.muted("Can't perform ATX click or operation was not completed")
async def __inner_click(self, name: str, pin: int, delay: float) -> None:
try:
gpio.write(pin, True)
await asyncio.sleep(delay)
finally:
try:
gpio.write(pin, False)
await asyncio.sleep(1)
finally:
self.__region.exit()
get_logger(0).info("Clicked ATX button %r", name)

40
kvmd/apps/kvmd/hid.py
View File

@ -169,13 +169,24 @@ class Hid(multiprocessing.Process): # pylint: disable=too-many-instance-attribu
prev_state = state prev_state = state
await asyncio.sleep(self.__state_poll) await asyncio.sleep(self.__state_poll)
@aiotools.tasked
@aiotools.atomic @aiotools.atomic
async def reset(self) -> None: async def reset(self) -> None:
async with self.__lock: async with aiotools.unlock_only_on_exception(self.__lock):
await self.__inner_reset()
@aiotools.tasked
@aiotools.muted("Can't reset HID or operation was not completed")
async def __inner_reset(self) -> None:
try:
gpio.write(self.__reset_pin, True) gpio.write(self.__reset_pin, True)
await asyncio.sleep(self.__reset_delay) await asyncio.sleep(self.__reset_delay)
gpio.write(self.__reset_pin, False) finally:
try:
gpio.write(self.__reset_pin, False)
await asyncio.sleep(1)
finally:
self.__lock.release()
get_logger(0).info("Reset HID performed")
async def send_key_event(self, key: str, state: bool) -> None: async def send_key_event(self, key: str, state: bool) -> None:
await self.__send_bool_event(_KeyEvent(key, state), self.__pressed_keys) await self.__send_bool_event(_KeyEvent(key, state), self.__pressed_keys)
@ -196,17 +207,20 @@ class Hid(multiprocessing.Process): # pylint: disable=too-many-instance-attribu
@aiotools.atomic @aiotools.atomic
async def cleanup(self) -> None: async def cleanup(self) -> None:
logger = get_logger(0)
async with self.__lock: async with self.__lock:
if self.is_alive(): try:
self.__unsafe_clear_events() if self.is_alive():
get_logger(0).info("Stopping HID daemon ...") self.__unsafe_clear_events()
self.__stop_event.set() logger.info("Stopping HID daemon ...")
else: self.__stop_event.set()
get_logger(0).warning("Emergency cleaning up HID events ...") else:
self.__emergency_clear_events() logger.warning("Emergency cleaning up HID events ...")
if self.exitcode is not None: self.__emergency_clear_events()
self.join() if self.exitcode is not None:
gpio.write(self.__reset_pin, False) self.join()
finally:
gpio.write(self.__reset_pin, False)
async def __send_bool_event(self, event: _BoolEvent, pressed: Set[str]) -> None: async def __send_bool_event(self, event: _BoolEvent, pressed: Set[str]) -> None:
if not self.__stop_event.is_set(): if not self.__stop_event.is_set():

62
kvmd/apps/kvmd/msd.py
View File

@ -58,27 +58,27 @@ class MsdOperationError(MsdError):
class MsdDisabledError(MsdOperationError): class MsdDisabledError(MsdOperationError):
def __init__(self) -> None: def __init__(self) -> None:
super().__init__("Mass-storage device is disabled") super().__init__("MSD is disabled")
class MsdOfflineError(MsdOperationError): class MsdOfflineError(MsdOperationError):
def __init__(self) -> None: def __init__(self) -> None:
super().__init__("Mass-storage device is not found") super().__init__("MSD is not found")
class MsdAlreadyOnServerError(MsdOperationError): class MsdAlreadyOnServerError(MsdOperationError):
def __init__(self) -> None: def __init__(self) -> None:
super().__init__("Mass-storage is already connected to Server") super().__init__("MSD is already connected to Server")
class MsdAlreadyOnKvmError(MsdOperationError): class MsdAlreadyOnKvmError(MsdOperationError):
def __init__(self) -> None: def __init__(self) -> None:
super().__init__("Mass-storage is already connected to KVM") super().__init__("MSD is already connected to KVM")
class MsdNotOnKvmError(MsdOperationError): class MsdNotOnKvmError(MsdOperationError):
def __init__(self) -> None: def __init__(self) -> None:
super().__init__("Mass-storage is not connected to KVM") super().__init__("MSD is not connected to KVM")
class MsdIsBusyError(MsdOperationError, aioregion.RegionIsBusyError): class MsdIsBusyError(MsdOperationError, aioregion.RegionIsBusyError):
@ -226,16 +226,16 @@ class MassStorageDevice: # pylint: disable=too-many-instance-attributes
logger = get_logger(0) logger = get_logger(0)
if self._enabled: if self._enabled:
logger.info("Using %r as mass-storage device", self.__device_path) logger.info("Using %r as MSD", self.__device_path)
try: try:
aiotools.run_sync(self.__load_device_info()) aiotools.run_sync(self.__load_device_info())
if self.__write_meta: if self.__write_meta:
logger.info("Enabled image metadata writing") logger.info("Enabled image metadata writing")
except Exception as err: except Exception as err:
log = (logger.error if isinstance(err, MsdError) else logger.exception) log = (logger.error if isinstance(err, MsdError) else logger.exception)
log("Mass-storage device is offline: %s", err) log("MSD is offline: %s", err)
else: else:
logger.info("Mass-storage device is disabled") logger.info("MSD is disabled")
def get_state(self) -> Dict: def get_state(self) -> Dict:
online = (self._enabled and bool(self._device_info)) online = (self._enabled and bool(self._device_info))
@ -282,7 +282,7 @@ class MassStorageDevice: # pylint: disable=too-many-instance-attributes
gpio.write(self.__target_pin, True) gpio.write(self.__target_pin, True)
raise raise
self.__on_kvm = True self.__on_kvm = True
get_logger().info("Mass-storage device switched to KVM: %s", self._device_info) get_logger().info("MSD switched to KVM: %s", self._device_info)
state = self.get_state() state = self.get_state()
return state return state
@ -303,7 +303,7 @@ class MassStorageDevice: # pylint: disable=too-many-instance-attributes
gpio.write(self.__target_pin, True) gpio.write(self.__target_pin, True)
self.__on_kvm = False self.__on_kvm = False
get_logger().info("Mass-storage device switched to Server") get_logger().info("MSD switched to Server")
state = self.get_state() state = self.get_state()
return state return state
@ -311,28 +311,34 @@ class MassStorageDevice: # pylint: disable=too-many-instance-attributes
if notify: if notify:
await self.__state_queue.put(state or self.get_state()) await self.__state_queue.put(state or self.get_state())
@aiotools.tasked
@aiotools.atomic @aiotools.atomic
async def reset(self) -> None: async def reset(self) -> None:
notify = False if not self._enabled:
raise MsdDisabledError()
with aiotools.unregion_only_on_exception(self.__region):
await self.__inner_reset()
@aiotools.tasked
@aiotools.muted("Can't reset MSD or operation was not completed")
async def __inner_reset(self) -> None:
try: try:
with self.__region: gpio.write(self.__reset_pin, True)
if not self._enabled: await asyncio.sleep(self.__reset_delay)
raise MsdDisabledError() gpio.write(self.__reset_pin, False)
notify = True
gpio.write(self.__reset_pin, True) gpio.write(self.__target_pin, False)
await asyncio.sleep(self.__reset_delay) self.__on_kvm = True
gpio.write(self.__target_pin, False)
self.__on_kvm = True
await asyncio.sleep(self.__reset_delay)
gpio.write(self.__reset_pin, False)
await self.__load_device_info() await self.__load_device_info()
get_logger(0).info("Mass-storage device reset has been successful") get_logger(0).info("MSD reset has been successful")
finally: finally:
if notify: try:
await self.__state_queue.put(self.get_state()) gpio.write(self.__reset_pin, False)
finally:
try:
await self.__state_queue.put(self.get_state())
finally:
self.__region.exit()
@_msd_working @_msd_working
@aiotools.atomic @aiotools.atomic
@ -392,12 +398,12 @@ class MassStorageDevice: # pylint: disable=too-many-instance-attributes
async def __close_device_file(self) -> None: async def __close_device_file(self) -> None:
try: try:
if self.__device_file: if self.__device_file:
get_logger().info("Closing mass-storage device file ...") get_logger().info("Closing device file ...")
await self.__device_file.close() await self.__device_file.close()
except asyncio.CancelledError: # pylint: disable=try-except-raise except asyncio.CancelledError: # pylint: disable=try-except-raise
raise raise
except Exception: except Exception:
get_logger().exception("Can't close mass-storage device file") get_logger().exception("Can't close device file")
finally: finally:
self.__device_file = None self.__device_file = None
self.__written = 0 self.__written = 0

15
kvmd/apps/kvmd/server.py
View File

@ -480,7 +480,7 @@ class Server: # pylint: disable=too-many-instance-attributes
data_field = await _get_multipart_field(reader, "image_data") data_field = await _get_multipart_field(reader, "image_data")
logger.info("Writing image %r to mass-storage device ...", image_name) logger.info("Writing image %r to MSD ...", image_name)
await self.__msd.write_image_info(image_name, False) await self.__msd.write_image_info(image_name, False)
while True: while True:
chunk = await data_field.read_chunk(self.__msd.chunk_size) chunk = await data_field.read_chunk(self.__msd.chunk_size)
@ -490,7 +490,7 @@ class Server: # pylint: disable=too-many-instance-attributes
await self.__msd.write_image_info(image_name, True) await self.__msd.write_image_info(image_name, True)
finally: finally:
if written != 0: if written != 0:
logger.info("Written %d bytes to mass-storage device", written) logger.info("Written %d bytes to MSD", written)
return _json({"written": written}) return _json({"written": written})
@_exposed("POST", "/msd/reset") @_exposed("POST", "/msd/reset")
@ -549,12 +549,14 @@ class Server: # pylint: disable=too-many-instance-attributes
logger = get_logger(0) logger = get_logger(0)
logger.info("Waiting short tasks ...") logger.info("Waiting short tasks ...")
await aiotools.gather_short_tasks() await asyncio.gather(*aiotools.get_short_tasks(), return_exceptions=True)
logger.info("Cancelling system tasks ...") logger.info("Cancelling system tasks ...")
for task in self.__system_tasks: for task in self.__system_tasks:
task.cancel() task.cancel()
await asyncio.gather(*self.__system_tasks)
logger.info("Waiting system tasks ...")
await asyncio.gather(*self.__system_tasks, return_exceptions=True)
logger.info("Disconnecting clients ...") logger.info("Disconnecting clients ...")
for ws in list(self.__sockets): for ws in list(self.__sockets):
@ -566,15 +568,14 @@ class Server: # pylint: disable=too-many-instance-attributes
self._auth_manager, self._auth_manager,
self.__streamer, self.__streamer,
self.__msd, self.__msd,
self.__atx,
self.__hid, self.__hid,
]: ]:
logger.info("Cleaning up %s ...", type(obj).__name__) logger.info("Cleaning up %s ...", type(obj).__name__)
try: try:
await obj.cleanup() # type: ignore await obj.cleanup() # type: ignore
except asyncio.CancelledError: # pylint: disable=try-except-raise
raise
except Exception: except Exception:
logger.exception("Cleanup error") logger.exception("Cleanup error on %s", type(obj).__name__)
async def __broadcast_event(self, event_type: _Events, event_attrs: Dict) -> None: async def __broadcast_event(self, event_type: _Events, event_attrs: Dict) -> None:
if self.__sockets: if self.__sockets:

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

@ -162,11 +162,14 @@ class Streamer: # pylint: disable=too-many-instance-attributes
@aiotools.atomic @aiotools.atomic
async def cleanup(self) -> None: async def cleanup(self) -> None:
if self.is_running(): try:
await self.stop() if self.is_running():
if self.__http_session: await self.stop()
await self.__http_session.close() if self.__http_session:
self.__http_session = None await self.__http_session.close()
self.__http_session = None
finally:
await self.__set_hw_enabled(False)
def __ensure_session(self) -> aiohttp.ClientSession: def __ensure_session(self) -> aiohttp.ClientSession:
if not self.__http_session: if not self.__http_session: