import logging import queue import sys import threading import traceback from functools import wraps from .base import AbstractSession from ..exceptions import SessionNotFoundException, SessionClosedException, SessionException from ..utils import random_str, LimitedSizeQueue, isgeneratorfunction, iscoroutinefunction, catch_exp_call, \ get_function_name logger = logging.getLogger(__name__) """ 基于线程的会话实现 主任务线程退出后，连接关闭，但不会清理主任务线程产生的其他线程 客户端连接关闭后，后端线程不会退出，但是再次调用输入输出函数会引发异常 todo: thread 重名 """ # todo 线程安全 class ThreadBasedSession(AbstractSession): thread2session = {} # thread_id -> session unhandled_task_mq_maxsize = 1000 event_mq_maxsize = 100 callback_mq_maxsize = 100 _active_session_cnt = 0 @classmethod def active_session_count(cls): return cls._active_session_cnt @classmethod def get_current_session(cls) -> "ThreadBasedSession": curr = id(threading.current_thread()) session = cls.thread2session.get(curr) if session is None: raise SessionNotFoundException( "Can't find current session. Maybe session closed. Did you forget to use `register_thread` ?") return session @classmethod def get_current_task_id(cls): return cls._get_task_id(threading.current_thread()) @staticmethod def _get_task_id(thread: threading.Thread): tname = getattr(thread, '_target', 'task') tname = getattr(tname, '__name__', tname) return '%s-%s' % (tname, id(thread)) def __init__(self, target, session_info, on_task_command=None, on_session_close=None, loop=None): """ :param target: 会话运行的函数 :param on_task_command: 当Task内发送Command给session的时候触发的处理函数 :param on_session_close: 会话结束的处理函数 :param loop: 事件循环。若 on_task_command 或者 on_session_close 中有调用使用asyncio事件循环的调用， 则需要事件循环实例来将回调在事件循环的线程中执行 """ assert (not iscoroutinefunction(target)) and (not isgeneratorfunction(target)), ValueError( "ThreadBasedSession only accept a simple function as task function, " "not coroutine function or generator function. ") ThreadBasedSession._active_session_cnt += 1 self.info = session_info self._on_task_command = on_task_command or (lambda _: None) self._on_session_close = on_session_close or (lambda: None) self._loop = loop # 会话结束时运行的函数 self.deferred_functions = [] self.threads = [] # 注册到当前会话的线程集合 self.unhandled_task_msgs = LimitedSizeQueue(maxsize=self.unhandled_task_mq_maxsize) self.task_mqs = {} # task_id -> event msg queue self._closed = False # 用于实现回调函数的注册 self.callback_mq = None self.callback_thread = None self.callbacks = {} # callback_id -> (callback_func, is_mutex) self._start_main_task(target) def _start_main_task(self, target): @wraps(target) def main_task(target): try: target() except Exception as e: if not isinstance(e, SessionException): self.on_task_exception() finally: for t in self.threads: if t.is_alive() and t is not threading.current_thread(): t.join() try: self.send_task_command(dict(command='close_session')) except SessionClosedException: pass self._trigger_close_event() self.close() thread = threading.Thread(target=main_task, kwargs=dict(target=target), daemon=True, name='main_task') self.register_thread(thread) thread.start() def send_task_command(self, command): """向会话发送来自协程内的消息 :param dict command: 消息 """ if self.closed(): raise SessionClosedException() self.unhandled_task_msgs.put(command) if self._loop: self._loop.call_soon_threadsafe(self._on_task_command, self) else: self._on_task_command(self) def next_client_event(self): # 函数开始不需要判断 self.closed() # 如果会话关闭，对 get_current_session().next_client_event() 的调用会抛出SessionNotFoundException task_id = self.get_current_task_id() event_mq = self.get_current_session().task_mqs.get(task_id) event = event_mq.get() if event is None: raise SessionClosedException return event def send_client_event(self, event): """向会话发送来自用户浏览器的事件️ :param dict event: 事件️消息 """ task_id = event['task_id'] mq = self.task_mqs.get(task_id) if not mq and task_id in self.callbacks: mq = self.callback_mq if not mq: logger.error('event_mqs not found, task_id:%s', task_id) return mq.put(event) def get_task_commands(self): return self.unhandled_task_msgs.get() def _trigger_close_event(self): """触发Backend on_session_close callback""" if self._loop: self._loop.call_soon_threadsafe(self._on_session_close) else: self._on_session_close() def _cleanup(self): self.unhandled_task_msgs.wait_empty(8) if not self.unhandled_task_msgs.empty(): logger.debug("Unhandled task messages when session close:%s", self.unhandled_task_msgs.get()) raise RuntimeError('There are unhandled task messages when session close!') for t in self.threads: del ThreadBasedSession.thread2session[id(t)] if self.callback_mq is not None: # 回调功能已经激活 self.callback_mq.put(None) # 结束回调线程 for mq in self.task_mqs.values(): mq.put(None) # 消费端接收到None消息会抛出SessionClosedException异常 self.task_mqs = {} ThreadBasedSession._active_session_cnt -= 1 def close(self): """关闭当前Session。由Backend调用""" # todo self._closed 会有竞争条件 if self._closed: return self._closed = True self._cleanup() self.deferred_functions.reverse() while self.deferred_functions: func = self.deferred_functions.pop() catch_exp_call(func, logger) def closed(self): return self._closed def on_task_exception(self): from ..output import put_markdown # todo logger.exception('Error in thread executing') type, value, tb = sys.exc_info() tb_len = len(list(traceback.walk_tb(tb))) lines = traceback.format_exception(type, value, tb, limit=1 - tb_len) traceback_msg = ''.join(lines) try: put_markdown("发生错误：\n```\n%s\n```" % traceback_msg) except Exception: pass def _activate_callback_env(self): """激活回调功能 ThreadBasedSession 的回调实现原理是：创建一个单独的线程用于接收回调事件，进而调用相关的回调函数。 当用户Task中并没有使用到回调功能时，不必开启此线程，可以节省资源 """ if self.callback_mq is not None: # 回调功能已经激活 return self.callback_mq = queue.Queue(maxsize=self.callback_mq_maxsize) self.callback_thread = threading.Thread(target=self._dispatch_callback_event, daemon=True, name='callback-' + random_str(10)) # self.register_thread(self.callback_thread) self.thread2session[id(self.callback_thread)] = self # 用于在线程内获取会话 self.callback_thread.start() logger.debug('Callback thread start') def _dispatch_callback_event(self): while not self.closed(): event = self.callback_mq.get() if event is None: # 结束信号 logger.debug('Callback thread exit') break callback_info = self.callbacks.get(event['task_id']) if not callback_info: logger.error("No callback for callback_id:%s", event['task_id']) return callback, mutex = callback_info @wraps(callback) def run(callback): try: callback(event['data']) except Exception: # 子类可能会重写 get_current_session ，所以不要用 ThreadBasedSession.get_current_session 来调用 self.get_current_session().on_task_exception() if mutex: run(callback) else: t = threading.Thread(target=run, kwargs=dict(callback=callback), daemon=True) self.register_thread(t) t.start() def register_callback(self, callback, serial_mode=False): """ 向Session注册一个回调函数，返回回调id Session需要保证当收到前端发送的事件消息 ``{event: "callback"，task_id: 回调id, data:...}`` 时， ``callback`` 回调函数被执行， 并传入事件消息中的 ``data`` 字段值作为参数 :param bool serial_mode: 串行模式模式。若为 ``True`` ，则对于同一组件的点击事件，串行执行其回调函数 """ assert (not iscoroutinefunction(callback)) and (not isgeneratorfunction(callback)), ValueError( "In ThreadBasedSession.register_callback, `callback` must be a simple function, " "not coroutine function or generator function. ") self._activate_callback_env() callback_id = 'CB-%s-%s' % (get_function_name(callback, 'callback'), random_str(10)) self.callbacks[callback_id] = (callback, serial_mode) return callback_id def register_thread(self, t: threading.Thread): """将线程注册到当前会话，以便在线程内调用 pywebio 交互函数。 会话会一直保持直到所有通过 `register_thread` 注册的线程以及当前会话的主任务线程退出 :param threading.Thread thread: 线程对象 """ self.threads.append(t) # 保存 registered thread，用于主任务线程退出后等待注册线程结束 self.thread2session[id(t)] = self # 用于在线程内获取会话 event_mq = queue.Queue(maxsize=self.event_mq_maxsize) # 线程内的用户事件队列 self.task_mqs[self._get_task_id(t)] = event_mq def defer_call(self, func): """设置会话结束时调用的函数。可以用于资源清理。""" self.deferred_functions.append(func) class ScriptModeSession(ThreadBasedSession): """Script mode的会话实现""" @classmethod def get_current_session(cls) -> "ScriptModeSession": if cls.instance is None: raise SessionNotFoundException("Can't find current session. It might be a bug.") if cls.instance.closed(): raise SessionClosedException() return cls.instance @classmethod def get_current_task_id(cls): task_id = super().get_current_task_id() session = cls.get_current_session() if task_id not in session.task_mqs: session.register_thread(threading.current_thread()) return task_id instance = None def __init__(self, thread, session_info, on_task_command=None, loop=None): """ :param thread: 第一次调用PyWebIO交互函数的线程 todo 貌似本参数并不必要 :param on_task_command: 会话结束的处理函数。后端Backend在相应on_session_close时关闭连接时， 需要保证会话内的所有消息都传送到了客户端 :param loop: 事件循环。若 on_task_command 或者on_session_close中有调用使用asyncio事件循环的调用， 则需要事件循环实例来将回调在事件循环的线程中执行 """ if ScriptModeSession.instance is not None: raise RuntimeError("ScriptModeSession can only be created once.") ScriptModeSession.instance = self ThreadBasedSession._active_session_cnt += 1 self.info = session_info self._on_task_command = on_task_command or (lambda _: None) self._on_session_close = lambda: None self._loop = loop # 会话结束时运行的函数 self.deferred_functions = [] self.threads = [] # 当前会话的线程 self.unhandled_task_msgs = LimitedSizeQueue(maxsize=self.unhandled_task_mq_maxsize) self.task_mqs = {} # task_id -> event msg queue self._closed = False # 用于实现回调函数的注册 self.callback_mq = None self.callback_thread = None self.callbacks = {} # callback_id -> (callback_func, is_mutex) tid = id(thread) event_mq = queue.Queue(maxsize=self.event_mq_maxsize) self.task_mqs[tid] = event_mq