Lib.utils.QQDataCacher

  1"""
  2QQ数据缓存
  3"""
  4import time
  5import threading
  6
  7from ..core import OnebotAPI, ConfigManager
  8from . import Logger
  9
 10NotFetched = type("NotFetched", (), {"__getattr__": lambda _, __: NotFetched,
 11                                     "__repr__": lambda _: "NotFetched",
 12                                     "__bool__": lambda _: False})
 13api = OnebotAPI.api
 14logger = Logger.get_logger()
 15
 16if ConfigManager.GlobalConfig().qq_data_cache.enable:
 17    expire_time = ConfigManager.GlobalConfig().qq_data_cache.expire_time
 18else:
 19    expire_time = 0
 20
 21
 22class QQDataItem:
 23    """
 24    QQ数据缓存类
 25    """
 26
 27    def __init__(self):
 28        self._data = NotFetched  # 数据
 29        self.last_update = time.time()  # 最后刷新时间
 30        self.last_use = -1  # 最后被使用时间（数据被使用）
 31
 32    def refresh_cache(self):
 33        """
 34        刷新缓存
 35        Returns:
 36            None
 37        """
 38        self.last_update = time.time()
 39
 40
 41class UserData(QQDataItem):
 42    """
 43    QQ用户数据缓存类
 44    """
 45
 46    def __init__(
 47            self,
 48            user_id: int,
 49            nickname: str = NotFetched,
 50            sex: str = NotFetched,
 51            age: int = NotFetched,
 52            is_friend: bool = NotFetched,
 53            remark: str | None = NotFetched  # 此值仅在是好友的时候会存在
 54    ):
 55        super().__init__()
 56        self._user_id = user_id
 57        self._data = {
 58            "user_id": user_id,
 59            "nickname": nickname,
 60            "sex": sex,
 61            "age": age,
 62            "is_friend": is_friend,
 63            "remark": remark
 64        }
 65
 66    def refresh_cache(self):
 67        """
 68        刷新缓存
 69        Returns:
 70            None
 71        """
 72        if int(self._user_id) <= 0:
 73            logger.warn(f"获取用户{self._user_id}缓存信息失败: user_id小于等于0")
 74            return
 75        try:
 76            data = api.get_stranger_info(self._user_id)
 77            for k in data:
 78                self._data[k] = data[k]
 79            self._data["is_friend"] = NotFetched
 80            self._data["remark"] = NotFetched
 81        except Exception as e:
 82            logger.warn(f"获取用户{self._user_id}缓存信息失败: {repr(e)}")
 83            return
 84
 85    def __getattr__(self, item):
 86        if item == "_data" or item == "data":
 87            return self._data
 88
 89        if item in ["remark", "is_friend"] and self._data.get(item) != NotFetched:
 90            try:
 91                res = api.get_friend_list()
 92                for friend in res:
 93                    if friend["user_id"] == self._user_id:
 94                        self._data["remark"] = friend["remark"]
 95                        self._data["is_friend"] = True
 96                        break
 97                else:
 98                    self._data["is_friend"] = False
 99                    self._data["remark"] = None
