speed up msd

2026-01-29 00:51:53 +08:00 · 2021-05-26 12:26:31 +03:00
parent 5da412ae7c
commit e480629724
10 changed files with 66 additions and 17 deletions
--- a/kvmd/plugins/msd/init.py
+++ b/kvmd/plugins/msd/init.py
@@ -118,6 +118,9 @@ class BaseMsd(BasePlugin):
            raise NotImplementedError()
        yield

+    def get_upload_chunk_size(self) -> int:
+        raise NotImplementedError()
+
    async def write_image_chunk(self, chunk: bytes) -> int:
        raise NotImplementedError()

--- a/kvmd/plugins/msd/disabled.py
+++ b/kvmd/plugins/msd/disabled.py
@@ -75,6 +75,9 @@ class Plugin(BaseMsd):
            raise MsdDisabledError()
        yield

+    def get_upload_chunk_size(self) -> int:
+        raise MsdDisabledError()
+
    async def write_image_chunk(self, chunk: bytes) -> int:
        raise MsdDisabledError()

--- a/kvmd/plugins/msd/otg/init.py
+++ b/kvmd/plugins/msd/otg/init.py
@@ -24,6 +24,7 @@ import os
 import asyncio
 import contextlib
 import dataclasses
+import functools
 import time

 from typing import List
@@ -42,6 +43,7 @@ from ....inotify import Inotify
 from ....yamlconf import Option

 from ....validators.basic import valid_bool
+from ....validators.basic import valid_number
 from ....validators.os import valid_abs_dir
 from ....validators.os import valid_printable_filename
 from ....validators.os import valid_command
@@ -135,6 +137,9 @@ class _State:
 class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
    def __init__(  # pylint: disable=super-init-not-called
        self,
+        upload_chunk_size: int,
+        sync_chunk_size: int,
+
        storage_path: str,

        remount_cmd: List[str],
@@ -145,6 +150,9 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
        gadget: str,  # XXX: Not from options, see /kvmd/apps/kvmd/__init__.py for details
    ) -> None:

+        self.__upload_chunk_size = upload_chunk_size
+        self.__sync_chunk_size = sync_chunk_size
+
        self.__storage_path = os.path.normpath(storage_path)
        self.__images_path = os.path.join(self.__storage_path, "images")
        self.__meta_path = os.path.join(self.__storage_path, "meta")
@@ -159,6 +167,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes

        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.__notifier = aiotools.AioNotifier()
