pikvm/pikvm#321: server-side uploading counters

8 files changed, 115 insertions, 95 deletions

diff --git a/kvmd/apps/kvmd/api/msd.py b/kvmd/apps/kvmd/api/msd.py
index 0b8fa6e2..315a86f1 100644
--- a/kvmd/apps/kvmd/api/msd.py
+++ b/kvmd/apps/kvmd/api/msd.py
@@ -28,10 +28,12 @@ from ....logging import get_logger
 from ....plugins.msd import BaseMsd
 
 from ....validators.basic import valid_bool
+from ....validators.basic import valid_int_f0
 from ....validators.kvm import valid_msd_image_name
 
 from ..http import exposed_http
 from ..http import make_json_response
+from ..http import get_field_value
 from ..http import get_multipart_field
 
 
@@ -71,12 +73,12 @@ class MsdApi:
         name = ""
         written = 0
         try:
-            name_field = await get_multipart_field(reader, "image")
-            name = valid_msd_image_name((await name_field.read()).decode("utf-8"))
+            name = valid_msd_image_name(await get_field_value(reader, "image"))
+            size = valid_int_f0(await get_field_value(reader, "size"))
 
             data_field = await get_multipart_field(reader, "data")
 
-            async with self.__msd.write_image(name):
+            async with self.__msd.write_image(name, size):
                 logger.info("Writing image %r to MSD ...", name)
                 while True:
                     chunk = await data_field.read_chunk(self.__msd.get_upload_chunk_size())
diff --git a/kvmd/apps/kvmd/http.py b/kvmd/apps/kvmd/http.py
index e46be2fa..941205b2 100644
--- a/kvmd/apps/kvmd/http.py
+++ b/kvmd/apps/kvmd/http.py
@@ -171,6 +171,11 @@ def make_json_exception(err: Exception, status: Optional[int]=None) -> aiohttp.w
 
 
 # =====
+async def get_field_value(reader: aiohttp.MultipartReader, name: str) -> str:
+    field = await get_multipart_field(reader, name)
+    return (await field.read()).decode("utf-8")
+
+
 async def get_multipart_field(reader: aiohttp.MultipartReader, name: str) -> aiohttp.BodyPartReader:
     field = await reader.next()
     if not isinstance(field, aiohttp.BodyPartReader):
diff --git a/kvmd/plugins/msd/__init__.py b/kvmd/plugins/msd/__init__.py
index c8015a2a..680ccece 100644
--- a/kvmd/plugins/msd/__init__.py
+++ b/kvmd/plugins/msd/__init__.py
@@ -20,6 +20,7 @@
 # ========================================================================== #
 
 
+import os
 import contextlib
 
 from typing import Dict
@@ -27,6 +28,11 @@ from typing import Type
 from typing import AsyncGenerator
 from typing import Optional
 
+import aiofiles
+import aiofiles.base
+
+from ... import aiofs
+
 from ...errors import OperationError
 from ...errors import IsBusyError
 
@@ -113,7 +119,7 @@ class BaseMsd(BasePlugin):
         raise NotImplementedError()
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str) -> AsyncGenerator[None, None]:  # pylint: disable=unused-argument
+    async def write_image(self, name: str, size: int) -> AsyncGenerator[None, None]:  # pylint: disable=unused-argument
         if self is not None:  # XXX: Vulture and pylint hack
             raise NotImplementedError()
         yield
@@ -128,6 +134,52 @@ class BaseMsd(BasePlugin):
         raise NotImplementedError()
 
 
+class MsdImageWriter:
+    def __init__(self, path: str, size: int, sync: int) -> None:
+        self.__name = os.path.basename(path)
+        self.__path = path
+        self.__size = size
+        self.__sync = sync
+
+        self.__file: Optional[aiofiles.base.AiofilesContextManager] = None
+        self.__written = 0
+        self.__unsynced = 0
+
+    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,
+            "written": self.__written,
+        }
+
+    async def open(self) -> "MsdImageWriter":
+        assert self.__file is None
+        self.__file = await aiofiles.open(self.__path, mode="w+b", buffering=0)  # type: ignore
+        return self
+
+    async def write(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:
+            await aiofs.afile_sync(self.__file)
+            self.__unsynced = 0
+
+        return self.__written
+
+    async def close(self) -> None:
+        assert self.__file is not None
+        await aiofs.afile_sync(self.__file)
+        await self.__file.close()  # type: ignore
+
+
 # =====
 def get_msd_class(name: str) -> Type[BaseMsd]:
     return get_plugin_class("msd", name)  # type: ignore
diff --git a/kvmd/plugins/msd/disabled.py b/kvmd/plugins/msd/disabled.py
index 0f524ce9..fc1979e0 100644
--- a/kvmd/plugins/msd/disabled.py
+++ b/kvmd/plugins/msd/disabled.py
@@ -70,7 +70,7 @@ class Plugin(BaseMsd):
         raise MsdDisabledError()
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str) -> AsyncGenerator[None, None]:
+    async def write_image(self, name: str, size: int) -> AsyncGenerator[None, None]:
         if self is not None:  # XXX: Vulture and pylint hack
             raise MsdDisabledError()
         yield
