/msd/write_remote handle

4 files changed, 18 insertions, 20 deletions

diff --git a/kvmd/plugins/msd/__init__.py b/kvmd/plugins/msd/__init__.py
index 680ccece..547d778b 100644
--- a/kvmd/plugins/msd/__init__.py
+++ b/kvmd/plugins/msd/__init__.py
@@ -31,6 +31,8 @@ from typing import Optional
 import aiofiles
 import aiofiles.base
 
+from ...logging import get_logger
+
 from ... import aiofs
 
 from ...errors import OperationError
@@ -119,13 +121,10 @@ class BaseMsd(BasePlugin):
         raise NotImplementedError()
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str, size: int) -> AsyncGenerator[None, None]:  # pylint: disable=unused-argument
+    async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:  # pylint: disable=unused-argument
         if self is not None:  # XXX: Vulture and pylint hack
             raise NotImplementedError()
-        yield
-
-    def get_upload_chunk_size(self) -> int:
-        raise NotImplementedError()
+        yield 1
 
     async def write_image_chunk(self, chunk: bytes) -> int:
         raise NotImplementedError()
@@ -158,6 +157,7 @@ class MsdImageWriter:
 
     async def open(self) -> "MsdImageWriter":
         assert self.__file is None
+        get_logger(1).info("Writing %r image (%d bytes) to MSD ...", self.__name, self.__size)
         self.__file = await aiofiles.open(self.__path, mode="w+b", buffering=0)  # type: ignore
         return self
 
@@ -176,6 +176,13 @@ class MsdImageWriter:
 
     async def close(self) -> None:
         assert self.__file is not None
+        if self.__written == self.__size:
+            (log, result) = (get_logger().info, "OK")
+        elif self.__written < self.__size:
+            (log, result) = (get_logger().error, "INCOMPLETE")
+        else:  # written > size
+            (log, result) = (get_logger().warning, "OVERFLOW")
+        log("Written %d of %d bytes to MSD image %r: %s", self.__written, self.__size, self.__name, result)
         await aiofs.afile_sync(self.__file)
         await self.__file.close()  # type: ignore
 
diff --git a/kvmd/plugins/msd/disabled.py b/kvmd/plugins/msd/disabled.py
index fc1979e0..220eaf7a 100644
--- a/kvmd/plugins/msd/disabled.py
+++ b/kvmd/plugins/msd/disabled.py
@@ -70,13 +70,10 @@ class Plugin(BaseMsd):
         raise MsdDisabledError()
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str, size: int) -> AsyncGenerator[None, None]:
+    async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:
         if self is not None:  # XXX: Vulture and pylint hack
             raise MsdDisabledError()
-        yield
-
-    def get_upload_chunk_size(self) -> int:
-        raise MsdDisabledError()
+        yield 1
 
     async def write_image_chunk(self, chunk: bytes) -> int:
         raise MsdDisabledError()
diff --git a/kvmd/plugins/msd/otg/__init__.py b/kvmd/plugins/msd/otg/__init__.py
index 03e06650..4289db2a 100644
--- a/kvmd/plugins/msd/otg/__init__.py
+++ b/kvmd/plugins/msd/otg/__init__.py
@@ -306,7 +306,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
             self.__state.vd.connected = connected
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str, size: int) -> AsyncGenerator[None, None]:
+    async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:
         try:
             async with self.__state._region:  # pylint: disable=protected-access
                 try:
@@ -328,7 +328,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                         self.__new_writer = await MsdImageWriter(path, size, self.__sync_chunk_size).open()
 
                     await self.__notifier.notify()
-                    yield
+                    yield self.__upload_chunk_size
                     self.__set_image_complete(name, True)
 
                 finally:
@@ -343,9 +343,6 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
             await self.__reload_state()
             await self.__notifier.notify()
 
-    def get_upload_chunk_size(self) -> int:
-        return self.__upload_chunk_size
-
     async def write_image_chunk(self, chunk: bytes) -> int:
         assert self.__new_writer
         written = await self.__new_writer.write(chunk)
diff --git a/kvmd/plugins/msd/relay/__init__.py b/kvmd/plugins/msd/relay/__init__.py
index 82de6596..a608e535 100644
--- a/kvmd/plugins/msd/relay/__init__.py
+++ b/kvmd/plugins/msd/relay/__init__.py
@@ -208,7 +208,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                 self.__connected = connected
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str, size: int) -> AsyncGenerator[None, None]:
+    async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:
         async with self.__working():
             async with self.__region:
                 try:
@@ -220,15 +220,12 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
                     await self.__write_image_info(False)
                     await self.__notifier.notify()
-                    yield
+                    yield self.__upload_chunk_size
                     await self.__write_image_info(True)
                 finally:
                     await self.__close_device_writer()
                     await self.__load_device_info()
 
-    def get_upload_chunk_size(self) -> int:
-        return self.__upload_chunk_size
-
     async def write_image_chunk(self, chunk: bytes) -> int:
         assert self.__device_writer
         return (await self.__device_writer.write(chunk))