partially asynchronous msd

3 files changed, 35 insertions, 28 deletions

diff --git a/kvmd/plugins/msd/__init__.py b/kvmd/plugins/msd/__init__.py
index 34efaafe..65da192c 100644
--- a/kvmd/plugins/msd/__init__.py
+++ b/kvmd/plugins/msd/__init__.py
@@ -27,6 +27,7 @@ import time
 from typing import AsyncGenerator
 
 import aiofiles
+import aiofiles.os
 import aiofiles.base
 
 from ...logging import get_logger
@@ -207,7 +208,7 @@ class MsdFileReader(BaseMsdReader):  # pylint: disable=too-many-instance-attribu
     async def open(self) -> "MsdFileReader":
         assert self.__file is None
         get_logger(1).info("Reading %r image from MSD ...", self.__name)
-        self.__file_size = os.stat(self.__path).st_size
+        self.__file_size = (await aiofiles.os.stat(self.__path)).st_size
         self.__file = await aiofiles.open(self.__path, mode="rb")  # type: ignore
         return self
 
@@ -269,7 +270,7 @@ class MsdFileWriter(BaseMsdWriter):  # pylint: disable=too-many-instance-attribu
     async def open(self) -> "MsdFileWriter":
         assert self.__file is None
         get_logger(1).info("Writing %r image (%d bytes) to MSD ...", self.__name, self.__file_size)
-        os.makedirs(os.path.dirname(self.__path), exist_ok=True)
+        await aiofiles.os.makedirs(os.path.dirname(self.__path), exist_ok=True)
         self.__file = await aiofiles.open(self.__path, mode="w+b", buffering=0)  # type: ignore
         return self
 
diff --git a/kvmd/plugins/msd/otg/__init__.py b/kvmd/plugins/msd/otg/__init__.py
index d1a9dc6a..51b9ca75 100644
--- a/kvmd/plugins/msd/otg/__init__.py
+++ b/kvmd/plugins/msd/otg/__init__.py
@@ -254,7 +254,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
             if name is not None:
                 if name:
-                    self.__state.vd.image = self.__STATE_get_storage_image(name)
+                    self.__state.vd.image = await self.__STATE_get_storage_image(name)
                 else:
                     self.__state.vd.image = None
 
@@ -278,7 +278,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                 if self.__state.vd.image is None:
                     raise MsdImageNotSelected()
 
-                if not self.__state.vd.image.exists():
+                if not (await self.__state.vd.image.exists()):
                     raise MsdUnknownImageError()
 
                 assert self.__state.vd.image.in_storage
@@ -304,7 +304,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                     async with self.__state._lock:  # pylint: disable=protected-access
                         self.__notifier.notify()
                         self.__STATE_check_disconnected()
