PyWebIO/pywebio/platform/adaptor/ws.py

import abc
import asyncio
import json
import logging
import time
import typing
from typing import Dict, Optional

from ...session import CoroutineBasedSession, Session, ThreadBasedSession
from ...utils import LRUDict, iscoroutinefunction, isgeneratorfunction, random_str
from ..utils import deserialize_binary_event

logger = logging.getLogger(__name__)


# used to store global state when reconnect enabled
class _reconnect_state:
    # used to clean up session
    detached_sessions = LRUDict()  # session_id -> detached timestamp. In increasing order of the time

    # unclosed and unexpired session
    # used to clean up session
    # used to retrieve session by id when new connection
    unclosed_sessions: Dict[str, Session] = {}  # session_id -> session

    # the messages that can't deliver to browser when session close due to connection lost
    session_will_messages: Dict[str, list] = {}  # session_id -> unhandled message list

    # used to get the active conn in session's callbacks
    active_connections: Dict[str, 'WebSocketConnection'] = {}  # session_id -> WSHandler

    expire_second = 0


def set_expire_second(sec):
    _reconnect_state.expire_second = max(_reconnect_state.expire_second, sec)


def clean_expired_sessions():
    while _reconnect_state.detached_sessions:
        session_id, detached_ts = _reconnect_state.detached_sessions.popitem(last=False)  # 弹出最早过期的session

        if time.time() < detached_ts + _reconnect_state.expire_second:
            # this session is not expired
            _reconnect_state.detached_sessions[session_id] = detached_ts  # restore
            _reconnect_state.detached_sessions.move_to_end(session_id, last=False)  # move to head
            break

        # clean this session
        logger.debug("session %s expired" % session_id)
        _reconnect_state.active_connections.pop(session_id, None)
        _reconnect_state.session_will_messages.pop(session_id, None)
        session = _reconnect_state.unclosed_sessions.pop(session_id, None)
        if session:
            session.close(nonblock=True)


_session_clean_task_started = False


async def session_clean_task():
    global _session_clean_task_started
    if _session_clean_task_started or not _reconnect_state.expire_second:
        return

    _session_clean_task_started = True
    logger.debug("Start session cleaning task")
    while True:
        try:
            clean_expired_sessions()
        except Exception:
            logger.exception("Error when clean expired sessions")

        await asyncio.sleep(_reconnect_state.expire_second // 2)


class WebSocketConnection(abc.ABC):
    @abc.abstractmethod
    def get_query_argument(self, name) -> typing.Optional[str]:
        pass

    @abc.abstractmethod
    def make_session_info(self) -> dict:
        pass

    @abc.abstractmethod
    def write_message(self, message: dict):
        pass

    @abc.abstractmethod
    def closed(self) -> bool:
        return False

    @abc.abstractmethod
    def close(self):
        pass


class WebSocketHandler:
    """
    hold by one connection,
    share one session with multiple connection in session lifetime, but one conn at a time
    """

    session_id: Optional[str] = None
    session: Optional[Session] = None  # the session that current connection attaches
    connection: WebSocketConnection
    reconnectable: bool

    def __init__(self, connection: WebSocketConnection, application, reconnectable: bool, ioloop=None):
        logger.debug("WebSocket opened")
        self.connection = connection
        self.reconnectable = reconnectable
        self.session_id = connection.get_query_argument('session')
        self.ioloop = ioloop or asyncio.get_event_loop()

        if self.session_id in ('NEW', None):  # 初始请求，创建新 Session
            self._init_session(application)
            if reconnectable:
                _reconnect_state.active_connections[self.session_id] = self.connection
                _reconnect_state.unclosed_sessions[self.session_id] = self.session
                # set session id to client, so the client can send it back to server to recover a session when it
                # resumes form a connection lost
                connection.write_message(dict(command='set_session_id', spec=self.session_id))
        elif self.session_id not in _reconnect_state.unclosed_sessions:  # session is expired
            bye_msg = dict(command='close_session')
            for m in _reconnect_state.session_will_messages.get(self.session_id, [bye_msg]):
                try:
                    connection.write_message(m)
                except Exception:
                    logger.exception("Error in sending message via websocket")
        else:  # resumes form a connection lost
            self.session = _reconnect_state.unclosed_sessions[self.session_id]
            _reconnect_state.detached_sessions.pop(self.session_id, None)
            _reconnect_state.active_connections[self.session_id] = connection
            # send the latest messages to client
            self._send_msg_to_client()

        logger.debug('session id: %s' % self.session_id)

    def _init_session(self, application):
        session_info = self.connection.make_session_info()
        self.session_id = random_str(24)

        if iscoroutinefunction(application) or isgeneratorfunction(application):
            self.session = CoroutineBasedSession(
                application, session_info=session_info,
                on_task_command=self._send_msg_to_client,
                on_session_close=self._close_from_session)
        else:
            self.session = ThreadBasedSession(
                application, session_info=session_info,
                on_task_command=self._send_msg_to_client,
                on_session_close=self._close_from_session,
                loop=self.ioloop)

    def _get_active_connection(self) -> Optional[WebSocketConnection]:
        # when reconnect enabled, the active connection for this session is in _reconnect_state.active_connections,
        # otherwise, it's self.connection.
        if self.reconnectable:
            conn = _reconnect_state.active_connections.get(self.session_id)
        else:
            conn = self.connection

        return conn

    def _send_msg_to_client(self, session: Session = None):
        conn = self._get_active_connection()
        session = session or self.session

        if not conn or conn.closed():
            return

        for msg in session.get_task_commands():
            try:
                conn.write_message(msg)
            except TypeError as e:
                logger.exception('Data serialization error: %s\n'
                                 'This may be because you pass the wrong type of parameter to the function'
                                 ' of PyWebIO.\nData content: %s', e, msg)
            except Exception:
                logger.exception("Error in sending message via websocket")

    def _close_from_session(self):
        conn = self._get_active_connection()
        if conn and not conn.closed():
            self._send_msg_to_client()
            conn.close()
        elif self.reconnectable:  # no active connection, and reconnect is enabled
            _reconnect_state.session_will_messages[self.session_id] = self.session.get_task_commands()
        self.session = None

    def send_client_data(self, data):
        if isinstance(data, bytes):
            event = deserialize_binary_event(data)
        else:
            event = json.loads(data)
        if event is None:
            return
        self.session.send_client_event(event)

    def notify_connection_lost(self):
        logger.debug("WebSocket closed")
        if not self.reconnectable and self.session:
            # when the connection lost is caused by `on_session_close()`, it's OK to close the session here though.
            # because the `session.close()` is reentrant
            self.session.close(nonblock=True)
            self.session = None  # reset the reference
            return

        _reconnect_state.active_connections.pop(self.session_id, None)
        if self.session_id in _reconnect_state.unclosed_sessions:
            _reconnect_state.detached_sessions[self.session_id] = time.time()