diff --git a/kvmd/plugins/msd/otg/__init__.py b/kvmd/plugins/msd/otg/__init__.py
index 8fae87bb..9a741a88 100644
--- a/kvmd/plugins/msd/otg/__init__.py
+++ b/kvmd/plugins/msd/otg/__init__.py
@@ -32,9 +32,6 @@ from typing import Dict
 from typing import AsyncGenerator
 from typing import Optional
 
-import aiofiles
-import aiofiles.base
-
 from ....logging import get_logger
 
 from ....inotify import InotifyMask
@@ -49,7 +46,6 @@ from ....validators.os import valid_printable_filename
 from ....validators.os import valid_command
 
 from .... import aiotools
-from .... import aiofs
 
 from .. import MsdError
 from .. import MsdIsBusyError
@@ -60,6 +56,7 @@ from .. import MsdImageNotSelected
 from .. import MsdUnknownImageError
 from .. import MsdImageExistsError
 from .. import BaseMsd
+from .. import MsdImageWriter
 
 from . import fs
 from . import helpers
@@ -165,10 +162,8 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
         self.__drive = Drive(gadget, instance=0, lun=0)
 
-        self.__new_file: Optional[aiofiles.base.AiofilesContextManager] = None
-        self.__new_file_written = 0
-        self.__new_file_unsynced = 0
-        self.__new_file_tick = 0.0
+        self.__new_writer: Optional[MsdImageWriter] = None
+        self.__new_writer_tick = 0.0
 
         self.__notifier = aiotools.AioNotifier()
         self.__state = _State(self.__notifier)
@@ -204,11 +199,14 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                     del storage["images"][name]["path"]
                     del storage["images"][name]["in_storage"]
 
-                storage["uploading"] = bool(self.__new_file)
-                if self.__new_file:  # При загрузке файла показываем размер вручную
+                if self.__new_writer:
+                    # При загрузке файла показываем актуальную статистику вручную
+                    storage["uploading"] = self.__new_writer.get_state()
                     space = fs.get_fs_space(self.__storage_path, fatal=False)
                     if space:
                         storage.update(dataclasses.asdict(space))
+                else:
+                    storage["uploading"] = None
 
             vd: Optional[Dict] = None
             if self.__state.vd:
@@ -253,7 +251,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
     @aiotools.atomic
     async def cleanup(self) -> None:
-        await self.__close_new_file()
+        await self.__close_new_writer()
 
     # =====
 
@@ -308,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) -> AsyncGenerator[None, None]:
+    async def write_image(self, name: str, size: int) -> AsyncGenerator[None, None]:
         try:
             async with self.__state._region:  # pylint: disable=protected-access
                 try:
@@ -326,16 +324,15 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
                         await self.__remount_storage(rw=True)
                         self.__set_image_complete(name, False)
-                        self.__new_file_written = 0
-                        self.__new_file_unsynced = 0
-                        self.__new_file = await aiofiles.open(path, mode="w+b", buffering=0)  # type: ignore
+
+                        self.__new_writer = await MsdImageWriter(path, size, self.__sync_chunk_size).open()
 
                     await self.__notifier.notify()
                     yield
                     self.__set_image_complete(name, True)
 
                 finally:
-                    await self.__close_new_file()
+                    await self.__close_new_writer()
                     try:
                         await self.__remount_storage(rw=False)
                     except Exception:
@@ -350,22 +347,14 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         return self.__upload_chunk_size
 
     async def write_image_chunk(self, chunk: bytes) -> int:
-        assert self.__new_file
-
-        await self.__new_file.write(chunk)  # type: ignore
-        self.__new_file_written += len(chunk)
-
-        self.__new_file_unsynced += len(chunk)
-        if self.__new_file_unsynced >= self.__sync_chunk_size:
-            await aiofs.afile_sync(self.__new_file)
-            self.__new_file_unsynced = 0
-
+        assert self.__new_writer
+        written = await self.__new_writer.write(chunk)
         now = time.monotonic()
