From 26d9edac3388fbb9773fc73b9d7a38daca4d1b03 Mon Sep 17 00:00:00 2001
From: host <host@test.local>
Date: Mon, 26 Oct 2026 06:19:01 +0800
Subject: [PATCH] Perf: light worker poll query + tasks(status) index

Worker polls the DB every interval (production 30s) and every
poll ran SELECT * FROM tasks: full rows including result_text
and error_text (up to 4000 chars each) for every task ever,
growing with history. load_tasks now selects only the columns
pick_next actually needs (id/project/prompt/status/depends/
priority/timestamps/executor) and only queued/completed/merged
rows; running/failed rows are neither ready nor completed, so
excluding them cannot change pick_next. get_task still does
SELECT * for the single claimed task. Plus CREATE INDEX IF NOT
EXISTS idx_tasks_status on tasks(status) in both SQLite
init_db and the Postgres schema, so the poll filter stops being
a full table scan on old databases too. Test: load_tasks returns
only queued+completed out of queued/completed/running/failed.
---
 cloud/db.py                        |  1 +
 cloud/server.py                    |  6 ++++++
 cloud/tests/test_fallback_retry.py | 17 +++++++++++++++++
 cloud/worker.py                    | 18 +++++++++++++++++-
 4 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/cloud/db.py b/cloud/db.py
index 5f62190..6c977dd 100644
--- a/cloud/db.py
+++ b/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, "
diff --git a/cloud/server.py b/cloud/server.py
index d7b18fd..b669f39 100644
--- a/cloud/server.py
+++ b/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()
diff --git a/cloud/tests/test_fallback_retry.py b/cloud/tests/test_fallback_retry.py
index d012dff..ada9285 100644
--- a/cloud/tests/test_fallback_retry.py
+++ b/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()
diff --git a/cloud/worker.py b/cloud/worker.py
index 4c92c6c..9f977ed 100644
--- a/cloud/worker.py
+++ b/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]