优化时间

2026-09-11 13:41:40 +08:00
parent 350223c5ad
commit fd8f556431
14 changed files with 1315 additions and 0 deletions
@@ -0,0 +1,187 @@
+"""SQLite task queue for one service process with bounded worker threads."""
+import json
+import logging
+import sqlite3
+import threading
+from contextlib import contextmanager
+from datetime import datetime
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+class QueueFull(Exception):
+    pass
+
+
+class TaskQueue:
+    def __init__(self, path, concurrency, max_pending, runner):
+        if concurrency < 1 or max_pending < 1:
+            raise ValueError("任务并发数和队列容量必须大于 0")
+        self.path = Path(path).resolve()
+        self.concurrency = concurrency
+        self.max_pending = max_pending
+        self.runner = runner
+        self.stop_event = threading.Event()
+        self.threads = []
+        self.lock_file = None
+
+    @contextmanager
+    def connect(self):
+        conn = sqlite3.connect(str(self.path), timeout=30)
+        conn.row_factory = sqlite3.Row
+        try:
+            with conn:
+                yield conn
+        finally:
+            conn.close()
+
+    def start(self):
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        # Hold a kernel lock for the full process lifecycle. A second API worker
+        # must not recover this process's running tasks or multiply the limit.
+        lock_file = open(str(self.path) + ".lock", "a+b")
+        try:
+            import os
+            if os.name == "nt":
+                import msvcrt
+                lock_file.seek(0, 2)
+                if lock_file.tell() == 0:
+                    lock_file.write(b"0")
+                    lock_file.flush()
+                lock_file.seek(0)
+                msvcrt.locking(lock_file.fileno(), msvcrt.LK_NBLCK, 1)
+            else:
+                import fcntl
+                fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except OSError as exc:
+            lock_file.close()
+            raise RuntimeError("任务数据库已被其他服务使用，请使用单个 API worker") from exc
+        self.lock_file = lock_file
+        try:
+            with self.connect() as db:
+                db.execute("PRAGMA journal_mode=WAL")
+                db.execute("""CREATE TABLE IF NOT EXISTS tasks (
+                    task_id TEXT PRIMARY KEY, status TEXT NOT NULL,
+                    created_at TEXT NOT NULL, completed_at TEXT, error TEXT,
+                    request TEXT, input_dir TEXT, callback_status INTEGER
+                )""")
+                db.execute("CREATE INDEX IF NOT EXISTS tasks_status_created ON tasks(status, created_at)")
+                db.execute("""UPDATE tasks SET status='failed', completed_at=?,
+                    error='服务中断，任务未完成，请重新提交' WHERE status='running'""",
+                    (datetime.now().isoformat(),))
+            self.stop_event.clear()
+            if hasattr(self.runner, 'start'):
+                self.runner.start()
+            for number in range(self.concurrency):
+                thread = threading.Thread(target=self.work, name=f"analysis-{number}")
+                self.threads.append(thread)
+                thread.start()
+        except BaseException:
+            self.stop()
+            raise
+
+    def stop(self):
+        self.stop_event.set()
+        for thread in self.threads:
+            if thread.ident is not None:
+                thread.join()
+        self.threads.clear()
+        if hasattr(self.runner, 'stop'):
+            self.runner.stop()
+        if self.lock_file:
+            self.lock_file.close()
+            self.lock_file = None
+
+    @staticmethod
+    def decode(row):
+        if row is None:
+            return None
+        task = dict(row)
+        task['request'] = json.loads(task['request']) if task['request'] else None
+        return task
+
+    def submit(self, task_id, request, input_dir):
+        with self.connect() as db:
+            db.execute("BEGIN IMMEDIATE")
+            row = db.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()
+            if row and row['status'] in ('pending', 'running'):
+                return self.decode(row), False
+            pending = db.execute("SELECT COUNT(*) FROM tasks WHERE status='pending'").fetchone()[0]
+            if pending >= self.max_pending:
+                raise QueueFull("等待队列已满，请稍后重试")
+            db.execute("""INSERT INTO tasks VALUES (?, 'pending', ?, NULL, NULL, ?, ?, NULL)
+                ON CONFLICT(task_id) DO UPDATE SET status='pending',
+                created_at=excluded.created_at, completed_at=NULL, error=NULL,
+                request=excluded.request, input_dir=excluded.input_dir, callback_status=NULL""",
+                (task_id, datetime.now().isoformat(), json.dumps(request, ensure_ascii=False), input_dir))
+            return self.decode(db.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()), True
+
+    def claim(self):
+        with self.connect() as db:
+            db.execute("BEGIN IMMEDIATE")
+            row = db.execute("SELECT * FROM tasks WHERE status='pending' ORDER BY created_at, rowid LIMIT 1").fetchone()
+            if row is None:
+                return None
+            db.execute("UPDATE tasks SET status='running' WHERE task_id=?", (row['task_id'],))
+            task = self.decode(row)
+            task['status'] = 'running'
+            return task
+
+    def finish(self, task_id, error=None):
+        with self.connect() as db:
+            # A callback cannot release a slot before the subprocess exits.
+            # Preserve a failure callback even when the script exits with zero.
+            db.execute("""UPDATE tasks SET
+                status=CASE WHEN ? IS NOT NULL OR callback_status=0 THEN 'failed' ELSE 'completed' END,
+                error=CASE WHEN ? IS NOT NULL THEN ? WHEN callback_status=0 THEN '分析任务回调失败' ELSE NULL END,
+                completed_at=? WHERE task_id=? AND status='running'""",
+                (error, error, error, datetime.now().isoformat(), task_id))
+
+    def work(self):
+        while not self.stop_event.is_set():
+            try:
+                task = self.claim()
+                if task is None:
+                    self.stop_event.wait(0.25)
+                    continue
+                error = None
+                try:
+                    self.runner(task, self.stop_event)
+                except Exception as exc:
+                    error = str(exc) or type(exc).__name__
+                    logger.exception("任务失败：%s", task['task_id'])
+                self.finish(task['task_id'], error)
+            except Exception:
+                logger.exception("任务队列异常")
+                self.stop_event.wait(1)
+
+    def get(self, task_id):
+        with self.connect() as db:
+            return self.decode(db.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone())
+
+    def list(self, status=None, limit=20):
+        with self.connect() as db:
+            if status:
+                rows = db.execute("SELECT * FROM tasks WHERE status=? ORDER BY created_at DESC LIMIT ?", (status, limit)).fetchall()
+            else:
+                rows = db.execute("SELECT * FROM tasks ORDER BY created_at DESC LIMIT ?", (limit,)).fetchall()
+            return [self.decode(row) for row in rows]
+
+    def counts(self):
+        counts = dict.fromkeys(('pending', 'running', 'completed', 'failed'), 0)
+        with self.connect() as db:
+            counts.update(dict(db.execute("SELECT status, COUNT(*) FROM tasks GROUP BY status").fetchall()))
+        counts['total'] = sum(counts.values())
+        return counts
+
+    def callback(self, task_id, status):
+        with self.connect() as db:
+            db.execute("BEGIN IMMEDIATE")
+            row = db.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()
+            if row is None:
+                return None
+            if row['status'] == 'running':
+                db.execute("UPDATE tasks SET callback_status=? WHERE task_id=?", (status, task_id))
+            # Late callbacks must not rewrite terminal state or consume pending work.
+            return row['status']