94 lines
2.6 KiB
Python
94 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import threading
|
|
from typing import Any
|
|
|
|
|
|
class AsyncLock:
|
|
def __init__(self):
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def __aenter__(self):
|
|
await self._lock.acquire()
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
self._lock.release()
|
|
|
|
@property
|
|
def locked(self) -> bool:
|
|
return self._lock.locked()
|
|
|
|
|
|
class StartupSerialLock:
|
|
def __init__(self):
|
|
self._lock = threading.Lock()
|
|
self._startup_in_progress = False
|
|
|
|
@property
|
|
def is_starting(self) -> bool:
|
|
return self._startup_in_progress
|
|
|
|
def acquire_startup(self) -> bool:
|
|
if self._lock.acquire(blocking=False):
|
|
self._startup_in_progress = True
|
|
return True
|
|
return False
|
|
|
|
def release_startup(self) -> None:
|
|
self._startup_in_progress = False
|
|
try:
|
|
self._lock.release()
|
|
except RuntimeError:
|
|
pass
|
|
|
|
|
|
class AccountDataWriteQueue:
|
|
def __init__(self, max_size: int = 500):
|
|
self._queue: asyncio.Queue[tuple[str, dict[str, Any]]] = asyncio.Queue(maxsize=max_size)
|
|
self._task: asyncio.Task | None = None
|
|
|
|
async def enqueue(self, account_type: str, data: dict[str, Any]) -> None:
|
|
await self._queue.put((account_type, data))
|
|
|
|
async def drain(self) -> list[tuple[str, dict[str, Any]]]:
|
|
items: list[tuple[str, dict[str, Any]]] = []
|
|
while not self._queue.empty():
|
|
try:
|
|
items.append(self._queue.get_nowait())
|
|
except asyncio.QueueEmpty:
|
|
break
|
|
return items
|
|
|
|
def start_consumer(
|
|
self,
|
|
write_fn,
|
|
interval: float = 1.0,
|
|
) -> asyncio.Task:
|
|
async def _consumer():
|
|
while True:
|
|
try:
|
|
batch = await self.drain()
|
|
if batch:
|
|
for account_type, data in batch:
|
|
try:
|
|
await write_fn(account_type, data)
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
await asyncio.sleep(interval)
|
|
|
|
self._task = asyncio.create_task(_consumer())
|
|
return self._task
|
|
|
|
async def stop_consumer(self) -> None:
|
|
if self._task and not self._task.done():
|
|
self._task.cancel()
|
|
try:
|
|
await self._task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
self._task = None
|