diff options
| -rw-r--r-- | kvmd/plugins/msd/__init__.py | 12 | ||||
| -rw-r--r-- | kvmd/plugins/msd/otg/__init__.py | 187 |
2 files changed, 111 insertions, 88 deletions
diff --git a/kvmd/plugins/msd/__init__.py b/kvmd/plugins/msd/__init__.py index 11cd8be5..8b5fa8e2 100644 --- a/kvmd/plugins/msd/__init__.py +++ b/kvmd/plugins/msd/__init__.py @@ -266,9 +266,6 @@ class MsdFileWriter(BaseMsdWriter): # pylint: disable=too-many-instance-attribu return self.__written - def is_complete(self) -> bool: - return (self.__written >= self.__file_size) - 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) @@ -276,6 +273,10 @@ class MsdFileWriter(BaseMsdWriter): # pylint: disable=too-many-instance-attribu self.__file = await aiofiles.open(self.__path, mode="w+b", buffering=0) # type: ignore return self + async def finish(self) -> bool: + await self.__sync() + return (self.__written >= self.__file_size) + async def close(self) -> None: assert self.__file is not None logger = get_logger() @@ -288,10 +289,7 @@ class MsdFileWriter(BaseMsdWriter): # pylint: disable=too-many-instance-attribu else: # written > size (log, result) = (logger.warning, "OVERFLOW") log("Written %d of %d bytes to MSD image %r: %s", self.__written, self.__file_size, self.__name, result) - try: - await self.__sync() - finally: - await self.__file.close() # type: ignore + await self.__file.close() # type: ignore except Exception: logger.exception("Can't close image writer") diff --git a/kvmd/plugins/msd/otg/__init__.py b/kvmd/plugins/msd/otg/__init__.py index 203e0b95..99de95e0 100644 --- a/kvmd/plugins/msd/otg/__init__.py +++ b/kvmd/plugins/msd/otg/__init__.py @@ -25,7 +25,6 @@ import contextlib import dataclasses import functools import copy -import time from typing import AsyncGenerator @@ -96,15 +95,17 @@ class _State: @contextlib.asynccontextmanager async def busy(self, check_online: bool=True) -> AsyncGenerator[None, None]: - with self._region: - async with self._lock: - self.__notifier.notify() - if check_online: - if self.vd is None: - raise MsdOfflineError() - assert self.storage - yield - self.__notifier.notify() + try: + with self._region: + async with self._lock: + self.__notifier.notify() + if check_online: + if self.vd is None: + raise MsdOfflineError() + assert self.storage + yield + finally: + self.__notifier.notify() def is_busy(self) -> bool: return self._region.is_busy() @@ -143,7 +144,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes logger = get_logger(0) logger.info("Using OTG gadget %r as MSD", gadget) - aiotools.run_sync(self.__reload_state(notify=False)) + aiotools.run_sync(self.__unsafe_reload_state()) @classmethod def get_plugin_options(cls) -> dict: @@ -163,13 +164,15 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes }, } + # ===== + async def get_state(self) -> dict: async with self.__state._lock: # pylint: disable=protected-access storage: (dict | None) = None if self.__state.storage: if self.__writer: # При загрузке файла показываем актуальную статистику вручную - await self.__storage.reload_parts_info() + await aiotools.shield_fg(self.__storage.reload_parts_info()) storage = dataclasses.asdict(self.__state.storage) for name in list(storage["images"]): @@ -212,9 +215,6 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes prev = copy.deepcopy(new) yield new - async def systask(self) -> None: - await self.__watch_inotify() - @aiotools.atomic_fg async def reset(self) -> None: async with self.__state.busy(check_online=False): @@ -226,11 +226,6 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes except Exception: get_logger(0).exception("Can't reset MSD properly") - @aiotools.atomic_fg - async def cleanup(self) -> None: - await self.__close_reader() - await self.__close_writer() - # ===== @aiotools.atomic_fg @@ -297,6 +292,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes async with self.__state._lock: # pylint: disable=protected-access self.__notifier.notify() self.__STATE_check_disconnected() + image = await self.__STATE_get_storage_image(name) self.__reader = await MsdFileReader( notifier=self.__notifier, @@ -304,7 +300,10 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes path=image.path, chunk_size=self.__read_chunk_size, ).open() + + self.__notifier.notify() yield self.__reader + finally: await aiotools.shield_fg(self.__close_reader()) finally: @@ -312,18 +311,37 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes @contextlib.asynccontextmanager async def write_image(self, name: str, size: int, remove_incomplete: (bool | None)) -> AsyncGenerator[MsdFileWriter, None]: + image: (Image | None) = None + complete = False + + async def finish_writing() -> None: + # Делаем под блокировкой, чтобы эвент айнотифи не был обработан + # до того, как мы не закончим все процедуры. + async with self.__state._lock: # pylint: disable=protected-access + try: + if image: + await image.set_complete(complete) + finally: + try: + if image and remove_incomplete and not complete: + await image.remove(fatal=False) + finally: + try: + await self.__close_writer() + finally: + if image: + await image.remount_rw(False, fatal=False) + try: with self.__state._region: # pylint: disable=protected-access - image: (Image | None) = None try: async with self.__state._lock: # pylint: disable=protected-access self.__notifier.notify() self.__STATE_check_disconnected() - image = await self.__STORAGE_create_new_image(name) + image = await self.__STORAGE_create_new_image(name) await image.remount_rw(True) await image.set_complete(False) - self.__writer = await MsdFileWriter( notifier=self.__notifier, name=image.name, @@ -335,22 +353,12 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes self.__notifier.notify() yield self.__writer - await image.set_complete(self.__writer.is_complete()) + complete = await self.__writer.finish() finally: - try: - if image and remove_incomplete and self.__writer and not self.__writer.is_complete(): - await image.remove(fatal=False) - finally: - try: - await aiotools.shield_fg(self.__close_writer()) - finally: - if image: - await aiotools.shield_fg(image.remount_rw(False, fatal=False)) + await aiotools.shield_fg(finish_writing()) finally: - # Между закрытием файла и эвентом айнотифи состояние может быть не обновлено, - # так что форсим обновление вручную, чтобы получить актуальное состояние. - await aiotools.shield_fg(self.__reload_state()) + self.__notifier.notify() @aiotools.atomic_fg async def remove(self, name: str) -> None: @@ -400,17 +408,26 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes async def __close_reader(self) -> None: if self.__reader: - await self.__reader.close() - self.__reader = None + try: + await self.__reader.close() + finally: + self.__reader = None async def __close_writer(self) -> None: if self.__writer: - await self.__writer.close() - self.__writer = None + try: + await self.__writer.close() + finally: + self.__writer = None # ===== - async def __watch_inotify(self) -> None: + @aiotools.atomic_fg + async def cleanup(self) -> None: + await self.__close_reader() + await self.__close_writer() + + async def systask(self) -> None: logger = get_logger(0) while True: try: @@ -425,7 +442,11 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes await inotify.watch_all_modify(*self.__storage.get_watchable_paths()) await inotify.watch_all_modify(*self.__drive.get_watchable_paths()) - # После установки вотчеров еще раз проверяем стейт, чтобы ничего не потерять + # После установки вотчеров еще раз проверяем стейт, + # чтобы не потерять состояние привода. + # Из-за гонки между первым релоадом и установкой вотчеров, + # мы можем потерять какие-то каталоги стораджа, но это допустимо, + # так как всегда есть ручной перезапуск. await self.__reload_state() while self.__state.vd: # Если живы после предыдущей проверки @@ -445,54 +466,58 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes await self.__reload_state() except Exception: logger.exception("Unexpected MSD watcher error") - time.sleep(1) + await asyncio.sleep(1) - async def __reload_state(self, notify: bool=True) -> None: - logger = get_logger(0) + async def __reload_state(self) -> None: async with self.__state._lock: # pylint: disable=protected-access - try: - path = self.__drive.get_image_path() - drive_state = _DriveState( - image=((await self.__storage.make_image_by_path(path)) if path else None), - cdrom=self.__drive.get_cdrom_flag(), - rw=self.__drive.get_rw_flag(), - ) - - await self.__storage.reload() - - if self.__state.vd is None and drive_state.image is None: - # Если только что включились и образ не подключен - попробовать - # перемонтировать хранилище (и создать images и meta). - logger.info("Probing to remount storage ...") - await self.__storage.remount_rw(True) - await self.__storage.remount_rw(False) - await self.__setup_initial() + await self.__unsafe_reload_state() + self.__notifier.notify() - except Exception: - logger.exception("Error while reloading MSD state; switching to offline") - self.__state.storage = None - self.__state.vd = None + # ===== Don't call this directly ==== + async def __unsafe_reload_state(self) -> None: + logger = get_logger(0) + try: + path = self.__drive.get_image_path() + drive_state = _DriveState( + image=((await self.__storage.make_image_by_path(path)) if path else None), + cdrom=self.__drive.get_cdrom_flag(), + rw=self.__drive.get_rw_flag(), + ) + + await self.__storage.reload() + + if self.__state.vd is None and drive_state.image is None: + # Если только что включились и образ не подключен - попробовать + # перемонтировать хранилище (и создать images и meta). + logger.info("Probing to remount storage ...") + await self.__storage.remount_rw(True) + await self.__storage.remount_rw(False) + await self.__unsafe_setup_initial() + + except Exception: + logger.exception("Error while reloading MSD state; switching to offline") + self.__state.storage = None + self.__state.vd = None + + else: + self.__state.storage = self.__storage + if drive_state.image: + # При подключенном образе виртуальный стейт заменяется реальным + self.__state.vd = _VirtualDriveState.from_drive_state(drive_state) else: - self.__state.storage = self.__storage - if drive_state.image: - # При подключенном образе виртуальный стейт заменяется реальным + if self.__state.vd is None: + # Если раньше MSD был отключен self.__state.vd = _VirtualDriveState.from_drive_state(drive_state) - else: - if self.__state.vd is None: - # Если раньше MSD был отключен - self.__state.vd = _VirtualDriveState.from_drive_state(drive_state) - image = self.__state.vd.image - if image and (not image.in_storage or not (await image.exists())): - # Если только что отключили ручной образ вне хранилища или ранее выбранный образ был удален - self.__state.vd.image = None + image = self.__state.vd.image + if image and (not image.in_storage or not (await image.exists())): + # Если только что отключили ручной образ вне хранилища или ранее выбранный образ был удален + self.__state.vd.image = None - self.__state.vd.connected = False - if notify: - self.__notifier.notify() + self.__state.vd.connected = False - async def __setup_initial(self) -> None: + async def __unsafe_setup_initial(self) -> None: if self.__initial_image: logger = get_logger(0) image = await self.__storage.make_image_by_name(self.__initial_image) |