188 lines
7.9 KiB
Python
188 lines
7.9 KiB
Python
"""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']
|