PyWebIO/pywebio/session/threadbased.py

import asyncio
import inspect
import logging
import queue
import sys
import threading
import traceback
from functools import wraps
from .base import AbstractSession
from ..exceptions import SessionNotFoundException, SessionClosedException
from ..utils import random_str, LimitedSizeQueue

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, 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 asyncio.iscoroutinefunction(target)) and (not inspect.isgeneratorfunction(target)), ValueError(
            "ThreadBasedSession only accept a simple function as task function, "
            "not coroutine function or generator function. ")

        ThreadBasedSession._active_session_cnt += 1

        self._on_task_command = on_task_command or (lambda _: None)
        self._on_session_close = on_session_close or (lambda: None)
        self._loop = loop

        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:
                self.on_task_exception()
            finally:
                for t in self.threads:
                    if t.is_alive() and t is not threading.current_thread():
                        t.join()
                self.send_task_command(dict(command='close_session'))
                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: 消息
        """
        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):
        task_id = self.get_current_task_id()
        event_mq = self.get_current_session().task_mqs.get(task_id)
        return event_mq.get()

    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.task_mqs = {}

        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)  # 结束回调线程

        ThreadBasedSession._active_session_cnt -= 1

    def close(self):
        """关闭当前Session。由Backend调用"""
        if self._closed:
            return
        self._closed = True
        self._cleanup()

    def closed(self):
        return self._closed

    def on_task_exception(self):
        from ..output import put_markdown  # todo
        logger.exception('Error in coroutine 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:
            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.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:  # 结束信号
                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:
                    # 子类可能会重写 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 asyncio.iscoroutinefunction(callback)) and (not inspect.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' % (getattr(callback, '__name__', ''), 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)
        self.thread2session[id(t)] = self
        event_mq = queue.Queue(maxsize=self.event_mq_maxsize)
        self.task_mqs[self._get_task_id(t)] = event_mq


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, on_task_command=None, loop=None):
        """
        :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._on_task_command = on_task_command or (lambda _: None)
        self._on_session_close = lambda: None
        self._loop = loop

        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