refactored msd writer api

5 files changed, 60 insertions, 56 deletions

diff --git a/kvmd/plugins/msd/__init__.py b/kvmd/plugins/msd/__init__.py
index dbe3ff90..e5bf2ab1 100644
--- a/kvmd/plugins/msd/__init__.py
+++ b/kvmd/plugins/msd/__init__.py
@@ -117,6 +117,17 @@ class BaseMsdReader:
         yield
 
 
+class BaseMsdWriter:
+    def get_state(self) -> Dict:
+        raise NotImplementedError()
+
+    def get_chunk_size(self) -> int:
+        raise NotImplementedError()
+
+    async def write_chunk(self, chunk: bytes) -> int:
+        raise NotImplementedError()
+
+
 class BaseMsd(BasePlugin):
     async def get_state(self) -> Dict:
         raise NotImplementedError()
@@ -154,16 +165,13 @@ class BaseMsd(BasePlugin):
         yield BaseMsdReader()
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
+    async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[BaseMsdWriter, None]:
         _ = name
         _ = size
         _ = remove_incomplete
         if self is not None:  # XXX: Vulture and pylint hack
             raise NotImplementedError()
-        yield 1
-
-    async def write_image_chunk(self, chunk: bytes) -> int:
-        raise NotImplementedError()
+        yield BaseMsdWriter()
 
     async def remove(self, name: str) -> None:
         raise NotImplementedError()
@@ -225,47 +233,38 @@ class MsdFileReader(BaseMsdReader):  # pylint: disable=too-many-instance-attribu
             logger.exception("Can't close image reader")
 
 
-class MsdImageWriter:  # pylint: disable=too-many-instance-attributes
-    def __init__(self, notifier: aiotools.AioNotifier, path: str, size: int, sync: int) -> None:
+class MsdFileWriter(BaseMsdWriter):  # pylint: disable=too-many-instance-attributes
+    def __init__(self, notifier: aiotools.AioNotifier, path: str, file_size: int, sync_size: int, chunk_size: int) -> None:
         self.__notifier = notifier
         self.__name = os.path.basename(path)
         self.__path = path
-        self.__size = size
-        self.__sync = sync
+        self.__file_size = file_size
+        self.__sync_size = sync_size
+        self.__chunk_size = chunk_size
 
         self.__file: Optional[aiofiles.base.AiofilesContextManager] = None
         self.__written = 0
         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
-
     def get_state(self) -> Dict:
         return {
             "name": self.__name,
-            "size": self.__size,
+            "size": self.__file_size,
             "written": self.__written,
         }
 
-    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
+    def get_chunk_size(self) -> int:
+        return self.__chunk_size
 
-    async def write(self, chunk: bytes) -> int:
+    async def write_chunk(self, chunk: bytes) -> int:
         assert self.__file is not None
 
         await self.__file.write(chunk)  # type: ignore
         self.__written += len(chunk)
 
         self.__unsynced += len(chunk)
-        if self.__unsynced >= self.__sync:
+        if self.__unsynced >= self.__sync_size:
             await aiofs.afile_sync(self.__file)
             self.__unsynced = 0
 
@@ -276,18 +275,31 @@ class MsdImageWriter:  # pylint: disable=too-many-instance-attributes
 
         return self.__written
 
+    def is_complete(self) -> bool:
+        return (self.__written >= self.__file_size)
+
+    def get_file(self) -> aiofiles.base.AiofilesContextManager:
+        assert self.__file is not None
+        return self.__file
+
+    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)
+        self.__file = await aiofiles.open(self.__path, mode="w+b", buffering=0)  # type: ignore
+        return self
+
     async def close(self) -> None:
         assert self.__file is not None
         logger = get_logger()
         logger.info("Closing image writer ...")
         try:
-            if self.__written == self.__size:
+            if self.__written == self.__file_size:
                 (log, result) = (logger.info, "OK")
-            elif self.__written < self.__size:
+            elif self.__written < self.__file_size:
                 (log, result) = (logger.error, "INCOMPLETE")
             else:  # written > size
                 (log, result) = (logger.warning, "OVERFLOW")
-            log("Written %d of %d bytes to MSD image %r: %s", self.__written, self.__size, self.__name, result)
+            log("Written %d of %d bytes to MSD image %r: %s", self.__written, self.__file_size, self.__name, result)
             try:
                 await aiofs.afile_sync(self.__file)
             finally:
diff --git a/kvmd/plugins/msd/disabled.py b/kvmd/plugins/msd/disabled.py
index 5317fb09..a1c2b752 100644
--- a/kvmd/plugins/msd/disabled.py
+++ b/kvmd/plugins/msd/disabled.py
@@ -30,6 +30,7 @@ from ... import aiotools
 
 from . import MsdOperationError
 from . import BaseMsdReader
+from . import BaseMsdWriter
 from . import BaseMsd
 
 
@@ -84,13 +85,10 @@ class Plugin(BaseMsd):
         yield BaseMsdReader()
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
+    async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[BaseMsdWriter, None]:
         if self is not None:  # XXX: Vulture and pylint hack
             raise MsdDisabledError()
-        yield 1
-
-    async def write_image_chunk(self, chunk: bytes) -> int:
-        raise MsdDisabledError()
+        yield BaseMsdWriter()
 
     async def remove(self, name: str) -> None:
         raise MsdDisabledError()
diff --git a/kvmd/plugins/msd/otg/__init__.py b/kvmd/plugins/msd/otg/__init__.py
index 6ff6e9ab..96ef9289 100644
--- a/kvmd/plugins/msd/otg/__init__.py
+++ b/kvmd/plugins/msd/otg/__init__.py
@@ -58,7 +58,7 @@ from .. import MsdUnknownImageError
 from .. import MsdImageExistsError
 from .. import BaseMsd
 from .. import MsdFileReader
-from .. import MsdImageWriter
+from .. import MsdFileWriter
 
 from . import fs
 
@@ -166,7 +166,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         self.__drive = Drive(gadget, instance=0, lun=0)
 
         self.__reader: Optional[MsdFileReader] = None
-        self.__writer: Optional[MsdImageWriter] = None
+        self.__writer: Optional[MsdFileWriter] = None
 
         self.__notifier = aiotools.AioNotifier()
         self.__state = _State(self.__notifier)
@@ -359,7 +359,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
             await self.__notifier.notify()
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
+    async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[MsdFileWriter, None]:
         try:
             async with self.__state._region:  # pylint: disable=protected-access
                 path: str = ""
@@ -379,15 +379,16 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                         await self.__remount_rw(True)
                         self.__set_image_complete(name, False)
 
-                        self.__writer = await MsdImageWriter(
+                        self.__writer = await MsdFileWriter(
                             notifier=self.__notifier,
                             path=path,
-                            size=size,
-                            sync=self.__sync_chunk_size,
+                            file_size=size,
+                            sync_size=self.__sync_chunk_size,
+                            chunk_size=self.__write_chunk_size,
                         ).open()
 
                     await self.__notifier.notify()
-                    yield self.__write_chunk_size
+                    yield self.__writer
                     self.__set_image_complete(name, self.__writer.is_complete())
 
                 finally:
@@ -405,10 +406,6 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
             await self.__reload_state()
             await self.__notifier.notify()
 
-    async def write_image_chunk(self, chunk: bytes) -> int:
-        assert self.__writer
-        return (await self.__writer.write(chunk))
-
     @aiotools.atomic
     async def remove(self, name: str) -> None:
         async with self.__state.busy():
diff --git a/kvmd/plugins/msd/relay/__init__.py b/kvmd/plugins/msd/relay/__init__.py
index e8dd995f..14f903f1 100644
--- a/kvmd/plugins/msd/relay/__init__.py
+++ b/kvmd/plugins/msd/relay/__init__.py
@@ -52,7 +52,7 @@ from .. import MsdCdromNotSupported
 from .. import MsdRwNotSupported
 from .. import BaseMsdReader
 from .. import BaseMsd
-from .. import MsdImageWriter
+from .. import MsdFileWriter
 
 from .gpio import Gpio
 
@@ -89,7 +89,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         self.__device_info: Optional[DeviceInfo] = None
         self.__connected = False
 
-        self.__device_writer: Optional[MsdImageWriter] = None
+        self.__device_writer: Optional[MsdFileWriter] = None
 
         self.__notifier = aiotools.AioNotifier()
         self.__region = aiotools.AioExclusiveRegion(MsdIsBusyError, self.__notifier)
@@ -226,7 +226,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         yield BaseMsdReader()
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
+    async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[MsdFileWriter, None]:
         async with self.__working():
             if remove_incomplete is not None:
                 raise MsdMultiNotSupported()
@@ -236,25 +236,22 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                     if self.__connected:
                         raise MsdConnectedError()
 
-                    self.__device_writer = await MsdImageWriter(
+                    self.__device_writer = await MsdFileWriter(
                         notifier=self.__notifier,
                         path=self.__device_info.path,
-                        size=size,
-                        sync=self.__sync_chunk_size,
+                        file_size=size,
+                        sync_size=self.__sync_chunk_size,
+                        chunk_size=self.__upload_chunk_size,
                     ).open()
 
                     await self.__write_image_info(False)
                     await self.__notifier.notify()
-                    yield self.__upload_chunk_size
+                    yield self.__device_writer
                     await self.__write_image_info(True)
                 finally:
                     await self.__close_device_writer()
                     await self.__load_device_info()
 
-    async def write_image_chunk(self, chunk: bytes) -> int:
-        assert self.__device_writer
-        return (await self.__device_writer.write(chunk))
-
     @aiotools.atomic
     async def remove(self, name: str) -> None:
         async with self.__working():
diff --git a/kvmd/plugins/msd/relay/drive.py b/kvmd/plugins/msd/relay/drive.py
index 2e95e837..e1301f99 100644
--- a/kvmd/plugins/msd/relay/drive.py
+++ b/kvmd/plugins/msd/relay/drive.py
@@ -32,7 +32,7 @@ from typing import Optional
 from .... import aiotools
 from .... import aiofs
 
-from .. import MsdImageWriter
+from .. import MsdFileWriter
 
 
 # =====
@@ -121,7 +121,7 @@ class DeviceInfo:
             image=image_info,
         )
 
-    async def write_image_info(self, device_writer: MsdImageWriter, complete: bool) -> bool:
+    async def write_image_info(self, device_writer: MsdFileWriter, complete: bool) -> bool:
         device_file = device_writer.get_file()
         state = device_writer.get_state()
         image_info = ImageInfo(state["name"], state["written"], complete)