"""Enqueue / claim biz_state collect batches (Postgres SKIP LOCKED + NE mutex).""" from __future__ import annotations from collections import defaultdict import logging from datetime import datetime, timedelta from typing import Any from uuid import uuid4 from sqlalchemy.orm import Session from ..config import settings from ..db import SessionLocal from ..models import BizStateBatch, BizStateTask, BizStateTaskItem from ..timeutil import utcnow_naive _log = logging.getLogger("netx.biz_state.claim") def _utcnow() -> datetime: return utcnow_naive() def max_concurrent_tasks() -> int: return max(1, int(getattr(settings, "biz_state_max_concurrent_tasks", 16) or 16)) def max_sessions_per_ne() -> int: """Concurrent SSH collections allowed per network element (>= 1).""" return max(1, int(getattr(settings, "biz_state_max_sessions_per_ne", 2) or 2)) def enqueue_collect(task_id: str, *, manual: bool = False) -> dict[str, Any]: """Create a queued batch and mark task collect_running. Does not open SSH. Returns keys: ok, queued, reason?, batch_id?, task_id """ tid = str(task_id or "").strip() if not tid: return {"ok": False, "queued": False, "reason": "missing_task_id", "task_id": ""} db = SessionLocal() try: task = db.get(BizStateTask, tid) if not task: return {"ok": False, "queued": False, "reason": "task_not_found", "task_id": tid} if task.collect_group_id: group_id = task.collect_group_id db.close() return enqueue_collect_group(group_id, manual=manual, requested_task_id=tid) if str(task.source or "").strip().lower() == "import": return { "ok": False, "queued": False, "reason": "import_offline_only", "task_id": tid, } if bool(task.collect_running): return { "ok": True, "queued": False, "reason": "already_collecting", "task_id": tid, } st = str(task.status or "").strip() if manual: if st in ("", "deleted"): return {"ok": False, "queued": False, "reason": "bad_status", "task_id": tid} else: if st != "running": return {"ok": False, "queued": False, "reason": "not_scheduled", "task_id": tid} items = ( db.query(BizStateTaskItem) .filter( BizStateTaskItem.task_id == tid, BizStateTaskItem.enabled.is_(True), ) .limit(1) .all() ) if not items: task.last_error = "no enabled task items" task.updated_at = _utcnow() db.commit() return { "ok": False, "queued": False, "reason": "no_enabled_items", "task_id": tid, } now = _utcnow() task.collect_running = True task.last_collect_started_at = now task.last_error = "" task.updated_at = now if hasattr(task, "collect_queued_at"): task.collect_queued_at = now batch = BizStateBatch( id=uuid4().hex, task_id=task.id, source=task.source, ne_id=task.ne_id, ne_name=task.ne_name, vendor=task.vendor, status="queued", started_at=now, queued_at=now, message="queued" if not manual else "queued_manual", ) db.add(batch) db.commit() return { "ok": True, "queued": True, "batch_id": batch.id, "task_id": tid, "manual": bool(manual), } except Exception: _log.exception("enqueue_collect failed task=%s", tid) try: db.rollback() except Exception: pass return {"ok": False, "queued": False, "reason": "enqueue_error", "task_id": tid} finally: db.close() def enqueue_collect_group(group_id: str, *, manual: bool = False, requested_task_id: str = "") -> dict[str, Any]: """Atomically enqueue both devices with one round identity, or neither.""" db = SessionLocal() try: tasks = db.query(BizStateTask).filter(BizStateTask.collect_group_id == group_id).order_by( BizStateTask.id).with_for_update().populate_existing().all() result: dict[str, Any] = {"ok": False, "queued": False, "task_id": requested_task_id, "collect_group_id": group_id} # One task = a single device used as both sides (same-device cutover); two tasks = a pair. if len(tasks) not in (1, 2) or len({t.ne_id for t in tasks}) != len(tasks): return {**result, "reason": "collect_pair_invalid"} if max_concurrent_tasks() < len(tasks): return {**result, "reason": "collect_pair_capacity"} from ..cli_budget import clamp_cli_workers worker_cap = clamp_cli_workers(int(getattr(settings, "biz_state_worker_collect_threads", None) or getattr(settings, "biz_state_dispatch_workers", 8) or 8)) if worker_cap < len(tasks): return {**result, "reason": "collect_pair_capacity"} if any(t.collect_running for t in tasks): return {**result, "reason": "collect_pair_busy"} for task in tasks: if task.source == "import" or (not manual and (task.status != "running" or task.collect_manual_only)): return {**result, "reason": "not_scheduled"} if task.status in ("", "deleted"): return {**result, "reason": "bad_status"} if not db.query(BizStateTaskItem.id).filter(BizStateTaskItem.task_id == task.id, BizStateTaskItem.enabled.is_(True)).first(): return {**result, "reason": "no_enabled_items"} now, round_id = _utcnow(), uuid4().hex batches = [] for task in tasks: task.collect_running = True task.collect_queued_at = now task.last_error = "" task.updated_at = now batch = BizStateBatch(id=uuid4().hex, task_id=task.id, source=task.source, ne_id=task.ne_id, ne_name=task.ne_name, vendor=task.vendor, status="queued", started_at=now, queued_at=now, collect_group_id=group_id, collect_round_id=round_id, collect_priority=10 if task.collect_manual_only else 20, message="queued_manual" if manual else "queued") db.add(batch) batches.append({"batch_id": batch.id, "task_id": task.id}) db.commit() selected = next((b for b in batches if b["task_id"] == requested_task_id), batches[0]) return {**result, **selected, "ok": True, "queued": True, "collect_round_id": round_id, "batches": batches, "manual": manual} except Exception: db.rollback() _log.exception("enqueue collect pair failed group=%s", group_id) return {"ok": False, "queued": False, "reason": "enqueue_error", "task_id": requested_task_id} finally: db.close() def _dialect_supports_skip_locked(db: Session) -> bool: try: bind = db.get_bind() name = str(getattr(getattr(bind, "dialect", None), "name", "") or "").lower() return name == "postgresql" except Exception: return False def count_running_batches(db: Session | None = None) -> int: own = db is None if own: db = SessionLocal() assert db is not None try: return int( db.query(BizStateBatch).filter(BizStateBatch.status == "running").count() ) finally: if own: db.close() def claim_queued_batches(limit: int | None = None) -> list[dict[str, Any]]: """Claim up to ``limit`` queued batches (SKIP LOCKED + same-NE mutex). Returns list of {batch_id, task_id, source, ne_id, vendor, device_type}. Global ceiling is always ``max_concurrent_tasks()``; ``limit`` only caps how many this caller wants in one call. """ from sqlalchemy import and_, case, text global_cap = max_concurrent_tasks() db = SessionLocal() try: pg = _dialect_supports_skip_locked(db) if pg: # A short transaction lock makes capacity reservation and pair claims # atomic across worker processes. Never held during device I/O. db.execute(text("SELECT pg_advisory_xact_lock(hashtext('biz-state-claim-capacity'))")) running_n = ( db.query(BizStateBatch).filter(BizStateBatch.status == "running").count() ) slots = max(0, global_cap - int(running_n)) if limit is not None: slots = min(slots, max(0, int(limit))) if slots <= 0: return [] # Running collections per NE; an NE accepts new ones until it reaches its session cap. per_ne_cap = max_sessions_per_ne() ne_running: dict[str, int] = defaultdict(int) for (running_ne,) in db.query(BizStateBatch.ne_id).filter(BizStateBatch.status == "running").all(): if str(running_ne or "").strip(): ne_running[str(running_ne).strip()] += 1 # Manual route rounds yield to fresh monitoring, but must eventually # run even when a busy device always has regular metrics queued. priority = case((and_(BizStateBatch.collect_priority == 10, BizStateBatch.queued_at <= _utcnow() - timedelta(seconds=300)), 30), else_=BizStateBatch.collect_priority) q = ( db.query(BizStateBatch) .filter(BizStateBatch.status == "queued") .order_by(priority.desc(), BizStateBatch.started_at.asc(), BizStateBatch.id.asc()) ) if pg: q = q.with_for_update(skip_locked=True) else: q = q.with_for_update() candidates = q.limit(max(slots * 8, 64)).all() claimed_rows: list[BizStateBatch] = [] visited: set[str] = set() for batch in candidates: if batch.id in visited: continue if len(claimed_rows) >= slots: break group = [batch] if batch.collect_round_id: peers = db.query(BizStateBatch).filter( BizStateBatch.collect_round_id == batch.collect_round_id, ).with_for_update(skip_locked=pg).all() expected = db.query(BizStateTask.id).filter( BizStateTask.collect_group_id == batch.collect_group_id).count() if batch.collect_group_id else 2 if len(peers) != max(1, expected) or any(p.status != "queued" for p in peers): continue group = peers visited.update(p.id for p in group) if len(group) > slots - len(claimed_rows): continue wanted: dict[str, int] = defaultdict(int) for p in group: if str(p.ne_id or "").strip(): wanted[str(p.ne_id).strip()] += 1 if any(ne_running[ne] + n > per_ne_cap for ne, n in wanted.items()): continue now = _utcnow() for member in group: member.status = "running" member.message = "" if member.collect_round_id: member.started_at = now task = db.get(BizStateTask, member.task_id) if task: task.last_collect_started_at = now claimed_rows.append(member) for ne, n in wanted.items(): ne_running[ne] += n if not claimed_rows: db.rollback() return [] # Autoflush so COUNT sees our running marks; trim if over global cap. try: db.flush() except Exception: pass n_running = int( db.query(BizStateBatch).filter(BizStateBatch.status == "running").count() ) while n_running > global_cap and claimed_rows: last = claimed_rows[-1] demoted = [b for b in claimed_rows if b.collect_round_id == last.collect_round_id] if last.collect_round_id else [last] for demote in demoted: claimed_rows.remove(demote) demote.status = "queued" demote.message = "queued" n_running -= 1 if not claimed_rows: db.rollback() return [] out: list[dict[str, Any]] = [] for batch in claimed_rows: task = db.get(BizStateTask, batch.task_id) out.append( { "batch_id": batch.id, "task_id": str(batch.task_id or ""), "source": str( batch.source or (task.source if task else "") or "managed" ), "ne_id": str(batch.ne_id or (task.ne_id if task else "") or ""), "vendor": str(batch.vendor or (task.vendor if task else "") or ""), "device_type": str((task.device_type if task else "") or ""), } ) db.commit() return out except Exception: _log.exception("claim_queued_batches failed") try: db.rollback() except Exception: pass return [] finally: db.close() def reclaim_stale_queued(*, max_age_sec: int = 3600) -> int: """Fail queued batches older than max_age_sec (no worker drained them).""" from datetime import timedelta age = max(60, int(max_age_sec)) cutoff = _utcnow() - timedelta(seconds=age) db = SessionLocal() n = 0 try: stale = ( db.query(BizStateBatch) .filter( BizStateBatch.status == "queued", BizStateBatch.started_at < cutoff, ) .all() ) for b in stale: b.status = "failed" b.message = f"stale_queued_timeout ({age}s)"[:1020] b.ended_at = _utcnow() task = db.get(BizStateTask, b.task_id) if task and bool(task.collect_running): # Only clear if no other running/queued for this task. other = ( db.query(BizStateBatch) .filter( BizStateBatch.task_id == b.task_id, BizStateBatch.id != b.id, BizStateBatch.status.in_(("queued", "running")), ) .first() ) if not other: task.collect_running = False if hasattr(task, "collect_queued_at"): task.collect_queued_at = None task.last_collect_ended_at = _utcnow() task.last_error = str(b.message)[:1020] n += 1 if n: db.commit() _log.warning("reclaimed stale queued batches n=%s age_sec=%s", n, age) return n except Exception: _log.exception("reclaim_stale_queued failed") try: db.rollback() except Exception: pass return 0 finally: db.close() def open_slots() -> int: """How many new running batches we can still start under the ceiling.""" return max(0, max_concurrent_tasks() - count_running_batches())