100            except Exception as e:
101                logger.warn(f"获取用户{self._user_id}是否为好友失败: {repr(e)}")
102                return None
103
104        if self._data.get(item) == NotFetched or time.time() - self.last_update > expire_time:
105            self.refresh_cache()
106
107        if self._data.get(item) == NotFetched:
108            return None
109
110        if item in self._data:
111            self.last_use = time.time()
112
113        return self._data.get(item)
114
115    def get_nickname(self) -> str:
116        """
117        获取昵称（如果有备注名优先返回备注名）
118        Returns:
119            昵称
120        """
121        return self.remark or self.nickname
122
123    def __repr__(self):
124        return f"UserData(user_id={self._user_id})"
125
126
127class GroupMemberData(QQDataItem):
128    """
129    QQ群成员数据缓存类
130    """
131
132    def __init__(
133            self,
134            group_id: int,
135            user_id: int,
136            nickname: str = NotFetched,
137            card: str = NotFetched,
138            sex: str = NotFetched,
139            age: int = NotFetched,
140            area: str = NotFetched,
141            join_time: int = NotFetched,
142            last_sent_time: int = NotFetched,
143            level: str = NotFetched,
144            role: str = NotFetched,
145            unfriendly: bool = NotFetched,
146            title: str = NotFetched,
147            title_expire_time: int = NotFetched,
148            card_changeable: bool = NotFetched,
149    ):
150        super().__init__()
151        self._group_id = group_id
152        self._user_id = user_id
153        self._data = {
154            "group_id": group_id,
155            "user_id": user_id,
156            "nickname": nickname,
157            "card": card,
158            "sex": sex,
159            "age": age,
160            "area": area,
161            "join_time": join_time,
162            "last_sent_time": last_sent_time,
163            "level": level,
164            "role": role,
165            "unfriendly": unfriendly,
166            "title": title,
167            "title_expire_time": title_expire_time,
168            "card_changeable": card_changeable,
169        }
170
171    def refresh_cache(self):
172        """
173        刷新缓存
174        Returns:
175            None
176        """
177        if int(self._group_id) <= 0 or int(self._user_id) <= 0:
178            logger.warn(f"获取群{self._group_id}中成员{self._user_id}缓存信息失败: group_id或user_id小于等于0")
179            return
180        try:
181            data = api.get_group_member_info(self._group_id, self._user_id, no_cache=True)
182            for k in data:
183                self._data[k] = data[k]
184        except Exception as e:
185            logger.warn(f"获取群{self._group_id}中成员{self._user_id}缓存信息失败: {repr(e)}")
186            user_data = get_user_info(self._user_id)
187            self._data["nickname"] = user_data.nickname if user_data.nickname else NotFetched
188            self._data["sex"] = user_data.sex if user_data.sex else NotFetched
189            self._data["age"] = user_data.age if user_data.age else NotFetched
190        super().refresh_cache()
191
192    def __getattr__(self, item):
193        if item == "_data" or item == "data":
194            return self._data
195
196        if self._data.get(item) == NotFetched or time.time() - self.last_update > expire_time:
197            self.refresh_cache()
198
199        if self._data.get(item) == NotFetched:
200            return None
201
202        if item in self._data:
203            self.last_use = time.time()
204
205        return self._data.get(item)
206
207    def __repr__(self):
208        return f"GroupMemberData(group_id={self._group_id}, user_id={self._user_id})"
209
210    def get_nickname(self):
211        """
212        获取群名片（如果有群名片优先返回群名片）
213        Returns:
214            群名片
215        """
216        return self.card or self.nickname
217
218
219class GroupData(QQDataItem):
220    """
221    QQ群数据缓存类
222    """
223
224    def __init__(
225            self,
226            group_id: int,
227            group_name: str = NotFetched,
228            member_count: int = NotFetched,
229            max_member_count: int = NotFetched
230    ):
231        super().__init__()
232        self._group_id = group_id
233        self._data = {
234            "group_id": group_id,
235            "group_name": group_name,
236            "member_count": member_count,
237            "max_member_count": max_member_count,
238            "group_member_list": NotFetched
239        }
240
241    def refresh_cache(self):
242        """
243        刷新缓存
244        Returns:
245            None
246        """
247        if int(self._group_id) <= 0:
248            logger.warn(f"获取群{self._group_id}缓存信息失败: group_id小于等于0")
249            return
250        try:
251            data = api.get_group_info(group_id=self._group_id, no_cache=True)
252            for k in data:
253                self._data[k] = data[k]
254            self._data["group_member_list"] = NotFetched
255        except Exception as e:
256            logger.warn(f"获取群{self._group_id}缓存信息失败: {repr(e)}")
257            return
258        super().refresh_cache()
259
260    def __getattr__(self, item):
261        if item == "_data" or item == "data":
262            return self._data
263
264        if item == "group_member_list" and self._data.get(item) == NotFetched:
265            try:
266                res = api.get_group_member_list(self._group_id)
267                member_list = [GroupMemberData(**{k: (v if v is not None else NotFetched)
268                                                  for k, v in member.items()})
269                               for member in res]
270                self._data[item] = member_list
271            except Exception as e:
272                logger.warn(f"获取群{self._group_id}成员列表信息失败: {repr(e)}")
273                return
274
275        if self._data.get(item) == NotFetched or time.time() - self.last_update > expire_time:
276            self.refresh_cache()
277
278        if self._data.get(item) == NotFetched:
279            return None
280
281        if item in self._data:
282            self.last_use = time.time()
283
284        return self._data.get(item)
285
286    def __repr__(self):
287        return f"GroupData(group_id={self._group_id})"
288
289
290group_info = {}
291group_member_info = {}
292user_info = {}
293
294group_info_lock = threading.Lock()
295group_member_info_lock = threading.Lock()
296user_info_lock = threading.Lock()
297
298max_cache_size = ConfigManager.GlobalConfig().qq_data_cache.max_cache_size
299expire_time = expire_time
300
301
302def get_user_info(user_id: int, *args, **kwargs) -> UserData:
303    """
304    获取用户信息
305    Args:
306        user_id: 用户ID
307    Returns:
308        None
309    """
310    with user_info_lock:
311        if user_id not in user_info:
312            data = UserData(user_id, *args, **kwargs)
313            user_info[user_id] = data
314
315        data = user_info[user_id]
316        return data
317
318
319def get_group_info(group_id: int, *args, **kwargs) -> GroupData:
320    """
321    获取群信息
322    Args:
323        group_id: 群号
324    Returns:
325        None
326    """
327    with group_info_lock:
328        if group_id not in group_info:
329            data = GroupData(group_id, *args, **kwargs)
330            group_info[group_id] = data
331
332        data = group_info[group_id]
333        return data
334
335
336def get_group_member_info(group_id: int, user_id: int, *args, **kwargs) -> GroupMemberData:
337    """
338    获取群成员信息
339    Args:
340        group_id: 群号
341        user_id: 用户ID
342    Returns:
343        None
344    """
345    with group_member_info_lock:
346        if group_id not in group_member_info:
347            group_member_info[group_id] = {}
348
349        if user_id not in group_member_info[group_id]:
350            data = GroupMemberData(group_id, user_id, *args, **kwargs)
351            group_member_info[group_id][user_id] = data
352
353        data = group_member_info[group_id][user_id]
354        return data
355
356
357def garbage_collection():
358    """
359    垃圾回收
360    Returns:
361        None
362    """
363    counter = 0
364
365    with group_member_info_lock:
366        for k in list(group_member_info.keys()):
367            group_member_items = list(zip(group_member_info[k].keys(), group_member_info[k].values()))
368            max_last_use_time = max([item[1].last_use for item in group_member_items])
369
370            if max_last_use_time < time.time() - expire_time * 2:
371                del group_member_info[k]
372                counter += 1
373                continue
374
375            group_member_items.sort(key=lambda x: x[1].last_use)
376            if len(group_member_items) > max_cache_size * (2 / 3):
377                for user_id, _ in group_member_items[:int(max_cache_size * (1 / 3))]:
378                    del group_member_info[k][user_id]
379                    counter += 1
380            del group_member_items, max_last_use_time
381
382    with group_info_lock:
383        group_items = list(zip(group_info.keys(), group_info.values()))
384        group_items.sort(key=lambda x: x[1].last_use)
385        if len(group_items) > max_cache_size * (2 / 3):
386            for group_id, _ in group_items[:int(max_cache_size * (1 / 3))]:
387                del group_info[group_id]
388                counter += 1
389        del group_items
390
391    with user_info_lock:
392        user_items = list(zip(user_info.keys(), user_info.values()))
393        user_items.sort(key=lambda x: x[1].last_use)
394        if len(user_items) > max_cache_size * (2 / 3):
395            for user_id, _ in user_items[:int(max_cache_size * (1 / 3))]:
396                del user_info[user_id]
397                counter += 1
398        del user_items
399
400    return counter
401
402
403def scheduled_garbage_collection():
404    """
405    定时垃圾回收
406    Returns:
407        None
408    """
409    t = 0
410    while True:
411        time.sleep(60)
412        t += 1
413        if (
414                t > 4 or (
415                t > 1 and (
416                len(group_info) > max_cache_size or
417                len(user_info) > max_cache_size or
418                len(group_member_info) > max_cache_size
419        )
420        )
421        ):
422            t = 0
423            logger.debug("QQ数据缓存清理开始...")
424            try:
425                counter = garbage_collection()
426                logger.debug(f"QQ数据缓存清理完成，共清理了 {counter} 项信息。")
427            except Exception as e:
428                logger.warn(f"QQ数据缓存清理时出现异常: {repr(e)}")
429
430
431# 启动垃圾回收线程
432threading.Thread(target=scheduled_garbage_collection, daemon=True).start()