feat: add New cache class
This commit is contained in:
parent
6cb4d05b8f
commit
f6c8f406cc
106
backend/MediaCacheManager.py
Normal file
106
backend/MediaCacheManager.py
Normal file
@ -0,0 +1,106 @@
|
||||
import functools
|
||||
import logging
|
||||
import collections
|
||||
import asyncio
|
||||
|
||||
import diskcache
|
||||
|
||||
logger = logging.getLogger(__file__.split("/")[-1])
|
||||
|
||||
@functools.total_ordering
|
||||
class MediaBlockHolder(object):
|
||||
waiters: collections.deque[asyncio.Future]
|
||||
chunk_id: int = 0
|
||||
|
||||
def __init__(self, chat_id: int, msg_id: int, start: int, target_len: int) -> None:
|
||||
self.chat_id = chat_id
|
||||
self.msg_id = msg_id
|
||||
self.start = start
|
||||
self.target_len = target_len
|
||||
self.mem = bytes()
|
||||
self.length = len(self.mem)
|
||||
self.waiters = collections.deque()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"MediaBlockHolder,id:{self.chat_id}-{self.msg_id},start:{self.start},len:{self.length}/{self.target_len}"
|
||||
|
||||
def __eq__(self, other: 'MediaBlockHolder'|int):
|
||||
if isinstance(other, int):
|
||||
return self.start == other
|
||||
return self.start == other.start
|
||||
|
||||
def __le__(self, other: 'MediaBlockHolder'|int):
|
||||
if isinstance(other, int):
|
||||
return self.start <= other
|
||||
return self.start <= other.start
|
||||
|
||||
def __gt__(self, other: 'MediaBlockHolder'|int):
|
||||
if isinstance(other, int):
|
||||
return self.start > other
|
||||
return self.start > other.start
|
||||
|
||||
def __add__(self, other: 'MediaBlockHolder'|bytes):
|
||||
if isinstance(other, bytes):
|
||||
self.append_mem(other)
|
||||
elif isinstance(other, MediaBlockHolder):
|
||||
self.append_mem(other.mem)
|
||||
else:
|
||||
raise RuntimeError(f"{self} can't add {type(other)}")
|
||||
|
||||
def is_completed(self) -> bool:
|
||||
return self.length >= self.target_len
|
||||
|
||||
def notify_waiters(self) -> None:
|
||||
while self.waiters:
|
||||
waiter = self.waiters.popleft()
|
||||
if not waiter.done():
|
||||
waiter.set_result(None)
|
||||
|
||||
def append_mem(self, mem: bytes) -> None:
|
||||
self.mem = self.mem + mem
|
||||
self.length = len(self.mem)
|
||||
self.notify_waiters()
|
||||
if self.length > self.target_len:
|
||||
logger.warning(f"MeidaBlock Overflow:{self}")
|
||||
|
||||
async def wait_update(self) -> None:
|
||||
if self.is_completed():
|
||||
return
|
||||
waiter = asyncio.Future()
|
||||
self.waiters.append(waiter)
|
||||
try:
|
||||
await waiter
|
||||
except:
|
||||
waiter.cancel()
|
||||
try:
|
||||
self.waiters.remove(waiter)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
@functools.total_ordering
|
||||
class BlockInfo(object):
|
||||
def __init__(self, hashid: int, offset: int, length: int, in_mem: bool) -> None:
|
||||
self.hashid = hashid
|
||||
self.offset = offset
|
||||
self.length = length
|
||||
self.in_mem = in_mem
|
||||
|
||||
def __eq__(self, other: 'BlockInfo'|int):
|
||||
if isinstance(other, int):
|
||||
return self.offset == other
|
||||
return self.offset == other.offset
|
||||
|
||||
def __le__(self, other: 'BlockInfo'|int):
|
||||
if isinstance(other, int):
|
||||
return self.offset <= other
|
||||
return self.offset <= other.offset
|
||||
|
||||
class MediaBlockHolderManager(object):
|
||||
|
||||
DEFAULT_MAX_CACHE_SIZE = 1024 * 1024 * 1024 # 1Gb
|
||||
# chat_id -> msg_id -> list[BlockInfo]
|
||||
chunk_cache: dict[int, dict[int, list[BlockInfo]]] = {}
|
||||
|
||||
def __init__(self, limit_size: int = DEFAULT_MAX_CACHE_SIZE, dir: str = 'cache') -> None:
|
||||
pass
|
||||
|
Loading…
x
Reference in New Issue
Block a user