Fixed pikvm/pikvm#783: remove incomplete msd image

kvmd/plugins/msd/__init__.py

@@ -144,9 +144,10 @@ class BaseMsd(BasePlugin):
         raise NotImplementedError()
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:
+    async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
         _ = name
         _ = size
+        _ = remove_incomplete
         if self is not None:  # XXX: Vulture and pylint hack
             raise NotImplementedError()
         yield 1
@@ -223,6 +224,9 @@ class MsdImageWriter:  # pylint: disable=too-many-instance-attributes
         self.__unsynced = 0
         self.__tick = 0.0
 
+    def is_complete(self) -> bool:
+        return (self.__written >= self.__size)
+
     def get_file(self) -> aiofiles.base.AiofilesContextManager:
         assert self.__file is not None
         return self.__file
@@ -252,7 +256,7 @@ class MsdImageWriter:  # pylint: disable=too-many-instance-attributes
             self.__unsynced = 0
 
         now = time.monotonic()
-        if self.__tick + 1 < now or self.__written == self.__size:
+        if self.__tick + 1 < now:
             self.__tick = now
             await self.__notifier.notify()
 

kvmd/plugins/msd/disabled.py

@@ -86,7 +86,7 @@ class Plugin(BaseMsd):
         raise MsdDisabledError()
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:
+    async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
         if self is not None:  # XXX: Vulture and pylint hack
             raise MsdDisabledError()
         yield 1

kvmd/plugins/msd/otg/__init__.py

@@ -363,9 +363,10 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         return (await self.__reader.read())
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:
+    async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
         try:
             async with self.__state._region:  # pylint: disable=protected-access
+                path: str = ""
                 try:
                     async with self.__state._lock:  # pylint: disable=protected-access
                         await self.__notifier.notify()
@@ -391,9 +392,15 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
                     await self.__notifier.notify()
                     yield self.__write_chunk_size
-                    self.__set_image_complete(name, True)
+                    self.__set_image_complete(name, self.__writer.is_complete())
 
                 finally:
+                    if remove_incomplete and self.__writer and not self.__writer.is_complete():
+                        # Можно сперва удалить файл, потом закрыть его
+                        try:
+                            os.remove(path)
+                        except Exception:
+                            pass
                     await self.__close_writer()
                     await self.__remount_rw(False, fatal=False)
         finally:

kvmd/plugins/msd/relay/__init__.py

@@ -229,8 +229,10 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
             raise MsdMultiNotSupported()
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:
+    async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
         async with self.__working():
+            if remove_incomplete is not None:
+                raise MsdMultiNotSupported()
             async with self.__region:
                 try:
                     assert self.__device_info