Merge pull request #122 from openqj/feat-pm-wait-deadlock

worker: always start capacity pool threads in run_all

cloud/tests/test_run_all_pool.py

@@ -0,0 +1,143 @@
+"""run_all 池线程数单测：轮中新建任务不能死锁。
+
+根因：run_all 曾按本轮开始时的就绪快照开线程（min(capacity,
+ready_count)）。本轮开始时若只有 1 个就绪任务（如刚认领的 PM），
+只开 1 个线程；该线程跑 PM 模型回合后进入标记等待（最长 3600 秒），
+期间 PM 新建/重试的子任务入队，却再无空闲线程认领它——run_loop
+又被本轮 run_all 的 join 挡住、开不了新一轮；若此时另一 worker
+已缩容，就只剩一个被等待占住的 worker，子任务等满等待超时也
+无人跑。修法：池线程数恒为 capacity。
+
+不跑真实模型、不联网：用 init_db 建的临时 SQLite 库 + 一个模拟
+run_one 的 Worker 子类——PM 任务先建子任务、再轮询等子任务完成；
+子任务任务一被认领就立刻完成。无修法时 PM 轮询 5 秒超时（子任务
+在本轮内无人认领，PM 失败后下一轮才跑子任务）；有修法时第二线程
+在本轮内认领子任务，PM 轮询成功，两个任务都在本轮内派发。
+"""
+
+import os
+import sqlite3
+import sys
+import tempfile
+import time
+import unittest
+from pathlib import Path
+
+try:
+    from cloud.server import init_db
+    from cloud.worker import Worker
+except ImportError:  # discover -s cloud/tests 时把项目根目录补进 sys.path
+    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+    from cloud.server import init_db
+    from cloud.worker import Worker
+
+FAKE_ENV = {"XCM_API_KEY": "test-placeholder"}
+
+
+class _WaitWorker(Worker):
+    """模拟 PM 等待：PM 建子任务后轮询等它完成，子任务立刻完成。"""
+
+    def _claim(self, task_id):
+        """模拟真实 run_one 的 CAS 认领：只有第一个线程能拿到任务。"""
+        conn = sqlite3.connect(self.db_path)
+        try:
+            cur = conn.execute(
+                "UPDATE tasks SET status = 'running', updated_at = ?"
+                " WHERE id = ? AND status = 'queued'",
+                (time.time(), task_id))
+            conn.commit()
+            return cur.rowcount == 1
+        finally:
+            conn.close()
+
+    def run_one(self, task=None):
+        if task is None:
+            return {"dispatched": False}
+        task_id = task["id"]
+        if not self._claim(task_id):
+            return {"dispatched": False}
+        if task_id == "CHILD":
+            conn = sqlite3.connect(self.db_path)
+            try:
+                conn.execute(
+                    "UPDATE tasks SET status = 'completed',"
+                    " updated_at = ? WHERE id = 'CHILD'",
+                    (time.time(),))
+                conn.commit()
+            finally:
+                conn.close()
+            return {"dispatched": True}
+        if task_id == "PM":
+            now = time.time()
+            conn = sqlite3.connect(self.db_path)
+            try:
+                conn.execute(
+                    "INSERT OR IGNORE INTO tasks(id, project_id, prompt,"
+                    " status, idempotency_key, depends_on, priority,"
+                    " created_by, created_at, updated_at, last_heartbeat)"
+                    " VALUES('CHILD', 'p1', 'do work', 'queued', NULL,"
+                    " '[]', 0, 'tester', ?, ?, NULL)",
+                    (now, now))
+                conn.commit()
+            finally:
+                conn.close()
+            deadline = time.time() + 5
+            while time.time() < deadline:
+                conn = sqlite3.connect(self.db_path)
+                try:
+                    row = conn.execute(
+                        "SELECT status FROM tasks WHERE id = 'CHILD'"
+                        ).fetchone()
+                finally:
+                    conn.close()
+                if row is not None and row[0] == "completed":
+                    return {"dispatched": True}
+                time.sleep(0.1)
+            # 超时 = 本轮内无人认领子任务，即死锁症状。
+            return {"dispatched": False}
+        return {"dispatched": False}
+
+
+class RunAllPoolTestCase(unittest.TestCase):
+    def setUp(self):
+        self.tmp = tempfile.TemporaryDirectory()
+        self.db_path = os.path.join(self.tmp.name, "control.db")
+        init_db(self.db_path)
+        now = time.time()
+        conn = sqlite3.connect(self.db_path)
+        try:
+            conn.execute(
+                "INSERT INTO tasks(id, project_id, prompt, status,"
+                " idempotency_key, depends_on, priority, created_by,"
+                " created_at, updated_at, last_heartbeat)"
+                " VALUES('PM', 'p1', 'do work', 'queued', NULL,"
+                " '[]', 0, 'tester', ?, ?, NULL)",
+                (now, now))
+            conn.commit()
+        finally:
+            conn.close()
+
+    def tearDown(self):
+        self.tmp.cleanup()
+
+    def test_mid_round_child_claimed_in_same_run_all(self):
+        worker = _WaitWorker(self.db_path, "/fake/host.py",
+                             env=FAKE_ENV, capacity=2)
+        # max_tasks=2：模拟任务本身不标 completed（真实 run_one 会标），
+        # 用派发上限让 run_all 在 PM 与子任务都派发后收尾，避免测试
+        # 因 PM 仍 queued 而被池线程反复认领。
+        results = worker.run_all(max_tasks=2)
+        # 两个任务都在本轮内派发成功：PM 轮询等到了子任务完成。
+        self.assertEqual(len(results), 2)
+        self.assertTrue(all(r["dispatched"] for r in results))
+        conn = sqlite3.connect(self.db_path)
+        try:
+            row = conn.execute(
+                "SELECT status FROM tasks WHERE id = 'CHILD'").fetchone()
+        finally:
+            conn.close()
+        self.assertEqual(row[0], "completed")
+
+
+if __name__ == "__main__":
+    unittest.main()

cloud/worker.py

@@ -4602,15 +4602,18 @@ class Worker:
                     with results_lock:
                         results.append(result)
 
-        try:
+        # 池线程数恒为 capacity，绝不按本轮开始时的就绪快照缩减。
-            ready_count = len(self._ready_tasks())
+        # 快照缩减是 PM 等待死锁的根因：本轮开始时若只有 1 个就绪
-        except Exception:
+        # 任务（如刚认领的 PM），只开 1 个线程；该线程跑 PM 模型回合
-            # 深度快照失败时退回开满线程，池线程自己的认领
+        # 后进入标记等待（最长 3600 秒），期间 PM 新建/重试的子任务
-            # 失败容错会处理，不让一次读库失败炸掉整轮。
+        # 入队，却再无空闲线程认领它——run_loop 又被本轮 run_all 的
-            ready_count = self.capacity
+        # join 挡住、开不了新一轮；若此时另一 worker 已缩容，就只剩
+        # 一个被等待占住的 worker，子任务等满等待超时也无人跑。
+        # 恒开 capacity 个线程没有空转代价：空闲线程在 task 为 None
+        # 且 active 为 0 时立刻退出，有在跑任务时才 0.2 秒轮询认领，
+        # 正好接住轮中新建的任务。
         threads = [threading.Thread(target=_pool_thread, daemon=True)
-                   for _ in range(min(self.capacity,
+                   for _ in range(self.capacity)]
-                                      max(1, ready_count)))]
         for t in threads:
             t.start()
         for t in threads: