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

refactored msd writer api

Browse Source
This commit is contained in:
Maxim Devaev 2022-08-04 04:04:10 +03:00
parent 9925198762
commit 9ee63aba3e
6 changed files with 66 additions and 60 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
kvmd/apps/kvmd/api/msd.py
View File

@ -100,12 +100,13 @@ class MsdApi:
size = valid_int_f0(request.content_length)
remove_incomplete = self.__get_remove_incomplete(request)
written = 0
async with self.__msd.write_image(name, size, remove_incomplete) as chunk_size:
async with self.__msd.write_image(name, size, remove_incomplete) as writer:
chunk_size = writer.get_chunk_size()
while True:
chunk = await request.content.read(chunk_size)
if not chunk:
break
written = await self.__msd.write_image_chunk(chunk)
written = await writer.write_chunk(chunk)
return make_json_response(self.__make_write_info(name, size, written))
@exposed_http("POST", "/msd/write_remote")
@ -139,12 +140,13 @@ class MsdApi:
size = valid_int_f0(remote.content_length)
get_logger(0).info("Downloading image %r as %r to MSD ...", url, name)
async with self.__msd.write_image(name, size, remove_incomplete) as chunk_size:
async with self.__msd.write_image(name, size, remove_incomplete) as writer:
chunk_size = writer.get_chunk_size()
response = await start_streaming(request, "application/x-ndjson")
await stream_write_info()
last_report_ts = 0
async for chunk in remote.content.iter_chunked(chunk_size):
written = await self.__msd.write_image_chunk(chunk)
written = await writer.write_chunk(chunk)
now = int(time.time())
if last_report_ts + 1 < now:
await stream_write_info()

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

