PyWebIO/pywebio/session/threadbased.py

import asyncio
import inspect
import logging
import queue
import sys
import threading
import traceback

from .base import AbstractSession
from ..exceptions import SessionNotFoundException
from ..utils import random_str

logger = logging.getLogger(__name__)

"""
基于线程的会话实现

主任务线程退出后，连接关闭，但不会清理主任务线程产生的其他线程

客户端连接关闭后，后端线程不会退出，但是再次调用输入输出函数会引发异常
todo: thread 重名
"""


# todo 线程安全
class ThreadBasedSession(AbstractSession):
    thread2session = {}  # thread_id -> session

    event_mq_maxsize = 100
    callback_mq_maxsize = 100

    @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')
        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事件循环的调用，
            则需要事件循环实例来将回调在事件循环的线程中执行
        """
        self._on_task_command = on_task_command or (lambda _: None)
        self._on_session_close = on_session_close or (lambda: None)
        self._loop = loop

        self._server_msg_lock = threading.Lock()
        self.threads = []  # 注册到当前会话的线程集合
        self.unhandled_task_msgs = []

        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):
        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. ")

        def thread_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=thread_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: 消息
        """
        with self._server_msg_lock:
            self.unhandled_task_msgs.append(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):
        with self._server_msg_lock:
            msgs = self.unhandled_task_msgs
            self.unhandled_task_msgs = []
        return msgs

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

        # Don't clean unhandled_task_msgs, it may not send to client
        # self.unhandled_task_msgs = []

        for t in self.threads:
            del ThreadBasedSession.thread2session[id(t)]

        if self.callback_mq is not None:  # 回调功能已经激活
            self.callback_mq.put(None)  # 结束回调线程

    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)
        put_markdown("发生错误：\n```\n%s\n```" % traceback_msg)

    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

            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.")
        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

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

        self._server_msg_lock = threading.Lock()
        self.threads = []  # 当前会话的线程
        self.unhandled_task_msgs = []

        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