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

pikvm/pikvm#321: server-side uploading counters

Browse Source
This commit is contained in:
Devaev Maxim
2021-06-08 03:12:24 +03:00
parent cf08c04e55
commit b5ab5699c4
8 changed files with 115 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -29,13 +29,9 @@ from typing import Dict
from typing import AsyncGenerator
from typing import Optional
import aiofiles
import aiofiles.base
from ....logging import get_logger
from .... import aiotools
from .... import aiofs
from ....yamlconf import Option
@@ -54,10 +50,10 @@ from .. import MsdDisconnectedError
from .. import MsdMultiNotSupported
from .. import MsdCdromNotSupported
from .. import BaseMsd
from .. import MsdImageWriter
from .gpio import Gpio
from .drive import ImageInfo
from .drive import DeviceInfo
@@ -91,9 +87,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
self.__device_info: Optional[DeviceInfo] = None
self.__connected = False
self.__device_file: Optional[aiofiles.base.AiofilesContextManager] = None
self.__written = 0
self.__unsynced = 0
self.__device_writer: Optional[MsdImageWriter] = None
self.__notifier = aiotools.AioNotifier()
self.__region = aiotools.AioExclusiveRegion(MsdIsBusyError, self.__notifier)
@@ -132,7 +126,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
storage = {
"size": self.__device_info.size,
"free": self.__device_info.free,
"uploading": bool(self.__device_file)
"uploading": (self.__device_writer.get_state() if self.__device_writer else None),
}
drive = {
"image": (self.__device_info.image and dataclasses.asdict(self.__device_info.image)),
@@ -177,7 +171,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
@aiotools.atomic
async def cleanup(self) -> None:
try:
await self.__close_device_file()
await self.__close_device_writer()
finally:
self.__gpio.close()
@@ -214,7 +208,7 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
self.__connected = connected
@contextlib.asynccontextmanager
async def write_image(self, name: str) -> AsyncGenerator[None, None]:
async def write_image(self, name: str, size: int) -> AsyncGenerator[None, None]:
async with self.__working():
async with self.__region:
try:
@@ -222,30 +216,22 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
if self.__connected:
raise MsdConnectedError()
self.__device_file = await aiofiles.open(self.__device_info.path, mode="w+b", buffering=0) # type: ignore
self.__written = 0
self.__unsynced = 0
self.__device_writer = await MsdImageWriter(self.__device_info.path, size, self.__sync_chunk_size).open()
await self.__write_image_info(name, complete=False)
await self.__write_image_info(False)
await self.__notifier.notify()
yield
await self.__write_image_info(name, complete=True)
await self.__write_image_info(True)
finally:
await self.__close_device_file()
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_file
await self.__device_file.write(chunk) # type: ignore
self.__written += len(chunk)
self.__unsynced += len(chunk)
if self.__unsynced >= self.__sync_chunk_size:
await aiofs.afile_sync(self.__device_file)
self.__unsynced = 0
return self.__written
assert self.__device_writer
return (await self.__device_writer.write(chunk))
@aiotools.atomic
async def remove(self, name: str) -> None:
@@ -262,27 +248,21 @@ class Plugin(BaseMsd): # pylint: disable=too-many-instance-attributes
# =====
async def __write_image_info(self, name: str, complete: bool) -> None:
assert self.__device_file
async def __write_image_info(self, complete: bool) -> None:
assert self.__device_writer
assert self.__device_info
if not (await self.__device_info.write_image_info(
device_file=self.__device_file,
image_info=ImageInfo(name, self.__written, complete),
)):
if not (await self.__device_info.write_image_info(self.__device_writer, complete)):
get_logger().error("Can't write image info because device is full")
async def __close_device_file(self) -> None:
async def __close_device_writer(self) -> None:
try:
if self.__device_file:
if self.__device_writer:
get_logger().info("Closing device file ...")
await aiofs.afile_sync(self.__device_file)
await self.__device_file.close() # type: ignore
await self.__device_writer.close() # type: ignore
except Exception:
get_logger().exception("Can't close device file")
finally:
self.__device_file = None
self.__written = 0
self.__unsynced = 0
self.__device_writer = None
async def __load_device_info(self) -> None:
retries = self.__init_retries

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

@@ -29,11 +29,11 @@ import dataclasses
from typing import IO
from typing import Optional
import aiofiles.base
from .... import aiotools
from .... import aiofs
from .. import MsdImageWriter
# =====
_IMAGE_INFO_SIZE = 4096
@@ -121,11 +121,10 @@ class DeviceInfo:
image=image_info,
)
async def write_image_info(
self,
device_file: aiofiles.base.AiofilesContextManager,
image_info: ImageInfo,
) -> bool:
async def write_image_info(self, device_writer: MsdImageWriter, complete: bool) -> bool:
device_file = device_writer.get_file()
state = device_writer.get_state()
image_info = ImageInfo(state["name"], state["written"], complete)
if self.size - image_info.size > _IMAGE_INFO_SIZE:
await device_file.seek(self.size - _IMAGE_INFO_SIZE) # type: ignore