@ -117,6 +117,17 @@ class BaseMsdReader:
yield
class BaseMsdWriter:
def get_state(self) -> Dict:
raise NotImplementedError()
def get_chunk_size(self) -> int:
raise NotImplementedError()
async def write_chunk(self, chunk: bytes) -> int:
raise NotImplementedError()
class BaseMsd(BasePlugin):
async def get_state(self) -> Dict:
raise NotImplementedError()
@ -154,16 +165,13 @@ class BaseMsd(BasePlugin):
yield BaseMsdReader()
@contextlib.asynccontextmanager
async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[BaseMsdWriter, None]:
_ = name
_ = size
_ = remove_incomplete
if self is not None: # XXX: Vulture and pylint hack
raise NotImplementedError()
yield 1
async def write_image_chunk(self, chunk: bytes) -> int:
raise NotImplementedError()
yield BaseMsdWriter()
async def remove(self, name: str) -> None:
raise NotImplementedError()
@ -225,47 +233,38 @@ class MsdFileReader(BaseMsdReader): # pylint: disable=too-many-instance-attribu
logger.exception("Can't close image reader")
class MsdImageWriter: # pylint: disable=too-many-instance-attributes
def __init__(self, notifier: aiotools.AioNotifier, path: str, size: int, sync: int) -> None:
class MsdFileWriter(BaseMsdWriter): # pylint: disable=too-many-instance-attributes
def __init__(self, notifier: aiotools.AioNotifier, path: str, file_size: int, sync_size: int, chunk_size: int) -> None:
self.__notifier = notifier
self.__name = os.path.basename(path)
self.__path = path
self.__size = size
self.__sync = sync
self.__file_size = file_size
self.__sync_size = sync_size
self.__chunk_size = chunk_size
self.__file: Optional[aiofiles.base.AiofilesContextManager] = None
self.__written = 0
self.__unsynced = 0
self.__tick = 0.0
def is_complete(self) -> bool:
return (self.__written >= self.__size)
def get_file(self) -> aiofiles.base.AiofilesContextManager:
assert self.__file is not None
return self.__file
def get_state(self) -> Dict:
return {
"name": self.__name,
"size": self.__size,
"size": self.__file_size,
"written": self.__written,
}
async def open(self) -> "MsdImageWriter":
assert self.__file is None
get_logger(1).info("Writing %r image (%d bytes) to MSD ...", self.__name, self.__size)
self.__file = await aiofiles.open(self.__path, mode="w+b", buffering=0) # type: ignore
return self
def get_chunk_size(self) -> int:
return self.__chunk_size
async def write(self, chunk: bytes) -> int:
async def write_chunk(self, chunk: bytes) -> int:
assert self.__file is not None
await self.__file.write(chunk) # type: ignore
self.__written += len(chunk)
self.__unsynced += len(chunk)
if self.__unsynced >= self.__sync:
if self.__unsynced >= self.__sync_size:
await aiofs.afile_sync(self.__file)
self.__unsynced = 0
@ -276,18 +275,31 @@ class MsdImageWriter: # pylint: disable=too-many-instance-attributes
return self.__written
def is_complete(self) -> bool:
return (self.__written >= self.__file_size)
def get_file(self) -> aiofiles.base.AiofilesContextManager:
assert self.__file is not None
return self.__file
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)
self.__file = await aiofiles.open(self.__path, mode="w+b", buffering=0) # type: ignore
return self
async def close(self) -> None:
assert self.__file is not None
logger = get_logger()
logger.info("Closing image writer ...")
try:
if self.__written == self.__size:
if self.__written == self.__file_size:
(log, result) = (logger.info, "OK")
elif self.__written < self.__size:
elif self.__written < self.__file_size:
(log, result) = (logger.error, "INCOMPLETE")
else: # written > size
(log, result) = (logger.warning, "OVERFLOW")
log("Written %d of %d bytes to MSD image %r: %s", self.__written, self.__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 aiofs.afile_sync(self.__file)
finally:

8
kvmd/plugins/msd/disabled.py
View File

@ -30,6 +30,7 @@ from ... import aiotools
from . import MsdOperationError
from . import BaseMsdReader
from . import BaseMsdWriter
from . import BaseMsd
@ -84,13 +85,10 @@ class Plugin(BaseMsd):
yield BaseMsdReader()
@contextlib.asynccontextmanager
async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[BaseMsdWriter, None]:
if self is not None: # XXX: Vulture and pylint hack
raise MsdDisabledError()
yield 1
async def write_image_chunk(self, chunk: bytes) -> int:
raise MsdDisabledError()
yield BaseMsdWriter()
async def remove(self, name: str) -> None:
raise MsdDisabledError()

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

@ -58,7 +58,7 @@ from .. import MsdUnknownImageError
from .. import MsdImageExistsError
from .. import BaseMsd
from .. import MsdFileReader
from .. import MsdImageWriter
from .. import MsdFileWriter
from . import fs
@ -166,7 +166,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
self.__drive = Drive(gadget, instance=0, lun=0)
self.__reader: Optional[MsdFileReader] = None
self.__writer: Optional[MsdImageWriter] = None
self.__writer: Optional[MsdFileWriter] = None
self.__notifier = aiotools.AioNotifier()
self.__state = _State(self.__notifier)
@ -359,7 +359,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
await self.__notifier.notify()
@contextlib.asynccontextmanager
async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[MsdFileWriter, None]:
try:
async with self.__state._region: # pylint: disable=protected-access
path: str = ""
@ -379,15 +379,16 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
await self.__remount_rw(True)
self.__set_image_complete(name, False)
self.__writer = await MsdImageWriter(
self.__writer = await MsdFileWriter(
notifier=self.__notifier,
path=path,
size=size,
sync=self.__sync_chunk_size,
file_size=size,
sync_size=self.__sync_chunk_size,
chunk_size=self.__write_chunk_size,
).open()
await self.__notifier.notify()
yield self.__write_chunk_size
yield self.__writer
self.__set_image_complete(name, self.__writer.is_complete())
finally:
@ -405,10 +406,6 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
await self.__reload_state()
await self.__notifier.notify()
async def write_image_chunk(self, chunk: bytes) -> int:
assert self.__writer
return (await self.__writer.write(chunk))
@aiotools.atomic
async def remove(self, name: str) -> None:
async with self.__state.busy():

19
kvmd/plugins/msd/relay/__init__.py
View File

@ -52,7 +52,7 @@ from .. import MsdCdromNotSupported
from .. import MsdRwNotSupported
from .. import BaseMsdReader
from .. import BaseMsd
from .. import MsdImageWriter
from .. import MsdFileWriter
from .gpio import Gpio
@ -89,7 +89,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
self.__device_info: Optional[DeviceInfo] = None
self.__connected = False
self.__device_writer: Optional[MsdImageWriter] = None
self.__device_writer: Optional[MsdFileWriter] = None
self.__notifier = aiotools.AioNotifier()
self.__region = aiotools.AioExclusiveRegion(MsdIsBusyError, self.__notifier)
@ -226,7 +226,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
yield BaseMsdReader()
@contextlib.asynccontextmanager
async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[MsdFileWriter, None]:
async with self.__working():
if remove_incomplete is not None:
raise MsdMultiNotSupported()
@ -236,25 +236,22 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
if self.__connected:
raise MsdConnectedError()
self.__device_writer = await MsdImageWriter(
self.__device_writer = await MsdFileWriter(
notifier=self.__notifier,
path=self.__device_info.path,
size=size,
sync=self.__sync_chunk_size,
file_size=size,
sync_size=self.__sync_chunk_size,
chunk_size=self.__upload_chunk_size,
).open()
await self.__write_image_info(False)
await self.__notifier.notify()
yield self.__upload_chunk_size
yield self.__device_writer
await self.__write_image_info(True)
finally:
await self.__close_device_writer()
await self.__load_device_info()
async def write_image_chunk(self, chunk: bytes) -> int:
assert self.__device_writer
return (await self.__device_writer.write(chunk))
@aiotools.atomic
async def remove(self, name: str) -> None:
async with self.__working():

4
kvmd/plugins/msd/relay/drive.py
View File

@ -32,7 +32,7 @@ from typing import Optional
from .... import aiotools
from .... import aiofs
from .. import MsdImageWriter
from .. import MsdFileWriter
# =====
@ -121,7 +121,7 @@ class DeviceInfo:
image=image_info,
)
async def write_image_info(self, device_writer: MsdImageWriter, complete: bool) -> bool:
async def write_image_info(self, device_writer: MsdFileWriter, complete: bool) -> bool:
device_file = device_writer.get_file()
state = device_writer.get_state()
image_info = ImageInfo(state["name"], state["written"], complete)