✨ 支持上报超时的任务

nonebot_bison/scheduler/scheduler.py

@@ -1,8 +1,10 @@
 from dataclasses import dataclass
 from collections import defaultdict
+from collections.abc import Callable
 
 from nonebot.log import logger
 from nonebot_plugin_apscheduler import scheduler
+from apscheduler.events import EVENT_JOB_MAX_INSTANCES
 from nonebot_plugin_saa.utils.exceptions import NoBotFound
 
 from nonebot_bison.utils import ClientManager
@@ -24,6 +26,15 @@ class Schedulable:
     use_batch: bool = False
 
 
+def handle_time_exceeded(event):
+    # event.job_id 是该任务在 apscheduler 的 id, 进而可以获得该任务的函数，再获取该函数绑定的对象
+    logger.warning(f"{scheduler.get_job(event.job_id).func.__self__.name} 抓取执行超时")
+    scheduler.get_job(event.job_id).func.__self__.metrics_report(False)
+
+
+scheduler.add_listener(handle_time_exceeded, EVENT_JOB_MAX_INSTANCES)
+
+
 class Scheduler:
     schedulable_list: list[Schedulable]  # for load weigth from db
     batch_api_target_cache: dict[str, dict[Target, list[Target]]]  # platform_name -> (target -> [target])
@@ -56,6 +67,7 @@ class Scheduler:
 
         self.platform_name_list = platform_name_list
         self.pre_weight_val = 0  # 轮调度中“本轮”增加权重和的初值
+        self.metrics_report: Callable[[bool], None] | None = None  # 作为函数变量，允许外部调用来上报此次抓取是否成功
         logger.info(
             f"register scheduler for {self.name} with "
             f"{self.scheduler_config.schedule_type} {self.scheduler_config.schedule_setting}"
@@ -97,6 +109,13 @@ class Scheduler:
 
         success_flag = False
         platform_obj = platform_manager[schedulable.platform_name](context)
+        # 通过闭包的形式，将此次抓取任务的信息保存为函数变量，允许在该任务无法正常结束时由外部上报
+        self.metrics_report = lambda x: request_counter.labels(
+            platform_name=schedulable.platform_name,
+            site_name=platform_obj.site.name,
+            target=schedulable.target,
+            success=x,
+        ).inc()
         try:
             with request_histogram.labels(
                 platform_name=schedulable.platform_name, site_name=platform_obj.site.name
@@ -123,12 +142,7 @@ class Scheduler:
             err.args += (records,)
             raise
 
-        request_counter.labels(
+        self.metrics_report(success_flag)
-            platform_name=schedulable.platform_name,
-            site_name=platform_obj.site.name,
-            target=schedulable.target,
-            success=success_flag,
-        ).inc()
         if not to_send:
             return
         sent_counter.labels(