pikvm/pikvm#321: server-side uploading counters

1 files changed, 18 insertions, 38 deletions

diff --git a/kvmd/plugins/msd/relay/__init__.py b/kvmd/plugins/msd/relay/__init__.py
index 659594d8..82de6596 100644
--- a/kvmd/plugins/msd/relay/__init__.py
+++ b/kvmd/plugins/msd/relay/__init__.py
@@ -29,13 +29,9 @@ from typing import Dict
 from typing import AsyncGenerator
 from typing import Optional
 
-import aiofiles
-import aiofiles.base
-
 from ....logging import get_logger
 
 from .... import aiotools
-from .... import aiofs
 
 from ....yamlconf import Option
 
@@ -54,10 +50,10 @@ from .. import MsdDisconnectedError
 from .. import MsdMultiNotSupported
 from .. import MsdCdromNotSupported
 from .. import BaseMsd
+from .. import MsdImageWriter
 
 from .gpio import Gpio
 
-from .drive import ImageInfo
 from .drive import DeviceInfo
 
 
@@ -91,9 +87,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         self.__device_info: Optional[DeviceInfo] = None
         self.__connected = False
 
-        self.__device_file: Optional[aiofiles.base.AiofilesContextManager] = None
-        self.__written = 0
-        self.__unsynced = 0
+        self.__device_writer: Optional[MsdImageWriter] = None
 
         self.__notifier = aiotools.AioNotifier()
         self.__region = aiotools.AioExclusiveRegion(MsdIsBusyError, self.__notifier)
@@ -132,7 +126,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
             storage = {
                 "size": self.__device_info.size,
                 "free": self.__device_info.free,
-                "uploading": bool(self.__device_file)
+                "uploading": (self.__device_writer.get_state() if self.__device_writer else None),
             }
             drive = {
                 "image": (self.__device_info.image and dataclasses.asdict(self.__device_info.image)),
@@ -177,7 +171,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
     @aiotools.atomic
     async def cleanup(self) -> None:
         try:
-            await self.__close_device_file()
+            await self.__close_device_writer()
         finally:
             self.__gpio.close()
 
@@ -214,7 +208,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                 self.__connected = connected
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str) -> AsyncGenerator[None, None]:
+    async def write_image(self, name: str, size: int) -> AsyncGenerator[None, None]:
         async with self.__working():
             async with self.__region:
                 try:
@@ -222,30 +216,22 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                     if self.__connected:
                         raise MsdConnectedError()
 
-                    self.__device_file = await aiofiles.open(self.__device_info.path, mode="w+b", buffering=0)  # type: ignore
-                    self.__written = 0
-                    self.__unsynced = 0
+                    self.__device_writer = await MsdImageWriter(self.__device_info.path, size, self.__sync_chunk_size).open()
 
-                    await self.__write_image_info(name, complete=False)
+                    await self.__write_image_info(False)
                     await self.__notifier.notify()
                     yield
-                    await self.__write_image_info(name, complete=True)
+                    await self.__write_image_info(True)
                 finally:
-                    await self.__close_device_file()
+                    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_file
-        await self.__device_file.write(chunk)  # type: ignore
-        self.__written += len(chunk)
-        self.__unsynced += len(chunk)
-        if self.__unsynced >= self.__sync_chunk_size:
-            await aiofs.afile_sync(self.__device_file)
-            self.__unsynced = 0
-        return self.__written
+        assert self.__device_writer
+        return (await self.__device_writer.write(chunk))
 
     @aiotools.atomic
     async def remove(self, name: str) -> None:
@@ -262,27 +248,21 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
     # =====
 
-    async def __write_image_info(self, name: str, complete: bool) -> None:
-        assert self.__device_file
+    async def __write_image_info(self, complete: bool) -> None:
+        assert self.__device_writer
         assert self.__device_info
-        if not (await self.__device_info.write_image_info(
-            device_file=self.__device_file,
-            image_info=ImageInfo(name, self.__written, complete),
-        )):
+        if not (await self.__device_info.write_image_info(self.__device_writer, complete)):
             get_logger().error("Can't write image info because device is full")
 
-    async def __close_device_file(self) -> None:
+    async def __close_device_writer(self) -> None:
         try:
-            if self.__device_file:
+            if self.__device_writer:
                 get_logger().info("Closing device file ...")
-                await aiofs.afile_sync(self.__device_file)
-                await self.__device_file.close()  # type: ignore
+                await self.__device_writer.close()  # type: ignore
         except Exception:
             get_logger().exception("Can't close device file")
         finally:
-            self.__device_file = None
-            self.__written = 0
-            self.__unsynced = 0
+            self.__device_writer = None
 
     async def __load_device_info(self) -> None:
         retries = self.__init_retries