refactoring

3 files changed, 24 insertions, 12 deletions

diff --git a/kvmd/plugins/msd/__init__.py b/kvmd/plugins/msd/__init__.py
index 1bb11413..b830b0a9 100644
--- a/kvmd/plugins/msd/__init__.py
+++ b/kvmd/plugins/msd/__init__.py
@@ -22,6 +22,7 @@
 
 import os
 import contextlib
+import time
 
 from typing import Dict
 from typing import Type
@@ -36,6 +37,7 @@ from ...logging import get_logger
 from ...errors import OperationError
 from ...errors import IsBusyError
 
+from ... import aiotools
 from ... import aiofs
 
 from .. import BasePlugin
@@ -187,8 +189,9 @@ class MsdImageReader:
             logger.exception("Can't close image reader")
 
 
-class MsdImageWriter:
-    def __init__(self, path: str, size: int, sync: int) -> None:
+class MsdImageWriter:  # pylint: disable=too-many-instance-attributes
+    def __init__(self, notifier: aiotools.AioNotifier, path: str, size: int, sync: int) -> None:
+        self.__notifier = notifier
         self.__name = os.path.basename(path)
         self.__path = path
         self.__size = size
@@ -197,6 +200,7 @@ class MsdImageWriter:
         self.__file: Optional[aiofiles.base.AiofilesContextManager] = None
         self.__written = 0
         self.__unsynced = 0
+        self.__tick = 0.0
 
     def get_file(self) -> aiofiles.base.AiofilesContextManager:
         assert self.__file is not None
@@ -226,6 +230,11 @@ class MsdImageWriter:
             await aiofs.afile_sync(self.__file)
             self.__unsynced = 0
 
+        now = time.monotonic()
+        if self.__tick + 1 < now or self.__written == self.__size:
+            self.__tick = now
+            await self.__notifier.notify()
+
         return self.__written
 
     async def close(self) -> None:
diff --git a/kvmd/plugins/msd/otg/__init__.py b/kvmd/plugins/msd/otg/__init__.py
index e26aac44..a98924ac 100644
--- a/kvmd/plugins/msd/otg/__init__.py
+++ b/kvmd/plugins/msd/otg/__init__.py
@@ -167,7 +167,6 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
         self.__reader: Optional[MsdImageReader] = None
         self.__writer: Optional[MsdImageWriter] = None
-        self.__writer_tick = 0.0
 
         self.__notifier = aiotools.AioNotifier()
         self.__state = _State(self.__notifier)
@@ -377,7 +376,12 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                         await self.__remount_rw(True)
                         self.__set_image_complete(name, False)
 
-                        self.__writer = await MsdImageWriter(path, size, self.__sync_chunk_size).open()
+                        self.__writer = await MsdImageWriter(
+                            notifier=self.__notifier,
+                            path=path,
+                            size=size,
+                            sync=self.__sync_chunk_size,
+                        ).open()
 
                     await self.__notifier.notify()
                     yield self.__write_chunk_size
@@ -394,13 +398,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
     async def write_image_chunk(self, chunk: bytes) -> int:
         assert self.__writer
-        written = await self.__writer.write(chunk)
-        now = time.monotonic()
-        if self.__writer_tick + 1 < now:
-            # Это нужно для ручного оповещения о свободном пространстве на диске, см. get_state()
-            self.__writer_tick = now
-            await self.__notifier.notify()
-        return written
+        return (await self.__writer.write(chunk))
 
     @aiotools.atomic
     async def remove(self, name: str) -> None:
diff --git a/kvmd/plugins/msd/relay/__init__.py b/kvmd/plugins/msd/relay/__init__.py
index 2b45bce3..c4891a96 100644
--- a/kvmd/plugins/msd/relay/__init__.py
+++ b/kvmd/plugins/msd/relay/__init__.py
@@ -237,7 +237,12 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                     if self.__connected:
                         raise MsdConnectedError()
 
-                    self.__device_writer = await MsdImageWriter(self.__device_info.path, size, self.__sync_chunk_size).open()
+                    self.__device_writer = await MsdImageWriter(
+                        notifier=self.__notifier,
+                        path=self.__device_info.path,
+                        size=size,
+                        sync=self.__sync_chunk_size,
+                    ).open()
 
                     await self.__write_image_info(False)
                     await self.__notifier.notify()