-                        image = self.__STATE_get_storage_image(name)
+                        image = await self.__STATE_get_storage_image(name)
                         self.__reader = await MsdFileReader(
                             notifier=self.__notifier,
                             name=image.name,
@@ -326,10 +326,10 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                     async with self.__state._lock:  # pylint: disable=protected-access
                         self.__notifier.notify()
                         self.__STATE_check_disconnected()
-                        image = self.__STORAGE_create_new_image(name)
+                        image = await self.__STORAGE_create_new_image(name)
 
                         await image.remount_rw(True)
-                        image.set_complete(False)
+                        await image.set_complete(False)
 
                         self.__writer = await MsdFileWriter(
                             notifier=self.__notifier,
@@ -342,16 +342,18 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
                     self.__notifier.notify()
                     yield self.__writer
-                    image.set_complete(self.__writer.is_complete())
+                    await image.set_complete(self.__writer.is_complete())
 
                 finally:
-                    if image and remove_incomplete and self.__writer and not self.__writer.is_complete():
-                        image.remove(fatal=False)
                     try:
-                        await aiotools.shield_fg(self.__close_writer())
+                        if image and remove_incomplete and self.__writer and not self.__writer.is_complete():
+                            await image.remove(fatal=False)
                     finally:
-                        if image:
-                            await aiotools.shield_fg(image.remount_rw(False, fatal=False))
+                        try:
+                            await aiotools.shield_fg(self.__close_writer())
+                        finally:
+                            if image:
+                                await aiotools.shield_fg(image.remount_rw(False, fatal=False))
         finally:
             # Между закрытием файла и эвентом айнотифи состояние может быть не обновлено,
             # так что форсим обновление вручную, чтобы получить актуальное состояние.
@@ -363,17 +365,17 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
             assert self.__state.storage
             assert self.__state.vd
             self.__STATE_check_disconnected()
-            image = self.__STATE_get_storage_image(name)
+            image = await self.__STATE_get_storage_image(name)
 
             if self.__state.vd.image == image:
                 self.__state.vd.image = None
 
             await image.remount_rw(True)
             try:
-                image.remove(fatal=True)
+                await image.remove(fatal=True)
                 del self.__state.storage.images[name]
             finally:
-                await image.remount_rw(False, fatal=False)
+                await aiotools.shield_fg(image.remount_rw(False, fatal=False))
 
     # =====
 
@@ -387,18 +389,18 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         if self.__state.vd.connected or self.__drive.get_image_path():
             raise MsdConnectedError()
 
-    def __STATE_get_storage_image(self, name: str) -> Image:  # pylint: disable=invalid-name
+    async def __STATE_get_storage_image(self, name: str) -> Image:  # pylint: disable=invalid-name
         assert self.__state.storage
         image = self.__state.storage.images.get(name)
-        if image is None or not image.exists():
+        if image is None or not (await image.exists()):
             raise MsdUnknownImageError()
         assert image.in_storage
         return image
 
-    def __STORAGE_create_new_image(self, name: str) -> Image:  # pylint: disable=invalid-name
+    async def __STORAGE_create_new_image(self, name: str) -> Image:  # pylint: disable=invalid-name
         assert self.__state.storage
         image = self.__storage.get_image_by_name(name)
-        if image.name in self.__state.storage.images or image.exists():
+        if image.name in self.__state.storage.images or (await image.exists()):
             raise MsdImageExistsError()
         return image
 
diff --git a/kvmd/plugins/msd/otg/storage.py b/kvmd/plugins/msd/otg/storage.py
index b6b6d52a..0bdbbaae 100644
--- a/kvmd/plugins/msd/otg/storage.py
+++ b/kvmd/plugins/msd/otg/storage.py
@@ -26,6 +26,9 @@ import dataclasses
 from typing import Generator
 from typing import Optional
 
+import aiofiles
+import aiofiles.os
+
 from ....logging import get_logger
 
 from .... import aiohelpers
@@ -85,32 +88,33 @@ class Image(_Image):
         except Exception:
             return 0.0
 
-    def exists(self) -> bool:
-        return os.path.exists(self.path)
+    async def exists(self) -> bool:
+        return (await aiofiles.os.path.exists(self.path))
 
     async def remount_rw(self, rw: bool, fatal: bool=True) -> None:
         assert self.__storage
         if not self.__adopted:
             await self.__storage.remount_rw(rw, fatal)
 
-    def remove(self, fatal: bool) -> None:
+    async def remove(self, fatal: bool) -> None:
         assert self.__storage
         try:
-            os.remove(self.path)
+            await aiofiles.os.remove(self.path)
         except FileNotFoundError:
             pass
         except Exception:
             if fatal:
                 raise
-        self.set_complete(False)
+        await self.set_complete(False)
 
-    def set_complete(self, flag: bool) -> None:
+    async def set_complete(self, flag: bool) -> None:
         assert self.__storage
         if flag:
-            open(self.__complete_path, "w").close()  # pylint: disable=consider-using-with
+            async with aiofiles.open(self.__complete_path, "w"):
+                pass
         else:
             try:
-                os.remove(self.__complete_path)
+                await aiofiles.os.remove(self.__complete_path)
             except FileNotFoundError:
                 pass