-        if self.__new_file_tick + 1 < now:
+        if self.__new_writer_tick + 1 < now:
             # Это нужно для ручного оповещения о свободном пространстве на диске, см. get_state()
-            self.__new_file_tick = now
+            self.__new_writer_tick = now
             await self.__notifier.notify()
-        return self.__new_file_written
+        return written
 
     @aiotools.atomic
     async def remove(self, name: str) -> None:
@@ -392,18 +381,15 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
     # =====
 
-    async def __close_new_file(self) -> None:
+    async def __close_new_writer(self) -> None:
         try:
-            if self.__new_file:
+            if self.__new_writer:
                 get_logger().info("Closing new image file ...")
-                await aiofs.afile_sync(self.__new_file)
-                await self.__new_file.close()  # type: ignore
+                await self.__new_writer.close()
         except Exception:
             get_logger().exception("Can't close image file")
         finally:
-            self.__new_file = None
-            self.__new_file_written = 0
-            self.__new_file_unsynced = 0
+            self.__new_writer = None
 
     # =====
 
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
diff --git a/kvmd/plugins/msd/relay/drive.py b/kvmd/plugins/msd/relay/drive.py
index e9b9f681..a85b3fa2 100644
--- a/kvmd/plugins/msd/relay/drive.py
+++ b/kvmd/plugins/msd/relay/drive.py
@@ -29,11 +29,11 @@ import dataclasses
 from typing import IO
 from typing import Optional
 
-import aiofiles.base
-
 from .... import aiotools
 from .... import aiofs
 
+from .. import MsdImageWriter
+
 
 # =====
 _IMAGE_INFO_SIZE = 4096
@@ -121,11 +121,10 @@ class DeviceInfo:
             image=image_info,
         )
 
-    async def write_image_info(
-        self,
-        device_file: aiofiles.base.AiofilesContextManager,
-        image_info: ImageInfo,
-    ) -> bool:
+    async def write_image_info(self, device_writer: MsdImageWriter, complete: bool) -> bool:
+        device_file = device_writer.get_file()
+        state = device_writer.get_state()
+        image_info = ImageInfo(state["name"], state["written"], complete)
 
         if self.size - image_info.size > _IMAGE_INFO_SIZE:
             await device_file.seek(self.size - _IMAGE_INFO_SIZE)  # type: ignore
diff --git a/web/share/js/kvm/msd.js b/web/share/js/kvm/msd.js
index 68548fdc..23ad1bea 100644
--- a/web/share/js/kvm/msd.js
+++ b/web/share/js/kvm/msd.js
@@ -101,13 +101,13 @@ export function Msd() {
 	var __clickUploadNewImageButton = function() {
 		let form_data = new FormData();
 		form_data.append("image", __image_file.name);
+		form_data.append("size", __image_file.size);
 		form_data.append("data", __image_file);
 
 		__upload_http = new XMLHttpRequest();
 		__upload_http.open("POST", "/api/msd/write", true);
 		__upload_http.upload.timeout = 15000;
 		__upload_http.onreadystatechange = __uploadStateChange;
-		__upload_http.upload.onprogress = __uploadProgress;
 		__upload_http.send(form_data);
 	};
 
@@ -123,16 +123,8 @@ export function Msd() {
 		}
 	};
 
-	var __uploadProgress = function(event) {
-		if(event.lengthComputable) {
-			let percent = Math.round((event.loaded * 100) / event.total);
-			tools.progressSetValue($("msd-uploading-progress"), `${percent}%`, percent);
-		}
-	};
-
 	var __clickAbortUploadingButton = function() {
 		__upload_http.onreadystatechange = null;
-		__upload_http.upload.onprogress = null;
 		__upload_http.abort();
 		__upload_http = null;
 		tools.progressSetValue($("msd-uploading-progress"), "Aborted", 0);
@@ -238,8 +230,12 @@ export function Msd() {
 			tools.hiddenSetVisible($("msd-submenu-new-image"), __image_file);
 			$("msd-new-image-name").innerHTML = (__image_file ? __image_file.name : "");
 			$("msd-new-image-size").innerHTML = (__image_file ? tools.formatSize(__image_file.size) : "");
+
 			if (!__upload_http) {
 				tools.progressSetValue($("msd-uploading-progress"), "Waiting for upload (press UPLOAD button) ...", 0);
+			} else if (__state.storage.uploading) {
+				let percent = Math.round(__state.storage.uploading.written * 100 / __state.storage.uploading.size);
+				tools.progressSetValue($("msd-uploading-progress"), `${percent}%`, percent);
 			}
 
 		} else {