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

/msd/write_remote handle

Browse Source
This commit is contained in:
Maxim Devaev
2021-07-27 05:25:54 +03:00
parent 3c421fa94c
commit 6b07a80834
9 changed files with 173 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -31,6 +31,8 @@ from typing import Optional
import aiofiles
import aiofiles.base
from ...logging import get_logger
from ... import aiofs
from ...errors import OperationError
@@ -119,13 +121,10 @@ class BaseMsd(BasePlugin):
raise NotImplementedError()
@contextlib.asynccontextmanager
async def write_image(self, name: str, size: int) -> AsyncGenerator[None, None]: # pylint: disable=unused-argument
async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]: # pylint: disable=unused-argument
if self is not None: # XXX: Vulture and pylint hack
raise NotImplementedError()
yield
def get_upload_chunk_size(self) -> int:
raise NotImplementedError()
yield 1
async def write_image_chunk(self, chunk: bytes) -> int:
raise NotImplementedError()
@@ -158,6 +157,7 @@ class MsdImageWriter:
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
@@ -176,6 +176,13 @@ class MsdImageWriter:
async def close(self) -> None:
assert self.__file is not None
if self.__written == self.__size:
(log, result) = (get_logger().info, "OK")
elif self.__written < self.__size:
(log, result) = (get_logger().error, "INCOMPLETE")
else: # written > size
(log, result) = (get_logger().warning, "OVERFLOW")
log("Written %d of %d bytes to MSD image %r: %s", self.__written, self.__size, self.__name, result)
await aiofs.afile_sync(self.__file)
await self.__file.close() # type: ignore

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

@@ -70,13 +70,10 @@ class Plugin(BaseMsd):
raise MsdDisabledError()
@contextlib.asynccontextmanager
async def write_image(self, name: str, size: int) -> AsyncGenerator[None, None]:
async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:
if self is not None: # XXX: Vulture and pylint hack
raise MsdDisabledError()
yield
def get_upload_chunk_size(self) -> int:
raise MsdDisabledError()
yield 1
async def write_image_chunk(self, chunk: bytes) -> int:
raise MsdDisabledError()

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

@@ -306,7 +306,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
self.__state.vd.connected = connected
@contextlib.asynccontextmanager
async def write_image(self, name: str, size: int) -> AsyncGenerator[None, None]:
async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:
try:
async with self.__state._region: # pylint: disable=protected-access
try:
@@ -328,7 +328,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
self.__new_writer = await MsdImageWriter(path, size, self.__sync_chunk_size).open()
await self.__notifier.notify()
yield
yield self.__upload_chunk_size
self.__set_image_complete(name, True)
finally:
@@ -343,9 +343,6 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
await self.__reload_state()
await self.__notifier.notify()
def get_upload_chunk_size(self) -> int:
return self.__upload_chunk_size
async def write_image_chunk(self, chunk: bytes) -> int:
assert self.__new_writer
written = await self.__new_writer.write(chunk)

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

@@ -208,7 +208,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
self.__connected = connected
@contextlib.asynccontextmanager
async def write_image(self, name: str, size: int) -> AsyncGenerator[None, None]:
async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:
async with self.__working():
async with self.__region:
try:
@@ -220,15 +220,12 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
await self.__write_image_info(False)
await self.__notifier.notify()
yield
yield self.__upload_chunk_size
await self.__write_image_info(True)
finally:
await self.__close_device_writer()
await self.__load_device_info()
def get_upload_chunk_size(self) -> int:
return self.__upload_chunk_size
async def write_image_chunk(self, chunk: bytes) -> int:
assert self.__device_writer
return (await self.__device_writer.write(chunk))