Skip to content

Commit

Permalink
refactor: Optimized auto reload task in JsonDataSource
Browse files Browse the repository at this point in the history
  • Loading branch information
Starlwr committed Mar 31, 2024
1 parent fcd556d commit 224819e
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 14 deletions.
5 changes: 3 additions & 2 deletions starbot_datasource/__init__.py
Original file line number Diff line number Diff line change
@@ -1,2 +1,3 @@
from .core.event import *
from .core.datasource import *
from .core.event import EventType, DataSourceEvent
from .core.datasource import DataSource, DictDataSource, JsonDataSource
from .core.config import config
46 changes: 34 additions & 12 deletions starbot_datasource/core/datasource.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
import abc
import asyncio
import json
import os
from json import JSONDecodeError
from typing import List, Dict, Union, Optional, NoReturn, Set

from loguru import logger
Expand Down Expand Up @@ -124,7 +126,6 @@ class JsonDataSource(DataSource):
def __init__(self, json_file: str, auto_reload: Optional[bool] = True, auto_reload_interval: Optional[int] = 5):
super().__init__()
self.__config = None
self.__content = ""

self.__json_file = json_file
self.__auto_reload = auto_reload
Expand All @@ -141,8 +142,9 @@ async def load(self) -> NoReturn:
logger.info("开始从 JSON 中初始化 Bot 配置")

try:
modify_time = os.path.getmtime(self.__json_file)
with open(self.__json_file, encoding="utf-8") as file:
self.__content = file.read()
content = file.read()
except FileNotFoundError:
raise DataSourceException("JSON 文件不存在, 请检查文件路径是否正确")
except UnicodeDecodeError:
Expand All @@ -151,7 +153,7 @@ async def load(self) -> NoReturn:
raise DataSourceException(f"读取 JSON 文件异常 {ex}")

try:
self.__config = json.loads(self.__content)
self.__config = json.loads(content)
except Exception:
raise DataSourceException("JSON 文件内容格式不正确")

Expand All @@ -167,21 +169,30 @@ async def load(self) -> NoReturn:
logger.success(f"成功从 JSON 中导入了 {len(self.ups)} 个 UP 主")

if self.__auto_reload:
executor.create_task(self.__auto_reload_task())
executor.create_task(self.__auto_reload_task(modify_time))

async def __auto_reload_task(self) -> NoReturn:
async def __auto_reload_task(self, modify_time: float) -> NoReturn:
"""
JSON 文件内容发生变化时自动重载配置
Args:
modify_time: 配置文件修改时间
"""
while True:
await asyncio.sleep(self.__auto_reload_interval)
try:
with open(self.__json_file, encoding="utf-8") as file:
content = file.read()
new_modify_time = os.path.getmtime(self.__json_file)

if content == self.__content:
if new_modify_time == modify_time:
continue

modify_time = new_modify_time

logger.info(f"数据源配置已更新, 开始重载配置")

with open(self.__json_file, encoding="utf-8") as file:
content = file.read()

conf = json.loads(content)
if isinstance(conf, dict):
conf = [conf]
Expand All @@ -194,17 +205,16 @@ async def __auto_reload_task(self) -> NoReturn:
if len(new_up_set) != len(new_ups):
raise DataSourceException("数据源中不可含有重复的主播")

self.__content = content
self.__config = conf

added_ups = new_up_set - self.ups
removed_ups = self.ups - new_up_set

tip = [f'{act}{len(ups)} 个主播' for act, ups in [('新增', added_ups), ('移除', removed_ups)] if ups]
if tip:
logger.info(f"检测到数据源配置中 {', '.join(tip)}")
logger.info(f"检测到 {', '.join(tip)}")
else:
logger.info(f"检测到数据源配置已更新, 开始重载配置")
logger.info("未检测到新增或移除的主播, 开始重载已存在主播的配置")

for up in new_up_set.union(self.ups):
if up in added_ups:
Expand All @@ -213,5 +223,17 @@ async def __auto_reload_task(self) -> NoReturn:
self.remove(up.uid)
else:
self.update(up)
except Exception:

logger.success("数据源配置重载成功")
except FileNotFoundError:
logger.error("数据源配置 JSON 文件不存在")
except UnicodeDecodeError:
logger.error("数据源配置 JSON 文件编码不正确, 请将其转换为 UTF-8 格式编码")
except JSONDecodeError:
logger.error("数据源配置 JSON 文件内容格式不正确")
except DataSourceException as ex:
logger.error(ex.msg)
except Exception as ex:
logger.error(f"数据源自动重载任务异常 {ex}")
finally:
continue

0 comments on commit 224819e

Please sign in to comment.