174 lines
8.5 KiB
Python
174 lines
8.5 KiB
Python
"""Durable synchronization queue, shared by API processes on the data volume."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from contextlib import closing
|
|
from pathlib import Path
|
|
from typing import Callable
|
|
|
|
from tools.rag.rebuild_index import ReindexInProgressError, acquire_reindex_lock, rebuild_index
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SyncJobs:
|
|
def __init__(self, data_dir: Path, *, runner: Callable | None = None):
|
|
self.data_dir = Path(data_dir)
|
|
self.data_dir.mkdir(parents=True, exist_ok=True)
|
|
self.path = self.data_dir / 'sync_jobs.sqlite'
|
|
self.runner = runner or rebuild_index
|
|
self.stopping = threading.Event()
|
|
self.thread: threading.Thread | None = None
|
|
with closing(self.connect()) as conn, conn:
|
|
conn.executescript('''
|
|
CREATE TABLE IF NOT EXISTS jobs (
|
|
id TEXT PRIMARY KEY, status TEXT NOT NULL, phase TEXT NOT NULL,
|
|
pull_git INTEGER NOT NULL, created_at REAL NOT NULL, updated_at REAL NOT NULL,
|
|
started_at REAL, finished_at REAL, requested_revision TEXT,
|
|
request_count INTEGER NOT NULL DEFAULT 1, attempts INTEGER NOT NULL DEFAULT 0,
|
|
available_at REAL NOT NULL DEFAULT 0, reason TEXT, result_json TEXT);
|
|
CREATE INDEX IF NOT EXISTS jobs_status ON jobs(status, created_at);
|
|
CREATE TABLE IF NOT EXISTS deliveries (
|
|
delivery_key TEXT PRIMARY KEY, job_id TEXT NOT NULL REFERENCES jobs(id));
|
|
''')
|
|
|
|
def connect(self):
|
|
conn = sqlite3.connect(self.path, timeout=15)
|
|
conn.row_factory = sqlite3.Row
|
|
conn.execute('PRAGMA foreign_keys=ON')
|
|
return conn
|
|
|
|
@staticmethod
|
|
def public(row):
|
|
if row is None:
|
|
return None
|
|
value = dict(row)
|
|
value['pull_git'] = bool(value['pull_git'])
|
|
result = value.pop('result_json')
|
|
value['result'] = json.loads(result) if result else None
|
|
return value
|
|
|
|
def enqueue(self, *, pull_git: bool, delivery_key: str | None = None,
|
|
requested_revision: str | None = None) -> dict:
|
|
now = time.time()
|
|
with closing(self.connect()) as conn, conn:
|
|
conn.execute('BEGIN IMMEDIATE')
|
|
if delivery_key:
|
|
duplicate = conn.execute('SELECT job_id FROM deliveries WHERE delivery_key=?',
|
|
(delivery_key,)).fetchone()
|
|
if duplicate:
|
|
return {'job': self.public(conn.execute('SELECT * FROM jobs WHERE id=?',
|
|
(duplicate['job_id'],)).fetchone()), 'coalesced': False, 'duplicate': True}
|
|
# Never merge a new request into a running job: its source snapshot
|
|
# may already exist. One queued follow-up absorbs further pushes.
|
|
queued = conn.execute("SELECT * FROM jobs WHERE status='queued' ORDER BY created_at LIMIT 1").fetchone()
|
|
if queued:
|
|
job_id = queued['id']
|
|
conn.execute('''UPDATE jobs SET pull_git=MAX(pull_git,?), updated_at=?,
|
|
requested_revision=COALESCE(?,requested_revision), request_count=request_count+1
|
|
WHERE id=?''', (int(pull_git), now, requested_revision, job_id))
|
|
else:
|
|
job_id = uuid.uuid4().hex
|
|
conn.execute('''INSERT INTO jobs(id,status,phase,pull_git,created_at,updated_at,requested_revision)
|
|
VALUES(?,'queued','queued',?,?,?,?)''',
|
|
(job_id, int(pull_git), now, now, requested_revision))
|
|
if delivery_key:
|
|
conn.execute('INSERT INTO deliveries VALUES(?,?)', (delivery_key, job_id))
|
|
job = self.public(conn.execute('SELECT * FROM jobs WHERE id=?', (job_id,)).fetchone())
|
|
return {'job': job, 'coalesced': queued is not None, 'duplicate': False}
|
|
|
|
def get(self, job_id: str):
|
|
with closing(self.connect()) as conn:
|
|
return self.public(conn.execute('SELECT * FROM jobs WHERE id=?', (job_id,)).fetchone())
|
|
|
|
def overview(self):
|
|
with closing(self.connect()) as conn:
|
|
active = [self.public(row) for row in conn.execute(
|
|
"SELECT * FROM jobs WHERE status IN ('queued','running') ORDER BY created_at")]
|
|
success = self.public(conn.execute(
|
|
"SELECT * FROM jobs WHERE status='succeeded' ORDER BY finished_at DESC LIMIT 1").fetchone())
|
|
failure = self.public(conn.execute(
|
|
"SELECT * FROM jobs WHERE status='failed' ORDER BY finished_at DESC LIMIT 1").fetchone())
|
|
return {'active': active, 'last_success': success, 'last_failure': failure}
|
|
|
|
def phase(self, job_id: str, phase: str):
|
|
with closing(self.connect()) as conn, conn:
|
|
conn.execute('UPDATE jobs SET phase=?,updated_at=? WHERE id=?', (phase, time.time(), job_id))
|
|
|
|
def run_once(self) -> bool:
|
|
"""The process lock covers recovery, claim and execution, including retries."""
|
|
try:
|
|
with acquire_reindex_lock(self.data_dir / '.sync-worker.lock'):
|
|
return self._run_locked()
|
|
except ReindexInProgressError:
|
|
return False
|
|
|
|
def _run_locked(self) -> bool:
|
|
now = time.time()
|
|
with closing(self.connect()) as conn, conn:
|
|
conn.execute('BEGIN IMMEDIATE')
|
|
# No other worker can be running while this process holds the lock.
|
|
# A persisted running row therefore belongs to an interrupted process.
|
|
conn.execute("""UPDATE jobs SET status='queued',phase='recovered',reason='interrupted',
|
|
updated_at=? WHERE status='running'""", (now,))
|
|
row = conn.execute("""SELECT * FROM jobs WHERE status='queued' AND available_at<=?
|
|
ORDER BY created_at LIMIT 1""", (now,)).fetchone()
|
|
if row is None:
|
|
return False
|
|
job_id = row['id']
|
|
conn.execute("""UPDATE jobs SET status='running',phase='starting',started_at=?,
|
|
updated_at=?,attempts=attempts+1,reason=NULL WHERE id=?""", (now, now, job_id))
|
|
try:
|
|
result = self.runner(pull_git=bool(row['pull_git']), data_dir=self.data_dir,
|
|
progress=lambda phase: self.phase(job_id, phase))
|
|
except ReindexInProgressError:
|
|
# A CLI rebuild may own the pipeline lock. Keep the HTTP request
|
|
# durable and try again; this is not a failed synchronization attempt.
|
|
with closing(self.connect()) as conn, conn:
|
|
conn.execute("""UPDATE jobs SET status='queued',phase='waiting_for_lock',
|
|
attempts=attempts-1,available_at=?,updated_at=? WHERE id=?""",
|
|
(time.time() + 1, time.time(), job_id))
|
|
except Exception:
|
|
logger.exception('Synchronizácia %s zlyhala.', job_id)
|
|
retry = row['attempts'] + 1 < 3
|
|
with closing(self.connect()) as conn, conn:
|
|
conn.execute('''UPDATE jobs SET status=?,phase=?,reason='sync_failed',
|
|
updated_at=?,finished_at=?,available_at=? WHERE id=?''',
|
|
('queued' if retry else 'failed', 'retry_wait' if retry else 'failed',
|
|
time.time(), None if retry else time.time(), time.time() + 5 * (row['attempts'] + 1), job_id))
|
|
else:
|
|
with closing(self.connect()) as conn, conn:
|
|
conn.execute("""UPDATE jobs SET status='succeeded',phase='complete',result_json=?,
|
|
updated_at=?,finished_at=?,reason=NULL WHERE id=?""",
|
|
(json.dumps(result, ensure_ascii=False), time.time(), time.time(), job_id))
|
|
return True
|
|
|
|
def _loop(self):
|
|
while not self.stopping.is_set():
|
|
try:
|
|
worked = self.run_once()
|
|
except Exception:
|
|
logger.exception('Chyba pracovníka synchronizačnej fronty.')
|
|
worked = False
|
|
if not worked:
|
|
self.stopping.wait(0.5)
|
|
|
|
def start(self):
|
|
if self.thread and self.thread.is_alive():
|
|
return
|
|
self.stopping.clear()
|
|
self.thread = threading.Thread(target=self._loop, name='index-sync', daemon=True)
|
|
self.thread.start()
|
|
|
|
def stop(self):
|
|
self.stopping.set()
|
|
if self.thread:
|
|
# Compose grants five minutes to finish publication before shutdown.
|
|
# An interrupted job remains durable and is recovered on startup.
|
|
self.thread.join(timeout=285)
|