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

refactored msd reader api

Browse Source
This commit is contained in:
Maxim Devaev 2022-08-03 19:44:08 +03:00
parent 0809daa199
commit 9925198762
5 changed files with 51 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -87,12 +87,10 @@ class MsdApi:
@exposed_http("GET", "/msd/read")
async def __read_handler(self, request: Request) -> StreamResponse:
name = valid_msd_image_name(request.query.get("image"))
async with self.__msd.read_image(name) as size:
async with self.__msd.read_image(name) as reader:
size = reader.get_total_size()
response = await start_streaming(request, "application/octet-stream", size, name)
while True:
chunk = await self.__msd.read_image_chunk()
if not chunk:
return response
async for chunk in reader.read_chunked():
await response.write(chunk)
return response

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

@ -104,6 +104,19 @@ class MsdRwNotSupported(MsdOperationError):
# =====
class BaseMsdReader:
def get_state(self) -> Dict:
raise NotImplementedError()
def get_total_size(self) -> int:
raise NotImplementedError()
async def read_chunked(self) -> AsyncGenerator[bytes, None]:
if self is not None: # XXX: Vulture and pylint hack
raise NotImplementedError()
yield
class BaseMsd(BasePlugin):
async def get_state(self) -> Dict:
raise NotImplementedError()
@ -134,14 +147,11 @@ class BaseMsd(BasePlugin):
raise NotImplementedError()
@contextlib.asynccontextmanager
async def read_image(self, name: str) -> AsyncGenerator[int, None]:
async def read_image(self, name: str) -> AsyncGenerator[BaseMsdReader, None]:
_ = name
if self is not None: # XXX: Vulture and pylint hack
raise NotImplementedError()
yield 1
async def read_image_chunk(self) -> bytes:
raise NotImplementedError()
yield BaseMsdReader()
@contextlib.asynccontextmanager
async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
@ -159,7 +169,7 @@ class BaseMsd(BasePlugin):
raise NotImplementedError()
class MsdImageReader: # pylint: disable=too-many-instance-attributes
class MsdFileReader(BaseMsdReader): # pylint: disable=too-many-instance-attributes
def __init__(self, notifier: aiotools.AioNotifier, path: str, chunk_size: int) -> None:
self.__notifier = notifier
self.__name = os.path.basename(path)
@ -171,10 +181,6 @@ class MsdImageReader: # pylint: disable=too-many-instance-attributes
self.__readed = 0
self.__tick = 0.0
def get_size(self) -> int:
assert self.__file is not None
return self.__file_size
def get_state(self) -> Dict:
return {
"name": self.__name,
@ -182,16 +188,17 @@ class MsdImageReader: # pylint: disable=too-many-instance-attributes
"readed": self.__readed,
}
async def open(self) -> "MsdImageReader":
assert self.__file is None
get_logger(1).info("Reading %r image from MSD ...", self.__name)
self.__file_size = os.stat(self.__path).st_size
self.__file = await aiofiles.open(self.__path, mode="rb") # type: ignore
return self
async def read(self) -> bytes:
def get_total_size(self) -> int:
assert self.__file is not None
return self.__file_size
async def read_chunked(self) -> AsyncGenerator[bytes, None]:
assert self.__file is not None
while True:
chunk = await self.__file.read(self.__chunk_size) # type: ignore
if not chunk:
break
self.__readed += len(chunk)
now = time.monotonic()
@ -199,7 +206,14 @@ class MsdImageReader: # pylint: disable=too-many-instance-attributes
self.__tick = now
await self.__notifier.notify()
return chunk
yield chunk
async def open(self) -> "MsdFileReader":
assert self.__file is None
get_logger(1).info("Reading %r image from MSD ...", self.__name)
self.__file_size = os.stat(self.__path).st_size
self.__file = await aiofiles.open(self.__path, mode="rb") # type: ignore
return self
async def close(self) -> None:
assert self.__file is not None

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

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

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

@ -57,7 +57,7 @@ from .. import MsdImageNotSelected
from .. import MsdUnknownImageError
from .. import MsdImageExistsError
from .. import BaseMsd
from .. import MsdImageReader
from .. import MsdFileReader
from .. import MsdImageWriter
from . import fs
@ -165,7 +165,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
self.__drive = Drive(gadget, instance=0, lun=0)
self.__reader: Optional[MsdImageReader] = None
self.__reader: Optional[MsdFileReader] = None
self.__writer: Optional[MsdImageWriter] = None
self.__notifier = aiotools.AioNotifier()
@ -329,7 +329,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
self.__state.vd.connected = connected
@contextlib.asynccontextmanager
async def read_image(self, name: str) -> AsyncGenerator[int, None]:
async def read_image(self, name: str) -> AsyncGenerator[MsdFileReader, None]:
try:
async with self.__state._region: # pylint: disable=protected-access
try:
@ -345,23 +345,19 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
if name not in self.__state.storage.images or not os.path.exists(path):
raise MsdUnknownImageError()
self.__reader = await MsdImageReader(
self.__reader = await MsdFileReader(
notifier=self.__notifier,
path=path,
chunk_size=self.__read_chunk_size,
).open()
yield self.__reader.get_size()
yield self.__reader
finally:
await self.__close_reader()
finally:
await self.__notifier.notify()
async def read_image_chunk(self) -> bytes:
assert self.__reader
return (await self.__reader.read())
@contextlib.asynccontextmanager
async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]:
try:

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

@ -50,6 +50,7 @@ from .. import MsdDisconnectedError
from .. import MsdMultiNotSupported
from .. import MsdCdromNotSupported
from .. import MsdRwNotSupported
from .. import BaseMsdReader
from .. import BaseMsd
from .. import MsdImageWriter
@ -218,15 +219,11 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
self.__connected = connected
@contextlib.asynccontextmanager
async def read_image(self, name: str) -> AsyncGenerator[int, None]:
async def read_image(self, name: str) -> AsyncGenerator[BaseMsdReader, None]:
async with self.__working():
if self is not None: # XXX: Vulture and pylint hack
raise MsdMultiNotSupported()
yield 1
async def read_image_chunk(self) -> bytes:
async with self.__working():
raise MsdMultiNotSupported()
yield BaseMsdReader()
@contextlib.asynccontextmanager
async def write_image(self, name: str, size: int, remove_incomplete: Optional[bool]) -> AsyncGenerator[int, None]: