shielded some tasks

This commit is contained in:
Maxim Devaev 2022-08-05 18:19:31 +03:00
parent eeaeebf7c7
commit d21e74700a
8 changed files with 30 additions and 20 deletions

View File

@ -63,7 +63,6 @@ class Stun:
self.__sock: Optional[socket.socket] = None self.__sock: Optional[socket.socket] = None
async def get_info(self, src_ip: str, src_port: int) -> Tuple[str, str]: async def get_info(self, src_ip: str, src_port: int) -> Tuple[str, str]:
(family, _, _, _, addr) = socket.getaddrinfo(src_ip, src_port, type=socket.SOCK_DGRAM)[0] (family, _, _, _, addr) = socket.getaddrinfo(src_ip, src_port, type=socket.SOCK_DGRAM)[0]
try: try:
with socket.socket(family, socket.SOCK_DGRAM) as self.__sock: with socket.socket(family, socket.SOCK_DGRAM) as self.__sock:

View File

@ -66,7 +66,7 @@ class ExtrasInfoSubmanager(BaseInfoSubmanager):
return None return None
finally: finally:
if sui is not None: if sui is not None:
await sui.close() await asyncio.shield(sui.close())
def __get_extras_path(self, *parts: str) -> str: def __get_extras_path(self, *parts: str) -> str:
return os.path.join(self.__global_config.kvmd.info.extras, *parts) return os.path.join(self.__global_config.kvmd.info.extras, *parts)

View File

@ -116,7 +116,7 @@ class Snapshoter: # pylint: disable=too-many-instance-attributes
logger.exception("Unhandled exception while taking snapshot") logger.exception("Unhandled exception while taking snapshot")
finally: finally:
self.__snapshoting = False self.__snapshoting = False
await notifier.notify() await asyncio.shield(notifier.notify())
async def __wakeup(self) -> None: async def __wakeup(self) -> None:
logger = get_logger(0) logger = get_logger(0)

View File

@ -93,8 +93,7 @@ class RfbClient(RfbClientStream): # pylint: disable=too-many-instance-attribute
# ===== # =====
async def _run(self, **coros: Coroutine) -> None: async def _run(self, **coros: Coroutine) -> None:
logger = get_logger(0) get_logger(0).info("%s [entry]: Starting client tasks ...", self._remote)
logger.info("%s [entry]: Starting client tasks ...", self._remote)
tasks = list(map(asyncio.create_task, [ # type: ignore # Я хз, почему github action фейлится здесь tasks = list(map(asyncio.create_task, [ # type: ignore # Я хз, почему github action фейлится здесь
self.__wrapper(name, coro) self.__wrapper(name, coro)
for (name, coro) in {"main": self.__main_task_loop(), **coros}.items() for (name, coro) in {"main": self.__main_task_loop(), **coros}.items()
@ -102,11 +101,14 @@ class RfbClient(RfbClientStream): # pylint: disable=too-many-instance-attribute
try: try:
await aiotools.wait_first(*tasks) await aiotools.wait_first(*tasks)
finally: finally:
await asyncio.shield(self.__cleanup(tasks))
async def __cleanup(self, tasks: List[asyncio.Task]) -> None:
for task in tasks: for task in tasks:
task.cancel() task.cancel()
await asyncio.gather(*tasks, return_exceptions=True) await asyncio.gather(*tasks, return_exceptions=True)
await self._close() await self._close()
logger.info("%s [entry]: Connection closed", self._remote) get_logger(0).info("%s [entry]: Connection closed", self._remote)
async def __wrapper(self, name: str, coro: Coroutine) -> None: async def __wrapper(self, name: str, coro: Coroutine) -> None:
logger = get_logger(0) logger = get_logger(0)

View File

@ -147,6 +147,9 @@ class _Client(RfbClient): # pylint: disable=too-many-instance-attributes
fb_sender=self.__fb_sender_task_loop(), fb_sender=self.__fb_sender_task_loop(),
) )
finally: finally:
await asyncio.shield(self.__cleanup())
async def __cleanup(self) -> None:
if self.__kvmd_session: if self.__kvmd_session:
await self.__kvmd_session.close() await self.__kvmd_session.close()
self.__kvmd_session = None self.__kvmd_session = None
@ -446,6 +449,10 @@ class VncServer: # pylint: disable=too-many-instance-attributes
shared_params = _SharedParams() shared_params = _SharedParams()
async def cleanup_client(writer: asyncio.StreamWriter) -> None:
if (await aiotools.close_writer(writer)):
get_logger(0).info("%s [entry]: Connection is closed in an emergency", rfb_format_remote(writer))
async def handle_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: async def handle_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
logger = get_logger(0) logger = get_logger(0)
remote = rfb_format_remote(writer) remote = rfb_format_remote(writer)
@ -491,8 +498,7 @@ class VncServer: # pylint: disable=too-many-instance-attributes
except Exception: except Exception:
logger.exception("%s [entry]: Unhandled exception in client task", remote) logger.exception("%s [entry]: Unhandled exception in client task", remote)
finally: finally:
if (await aiotools.close_writer(writer)): await asyncio.shield(cleanup_client(writer))
logger.info("%s [entry]: Connection is closed in an emergency", remote)
self.__handle_client = handle_client self.__handle_client = handle_client

View File

@ -170,9 +170,10 @@ class KvmdClientWs:
if writer_task: if writer_task:
writer_task.cancel() writer_task.cancel()
try: try:
await self.__ws.close() await asyncio.shield(self.__ws.close())
except Exception: except Exception:
pass pass
finally:
self.__communicated = False self.__communicated = False
async def send_key_event(self, key: str, state: bool) -> None: async def send_key_event(self, key: str, state: bool) -> None:

View File

@ -352,7 +352,7 @@ class HttpServer:
await self._on_ws_opened() await self._on_ws_opened()
yield ws yield ws
finally: finally:
await self.__close_ws(ws) await asyncio.shield(self.__close_ws(ws))
async def _ws_loop(self, ws: WsSession) -> WebSocketResponse: async def _ws_loop(self, ws: WsSession) -> WebSocketResponse:
logger = get_logger() logger = get_logger()

View File

@ -123,7 +123,9 @@ class Plugin(BaseUserGpioDriver):
else: else:
os.unlink(os.path.join(self.__profile_path, pin)) os.unlink(os.path.join(self.__profile_path, pin))
finally: finally:
try:
await asyncio.sleep(self.__init_delay) await asyncio.sleep(self.__init_delay)
finally:
self.__set_udc_enabled(True) self.__set_udc_enabled(True)
def __set_udc_enabled(self, enabled: bool) -> None: def __set_udc_enabled(self, enabled: bool) -> None: