"""Config sync worker: claim tasks, collect vendor configs, store zlib snapshots.""" from __future__ import annotations import io import logging import time from concurrent.futures import ThreadPoolExecutor from datetime import datetime from threading import Lock from typing import Any from uuid import uuid4 from fastapi import HTTPException from .cli_resolve import resolve_cli_target from .config import settings from .config_sync_codec import compress_text from .config_sync_commands import command_list, commands_for_vendor from .config_sync_service import finalize_cycle, sync_cycle_progress from .db import SessionLocal from .models import ConfigSyncCycle, ConfigSyncPolicy, ConfigSyncTask, NeConfigHistory, NeConfigSnapshot from .ne_cli_errors import format_cli_failure, session_log_text from .ne_netmiko import send_show_command from .ne_session_factory import close_netmiko_connection, open_netmiko_connection _log = logging.getLogger("netx.config_sync.runner") _pools: dict[str, ThreadPoolExecutor] = {} _pools_lock = Lock() def _utcnow() -> datetime: return datetime.utcnow() def _format_error(exc: BaseException) -> str: return format_cli_failure(exc, limit=1020) def _pool_for_cycle(cycle_id: str, concurrency: int) -> ThreadPoolExecutor: from .cli_budget import clamp_cli_workers with _pools_lock: pool = _pools.get(cycle_id) if pool is None: workers = clamp_cli_workers(int(concurrency or 5)) pool = ThreadPoolExecutor(max_workers=workers, thread_name_prefix=f"cfg-sync-{cycle_id[:8]}") _pools[cycle_id] = pool return pool def _release_pool(cycle_id: str) -> None: with _pools_lock: pool = _pools.pop(cycle_id, None) if pool is not None: try: pool.shutdown(wait=False, cancel_futures=False) except TypeError: pool.shutdown(wait=False) except Exception: _log.exception("config_sync pool shutdown failed cycle=%s", cycle_id) def shutdown_config_sync_pools(*, wait: bool = False) -> None: with _pools_lock: pools = list(_pools.items()) _pools.clear() for cycle_id, pool in pools: try: pool.shutdown(wait=wait, cancel_futures=True) except TypeError: pool.shutdown(wait=wait) except Exception: _log.exception("config_sync pool shutdown failed cycle=%s", cycle_id) def _update_task(task_id: str, **fields: Any) -> None: db = SessionLocal() try: row = db.get(ConfigSyncTask, task_id) if not row: return for key, val in fields.items(): setattr(row, key, val) db.commit() finally: db.close() def _claim_task(cycle_id: str, task_id: str) -> bool: for attempt in range(10): db = SessionLocal() try: task = db.get(ConfigSyncTask, task_id) cycle = db.get(ConfigSyncCycle, cycle_id) if not task or not cycle: time.sleep(0.05 * (attempt + 1)) continue cycle_status = str(cycle.status or "") task_status = str(task.status or "") if cycle_status == "paused": return False if cycle_status != "running": return False if task_status == "running": return False if task_status != "pending": return False task.status = "running" task.message = "collecting" task.started_at = _utcnow() db.commit() return True finally: db.close() time.sleep(0.05 * (attempt + 1)) return False def _collect_commands( creds: dict[str, Any], commands: list[str], *, conn_holder: dict[str, Any] | None = None, ) -> list[str]: per_cmd = int(settings.ne_collect_read_timeout_sec or 120) session_timeout = per_cmd * max(1, len(commands)) + 60 log_buf = io.BytesIO() try: conn = open_netmiko_connection(creds, session_timeout=session_timeout, session_log=log_buf) except Exception as exc: raise RuntimeError(format_cli_failure(exc, session_log_text(log_buf))) from exc if conn_holder is not None: conn_holder["conn"] = conn try: outputs: list[str] = [] for command in commands: if conn_holder is not None and conn_holder.get("timed_out"): raise TimeoutError("config_sync_aborted") try: out = send_show_command(conn, command, read_timeout=per_cmd) except Exception as exc: raise RuntimeError(format_cli_failure(exc, session_log_text(log_buf))) from exc outputs.append(str(out or "")) return outputs finally: if conn_holder is not None: conn_holder.pop("conn", None) close_netmiko_connection(conn) def _collect_with_timeout(creds: dict[str, Any], commands: list[str]) -> list[str]: from .cli_timeout import run_cli_with_timeout per_cmd = int(settings.ne_collect_read_timeout_sec or 120) cap = int(settings.ne_collect_run_timeout_cap_sec or 600) budget = min(cap, per_cmd * max(1, len(commands)) + 90) holder: dict[str, Any] = {} return run_cli_with_timeout( lambda: _collect_commands(creds, commands, conn_holder=holder), timeout_sec=budget, conn_holder=holder, label="config_sync", acquire_budget=True, ) def _history_keep(db) -> int: pol = db.get(ConfigSyncPolicy, 1) if not pol: return 3 return max(0, min(30, int(pol.history_keep if pol.history_keep is not None else 3))) def _save_success_snapshot( *, source: str, target_id: str, vendor: str, device_type: str, ne_name: str, ne_ip: str, primary_text: str, alt_text: str | None, commands: list[str], cycle_id: str, task_id: str, ) -> None: primary_blob, primary_sha, plain_size, zlib_size = compress_text(primary_text) alt_blob = None alt_sha = "" plain_alt = 0 zlib_alt = 0 if alt_text is not None: alt_blob, alt_sha, plain_alt, zlib_alt = compress_text(alt_text) db = SessionLocal() try: existing = db.get(NeConfigSnapshot, {"source": source, "target_id": target_id}) changed = True if existing is not None: changed = ( str(existing.config_sha256 or "") != primary_sha or str(existing.config_alt_sha256 or "") != alt_sha ) if changed: keep = _history_keep(db) if keep > 0: # Archive previous successful snapshot before overwrite. db.add( NeConfigHistory( id=uuid4().hex, source=existing.source, target_id=existing.target_id, vendor=str(existing.vendor or ""), device_type=str(existing.device_type or ""), ne_name=str(existing.ne_name or ""), ne_ip=str(existing.ne_ip or ""), config_zlib=existing.config_zlib or b"", config_alt_zlib=existing.config_alt_zlib, config_sha256=str(existing.config_sha256 or ""), config_alt_sha256=str(existing.config_alt_sha256 or ""), plain_size=int(existing.plain_size or 0), plain_alt_size=int(existing.plain_alt_size or 0), zlib_size=int(existing.zlib_size or 0), zlib_alt_size=int(existing.zlib_alt_size or 0), commands_json=existing.commands_json if isinstance(existing.commands_json, list) else [], collected_at=existing.collected_at or _utcnow(), cycle_id=str(existing.last_cycle_id or ""), task_id=str(existing.last_task_id or ""), ) ) db.flush() old_rows = ( db.query(NeConfigHistory) .filter(NeConfigHistory.source == source, NeConfigHistory.target_id == target_id) .order_by(NeConfigHistory.collected_at.desc()) .all() ) for stale in old_rows[keep:]: db.delete(stale) existing.vendor = vendor existing.device_type = device_type existing.ne_name = ne_name existing.ne_ip = ne_ip existing.config_zlib = primary_blob existing.config_alt_zlib = alt_blob existing.config_sha256 = primary_sha existing.config_alt_sha256 = alt_sha existing.plain_size = plain_size existing.plain_alt_size = plain_alt existing.zlib_size = zlib_size existing.zlib_alt_size = zlib_alt existing.commands_json = list(commands) existing.collected_at = _utcnow() existing.last_cycle_id = cycle_id existing.last_task_id = task_id else: db.add( NeConfigSnapshot( source=source, target_id=target_id, vendor=vendor, device_type=device_type, ne_name=ne_name, ne_ip=ne_ip, config_zlib=primary_blob, config_alt_zlib=alt_blob, config_sha256=primary_sha, config_alt_sha256=alt_sha, plain_size=plain_size, plain_alt_size=plain_alt, zlib_size=zlib_size, zlib_alt_size=zlib_alt, commands_json=list(commands), collected_at=_utcnow(), last_cycle_id=cycle_id, last_task_id=task_id, ) ) db.commit() finally: db.close() def _run_single(cycle_id: str, task_id: str) -> None: if not _claim_task(cycle_id, task_id): db = SessionLocal() try: sync_cycle_progress(db, cycle_id) finalize_cycle(db, cycle_id) finally: db.close() return db = SessionLocal() try: task = db.get(ConfigSyncTask, task_id) if not task: return source = str(task.source or "").strip().lower() target_id = str(task.target_id or "").strip() vendor_hint = str(task.vendor or "") try: if source == "managed": creds, device = resolve_cli_target(db, managed_ne_id=target_id) elif source == "ume": creds, device = resolve_cli_target(db, ume_ne_id=target_id) else: _update_task(task_id, status="fail", message="invalid_source", ended_at=_utcnow()) return except HTTPException as exc: detail = str(exc.detail) if exc.detail else "resolve_failed" _update_task(task_id, status="fail", message=detail[:1020], ended_at=_utcnow()) return except Exception as exc: _update_task(task_id, status="fail", message=_format_error(exc), ended_at=_utcnow()) return vendor = str(device.get("vendor") or vendor_hint or "") device_type = str(device.get("device_type") or "") ne_name = str(device.get("name") or task.ne_name or "") ne_ip = str(device.get("ip_address") or task.ne_ip or "") cmds = commands_for_vendor(vendor, device_type) if cmds is None: _update_task( task_id, status="fail", message="unsupported_vendor", vendor=vendor, ne_name=ne_name, ne_ip=ne_ip, ended_at=_utcnow(), ) return cmd_names = command_list(cmds) try: outputs = _collect_with_timeout(creds, cmd_names) except Exception as exc: _log.warning("config_sync collect failed task=%s: %s", task_id, _format_error(exc)) _update_task( task_id, status="fail", message=_format_error(exc), vendor=vendor, ne_name=ne_name, ne_ip=ne_ip, ended_at=_utcnow(), ) return if not outputs or not str(outputs[0] or "").strip(): _update_task( task_id, status="fail", message="empty_config_output", vendor=vendor, ne_name=ne_name, ne_ip=ne_ip, ended_at=_utcnow(), ) return primary = outputs[0] alt = outputs[1] if cmds.alt and len(outputs) > 1 else None try: _save_success_snapshot( source=source, target_id=target_id, vendor=vendor, device_type=device_type, ne_name=ne_name, ne_ip=ne_ip, primary_text=primary, alt_text=alt, commands=cmd_names, cycle_id=cycle_id, task_id=task_id, ) except Exception as exc: _log.exception("config_sync snapshot save failed task=%s", task_id) _update_task(task_id, status="fail", message=_format_error(exc), ended_at=_utcnow()) return _update_task( task_id, status="success", message="synced", vendor=vendor, ne_name=ne_name, ne_ip=ne_ip, ended_at=_utcnow(), ) finally: db.close() db2 = SessionLocal() try: sync_cycle_progress(db2, cycle_id) finalize_cycle(db2, cycle_id) cycle = db2.get(ConfigSyncCycle, cycle_id) if cycle and str(cycle.status or "") in ("success", "fail", "cancelled"): _release_pool(cycle_id) finally: db2.close() def _run_safe(cycle_id: str, task_id: str) -> None: try: _run_single(cycle_id, task_id) except Exception: _log.exception("config_sync worker crashed cycle=%s task=%s", cycle_id, task_id) _update_task( task_id, status="fail", message="config_sync_worker_crashed", ended_at=_utcnow(), ) db = SessionLocal() try: sync_cycle_progress(db, cycle_id) finalize_cycle(db, cycle_id) finally: db.close() def schedule_cycle_tasks(cycle_id: str, task_ids: list[str], concurrency: int) -> int: if not task_ids: db = SessionLocal() try: finalize_cycle(db, cycle_id) finally: db.close() return 0 pool = _pool_for_cycle(cycle_id, concurrency) submitted = 0 for tid in task_ids: pool.submit(_run_safe, cycle_id, str(tid)) submitted += 1 _log.info("scheduled config_sync cycle=%s tasks=%s concurrency=%s", cycle_id, submitted, concurrency) return submitted def dispatch_cycle(cycle_id: str) -> int: """Load pending tasks for a cycle and schedule workers.""" db = SessionLocal() try: cycle = db.get(ConfigSyncCycle, cycle_id) if not cycle: return 0 if str(cycle.status or "") not in ("running", "pending"): return 0 if str(cycle.status or "") == "pending": cycle.status = "running" if not cycle.started_at: cycle.started_at = _utcnow() db.commit() pending = ( db.query(ConfigSyncTask) .filter(ConfigSyncTask.cycle_id == cycle_id, ConfigSyncTask.status == "pending") .all() ) task_ids = [str(t.id) for t in pending] concurrency = max(1, min(32, int(cycle.concurrency or 5))) finally: db.close() return schedule_cycle_tasks(cycle_id, task_ids, concurrency)