shielded some tasks

5 files changed, 22 insertions, 15 deletions

diff --git a/kvmd/apps/janus/stun.py b/kvmd/apps/janus/stun.py
index a71cbd43..dc41b961 100644
--- a/kvmd/apps/janus/stun.py
+++ b/kvmd/apps/janus/stun.py
@@ -63,7 +63,6 @@ class Stun:
         self.__sock: Optional[socket.socket] = None
 
     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]
         try:
             with socket.socket(family, socket.SOCK_DGRAM) as self.__sock:
diff --git a/kvmd/apps/kvmd/info/extras.py b/kvmd/apps/kvmd/info/extras.py
index 37bf3ff7..3a1721c8 100644
--- a/kvmd/apps/kvmd/info/extras.py
+++ b/kvmd/apps/kvmd/info/extras.py
@@ -66,7 +66,7 @@ class ExtrasInfoSubmanager(BaseInfoSubmanager):
             return None
         finally:
             if sui is not None:
-                await sui.close()
+                await asyncio.shield(sui.close())
 
     def __get_extras_path(self, *parts: str) -> str:
         return os.path.join(self.__global_config.kvmd.info.extras, *parts)
diff --git a/kvmd/apps/kvmd/snapshoter.py b/kvmd/apps/kvmd/snapshoter.py
index 65c59eaa..0a8874a1 100644
--- a/kvmd/apps/kvmd/snapshoter.py
+++ b/kvmd/apps/kvmd/snapshoter.py
@@ -116,7 +116,7 @@ class Snapshoter:  # pylint: disable=too-many-instance-attributes
             logger.exception("Unhandled exception while taking snapshot")
         finally:
             self.__snapshoting = False
-            await notifier.notify()
+            await asyncio.shield(notifier.notify())
 
     async def __wakeup(self) -> None:
         logger = get_logger(0)
diff --git a/kvmd/apps/vnc/rfb/__init__.py b/kvmd/apps/vnc/rfb/__init__.py
index 5d7847a9..172d52ba 100644
--- a/kvmd/apps/vnc/rfb/__init__.py
+++ b/kvmd/apps/vnc/rfb/__init__.py
@@ -93,8 +93,7 @@ class RfbClient(RfbClientStream):  # pylint: disable=too-many-instance-attribute
     # =====
 
     async def _run(self, **coros: Coroutine) -> None:
-        logger = get_logger(0)
-        logger.info("%s [entry]: Starting client tasks ...", self._remote)
+        get_logger(0).info("%s [entry]: Starting client tasks ...", self._remote)
         tasks = list(map(asyncio.create_task, [  # type: ignore  # Я хз, почему github action фейлится здесь
             self.__wrapper(name, coro)
             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:
             await aiotools.wait_first(*tasks)
         finally:
-            for task in tasks:
-                task.cancel()
-            await asyncio.gather(*tasks, return_exceptions=True)
-            await self._close()
-            logger.info("%s [entry]: Connection closed", self._remote)
+            await asyncio.shield(self.__cleanup(tasks))
+
+    async def __cleanup(self, tasks: List[asyncio.Task]) -> None:
+        for task in tasks:
+            task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+        await self._close()
+        get_logger(0).info("%s [entry]: Connection closed", self._remote)
 
     async def __wrapper(self, name: str, coro: Coroutine) -> None:
         logger = get_logger(0)
diff --git a/kvmd/apps/vnc/server.py b/kvmd/apps/vnc/server.py
index 2f822ab9..54974bfb 100644
--- a/kvmd/apps/vnc/server.py
+++ b/kvmd/apps/vnc/server.py
@@ -147,9 +147,12 @@ class _Client(RfbClient):  # pylint: disable=too-many-instance-attributes
                 fb_sender=self.__fb_sender_task_loop(),
             )
         finally:
-            if self.__kvmd_session:
-                await self.__kvmd_session.close()
-                self.__kvmd_session = None
+            await asyncio.shield(self.__cleanup())
+
+    async def __cleanup(self) -> None:
+        if self.__kvmd_session:
+            await self.__kvmd_session.close()
+            self.__kvmd_session = None
 
     # =====
 
@@ -446,6 +449,10 @@ class VncServer:  # pylint: disable=too-many-instance-attributes
 
         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:
             logger = get_logger(0)
             remote = rfb_format_remote(writer)
@@ -491,8 +498,7 @@ class VncServer:  # pylint: disable=too-many-instance-attributes
             except Exception:
                 logger.exception("%s [entry]: Unhandled exception in client task", remote)
             finally:
-                if (await aiotools.close_writer(writer)):
-                    logger.info("%s [entry]: Connection is closed in an emergency", remote)
+                await asyncio.shield(cleanup_client(writer))
 
         self.__handle_client = handle_client