Merge pull request #20 from openqj/worker-poll-perf

Perf: light worker poll query + tasks(status) index

cloud/db.py

@@ -807,6 +807,7 @@ POSTGRES_SCHEMA = (
     "pr_url TEXT, "
     "pr_number INTEGER, "
     "UNIQUE(project_id, idempotency_key))",
+    "CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status)",
     "CREATE TABLE IF NOT EXISTS summaries("
     "id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, "
     "project_id TEXT, "

cloud/server.py

@@ -305,6 +305,12 @@ def init_db(db_path):
         )
         migrate_tasks_table(conn)
         migrate_projects_table(conn)
+        # worker 每轮 poll 都按 status 取任务（见 load_tasks 的
+        # WHERE status IN (...)），无索引时是全表扫描；任务表随
+        # 时间变大后索引是必须的。IF NOT EXISTS 保证老库幂等。
+        conn.execute(
+            "CREATE INDEX IF NOT EXISTS idx_tasks_status"
+            " ON tasks(status)")
         conn.commit()
     finally:
         conn.close()

cloud/tests/test_fallback_retry.py

@@ -384,6 +384,23 @@ class FallbackRetryTestCase(unittest.TestCase):
         self.assertEqual(len(self.calls), 1)
         self.assertEqual(result["events"], [])
 
+    def test_load_tasks_only_picks_relevant_statuses(self):
+        # load_tasks 只取 queued/completed/merged：running/failed
+        # 的任务对 pick_next 既不是就绪也不是已完成，排除它们
+        # 不影响结果，但省掉每轮 poll 全表 + 重字段的开销。
+        for tid, status in [("Q", "queued"), ("C", "completed"),
+                            ("R", "running"), ("F", "failed")]:
+            _insert_task(self.db_path, tid)
+            import sqlite3
+            conn = sqlite3.connect(self.db_path)
+            conn.execute(
+                "UPDATE tasks SET status = ? WHERE id = ?", (status, tid))
+            conn.commit()
+            conn.close()
+        worker = self._worker(runner=lambda cmd: None)
+        got = {t["id"] for t in worker.load_tasks()}
+        self.assertEqual(got, {"Q", "C"})
+
 
 if __name__ == "__main__":
     unittest.main()

cloud/worker.py

@@ -365,10 +365,26 @@ class Worker:
 
     # -- DB 读写 ------------------------------------------------------
     def load_tasks(self):
+        """pick_next 的输入：只取 queued/completed/merged 三类任务。
+
+        旧实现是 SELECT * 全表：prompt/result_text/error_text
+        （后两者各可达 4000 字）全量进内存，且每轮 poll 都来一次；
+        任务表随时间变大后这是纯浪费。pick_next 只需要 queued
+        任务和已完成 id（依赖判断），running/failed/interrupted
+        的任务对它既不是就绪、也不是已完成，排除它们不影响结果。
+        列清单必须覆盖 _task_from_row 直接取的全部字段（prompt、
+        idempotency_key 等），缺一列就会 KeyError；result_text、
+        error_text 等重字段不在此列，真正需要它们的 get_task
+        仍是 SELECT * 单行取。
+        """
         conn = _db_connect(self.db_path)
         try:
             rows = conn.execute(
-                "SELECT * FROM tasks ORDER BY created_at, id").fetchall()
+                "SELECT id, project_id, prompt, status, idempotency_key,"
+                " depends_on, priority, created_by, created_at, updated_at,"
+                " last_heartbeat, executor FROM tasks"
+                " WHERE status IN ('queued', 'completed', 'merged')"
+                " ORDER BY created_at, id").fetchall()
         finally:
             conn.close()
         return [_task_from_row(r) for r in rows]