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

partially asynchronous msd

Browse Source
This commit is contained in:
Maxim Devaev 2023-03-16 01:32:05 +02:00
parent baf454bef9
commit 4d6da37f40
3 changed files with 35 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -27,6 +27,7 @@ import time
from typing import AsyncGenerator from typing import AsyncGenerator
import aiofiles import aiofiles
import aiofiles.os
import aiofiles.base import aiofiles.base
from ...logging import get_logger from ...logging import get_logger
@ -207,7 +208,7 @@ class MsdFileReader(BaseMsdReader): # pylint: disable=too-many-instance-attribu
async def open(self) -> "MsdFileReader": async def open(self) -> "MsdFileReader":
assert self.__file is None assert self.__file is None
get_logger(1).info("Reading %r image from MSD ...", self.__name) get_logger(1).info("Reading %r image from MSD ...", self.__name)
self.__file_size = os.stat(self.__path).st_size self.__file_size = (await aiofiles.os.stat(self.__path)).st_size
self.__file = await aiofiles.open(self.__path, mode="rb") # type: ignore self.__file = await aiofiles.open(self.__path, mode="rb") # type: ignore
return self return self
@ -269,7 +270,7 @@ class MsdFileWriter(BaseMsdWriter): # pylint: disable=too-many-instance-attribu
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)
os.makedirs(os.path.dirname(self.__path), exist_ok=True) await aiofiles.os.makedirs(os.path.dirname(self.__path), exist_ok=True)
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

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

@ -254,7 +254,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
if name is not None: if name is not None:
if name: if name:
self.__state.vd.image = self.__STATE_get_storage_image(name) self.__state.vd.image = await self.__STATE_get_storage_image(name)
else: else:
self.__state.vd.image = None self.__state.vd.image = None
@ -278,7 +278,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
if self.__state.vd.image is None: if self.__state.vd.image is None:
raise MsdImageNotSelected() raise MsdImageNotSelected()
if not self.__state.vd.image.exists(): if not (await self.__state.vd.image.exists()):
raise MsdUnknownImageError() raise MsdUnknownImageError()
assert self.__state.vd.image.in_storage assert self.__state.vd.image.in_storage
@ -304,7 +304,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 = 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,
name=image.name, name=image.name,
@ -326,10 +326,10 @@ 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 = self.__STORAGE_create_new_image(name) image = await self.__STORAGE_create_new_image(name)
await image.remount_rw(True) await image.remount_rw(True)
image.set_complete(False) await image.set_complete(False)
self.__writer = await MsdFileWriter( self.__writer = await MsdFileWriter(
notifier=self.__notifier, notifier=self.__notifier,
@ -342,11 +342,13 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
self.__notifier.notify() self.__notifier.notify()
yield self.__writer yield self.__writer
image.set_complete(self.__writer.is_complete()) await image.set_complete(self.__writer.is_complete())
finally: finally:
try:
if image and remove_incomplete and self.__writer and not self.__writer.is_complete(): if image and remove_incomplete and self.__writer and not self.__writer.is_complete():
image.remove(fatal=False) await image.remove(fatal=False)
finally:
try: try:
await aiotools.shield_fg(self.__close_writer()) await aiotools.shield_fg(self.__close_writer())
finally: finally:
@ -363,17 +365,17 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
assert self.__state.storage assert self.__state.storage
assert self.__state.vd assert self.__state.vd
self.__STATE_check_disconnected() self.__STATE_check_disconnected()
image = self.__STATE_get_storage_image(name) image = await self.__STATE_get_storage_image(name)
if self.__state.vd.image == image: if self.__state.vd.image == image:
self.__state.vd.image = None self.__state.vd.image = None
await image.remount_rw(True) await image.remount_rw(True)
try: try:
image.remove(fatal=True) await image.remove(fatal=True)
del self.__state.storage.images[name] del self.__state.storage.images[name]
finally: finally:
await image.remount_rw(False, fatal=False) await aiotools.shield_fg(image.remount_rw(False, fatal=False))
# ===== # =====
@ -387,18 +389,18 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
if self.__state.vd.connected or self.__drive.get_image_path(): if self.__state.vd.connected or self.__drive.get_image_path():
raise MsdConnectedError() raise MsdConnectedError()
def __STATE_get_storage_image(self, name: str) -> Image: # pylint: disable=invalid-name async def __STATE_get_storage_image(self, name: str) -> Image: # pylint: disable=invalid-name
assert self.__state.storage assert self.__state.storage
image = self.__state.storage.images.get(name) image = self.__state.storage.images.get(name)
if image is None or not image.exists(): if image is None or not (await image.exists()):
raise MsdUnknownImageError() raise MsdUnknownImageError()
assert image.in_storage assert image.in_storage
return image return image
def __STORAGE_create_new_image(self, name: str) -> Image: # pylint: disable=invalid-name async def __STORAGE_create_new_image(self, name: str) -> Image: # pylint: disable=invalid-name
assert self.__state.storage assert self.__state.storage
image = self.__storage.get_image_by_name(name) image = self.__storage.get_image_by_name(name)
if image.name in self.__state.storage.images or image.exists(): if image.name in self.__state.storage.images or (await image.exists()):
raise MsdImageExistsError() raise MsdImageExistsError()
return image return image

20
kvmd/plugins/msd/otg/storage.py
View File

@ -26,6 +26,9 @@ import dataclasses
from typing import Generator from typing import Generator
from typing import Optional from typing import Optional
import aiofiles
import aiofiles.os
from ....logging import get_logger from ....logging import get_logger
from .... import aiohelpers from .... import aiohelpers
@ -85,32 +88,33 @@ class Image(_Image):
except Exception: except Exception:
return 0.0 return 0.0
def exists(self) -> bool: async def exists(self) -> bool:
return os.path.exists(self.path) return (await aiofiles.os.path.exists(self.path))
async def remount_rw(self, rw: bool, fatal: bool=True) -> None: async def remount_rw(self, rw: bool, fatal: bool=True) -> None:
assert self.__storage assert self.__storage
if not self.__adopted: if not self.__adopted:
await self.__storage.remount_rw(rw, fatal) await self.__storage.remount_rw(rw, fatal)
def remove(self, fatal: bool) -> None: async def remove(self, fatal: bool) -> None:
assert self.__storage assert self.__storage
try: try:
os.remove(self.path) await aiofiles.os.remove(self.path)
except FileNotFoundError: except FileNotFoundError:
pass pass
except Exception: except Exception:
if fatal: if fatal:
raise raise
self.set_complete(False) await self.set_complete(False)
def set_complete(self, flag: bool) -> None: async def set_complete(self, flag: bool) -> None:
assert self.__storage assert self.__storage
if flag: if flag:
open(self.__complete_path, "w").close() # pylint: disable=consider-using-with async with aiofiles.open(self.__complete_path, "w"):
pass
else: else:
try: try:
os.remove(self.__complete_path) await aiofiles.os.remove(self.__complete_path)
except FileNotFoundError: except FileNotFoundError:
pass pass