方案修改为 多线程 滚动

app/services/company_clean_service.py

@@ -1,4 +1,8 @@
-"""公司数据补充服务（协程版）"""
+"""公司数据补充服务（持续消费模型）
+
+启动 N 个 worker 协程，每个 worker 循环：取一条 → 处理 → 取下一条。
+没数据时短暂休眠后重试。
+"""
 
 import asyncio
 from datetime import datetime
@@ -13,47 +17,65 @@ from app.ai.prompts import COMPANY_ENRICH_SYSTEM
 from app.services.ai_tool import ai_chat_json
 from app.services.dict_cache_service import dict_cache
 
+# 停止信号
+_stop_event = asyncio.Event()
+
+
+def stop_company_clean():
+    """外部调用，通知所有 worker 停止"""
+    _stop_event.set()
+
 
 async def run_company_clean() -> None:
-    """一次批量公司补充任务"""
-    # 锁定一批待完善公司
+    """启动 N 个 worker 协程持续消费"""
+    _stop_event.clear()
+    worker_count = settings.company_worker_count
+    log.info("公司补充：启动 {} 个 worker 协程", worker_count)
+
+    workers = [asyncio.create_task(_worker(i)) for i in range(worker_count)]
+    await asyncio.gather(*workers)
+
+    log.info("公司补充：所有 worker 已退出")
+
+
+async def _worker(worker_id: int) -> None:
+    """单个 worker：循环取一条、处理一条"""
+    while not _stop_event.is_set():
+        data = await _fetch_one()
+        if data is None:
+            await asyncio.sleep(settings.company_idle_sleep)
+            continue
+
+        try:
+            await _do_clean(data)
+        except Exception as e:
+            log.error("[worker-{}] 公司补充异常, id={}, shortName={}: {}",
+                      worker_id, data["id"], data.get("short_name"), e)
+
+
+async def _fetch_one() -> dict | None:
+    """从 MySQL 锁定一条待完善公司并标记为 status=3"""
     async with MysqlSession() as mysql:
+        # MySQL 不支持 UPDATE ... RETURNING，分两步
         result = await mysql.execute(
             text("""
                 SELECT * FROM bg_company
                 WHERE status = 0
-                LIMIT :limit
+                LIMIT 1
                 FOR UPDATE SKIP LOCKED
             """),
-            {"limit": settings.company_batch_size},
         )
-        rows = result.mappings().all()
-        if not rows:
-            return
+        row = result.mappings().first()
+        if not row:
+            return None
 
-        ids = [r["id"] for r in rows]
-        # MySQL 批量 IN 用 format 拼接（id 是 bigint，安全）
-        ids_str = ",".join(str(i) for i in ids)
+        company_id = row["id"]
         await mysql.execute(
-            text(f"UPDATE bg_company SET status = 3, update_time = NOW() WHERE id IN ({ids_str})"),
+            text("UPDATE bg_company SET status = 3, update_time = NOW() WHERE id = :id"),
+            {"id": company_id},
         )
         await mysql.commit()
-
-    log.info("公司补充：锁定{}条数据", len(rows))
-
-    # 协程并发，信号量限流
-    sem = asyncio.Semaphore(settings.company_concurrency)
-    tasks = [_clean_one(sem, dict(r)) for r in rows]
-    await asyncio.gather(*tasks, return_exceptions=True)
-
-
-async def _clean_one(sem: asyncio.Semaphore, company: dict) -> None:
-    """单条公司补充"""
-    async with sem:
-        try:
-            await _do_clean(company)
-        except Exception as e:
-            log.error("公司补充异常, id={}, shortName={}: {}", company["id"], company.get("short_name"), e)
+        return dict(row)
 
 
 async def _do_clean(company: dict) -> None: