refactor: harden buffer flush state

Author: Boris-code

feapder/buffer/item_buffer.py

@@ -299,133 +299,134 @@ def __export_to_db(self, table, datas, is_update=False, update_keys=(), used_pip
     def __add_item_to_db(
             self, items, update_items, requests, callbacks, items_fingerprints
     ):
-        export_success = True
         self._is_adding_to_db = True
+        try:
+            export_success = True
 
-        # 去重
-        if setting.ITEM_FILTER_ENABLE:
-            items, items_fingerprints = self.__dedup_items(items, items_fingerprints)
+            # 去重
+            if setting.ITEM_FILTER_ENABLE:
+                items, items_fingerprints = self.__dedup_items(items, items_fingerprints)
 
-        # 分捡（返回值包含 pipelines_dict）
-        items_dict = self.__pick_items(items)
-        update_items_dict = self.__pick_items(update_items, is_update_item=True)
+            # 分捡（返回值包含 pipelines_dict）
+            items_dict = self.__pick_items(items)
+            update_items_dict = self.__pick_items(update_items, is_update_item=True)
 
-        # item批量入库
-        failed_items = {"add": [], "update": [], "requests": []}
-        while items_dict:
-            table, datas = items_dict.popitem()
-            used_pipelines = self._item_pipelines.get(table)
+            # item批量入库
+            failed_items = {"add": [], "update": [], "requests": []}
+            while items_dict:
+                table, datas = items_dict.popitem()
+                used_pipelines = self._item_pipelines.get(table)
 
-            log.debug(
-                """
+                log.debug(
+                    """
                 -------------- item 批量入库 --------------
                 表名: %s
                 datas: %s
                     """
-                % (table, tools.dumps_json(datas, indent=16))
-            )
+                    % (table, tools.dumps_json(datas, indent=16))
+                )
 
-            if not self.__export_to_db(table, datas, used_pipelines=used_pipelines):
-                export_success = False
-                failed_items["add"].append({"table": table, "datas": datas})
+                if not self.__export_to_db(table, datas, used_pipelines=used_pipelines):
+                    export_success = False
+                    failed_items["add"].append({"table": table, "datas": datas})
 
-        # 执行批量update
-        while update_items_dict:
-            table, datas = update_items_dict.popitem()
-            used_pipelines = self._item_pipelines.get(table)
+            # 执行批量update
+            while update_items_dict:
+                table, datas = update_items_dict.popitem()
+                used_pipelines = self._item_pipelines.get(table)
 
-            log.debug(
-                """
+                log.debug(
+                    """
                 -------------- item 批量更新 --------------
                 表名: %s
                 datas: %s
                     """
-                % (table, tools.dumps_json(datas, indent=16))
-            )
-
-            update_keys = self._item_update_keys.get(table)
-            if not self.__export_to_db(
-                    table, datas, is_update=True, update_keys=update_keys, used_pipelines=used_pipelines
-            ):
-                export_success = False
-                failed_items["update"].append(
-                    {"table": table, "datas": datas, "update_keys": update_keys}
+                    % (table, tools.dumps_json(datas, indent=16))
                 )
 
-        if export_success:
-            # 执行回调
-            while callbacks:
-                try:
-                    callback = callbacks.pop(0)
-                    callback()
-                except Exception as e:
-                    log.exception(e)
+                update_keys = self._item_update_keys.get(table)
+                if not self.__export_to_db(
+                        table, datas, is_update=True, update_keys=update_keys, used_pipelines=used_pipelines
+                ):
+                    export_success = False
+                    failed_items["update"].append(
+                        {"table": table, "datas": datas, "update_keys": update_keys}
+                    )
 
-            # 删除做过的request
-            if requests:
-                self.redis_db.zrem(self._table_request, requests)
+            if export_success:
+                # 执行回调
+                while callbacks:
+                    try:
+                        callback = callbacks.pop(0)
+                        callback()
+                    except Exception as e:
+                        log.exception(e)
 
-            # 去重入库
-            if setting.ITEM_FILTER_ENABLE:
-                if items_fingerprints:
-                    self.__class__.dedup.add(items_fingerprints, skip_check=True)
-        else:
-            failed_items["requests"] = requests
+                # 删除做过的request
+                if requests:
+                    self.redis_db.zrem(self._table_request, requests)
 
-            if self.export_retry_times > setting.EXPORT_DATA_MAX_RETRY_TIMES:
-                if self._redis_key != "air_spider":
-                    # 失败的item记录到redis
-                    self.redis_db.sadd(self._table_failed_items, failed_items)
+                # 去重入库
+                if setting.ITEM_FILTER_ENABLE:
+                    if items_fingerprints:
+                        self.__class__.dedup.add(items_fingerprints, skip_check=True)
+            else:
+                failed_items["requests"] = requests
 
-                    # 删除做过的request
-                    if requests:
-                        self.redis_db.zrem(self._table_request, requests)
+                if self.export_retry_times > setting.EXPORT_DATA_MAX_RETRY_TIMES:
+                    if self._redis_key != "air_spider":
+                        # 失败的item记录到redis
+                        self.redis_db.sadd(self._table_failed_items, failed_items)
 
-                    log.error(
-                        "入库超过最大重试次数，不再重试，数据记录到redis，items:\n {}".format(
-                            tools.dumps_json(failed_items)
-                        )
-                    )
-                self.export_retry_times = 0
+                        # 删除做过的request
+                        if requests:
+                            self.redis_db.zrem(self._table_request, requests)
 
-            else:
-                tip = ["入库不成功"]
-                if callbacks:
-                    tip.append("不执行回调")
-                if requests:
-                    tip.append("不删除任务")
-                    exists = self.redis_db.zexists(self._table_request, requests)
-                    for exist, request in zip(exists, requests):
-                        if exist:
-                            self.redis_db.zadd(self._table_request, requests, 300)
+                        log.error(
+                            "入库超过最大重试次数，不再重试，数据记录到redis，items:\n {}".format(
+                                tools.dumps_json(failed_items)
+                            )
+                        )
+                    self.export_retry_times = 0
 
-                if setting.ITEM_FILTER_ENABLE:
-                    tip.append("数据不入去重库")
+                else:
+                    tip = ["入库不成功"]
+                    if callbacks:
+                        tip.append("不执行回调")
+                    if requests:
+                        tip.append("不删除任务")
+                        exists = self.redis_db.zexists(self._table_request, requests)
+                        for exist, request in zip(exists, requests):
+                            if exist:
+                                self.redis_db.zadd(self._table_request, requests, 300)
 
-                if self._redis_key != "air_spider":
-                    tip.append("将自动重试")
+                    if setting.ITEM_FILTER_ENABLE:
+                        tip.append("数据不入去重库")
 
-                tip.append("失败items:\n {}".format(tools.dumps_json(failed_items)))
-                log.error("，".join(tip))
+                    if self._redis_key != "air_spider":
+                        tip.append("将自动重试")
 
-                self.export_falied_times += 1
+                    tip.append("失败items:\n {}".format(tools.dumps_json(failed_items)))
+                    log.error("，".join(tip))
 
-                if self._redis_key != "air_spider":
-                    self.export_retry_times += 1
+                    self.export_falied_times += 1
 
-            if self.export_falied_times > setting.EXPORT_DATA_MAX_FAILED_TIMES:
-                # 报警
-                msg = "《{}》爬虫导出数据失败，失败次数：{}，请检查爬虫是否正常".format(
-                    self._redis_key, self.export_falied_times
-                )
-                log.error(msg)
-                tools.send_msg(
-                    msg=msg,
-                    level="error",
-                    message_prefix="《%s》爬虫导出数据失败" % (self._redis_key),
-                )
+                    if self._redis_key != "air_spider":
+                        self.export_retry_times += 1
 
-        self._is_adding_to_db = False
+                if self.export_falied_times > setting.EXPORT_DATA_MAX_FAILED_TIMES:
+                    # 报警
+                    msg = "《{}》爬虫导出数据失败，失败次数：{}，请检查爬虫是否正常".format(
+                        self._redis_key, self.export_falied_times
+                    )
+                    log.error(msg)
+                    tools.send_msg(
+                        msg=msg,
+                        level="error",
+                        message_prefix="《%s》爬虫导出数据失败" % (self._redis_key),
+                    )
+        finally:
+            self._is_adding_to_db = False
 
     def metric_datas(self, table, datas):
         """

tests/test_runtime_status.py

@@ -6,6 +6,7 @@
 from feapder.core.collector import Collector
 from feapder.core.scheduler import Scheduler
 from feapder.core.runtime_state import RuntimeState
+from feapder.network.item import Item
 
 
 class FakeRedis:
@@ -185,3 +186,55 @@ def popleft(self):
     assert buffer._requests_deque.busy_observations
     assert all(buffer._requests_deque.busy_observations)
     assert buffer.is_idle() is True
+
+
+def test_request_buffer_flush_resets_flag_when_zadd_raises():
+    class ExplodingDB:
+        def zadd(self, table, values, prioritys=0):
+            raise RuntimeError("write failed")
+
+    request = type(
+        "FakeRequest",
+        (),
+        {
+            "priority": 300,
+            "filter_repeat": False,
+            "url": "https://example.com",
+            "to_dict": {"url": "https://example.com"},
+        },
+    )()
+    buffer = build_request_buffer(request_count=0, delete_count=0, flushing=False)
+    buffer._db = ExplodingDB()
+    buffer._table_request = "test:z_requests"
+    buffer._table_failed_request = "test:z_failed_requests"
+    buffer._requests_deque.append(request)
+
+    buffer.flush()
+
+    assert buffer.is_adding_to_db() is False
+
+
+def test_item_buffer_flush_resets_flag_when_export_raises(monkeypatch):
+    monkeypatch.setattr("feapder.buffer.item_buffer.setting.ITEM_FILTER_ENABLE", False)
+    item = Item(title="boom")
+    buffer = build_item_buffer(item_count=0, flushing=False)
+    buffer._items_queue.put(item)
+    buffer._redis_key = "test"
+    buffer._task_table = None
+    buffer._item_tables = {}
+    buffer._item_update_keys = {}
+    buffer._item_pipelines = {}
+    buffer._pipelines = []
+    buffer._have_mysql_pipeline = True
+    buffer._mysql_pipeline = None
+    buffer.export_retry_times = 0
+    buffer.export_falied_times = 0
+
+    def raise_export(table, datas, is_update=False, update_keys=(), used_pipelines=None):
+        raise RuntimeError("export failed")
+
+    buffer._ItemBuffer__export_to_db = raise_export
+
+    buffer.flush()
+
+    assert buffer.is_adding_to_db() is False