"""Abstract classes for managing threads and message queue"""
from __future__ import annotations
import asyncio
import contextlib
import typing
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass
import threading
import queue
from datetime import datetime
from enum import IntEnum, auto
from ..common import raise_exception_in_thread
from ..settings import config, logger
if typing.TYPE_CHECKING:
from starlette.websockets import WebSocket
from ..session import Session
from ..components.render.render_node import RenderNode
class ThreadMode(IntEnum):
NORMAL = auto()
REDUCED = auto()
@dataclass
class WorkerStat:
active: bool = False
last_tick: float = 0
mode: ThreadMode = ThreadMode.NORMAL
thread: threading.Thread = None
task_info: str = ""
[docs]
class BaseWorkerServer(ABC):
"""Abstract thread management and communication server class
It has basic thread manager implementation and abstract communication methods.
Read :doc:`threads` and :doc:`message_queue` for more information"""
run_with_web: typing.ClassVar[bool] = False
[docs]
class Listener(ABC):
"""Abstract communication listener class"""
[docs]
@abstractmethod
async def send(self, session_id: str, data: bytes):
"""Send packed message to client by `session_id`"""
[docs]
@abstractmethod
async def receive(self) -> tuple[str, bytes]:
"""Read packed message from client and return tuple (session_id, data)"""
[docs]
@abstractmethod
async def close(self):
"""Close active connection"""
workers: typing.ClassVar[dict[str, WorkerStat]] = {}
task_queue: typing.ClassVar[queue.Queue | None] = queue.Queue()
thread_counter: typing.ClassVar[int] = 0
listener: Listener
async_loop: asyncio.AbstractEventLoop
[docs]
@classmethod
def task_processor(cls):
"""Single task thread processor.
It extracts new function call with args from `task_queue` and let it run in this thread"""
try:
ident = threading.current_thread().name
logger.info('Task thread started')
while True:
try:
func, args, kwargs = cls.task_queue.get(timeout=5)
except queue.Empty:
if cls.workers[ident].mode == ThreadMode.REDUCED:
break
continue
if func is None: break
cls.workers[ident].last_tick = time.perf_counter()
cls.workers[ident].active = True
cls.workers[ident].task_info = func.__name__
func(*args, **kwargs)
if ident not in cls.workers:
break
cls.workers[ident].last_tick = time.perf_counter()
cls.workers[ident].active = False
cls.workers[ident].task_info = ''
except SystemExit:
logger.info('Task thread got exit signal')
@staticmethod
@contextlib.contextmanager
def wrap_session_task(node: RenderNode, func: typing.Callable):
from ..session import SessionTask
session = node.context.session
func_name = f'{node.context.oid}#{func.__name__}'
session.tasks[func_name] = SessionTask(threading.current_thread(), func)
yield
if func_name in session.tasks: # other thread is stopped this already, or we have similar funcs names
del session.tasks[func_name]
@staticmethod
def run_coroutine(node: RenderNode, func: typing.Callable, coro: typing.Coroutine):
from ..session import SessionTask
session = node.context.session
func_name = f'{node.oid}#{func.__name__}'
task = asyncio.run_coroutine_threadsafe(coro, session.server_worker.async_loop)
def on_done(future):
del session.tasks[func_name]
session.tasks[func_name] = SessionTask(task, func)
task.add_done_callback(on_done)
[docs]
@classmethod
def tasks_controller(cls):
"""Common task controller
It supervises all task threads life cycle. See also :doc:`threads`.
"""
while True:
time.sleep(1)
tick = time.perf_counter()
last_tick = 0
for k, v in list(cls.workers.items()): # type: str, WorkerStat
if not v.thread.is_alive():
logger.warning(f"Thread killed `{k}` ({v.task_info})")
del cls.workers[k]
elif v.active and config.THREAD_TIMEOUT and tick - v.last_tick > config.THREAD_TIMEOUT:
logger.warning(f"Thread timeout `{k}` ({v.task_info})")
raise_exception_in_thread(v.thread.native_id)
del cls.workers[k]
elif (config.MAX_TASK_THREADS and config.KILL_THREAD_LAG
and len(cls.workers) > config.MAX_TASK_THREADS
and not v.active and v.last_tick
and tick - v.last_tick > config.KILL_THREAD_LAG):
logger.warning(f"Thread killing `{k}` ({v.task_info})")
v.mode = ThreadMode.REDUCED
elif v.active and v.last_tick > last_tick:
last_tick = v.last_tick
if (config.MIN_TASK_THREADS and config.CREATE_THREAD_LAG
and len(cls.workers) < config.MIN_TASK_THREADS
or not cls.task_queue.empty() and last_tick
and tick - last_tick > config.CREATE_THREAD_LAG):
cls.thread_counter += 1
thread_name = f'X#{cls.thread_counter}'
thread = threading.Thread(target=cls.task_processor, name=thread_name, daemon=True)
logger.warning(f'New thread created `{thread_name}`')
cls.workers[thread_name] = WorkerStat(thread=thread)
thread.start()
[docs]
@staticmethod
def session_killer():
"""Session TTL controller"""
from ..session import Session
while True:
time.sleep(5)
now = datetime.now()
for session_id in frozenset(Session.sessions):
session = Session.sessions[session_id]
if not getattr(session, "just_connected", True) and (now - session.last_touch).seconds >= config.SESSION_TTL:
logger.warning(f'Session {session_id} killed by TTL limit {config.SESSION_TTL} seconds')
for task in list(session.tasks.keys()):
session.kill_task(task)
del Session.sessions[session_id]
@classmethod
def start_task_workers(cls):
logger.info("Starting task workers")
BaseWorkerServer.task_queue = queue.Queue()
for i in range(config.MIN_TASK_THREADS):
thread = threading.Thread(target=cls.task_processor, name=f'#{i}', daemon=True)
BaseWorkerServer.workers[thread.name] = WorkerStat(thread=thread)
thread.start()
threading.Thread(target=cls.tasks_controller, daemon=True).start()
threading.Thread(target=cls.session_killer, daemon=True).start()
[docs]
@abstractmethod
def start_listener(self):
"""Main entrypoint for listener subclass initialization
it should end by::
self.listener = WorkerServer.Listener(...)
"""
async def run_processor(self):
from ..session import Session
from ..serializer import serializer
from ..protocol import process_message
logger.info("Starting task processor")
self.async_loop = asyncio.get_running_loop()
while True:
session_id, message = await self.listener.receive()
data = serializer.decode(message)
if data['C'] == "SESSION":
Session(session_id, data['app'], data['lang'], data['params'])
else:
if (session:=Session.sessions.get(session_id, None)) is None:
from ..protocol import Messages
code = serializer.encode(Messages.reconnect())
await Session.server_worker.listener.send(session_id, code)
else:
session.last_touch = datetime.now()
await process_message(session, data)
[docs]
class BaseWorkerClient(ABC):
"""Abstract thread management and communication client class
It contains abstract communication methods."""
[docs]
class Connection(ABC):
"""Abstract communication connection class"""
[docs]
@abstractmethod
async def send(self, message: bytes):
"""Send a packed message to the server"""
[docs]
@abstractmethod
async def receive(self) -> bytes:
"""Receive a packed message from the server"""
[docs]
@abstractmethod
def close(self):
"""Close the connection"""
connection: Connection
sync_task: asyncio.Task
last_touch: datetime
[docs]
@abstractmethod
def open_connection(self, session_id: str):
"""Main entrypoint for connection subclass initialization
it should end by::
self.connection = WorkerClient.Connection(...)
"""
async def stop_websocket_binding(self):
self.sync_task.cancel()
try:
await self.sync_task
except asyncio.CancelledError:
pass
async def close_connection(self):
await self.stop_websocket_binding()
self.connection.close()
def bind_to_websocket(self, ws: WebSocket):
async def task():
while True:
message = await self.connection.receive()
self.last_touch = datetime.now()
await ws.send_bytes(message) #, compress=len(message)>=1000)
self.sync_task = asyncio.create_task(task())
async def connect_session(self, session_id: str, app: str, lang: list[str], params: dict[str, str]):
from ..serializer import serializer
logger.debug("Initiate session")
data = {
"C": "SESSION",
"app": app,
"lang": lang,
"params": params,
}
message = serializer.encode(data)
await self.connection.send(message)
@typing.overload
def __init__(self, session_id: str, ws: WebSocket, app: str, lang: list[str], params: dict[str, str]): ...
def __init__(self, *args):
self.args = args
self.last_touch = datetime.now()
async def __aenter__(self):
from ..session import Session
logger.debug("Connecting to task processor")
self.open_connection(self.args[0] + '/' + self.args[2])
await Session.remind_errors_client(self.args[1])
self.bind_to_websocket(self.args[1])
await self.connect_session(self.args[0], self.args[2], self.args[3], self.args[4])
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
await self.close_connection()