GithubFollow/One-KVM
1
0
Fork 0
mirror of https://github.com/mofeng-git/One-KVM.git synced 2025-12-12 01:00:29 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

refactoring

Browse Source
This commit is contained in:
Maxim Devaev 2024-10-28 10:46:12 +02:00
parent a84242c9bc
commit 60f413c1f4
2 changed files with 109 additions and 86 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
kvmd/plugins/msd/__init__.py
View File

@ -266,9 +266,6 @@ class MsdFileWriter(BaseMsdWriter): # pylint: disable=too-many-instance-attribu
return self.__written return self.__written
def is_complete(self) -> bool:
return (self.__written >= self.__file_size)
async def open(self) -> "MsdFileWriter": async def open(self) -> "MsdFileWriter":
assert self.__file is None assert self.__file is None
get_logger(1).info("Writing %r image (%d bytes) to MSD ...", self.__name, self.__file_size) 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 self.__file = await aiofiles.open(self.__path, mode="w+b", buffering=0) # type: ignore
return self return self
async def finish(self) -> bool:
await self.__sync()
return (self.__written >= self.__file_size)
async def close(self) -> None: async def close(self) -> None:
assert self.__file is not None assert self.__file is not None
logger = get_logger() logger = get_logger()
@ -288,9 +289,6 @@ class MsdFileWriter(BaseMsdWriter): # pylint: disable=too-many-instance-attribu
else: # written > size else: # written > size
(log, result) = (logger.warning, "OVERFLOW") (log, result) = (logger.warning, "OVERFLOW")
log("Written %d of %d bytes to MSD image %r: %s", self.__written, self.__file_size, self.__name, result) 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: except Exception:
logger.exception("Can't close image writer") logger.exception("Can't close image writer")

97
kvmd/plugins/msd/otg/__init__.py
View File