@@ -172,9 +181,14 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
    def get_plugin_options(cls) -> Dict:
        sudo = ["/usr/bin/sudo", "--non-interactive"]
        return {
-            "storage":      Option("/var/lib/kvmd/msd", type=valid_abs_dir, unpack_as="storage_path"),
-            "remount_cmd":  Option([*sudo, "/usr/bin/kvmd-helper-otgmsd-remount", "{mode}"], type=valid_command),
-            "unlock_cmd":   Option([*sudo, "/usr/bin/kvmd-helper-otgmsd-unlock", "unlock"],  type=valid_command),
+            "upload_chunk_size": Option(65536,   type=functools.partial(valid_number, min=1024)),
+            "sync_chunk_size":   Option(4194304, type=functools.partial(valid_number, min=1024)),
+
+            "storage": Option("/var/lib/kvmd/msd", type=valid_abs_dir, unpack_as="storage_path"),
+
+            "remount_cmd": Option([*sudo, "/usr/bin/kvmd-helper-otgmsd-remount", "{mode}"], type=valid_command),
+            "unlock_cmd":  Option([*sudo, "/usr/bin/kvmd-helper-otgmsd-unlock", "unlock"],  type=valid_command),
+
            "initial": {
                "image": Option("",    type=(lambda arg: (valid_printable_filename(arg) if arg else ""))),
                "cdrom": Option(False, type=valid_bool),
@@ -313,6 +327,7 @@ 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

                    await self.__notifier.notify()
@@ -331,10 +346,20 @@ 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_file
-        await aiofs.afile_write_now(self.__new_file, chunk)
+
+        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
+
        now = time.monotonic()
        if self.__new_file_tick + 1 < now:
            # Это нужно для ручного оповещения о свободном пространстве на диске, см. get_state()
@@ -377,6 +402,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
        finally:
            self.__new_file = None
            self.__new_file_written = 0
+            self.__new_file_unsynced = 0

    # =====

--- a/kvmd/plugins/msd/relay/init.py
+++ b/kvmd/plugins/msd/relay/init.py
@@ -23,6 +23,7 @@
 import asyncio
 import contextlib
 import dataclasses
+import functools

 from typing import Dict
 from typing import AsyncGenerator
@@ -39,6 +40,7 @@ from .... import aiofs
 from ....yamlconf import Option

 from ....validators.basic import valid_bool
+from ....validators.basic import valid_number
 from ....validators.basic import valid_int_f1
 from ....validators.basic import valid_float_f01
 from ....validators.os import valid_abs_path
@@ -61,8 +63,11 @@ from .drive import DeviceInfo

 # =====
 class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
-    def __init__(  # pylint: disable=super-init-not-called
+    def __init__(  # pylint: disable=super-init-not-called,too-many-arguments
        self,
+        upload_chunk_size: int,
+        sync_chunk_size: int,
+
        gpio_device_path: str,
        target_pin: int,
        reset_inverted: bool,
@@ -74,6 +79,9 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
        reset_delay: float,
    ) -> None:

+        self.__upload_chunk_size = upload_chunk_size
+        self.__sync_chunk_size = sync_chunk_size
+
        self.__device_path = device_path
        self.__init_delay = init_delay
        self.__init_retries = init_retries
@@ -85,6 +93,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes

        self.__device_file: Optional[aiofiles.base.AiofilesContextManager] = None
        self.__written = 0
+        self.__unsynced = 0

        self.__notifier = aiotools.AioNotifier()
        self.__region = aiotools.AioExclusiveRegion(MsdIsBusyError, self.__notifier)
@@ -92,6 +101,9 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
    @classmethod
    def get_plugin_options(cls) -> Dict:
        return {
+            "upload_chunk_size": Option(65536,   type=functools.partial(valid_number, min=1024)),
+            "sync_chunk_size":   Option(4194304, type=functools.partial(valid_number, min=1024)),
+
            "gpio_device":    Option("/dev/gpiochip0", type=valid_abs_path, unpack_as="gpio_device_path"),
            "target_pin":     Option(-1, type=valid_gpio_pin),
            "reset_pin":      Option(-1, type=valid_gpio_pin),
@@ -212,6 +224,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes

                    self.__device_file = await aiofiles.open(self.__device_info.path, mode="w+b", buffering=0)  # type: ignore
                    self.__written = 0
+                    self.__unsynced = 0

                    await self.__write_image_info(name, complete=False)
                    await self.__notifier.notify()
@@ -221,10 +234,17 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                    await self.__close_device_file()
                    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 aiofs.afile_write_now(self.__device_file, chunk)
+        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

    @aiotools.atomic
@@ -261,6 +281,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
        finally:
            self.__device_file = None
            self.__written = 0
+            self.__unsynced = 0

    async def __load_device_info(self) -> None:
        retries = self.__init_retries
--- a/kvmd/plugins/msd/relay/drive.py
+++ b/kvmd/plugins/msd/relay/drive.py
@@ -129,7 +129,8 @@ class DeviceInfo:

        if self.size - image_info.size > _IMAGE_INFO_SIZE:
            await device_file.seek(self.size - _IMAGE_INFO_SIZE)  # type: ignore
-            await aiofs.afile_write_now(device_file, image_info.to_bytes())
+            await device_file.write(image_info.to_bytes())  # type: ignore
+            await aiofs.afile_sync(device_file)
            await device_file.seek(0)  # type: ignore
            return True
        return False  # Device is full