async def consume(self) -> None:
"""
This function is a real consumer, get event message from event queue
and handle it
"""
if self.__queue is None:
raise RuntimeError("missing event queue for Handler")
while True:
if self.__queue.empty():
# If this consumer is idle too long, will be exited
if self.__pending_exit:
break
await asyncio.sleep(1)
continue
try:
item = self.__queue.get_nowait()
self.__active_time = time.time()
self.__task_count = self.__task_count + 1
await self.handle(item=item)
self.__task_count = self.__task_count - 1
except ValueError as err:
Logger.error(RuntimeError("Event has closed unexpectedly, consumer cannot get anything"))
except queue.Empty as err:
await asyncio.sleep(0.5)
def subprocess_initializer(log_queue, name="worker"):
"""
Initializer of worker subprocess
:param log_queue: Log queue
:param name: process name
"""
Logger.worker_configure(_queue=log_queue)
multiprocessing.current_process().name = name
Logger.info(f"active new worker: Process-{os.getpid()} {name}")
async def check_timeout(self):
"""
Check the idle time of consumer
"""
while True:
if time.time() - self.__active_time >= self.__timeout and self.__task_count == 0:
Logger.warning(
f"consumers on [Process-{os.getpid()} {multiprocessing.current_process().name}] "
f"is pending to exit, because it's idle too long")
self.__pending_exit = True
await asyncio.sleep(self.__timeout)
async def check_task() -> None:
while True:
done_or_canceled_count = 0
for task in tasks:
try:
if task.done() or task.cancelled():
done_or_canceled_count = done_or_canceled_count + 1
continue
exception = task.exception()
if exception is not None:
Logger.error(exception)
except asyncio.CancelledError as err:
Logger.error(
RuntimeError(
f"a consumer on [Process-{os.getpid()} {multiprocessing.current_process().name}] "
f"has cancelled unexpectedly"))
Logger.error(err)
except asyncio.InvalidStateError:
pass
if done_or_canceled_count == consumer_number:
Logger.warning(
f"consumers on [Process-{os.getpid()} {multiprocessing.current_process().name}] "
f"have all done or canceled, worker stopped")
return
await asyncio.sleep(2)
async def handle_command(self, text: str, event, **kwargs) -> None:
"""
Parse commands and its parameters
:param text: The text contains commands and its parameters
:param event: KHL Event object
:param kwargs: extra data
:return: None
"""
if len(self.get_commands()) == 0:
return
for commands in self.get_commands():
for command in commands:
pos = text.find(command)
if pos != -1:
params = []
for _ in range(commands[command][
CONFIG.COMMANDER_KEY_PARAM_NUMBER]):
next_space = text.find(' ', pos + len(command) + 1)
if next_space != -1:
params.append(text[pos + len(command) +
1:next_space])
else:
params.append(text[pos + len(command) + 1:])
if len(params[-1]) == 0:
params = []
break
pos = next_space + 1
if len(params) == commands[command][
CONFIG.COMMANDER_KEY_PARAM_NUMBER]:
func = commands[command][CONFIG.COMMANDER_KEY_HANDLE]
kwargs[CONFIG.BOT_KEY_EVENT] = event
try:
await func(*params, **kwargs)
except Exception as err:
if isinstance(func, functools.partial):
func_name = func.func.__name__
else:
func_name = func.__name__
Logger.warning(
f"command [{command}] and it's handle function [{func_name}] raise some error"
)
Logger.error(err)
break
def get_gateway(cls, token, url) -> {str, None}:
"""
Get gateway address of khl websocket
:param url: url of khl gateway api
:param token: khl bot token
:return: khl websocket address
"""
headers = {
"Authorization": f"Bot {token}"
}
try:
response = requests.get(url, headers=headers)
json_data = response.json()
if json_data["code"] != 0:
return None
return json_data["data"]["url"]
except Exception as e:
Logger.error(e)
async def get_from_khl_api(url, body, token, success=None, failed=None):
headers = {"Authorization": f"Bot {token}"}
async with aiohttp.ClientSession() as session:
try:
async with session.get(url=url, params=body,
headers=headers) as response:
json_rep = await response.json()
if response.status == 200 and json_rep["code"] == 0:
if success is not None:
success(json_rep)
else:
if failed is not None:
failed(json_rep)
return json_rep
except Exception as e:
Logger.error(e)
if failed is not None:
failed({"code": -1, "message": str(e)})
async def handle_subscribe(self, _type, condition: dict, handle,
event: Event) -> None:
"""
Handle subscribe events
:param _type: KHL Event type
:param condition: Conditions to filter event
:param handle: Handle function
:param event: Event Object
:return: None
"""
can_run = True
if condition is not None:
for item in condition:
keys = str(item).strip().split('.')
step = None
for key in keys:
if step is None:
step = condition[key]
else:
step = step[key]
if step != condition[item]:
can_run = False
break
if can_run:
kwargs = {CONFIG.COMMANDER_KEY_EVENT: event}
try:
await handle(**kwargs)
except Exception as err:
if isinstance(handle, functools.partial):
func_name = handle.func.__name__
else:
func_name = handle.__name__
Logger.warning(
f"subscribe event function [{func_name}] raises an error")
Logger.error(err)
def run(self) -> None:
"""
The entry for launch bot, in this function, bot will connect several components:
Handler(Consumer), KHLWss(Producer), Commander(Task)
"""
if self.__handler is None:
Logger.error(Exception("Please set messages handler."))
return
try:
multiprocessing.current_process(
).name = "Bot" # set main process name
atexit.register(self.__exit_handler) # handle exit event
# handle some signal
signal.signal(signal.SIGINT, Bot.signal_handler)
signal.signal(signal.SIGTERM, Bot.signal_handler)
# set up logger
self.__log_listener = Logger.listener_configure(
_queue=self.__log_queue)
self.__log_listener.start()
self.__wss.event_queue = self.__queue
self.__handler.event_queue = self.__queue
for commander in self.__commanders:
self.__handler.add_commands(commander.get_commands())
self.__handler.add_subscribes(commander.get_subscribes())
loop = asyncio.get_event_loop()
self.__loop = loop
loop.create_task(self.__wss.start())
loop.create_task(self.__check_queue_size())
self.__launch_subprocess(is_leader=True)
self.__running_process_count = self.__running_process_count + 1
# launch interval tasks
self.__launch_interval()
loop.run_forever()
except Exception as e:
Logger.error(e)
Logger.error(RuntimeError("fatal error, bot will be terminated"))
async def handle(self, item) -> None:
"""
Same as Handler.handle
"""
if item[CONFIG.BOT_KEY_MESSAGE_TYPE] == CONFIG.BOT_MESSAGE_TYPE_EVENT:
try:
item = item[CONFIG.BOT_KEY_MESSAGE_DATA]
event = Event(item)
if event.type == CONFIG.KHL_MSG_TEXT:
await self.handle_command(item["content"], event=event)
elif event.type == CONFIG.KHL_MSG_SYSTEM:
system_event_type = event.extra.type
_subscribes = self.get_subscribes()
if system_event_type is not None and system_event_type in _subscribes:
for item in _subscribes[system_event_type]:
await self.handle_subscribe(
_type=system_event_type,
condition=item[
CONFIG.COMMANDER_KEY_CONDITIONS],
handle=item[CONFIG.COMMANDER_KEY_HANDLE],
event=event)
except ValueError as err:
Logger.warning("Please check commands configuration")
Logger.error(err)
elif item[CONFIG.
BOT_KEY_MESSAGE_TYPE] == CONFIG.BOT_MESSAGE_TYPE_INTERVAL:
try:
await item[CONFIG.COMMANDER_KEY_HANDLE]()
except Exception as err:
func = item[CONFIG.COMMANDER_KEY_HANDLE]
if isinstance(func, functools.partial):
func_name = func.func.__name__
else:
func_name = func.__name__
Logger.warning(
f"interval function [{func_name}] raises an error")
Logger.error(err)
def __launch_subprocess(self, is_leader: bool = False) -> None:
"""
Driver function for starting new subprocess
:param is_leader: If True, start a leader subprocess for consumers,
it means the first subprocess for consumers, will not be exited when the process is idle
:return: None
"""
def process_count(result=None):
self.__running_process_count = self.__running_process_count - 1
def process_error(err):
Logger.warning("Have a error when launch worker")
Logger.error(err)
process_count()
self.__pool.apply_async(
Bot.subprocess_consumer,
(self.__handler, self.__config[CONFIG.MAX_CONSUMER_NUMBER],
is_leader),
error_callback=lambda err: Logger.error(err),
callback=process_count)
async def start(self) -> None:
"""
Launch websocket, and do some important operations
:return: None
"""
Logger.info("Getting KHL Websocket address")
uri = KHLWss.get_gateway(token=self.__token, url=KHL_API_BASEURL + KHL_API_GATEWAY)
if uri is None:
Logger.error(Exception("Failed to get websocket address"))
return
Logger.info(f"Success to get websocket address")
async with websockets.client.connect(uri) as ws_connection:
Logger.info("Connect to websocket success, launch bot")
asyncio.create_task(self.heartbeat(ws_connection))
while True:
try:
if ws_connection is None or not ws_connection.open:
Logger.warning("Reconnecting to khl websocket")
uri = KHLWss.get_gateway(token=self.__token, url=KHL_API_BASEURL + KHL_API_GATEWAY) # uri可能已经更新
ws_connection = await websockets.client.connect(uri)
Logger.warning("Reconnect to khl websocket success")
msg = await ws_connection.recv()
json_rep = decompress_to_json(msg)
if json_rep['s'] == 0:
self.latest_sn = json_rep["sn"]
if self.__event_queue is not None:
try:
self.__event_queue.put_nowait({
CONFIG.BOT_KEY_MESSAGE_TYPE: CONFIG.BOT_MESSAGE_TYPE_EVENT,
CONFIG.BOT_KEY_MESSAGE_DATA: json_rep['d']
})
except ValueError as err:
Logger.error(RuntimeError("Event queue has closed unexpectedly"))
Logger.error(err)
except queue.Full:
Logger.warning("Event queue is full")
except (websockets.ConnectionClosedError, websockets.ConnectionClosed) as e:
Logger.warning("The websockets have closed unexpectedly")
Logger.error(e)
# Try reconnection after 3 seconds, avoid too many times to reconnection on a short time
await asyncio.sleep(3)
except Exception as e:
Logger.warning("The websockets have closed unexpectedly")
Logger.error(e)
await asyncio.sleep(3)
def signal_handler(sig=None, frame=None):
Logger.warning("khlbot will be exited")
exit(0)
def process_error(err):
Logger.warning("Have a error when launch worker")
Logger.error(err)
process_count()