@ -25,7 +25,6 @@ import contextlib
import dataclasses import dataclasses
import functools import functools
import copy import copy
import time
from typing import AsyncGenerator from typing import AsyncGenerator
@ -96,6 +95,7 @@ class _State:
@contextlib.asynccontextmanager @contextlib.asynccontextmanager
async def busy(self, check_online: bool=True) -> AsyncGenerator[None, None]: async def busy(self, check_online: bool=True) -> AsyncGenerator[None, None]:
try:
with self._region: with self._region:
async with self._lock: async with self._lock:
self.__notifier.notify() self.__notifier.notify()
@ -104,6 +104,7 @@ class _State:
raise MsdOfflineError() raise MsdOfflineError()
assert self.storage assert self.storage
yield yield
finally:
self.__notifier.notify() self.__notifier.notify()
def is_busy(self) -> bool: def is_busy(self) -> bool:
@ -143,7 +144,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
logger = get_logger(0) logger = get_logger(0)
logger.info("Using OTG gadget %r as MSD", gadget) 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 @classmethod
def get_plugin_options(cls) -> dict: 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 def get_state(self) -> dict:
async with self.__state._lock: # pylint: disable=protected-access async with self.__state._lock: # pylint: disable=protected-access
storage: (dict | None) = None storage: (dict | None) = None
if self.__state.storage: if self.__state.storage:
if self.__writer: if self.__writer:
# При загрузке файла показываем актуальную статистику вручную # При загрузке файла показываем актуальную статистику вручную
await self.__storage.reload_parts_info() await aiotools.shield_fg(self.__storage.reload_parts_info())
storage = dataclasses.asdict(self.__state.storage) storage = dataclasses.asdict(self.__state.storage)
for name in list(storage["images"]): for name in list(storage["images"]):
@ -212,9 +215,6 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
prev = copy.deepcopy(new) prev = copy.deepcopy(new)
yield new yield new
async def systask(self) -> None:
await self.__watch_inotify()
@aiotools.atomic_fg @aiotools.atomic_fg
async def reset(self) -> None: async def reset(self) -> None:
async with self.__state.busy(check_online=False): async with self.__state.busy(check_online=False):
@ -226,11 +226,6 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
except Exception: except Exception:
get_logger(0).exception("Can't reset MSD properly") 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 @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 async with self.__state._lock: # pylint: disable=protected-access
self.__notifier.notify() self.__notifier.notify()
self.__STATE_check_disconnected() self.__STATE_check_disconnected()
image = await self.__STATE_get_storage_image(name) image = await self.__STATE_get_storage_image(name)
self.__reader = await MsdFileReader( self.__reader = await MsdFileReader(
notifier=self.__notifier, notifier=self.__notifier,
@ -304,7 +300,10 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
path=image.path, path=image.path,
chunk_size=self.__read_chunk_size, chunk_size=self.__read_chunk_size,
).open() ).open()
self.__notifier.notify()
yield self.__reader yield self.__reader
finally: finally:
await aiotools.shield_fg(self.__close_reader()) await aiotools.shield_fg(self.__close_reader())
finally: finally:
@ -312,18 +311,37 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
@contextlib.asynccontextmanager @contextlib.asynccontextmanager
async def write_image(self, name: str, size: int, remove_incomplete: (bool | None)) -> AsyncGenerator[MsdFileWriter, None]: 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: try:
with self.__state._region: # pylint: disable=protected-access with self.__state._region: # pylint: disable=protected-access
image: (Image | None) = None
try: try:
async with self.__state._lock: # pylint: disable=protected-access async with self.__state._lock: # pylint: disable=protected-access
self.__notifier.notify() self.__notifier.notify()
self.__STATE_check_disconnected() 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.remount_rw(True)
await image.set_complete(False) await image.set_complete(False)
self.__writer = await MsdFileWriter( self.__writer = await MsdFileWriter(
notifier=self.__notifier, notifier=self.__notifier,
name=image.name, name=image.name,
@ -335,22 +353,12 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
self.__notifier.notify() self.__notifier.notify()
yield self.__writer yield self.__writer
await image.set_complete(self.__writer.is_complete()) complete = await self.__writer.finish()
finally: finally:
try: await aiotools.shield_fg(finish_writing())
if image and remove_incomplete and self.__writer and not self.__writer.is_complete():
await image.remove(fatal=False)
finally: finally:
try: self.__notifier.notify()
await aiotools.shield_fg(self.__close_writer())
finally:
if image:
await aiotools.shield_fg(image.remount_rw(False, fatal=False))
finally:
# Между закрытием файла и эвентом айнотифи состояние может быть не обновлено,
# так что форсим обновление вручную, чтобы получить актуальное состояние.
await aiotools.shield_fg(self.__reload_state())
@aiotools.atomic_fg @aiotools.atomic_fg
async def remove(self, name: str) -> None: 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: async def __close_reader(self) -> None:
if self.__reader: if self.__reader:
try:
await self.__reader.close() await self.__reader.close()
finally:
self.__reader = None self.__reader = None
async def __close_writer(self) -> None: async def __close_writer(self) -> None:
if self.__writer: if self.__writer:
try:
await self.__writer.close() await self.__writer.close()
finally:
self.__writer = None 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) logger = get_logger(0)
while True: while True:
try: 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.__storage.get_watchable_paths())
await inotify.watch_all_modify(*self.__drive.get_watchable_paths()) await inotify.watch_all_modify(*self.__drive.get_watchable_paths())
# После установки вотчеров еще раз проверяем стейт, чтобы ничего не потерять # После установки вотчеров еще раз проверяем стейт,
# чтобы не потерять состояние привода.
# Из-за гонки между первым релоадом и установкой вотчеров,
# мы можем потерять какие-то каталоги стораджа, но это допустимо,
# так как всегда есть ручной перезапуск.
await self.__reload_state() await self.__reload_state()
while self.__state.vd: # Если живы после предыдущей проверки while self.__state.vd: # Если живы после предыдущей проверки
@ -445,11 +466,17 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
await self.__reload_state() await self.__reload_state()
except Exception: except Exception:
logger.exception("Unexpected MSD watcher error") logger.exception("Unexpected MSD watcher error")
time.sleep(1) await asyncio.sleep(1)
async def __reload_state(self, notify: bool=True) -> None: async def __reload_state(self) -> None:
logger = get_logger(0)
async with self.__state._lock: # pylint: disable=protected-access async with self.__state._lock: # pylint: disable=protected-access
await self.__unsafe_reload_state()
self.__notifier.notify()
# ===== Don't call this directly ====
async def __unsafe_reload_state(self) -> None:
logger = get_logger(0)
try: try:
path = self.__drive.get_image_path() path = self.__drive.get_image_path()
drive_state = _DriveState( drive_state = _DriveState(
@ -466,7 +493,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
logger.info("Probing to remount storage ...") logger.info("Probing to remount storage ...")
await self.__storage.remount_rw(True) await self.__storage.remount_rw(True)
await self.__storage.remount_rw(False) await self.__storage.remount_rw(False)
await self.__setup_initial() await self.__unsafe_setup_initial()
except Exception: except Exception:
logger.exception("Error while reloading MSD state; switching to offline") logger.exception("Error while reloading MSD state; switching to offline")
@ -489,10 +516,8 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
self.__state.vd.image = None self.__state.vd.image = None
self.__state.vd.connected = False self.__state.vd.connected = False
if notify:
self.__notifier.notify()
async def __setup_initial(self) -> None: async def __unsafe_setup_initial(self) -> None:
if self.__initial_image: if self.__initial_image:
logger = get_logger(0) logger = get_logger(0)
image = await self.__storage.make_image_by_name(self.__initial_image) image = await self.__storage.make_image_by_name(self.__initial_image)