using shield_fg()

7 files changed, 16 insertions, 30 deletions

diff --git a/kvmd/apps/kvmd/info/extras.py b/kvmd/apps/kvmd/info/extras.py
index 3a1721c8..9e70d05c 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 asyncio.shield(sui.close())
+                await aiotools.shield_fg(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/vnc/rfb/__init__.py b/kvmd/apps/vnc/rfb/__init__.py
index 172d52ba..059d4d1d 100644
--- a/kvmd/apps/vnc/rfb/__init__.py
+++ b/kvmd/apps/vnc/rfb/__init__.py
@@ -101,7 +101,7 @@ class RfbClient(RfbClientStream):  # pylint: disable=too-many-instance-attribute
         try:
             await aiotools.wait_first(*tasks)
         finally:
-            await asyncio.shield(self.__cleanup(tasks))
+            await aiotools.shield_fg(self.__cleanup(tasks))
 
     async def __cleanup(self, tasks: List[asyncio.Task]) -> None:
         for task in tasks:
diff --git a/kvmd/apps/vnc/server.py b/kvmd/apps/vnc/server.py
index 151f6374..7fac9aa4 100644
--- a/kvmd/apps/vnc/server.py
+++ b/kvmd/apps/vnc/server.py
@@ -147,7 +147,7 @@ class _Client(RfbClient):  # pylint: disable=too-many-instance-attributes
                 fb_sender=self.__fb_sender_task_loop(),
             )
         finally:
-            await asyncio.shield(self.__cleanup())
+            await aiotools.shield_fg(self.__cleanup())
 
     async def __cleanup(self) -> None:
         if self.__kvmd_session:
@@ -498,7 +498,7 @@ class VncServer:  # pylint: disable=too-many-instance-attributes
             except Exception:
                 logger.exception("%s [entry]: Unhandled exception in client task", remote)
             finally:
-                await asyncio.shield(cleanup_client(writer))
+                await aiotools.shield_fg(cleanup_client(writer))
 
         self.__handle_client = handle_client
 
diff --git a/kvmd/clients/kvmd.py b/kvmd/clients/kvmd.py
index 3f258a3b..f8b0d811 100644
--- a/kvmd/clients/kvmd.py
+++ b/kvmd/clients/kvmd.py
@@ -170,7 +170,7 @@ class KvmdClientWs:
             if writer_task:
                 writer_task.cancel()
             try:
-                await asyncio.shield(self.__ws.close())
+                await aiotools.shield_fg(self.__ws.close())
             except Exception:
                 pass
             finally:
diff --git a/kvmd/htserver.py b/kvmd/htserver.py
index 16d152b3..a83fc17c 100644
--- a/kvmd/htserver.py
+++ b/kvmd/htserver.py
@@ -61,6 +61,8 @@ from .errors import IsBusyError
 
 from .validators import ValidatorError
 
+from . import aiotools
+
 
 # =====
 class HttpError(Exception):
@@ -352,7 +354,7 @@ class HttpServer:
             await self._on_ws_opened()
             yield ws
         finally:
-            await asyncio.shield(self.__close_ws(ws))
+            await aiotools.shield_fg(self.__close_ws(ws))
 
     async def _ws_loop(self, ws: WsSession) -> WebSocketResponse:
         logger = get_logger()
diff --git a/kvmd/plugins/msd/otg/__init__.py b/kvmd/plugins/msd/otg/__init__.py
index 8d1ee8a5..cb781fae 100644
--- a/kvmd/plugins/msd/otg/__init__.py
+++ b/kvmd/plugins/msd/otg/__init__.py
@@ -354,12 +354,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                     yield self.__reader
 
                 finally:
-                    # FIXME: Перехват важен потому что по какой-то причине await
-                    # во вложенных finally путаются и выполняются не по порядку
-                    try:
-                        await asyncio.shield(self.__close_reader())
-                    except asyncio.CancelledError:
-                        pass
+                    await aiotools.shield_fg(self.__close_reader())
         finally:
             self.__notifier.notify()
 
@@ -404,20 +399,13 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                         except Exception:
                             pass
                     try:
-                        await asyncio.shield(self.__close_writer())
-                    except asyncio.CancelledError:
-                        pass
-                    try:
-                        await asyncio.shield(self.__remount_rw(False, fatal=False))
-                    except asyncio.CancelledError:
-                        pass
+                        await aiotools.shield_fg(self.__close_writer())
+                    finally:
+                        await aiotools.shield_fg(self.__remount_rw(False, fatal=False))
         finally:
             # Между закрытием файла и эвентом айнотифи состояние может быть не обновлено,
             # так что форсим обновление вручную, чтобы получить актуальное состояние.
-            try:
-                await asyncio.shield(self.__reload_state())
-            except asyncio.CancelledError:
-                pass
+            await aiotools.shield_fg(self.__reload_state())
 
     @aiotools.atomic
     async def remove(self, name: str) -> None:
diff --git a/kvmd/plugins/msd/relay/__init__.py b/kvmd/plugins/msd/relay/__init__.py
index f6b13dc5..febc3750 100644
--- a/kvmd/plugins/msd/relay/__init__.py
+++ b/kvmd/plugins/msd/relay/__init__.py
@@ -250,13 +250,9 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                     await self.__write_image_info(True)
                 finally:
                     try:
-                        await asyncio.shield(self.__close_device_writer())
-                    except asyncio.CancelledError:
-                        pass
-                    try:
-                        await asyncio.shield(self.__load_device_info())
-                    except asyncio.CancelledError:
-                        pass
+                        await aiotools.shield_fg(self.__close_device_writer())
+                    finally:
+                        await aiotools.shield_fg(self.__load_device_info())
 
     @aiotools.atomic
     async def remove(self, name: str) -> None: