|
|
# -*- coding: utf-8 -*-
|
|
|
import asyncio
|
|
|
import logging
|
|
|
from typing import *
|
|
|
|
|
|
import aiohttp
|
|
|
|
|
|
from . import handlers
|
|
|
from . import models
|
|
|
|
|
|
__all__ = (
|
|
|
'BlcPluginClient',
|
|
|
)
|
|
|
|
|
|
logger = logging.getLogger('blcsdk')
|
|
|
|
|
|
|
|
|
class BlcPluginClient:
|
|
|
"""
|
|
|
blivechat插件服务的客户端
|
|
|
|
|
|
:param ws_url: blivechat消息转发服务WebSocket地址
|
|
|
:param session: 连接池
|
|
|
:param heartbeat_interval: 发送心跳包的间隔时间（秒）
|
|
|
"""
|
|
|
|
|
|
def __init__(
|
|
|
self,
|
|
|
ws_url: str,
|
|
|
*,
|
|
|
session: Optional[aiohttp.ClientSession] = None,
|
|
|
heartbeat_interval: float = 30,
|
|
|
):
|
|
|
self._ws_url = ws_url
|
|
|
|
|
|
if session is None:
|
|
|
self._session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=10))
|
|
|
self._own_session = True
|
|
|
else:
|
|
|
self._session = session
|
|
|
self._own_session = False
|
|
|
assert self._session.loop is asyncio.get_event_loop() # noqa
|
|
|
|
|
|
self._heartbeat_interval = heartbeat_interval
|
|
|
|
|
|
self._handler: Optional[handlers.HandlerInterface] = None
|
|
|
"""消息处理器"""
|
|
|
|
|
|
# 在运行时初始化的字段
|
|
|
self._websocket: Optional[aiohttp.ClientWebSocketResponse] = None
|
|
|
"""WebSocket连接"""
|
|
|
self._network_future: Optional[asyncio.Future] = None
|
|
|
"""网络协程的future"""
|
|
|
self._heartbeat_timer_handle: Optional[asyncio.TimerHandle] = None
|
|
|
"""发心跳包定时器的handle"""
|
|
|
|
|
|
@property
|
|
|
def is_running(self) -> bool:
|
|
|
"""本客户端正在运行，注意调用stop后还没完全停止也算正在运行"""
|
|
|
return self._network_future is not None
|
|
|
|
|
|
def set_handler(self, handler: Optional['handlers.HandlerInterface']):
|
|
|
"""
|
|
|
设置消息处理器
|
|
|
|
|
|
注意消息处理器和网络协程运行在同一个协程，如果处理消息耗时太长会阻塞接收消息。如果是CPU密集型的任务，建议将消息推到线程池处理；
|
|
|
如果是IO密集型的任务，应该使用async函数，并且在handler里使用create_task创建新的协程
|
|
|
|
|
|
:param handler: 消息处理器
|
|
|
"""
|
|
|
self._handler = handler
|
|
|
|
|
|
def start(self):
|
|
|
"""启动本客户端"""
|
|
|
if self.is_running:
|
|
|
logger.warning('Plugin client is running, cannot start() again')
|
|
|
return
|
|
|
|
|
|
self._network_future = asyncio.create_task(self._network_coroutine_wrapper())
|
|
|
|
|
|
def stop(self):
|
|
|
"""停止本客户端"""
|
|
|
if not self.is_running:
|
|
|
logger.warning('Plugin client is stopped, cannot stop() again')
|
|
|
return
|
|
|
|
|
|
self._network_future.cancel()
|
|
|
|
|
|
async def stop_and_close(self):
|
|
|
"""便利函数，停止本客户端并释放本客户端的资源，调用后本客户端将不可用"""
|
|
|
if self.is_running:
|
|
|
self.stop()
|
|
|
await self.join()
|
|
|
await self.close()
|
|
|
|
|
|
async def join(self):
|
|
|
"""等待本客户端停止"""
|
|
|
if not self.is_running:
|
|
|
logger.warning('Plugin client is stopped, cannot join()')
|
|
|
return
|
|
|
|
|
|
await asyncio.shield(self._network_future)
|
|
|
|
|
|
async def close(self):
|
|
|
"""释放本客户端的资源，调用后本客户端将不可用"""
|
|
|
if self.is_running:
|
|
|
logger.warning('Plugin is calling close(), but client is running')
|
|
|
|
|
|
# 如果session是自己创建的则关闭session
|
|
|
if self._own_session:
|
|
|
await self._session.close()
|
|
|
|
|
|
async def send_cmd_data(self, cmd: models.Command, data: dict):
|
|
|
"""
|
|
|
发送消息给服务器
|
|
|
|
|
|
:param cmd: 消息类型，见Command
|
|
|
:param data: 消息体JSON数据
|
|
|
"""
|
|
|
if self._websocket is None or self._websocket.closed:
|
|
|
raise ConnectionResetError('websocket is closed')
|
|
|
|
|
|
body = {'cmd': cmd, 'data': data}
|
|
|
await self._websocket.send_json(body)
|
|
|
|
|
|
async def _network_coroutine_wrapper(self):
|
|
|
"""负责处理网络协程的异常，网络协程具体逻辑在_network_coroutine里"""
|
|
|
exc = None
|
|
|
try:
|
|
|
await self._network_coroutine()
|
|
|
except asyncio.CancelledError:
|
|
|
# 正常停止
|
|
|
pass
|
|
|
except Exception as e:
|
|
|
logger.exception('_network_coroutine() finished with exception:')
|
|
|
exc = e
|
|
|
finally:
|
|
|
logger.debug('_network_coroutine() finished')
|
|
|
self._network_future = None
|
|
|
|
|
|
if self._handler is not None:
|
|
|
self._handler.on_client_stopped(self, exc)
|
|
|
|
|
|
async def _network_coroutine(self):
|
|
|
"""网络协程，负责连接服务器、接收消息、解包"""
|
|
|
try:
|
|
|
# 连接
|
|
|
async with self._session.ws_connect(
|
|
|
self._ws_url,
|
|
|
receive_timeout=self._heartbeat_interval + 5,
|
|
|
) as websocket:
|
|
|
self._websocket = websocket
|
|
|
await self._on_ws_connect()
|
|
|
|
|
|
# 处理消息
|
|
|
message: aiohttp.WSMessage
|
|
|
async for message in websocket:
|
|
|
self._on_ws_message(message)
|
|
|
finally:
|
|
|
self._websocket = None
|
|
|
await self._on_ws_close()
|
|
|
# 插件消息都是本地通信的，这里不可能是因为网络问题而掉线，所以不尝试重连
|
|
|
|
|
|
async def _on_ws_connect(self):
|
|
|
"""WebSocket连接成功"""
|
|
|
self._heartbeat_timer_handle = asyncio.get_running_loop().call_later(
|
|
|
self._heartbeat_interval, self._on_send_heartbeat
|
|
|
)
|
|
|
|
|
|
async def _on_ws_close(self):
|
|
|
"""WebSocket连接断开"""
|
|
|
if self._heartbeat_timer_handle is not None:
|
|
|
self._heartbeat_timer_handle.cancel()
|
|
|
self._heartbeat_timer_handle = None
|
|
|
|
|
|
def _on_send_heartbeat(self):
|
|
|
"""定时发送心跳包的回调"""
|
|
|
if self._websocket is None or self._websocket.closed:
|
|
|
self._heartbeat_timer_handle = None
|
|
|
return
|
|
|
|
|
|
self._heartbeat_timer_handle = asyncio.get_running_loop().call_later(
|
|
|
self._heartbeat_interval, self._on_send_heartbeat
|
|
|
)
|
|
|
asyncio.create_task(self._send_heartbeat())
|
|
|
|
|
|
async def _send_heartbeat(self):
|
|
|
"""发送心跳包"""
|
|
|
try:
|
|
|
await self.send_cmd_data(models.Command.HEARTBEAT, {})
|
|
|
except (ConnectionResetError, aiohttp.ClientConnectionError) as e:
|
|
|
logger.warning('Plugin client _send_heartbeat() failed: %r', e)
|
|
|
except Exception: # noqa
|
|
|
logger.exception('Plugin client _send_heartbeat() failed:')
|
|
|
|
|
|
def _on_ws_message(self, message: aiohttp.WSMessage):
|
|
|
"""
|
|
|
收到WebSocket消息
|
|
|
|
|
|
:param message: WebSocket消息
|
|
|
"""
|
|
|
if message.type != aiohttp.WSMsgType.TEXT:
|
|
|
logger.warning('Unknown websocket message type=%s, data=%s', message.type, message.data)
|
|
|
return
|
|
|
|
|
|
try:
|
|
|
body = message.json()
|
|
|
self._handle_command(body)
|
|
|
except Exception:
|
|
|
logger.error('body=%s', message.data)
|
|
|
raise
|
|
|
|
|
|
def _handle_command(self, command: dict):
|
|
|
"""
|
|
|
处理业务消息
|
|
|
|
|
|
:param command: 业务消息
|
|
|
"""
|
|
|
if self._handler is not None:
|
|
|
try:
|
|
|
self._handler.handle(self, command)
|
|
|
except Exception as e:
|
|
|
logger.exception('Plugin client _handle_command() failed, command=%s', command, exc_info=e)
|