"""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']