"""Background server batch jobs (category, health check, Agent ops, path detect).""" from __future__ import annotations import asyncio import logging import shlex from typing import Any, Optional from server.application.server_batch_common import ( agent_install_skip_result, batch_server_display_name, detect_and_save_target_path, ensure_agent_api_key, install_error_msg, mark_agent_uninstalled, prefetch_batch_context, result_item_dict, sudo_wrap, ) from server.utils.agent_version import get_central_agent_version from server.config import settings from server.domain.models import AuditLog, Server from server.infrastructure.database.audit_log_repo import AuditLogRepositoryImpl from server.infrastructure.database.server_repo import ServerRepositoryImpl from server.infrastructure.database.session import AsyncSessionLocal from server.infrastructure.redis import server_batch_store as sbs from server.infrastructure.redis.client import get_redis from server.infrastructure.ssh.asyncssh_pool import exec_ssh_command logger = logging.getLogger("nexus.server_batch_service") _BG_TASKS: set[asyncio.Task] = set() CONCURRENCY = 5 OP_LABELS = { "category": "批量修改分类", "health-check": "健康检查", "detect-path": "自动检测目标路径", "onboard": "新服务器初始化", "install-agent": "批量安装 Agent", "upgrade-agent": "批量升级 Agent", "uninstall-agent": "批量卸载 Agent", } AUDIT_ACTIONS = { "category": "batch_update_category", "health-check": "batch_health_check", "detect-path": "batch_detect_path", "onboard": "batch_server_onboard", "install-agent": "batch_install_agent", "upgrade-agent": "batch_upgrade_agent", "uninstall-agent": "batch_uninstall_agent", } HEARTBEAT_PREFIX = "heartbeat:" async def _notify_progress(live: dict[str, Any]) -> None: try: from server.api.websocket import broadcast_script_progress stats = sbs.progress_stats(live) await broadcast_script_progress({ "execution_id": int(live["id"]), "task_kind": "server_batch", "op": live.get("op") or "", "label": live.get("label") or "", "status": live.get("status") or "running", "long_task": False, "operator": live.get("operator") or "", **stats, }) except Exception as e: logger.warning("server batch progress notify failed: %s", e) async def _notify_complete(live: dict[str, Any]) -> None: try: from server.api.websocket import broadcast_script_complete stats = sbs.progress_stats(live) await broadcast_script_complete({ "execution_id": int(live["id"]), "task_kind": "server_batch", "op": live.get("op") or "", "label": live.get("label") or "", "status": live.get("status") or "completed", "long_task": False, "operator": live.get("operator") or "", **stats, }) except Exception as e: logger.warning("server batch complete notify failed: %s", e) async def _reconcile_live_job( live: dict[str, Any], *, error: str, auto_stuck: bool, ) -> bool: """Fill pending servers as failed, finalize, persist, audit, notify. Returns True if reconciled.""" if live.get("status") != "running": return False pending = sbs.pending_server_ids(live) if not pending and sbs.progress_stats(live)["remaining"] == 0: sbs.finalize_status(live) await sbs.save_live(live) else: added = sbs.apply_pending_failures( live, server_ids=pending, error=error, auto_stuck=auto_stuck, ) if added == 0 and not pending: return False sbs.finalize_status(live) await sbs.save_live(live) await _persist_job_finalize(live) await _write_audit(live) await _notify_complete(live) job_id = int(live["id"]) logger.info( "Server batch job reconciled job_id=%s op=%s status=%s auto_stuck=%s", job_id, live.get("op"), live.get("status"), auto_stuck, ) return True async def recover_orphaned_batch_jobs() -> int: """On API startup: finalize batch jobs left running after process restart.""" from server.infrastructure.database.server_batch_job_repo import ( ServerBatchJobRepositoryImpl, ) reconciled_ids: set[int] = set() count = 0 for live in await sbs.list_running_live_jobs(): job_id = int(live["id"]) if job_id in reconciled_ids: continue if await _reconcile_live_job( live, error=sbs.INTERRUPTED_MESSAGE, auto_stuck=False, ): reconciled_ids.add(job_id) count += 1 async with AsyncSessionLocal() as session: repo = ServerBatchJobRepositoryImpl(session) for row in await repo.list_running(): job_id = int(row.id) if job_id in reconciled_ids: continue live = await sbs.get_live(job_id) if live is None: live = sbs.live_from_mysql_row(row) elif live.get("status") != "running": continue if await _reconcile_live_job( live, error=sbs.INTERRUPTED_MESSAGE, auto_stuck=False, ): reconciled_ids.add(job_id) count += 1 if count: logger.info("Recovered %s orphaned server batch job(s) after startup", count) return count async def detect_stuck_batch_jobs() -> int: """Mark running batch jobs with no progress for STUCK_NO_PROGRESS_SECONDS.""" count = 0 for live in await sbs.list_running_live_jobs(): if not sbs.is_stale_running(live): continue if await _reconcile_live_job( live, error=sbs.STUCK_MESSAGE, auto_stuck=True, ): count += 1 if count: logger.info("Auto-reconciled %s stuck server batch job(s)", count) return count async def start_batch_job( *, op: str, server_ids: list[int], operator: str, params: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: if op not in OP_LABELS: raise ValueError(f"未知的批量操作: {op}") if not server_ids: raise ValueError("server_ids 不能为空") job_id = await sbs.next_job_id() label = OP_LABELS[op] if op == "category" and params and params.get("category"): label = f"{label} → {params['category']}" live = sbs.init_live( job_id=job_id, op=op, label=label, server_ids=server_ids, operator=operator, params=params, ) await sbs.save_live(live) await _persist_job_start(live) await _notify_progress(live) task = asyncio.create_task(_run_background(job_id)) _BG_TASKS.add(task) task.add_done_callback(_BG_TASKS.discard) stats = sbs.progress_stats(live) return { "job_id": job_id, "status": "running", "label": label, "op": op, **stats, } async def get_batch_job(job_id: int) -> Optional[dict[str, Any]]: live = await sbs.get_live(job_id) if live: return sbs.live_to_api(live) return await _get_batch_job_from_mysql(job_id) async def list_batch_jobs( *, limit: int = 50, offset: int = 0, op: Optional[str] = None, ) -> dict[str, Any]: from server.infrastructure.database.server_batch_job_repo import ( ServerBatchJobRepositoryImpl, job_row_to_api, ) async with AsyncSessionLocal() as session: repo = ServerBatchJobRepositoryImpl(session) rows, total = await repo.list_recent(limit=limit, offset=offset, op=op) items = [job_row_to_api(row) for row in rows] return {"items": items, "total": total, "limit": limit, "offset": offset} async def _persist_job_start(live: dict[str, Any]) -> None: from server.infrastructure.database.server_batch_job_repo import ServerBatchJobRepositoryImpl try: async with AsyncSessionLocal() as session: repo = ServerBatchJobRepositoryImpl(session) await repo.upsert_running( job_id=int(live["id"]), op=str(live.get("op") or ""), label=str(live.get("label") or ""), server_ids=list(live.get("server_ids") or []), operator=str(live.get("operator") or ""), params=live.get("params") if isinstance(live.get("params"), dict) else {}, ) except Exception as e: logger.warning("Server batch MySQL create failed job_id=%s: %s", live.get("id"), e) async def _persist_job_finalize(live: dict[str, Any]) -> None: from server.infrastructure.database.server_batch_job_repo import ServerBatchJobRepositoryImpl try: api = sbs.live_to_api(live) async with AsyncSessionLocal() as session: repo = ServerBatchJobRepositoryImpl(session) await repo.finalize( int(live["id"]), status=str(live.get("status") or "completed"), results=list(api.get("results") or []), ) except Exception as e: logger.warning("Server batch MySQL finalize failed job_id=%s: %s", live.get("id"), e) async def _get_batch_job_from_mysql(job_id: int) -> Optional[dict[str, Any]]: from server.infrastructure.database.server_batch_job_repo import ( ServerBatchJobRepositoryImpl, job_row_to_api, ) try: async with AsyncSessionLocal() as session: repo = ServerBatchJobRepositoryImpl(session) row = await repo.get_by_id(job_id) if not row: return None return job_row_to_api(row) except Exception as e: logger.warning("Server batch MySQL read failed job_id=%s: %s", job_id, e) return None async def _run_background(job_id: int) -> None: live = await sbs.get_live(job_id) if not live: return op = live.get("op") or "" try: if op == "category": await _run_category(live) elif op == "health-check": await _run_health_check(live) elif op == "detect-path": await _run_detect_path(live) elif op == "onboard": await _run_onboard(live) elif op == "install-agent": await _run_install_agent(live) elif op == "upgrade-agent": await _run_upgrade_agent(live) elif op == "uninstall-agent": await _run_uninstall_agent(live) else: live["status"] = "failed" await sbs.save_live(live) return live = await sbs.get_live(job_id) if live: stats = sbs.progress_stats(live) if live.get("status") == "running" or stats["remaining"] == 0: sbs.finalize_status(live) await sbs.save_live(live) except Exception: logger.exception("Server batch job failed job_id=%s op=%s", job_id, op) live = await sbs.get_live(job_id) if live: live["status"] = "failed" await sbs.save_live(live) live = await sbs.get_live(job_id) if live: if live.get("status") == "running": sbs.finalize_status(live) await sbs.save_live(live) await _persist_job_finalize(live) await _write_audit(live) await _notify_complete(live) async def _write_audit(live: dict[str, Any]) -> None: import json op = live.get("op") or "" action = AUDIT_ACTIONS.get(op, "batch_server_op") stats = sbs.progress_stats(live) api = sbs.live_to_api(live) failures = [ {"server_name": (e.get("server_name") or "").strip() or "未知服务器", "error": (e.get("error") or "")[:300]} for e in (api.get("results") or []) if not e.get("success") ] summary_map = { "category": f"批量修改分类 更新{stats['success']}台", "health-check": f"批量健康检查: {stats['success']}/{stats['total']} 在线", "detect-path": f"批量检测路径: {stats['success']}/{stats['total']} 成功", "onboard": f"新服务器初始化: {stats['success']}/{stats['total']} 成功", "install-agent": f"批量安装Agent: {stats['success']}/{stats['total']} 成功", "upgrade-agent": f"批量升级Agent: {stats['success']}/{stats['total']} 成功", "uninstall-agent": f"批量卸载Agent: {stats['success']}/{stats['total']} 成功", } summary = summary_map.get(op, f"{live.get('label')} {stats['success']}/{stats['total']}") detail = json.dumps({ "job_id": int(live["id"]), "summary": summary, "success": stats["success"], "failed": stats["failed"], "total": stats["total"], "failures": failures[:30], }, ensure_ascii=False) try: async with AsyncSessionLocal() as session: audit_repo = AuditLogRepositoryImpl(session) await audit_repo.create(AuditLog( admin_username=live.get("operator") or "admin", action=action, target_type="server_batch_job", target_id=int(live["id"]), detail=detail, ip_address="", )) await session.commit() except Exception as e: logger.warning("Server batch audit failed job_id=%s: %s", live.get("id"), e) async def _finish_one(live: dict[str, Any], item: dict[str, Any]) -> None: sbs.record_result(live, item) await sbs.save_live(live) await _notify_progress(live) async def _run_with_semaphore(live: dict[str, Any], handler) -> None: from server.application.services.server_service import ServerService from server.infrastructure.database.server_repo import ServerRepositoryImpl from server.infrastructure.database.sync_log_repo import SyncLogRepositoryImpl server_ids = live.get("server_ids") or [] sem = asyncio.Semaphore(CONCURRENCY) async with AsyncSessionLocal() as session: service = ServerService( server_repo=ServerRepositoryImpl(session), sync_log_repo=SyncLogRepositoryImpl(session), ) server_map, labels = await prefetch_batch_context(service, server_ids) async def _one(sid: int) -> None: async with sem: item = await handler(sid, server_map, labels, live) await _finish_one(live, item) await asyncio.gather(*[_one(sid) for sid in server_ids]) async def _run_category(live: dict[str, Any]) -> None: params = live.get("params") or {} category = (params.get("category") or "").strip() or None server_ids = live.get("server_ids") or [] async with AsyncSessionLocal() as session: repo = ServerRepositoryImpl(session) servers = await repo.get_by_ids(server_ids) found_ids = {s.id for s in servers} labels = {s.id: batch_server_display_name(s) for s in servers} for sid in server_ids: labels.setdefault(sid, "未知服务器") for sid in server_ids: if sid not in found_ids: await _finish_one(live, result_item_dict( server_id=sid, server_name=labels[sid], success=False, error="服务器不存在", )) if found_ids: await repo.update_category_batch(list(found_ids), category) await session.commit() cat_label = category or "(清除)" for sid in found_ids: await _finish_one(live, result_item_dict( server_id=sid, server_name=labels[sid], success=True, stdout=f"分类 → {cat_label}", )) async def _run_health_check(live: dict[str, Any]) -> None: from server.infrastructure.ssh.remote_probe import ssh_health_probe async def handler(sid: int, server_map: dict[int, Server], labels: dict[int, str], _live) -> dict: server = server_map.get(sid) label = labels[sid] if not server: return result_item_dict(server_id=sid, server_name=label, success=False, error="服务器不存在") try: probe = await ssh_health_probe(server, timeout=15) online = probe.get("status") == "online" detail = probe.get("error") or probe.get("channel") or probe.get("status") or "" return result_item_dict( server_id=sid, server_name=label, success=online, stdout=str(detail)[:500] if online else "", error="" if online else str(probe.get("error") or "离线")[:300], ) except Exception as e: return result_item_dict(server_id=sid, server_name=label, success=False, error=str(e)[:300]) await _run_with_semaphore(live, handler) async def _run_detect_path(live: dict[str, Any]) -> None: async def handler(sid: int, server_map: dict[int, Server], labels: dict[int, str], _live) -> dict: server = server_map.get(sid) label = labels[sid] if not server: return result_item_dict(server_id=sid, server_name=label, success=False, error="服务器不存在") try: outcome = await detect_and_save_target_path(server) if outcome.get("success"): return result_item_dict( server_id=sid, server_name=label, success=True, stdout=str(outcome.get("stdout") or ""), ) return result_item_dict( server_id=sid, server_name=label, success=False, error=str(outcome.get("error") or "探测失败"), ) except Exception as e: return result_item_dict(server_id=sid, server_name=label, success=False, error=str(e)[:300]) await _run_with_semaphore(live, handler) async def _run_onboard(live: dict[str, Any]) -> None: from server.application.services.files_sudoers_service import install_nexus_files_sudoers from server.utils.posix_paths import is_unset_target_path async def handler(sid: int, server_map: dict[int, Server], labels: dict[int, str], _live) -> dict: server = server_map.get(sid) label = labels[sid] if not server: return result_item_dict(server_id=sid, server_name=label, success=False, error="服务器不存在") steps: list[str] = [] try: ssh_user = (server.username or "root").strip() or "root" if ssh_user == "root": steps.append("sudo: skip(root)") else: sudo_outcome = await install_nexus_files_sudoers(server) action = sudo_outcome.get("action") or "unknown" if not sudo_outcome.get("ok"): detail = sudo_outcome.get("detail") or action return result_item_dict( server_id=sid, server_name=label, success=False, error=f"sudo 配置失败: {detail}", stdout="; ".join(steps), ) steps.append(f"sudo: {action}") if is_unset_target_path(server.target_path): detect_outcome = await detect_and_save_target_path(server) if detect_outcome.get("success"): steps.append(str(detect_outcome.get("stdout") or "path: ok")) else: steps.append(f"path: {detect_outcome.get('error') or 'fail'}（推送仍可用 /www/wwwroot）") else: steps.append("path: already_set") return result_item_dict( server_id=sid, server_name=label, success=True, stdout="; ".join(steps), ) except Exception as e: return result_item_dict( server_id=sid, server_name=label, success=False, error=str(e)[:300], stdout="; ".join(steps), ) await _run_with_semaphore(live, handler) async def _run_install_agent(live: dict[str, Any]) -> None: base_url = (settings.API_BASE_URL or "").strip().rstrip("/") if not base_url: for sid in live.get("server_ids") or []: await _finish_one(live, result_item_dict( server_id=sid, server_name=str(sid), success=False, error="NEXUS_API_BASE_URL 未配置", )) return central_version = get_central_agent_version() async def handler(sid: int, server_map: dict[int, Server], labels: dict[int, str], _live) -> dict: server = server_map.get(sid) label = labels[sid] if not server: return result_item_dict(server_id=sid, server_name=label, success=False, error="服务器不存在") try: skip = await agent_install_skip_result( server=server, server_id=sid, server_name=label, central_version=central_version, ) if skip is not None: return skip api_key = await ensure_agent_api_key(server) ssh_user = (server.username or "root").strip() install_cmd = ( f"TMP=$(mktemp) && curl -fsSL {shlex.quote(base_url)}/agent/install.sh -o \"$TMP\" && bash \"$TMP\" -- " f"--url {shlex.quote(base_url)} --key {shlex.quote(api_key)} " f"--id {sid} --port {shlex.quote(str(server.agent_port or 8601))}" f"; RC=$?; rm -f \"$TMP\"; exit $RC" ) install_cmd = sudo_wrap(install_cmd, ssh_user) r = await exec_ssh_command(server, install_cmd, timeout=180) ok = r["exit_code"] == 0 if ok and "Status: FAILED" in r.get("stdout", ""): ok = False return result_item_dict( server_id=sid, server_name=label, success=ok, stdout=r["stdout"][:2000] if ok else r["stdout"][:1000], error="" if ok else install_error_msg(r.get("stdout", ""), r.get("stderr", ""), r["exit_code"]), ) except Exception as e: return result_item_dict(server_id=sid, server_name=label, success=False, error=str(e)[:300]) await _run_with_semaphore(live, handler) async def _run_upgrade_agent(live: dict[str, Any]) -> None: base_url = (settings.API_BASE_URL or "").strip().rstrip("/") if not base_url: for sid in live.get("server_ids") or []: await _finish_one(live, result_item_dict( server_id=sid, server_name=str(sid), success=False, error="NEXUS_API_BASE_URL 未配置", )) return async def handler(sid: int, server_map: dict[int, Server], labels: dict[int, str], _live) -> dict: server = server_map.get(sid) label = labels[sid] if not server: return result_item_dict(server_id=sid, server_name=label, success=False, error="服务器不存在") try: from server.utils.agent_deploy import ( AGENT_INSTALL_DIR, build_agent_upgrade_cmd, build_agent_upgrade_rollback_cmd, ) install_dir = AGENT_INSTALL_DIR ssh_user = (server.username or "root").strip() agent_port = server.agent_port or 8601 upgrade_cmd = build_agent_upgrade_cmd( base_url=base_url, ssh_user=ssh_user, install_dir=install_dir, agent_port=agent_port, ) r = await exec_ssh_command(server, upgrade_cmd, timeout=120) ok = r["exit_code"] == 0 and "upgrade_ok" in r["stdout"] if not ok: try: rollback = build_agent_upgrade_rollback_cmd( ssh_user=ssh_user, install_dir=install_dir, ) await exec_ssh_command(server, rollback, timeout=30) except Exception as rb_err: logger.warning("Rollback failed for server %s: %s", sid, rb_err) return result_item_dict( server_id=sid, server_name=label, success=ok, stdout=r["stdout"][:500] if ok else "", error=install_error_msg(r.get("stdout", ""), r.get("stderr", ""), r["exit_code"]) if not ok else "", ) except Exception as e: return result_item_dict(server_id=sid, server_name=label, success=False, error=str(e)[:200]) await _run_with_semaphore(live, handler) async def _run_uninstall_agent(live: dict[str, Any]) -> None: base_url = (settings.API_BASE_URL or "").strip().rstrip("/") if not base_url: for sid in live.get("server_ids") or []: await _finish_one(live, result_item_dict( server_id=sid, server_name=str(sid), success=False, error="NEXUS_API_BASE_URL 未配置", )) return async def handler(sid: int, server_map: dict[int, Server], labels: dict[int, str], _live) -> dict: server = server_map.get(sid) label = labels[sid] if not server: return result_item_dict(server_id=sid, server_name=label, success=False, error="服务器不存在") try: ssh_user = (server.username or "root").strip() uninstall_cmd = ( f"TMP=$(mktemp) && curl -fsSL {shlex.quote(base_url)}/agent/uninstall.sh -o \"$TMP\" " f"&& bash \"$TMP\" ; RC=$?; rm -f \"$TMP\"; exit $RC" ) uninstall_cmd = sudo_wrap(uninstall_cmd, ssh_user) r = await exec_ssh_command(server, uninstall_cmd, timeout=60) ok = r["exit_code"] == 0 if ok: try: redis = get_redis() await redis.delete(f"{HEARTBEAT_PREFIX}{sid}") except Exception as redis_err: logger.warning("Failed to clear Redis heartbeat for server %s: %s", sid, redis_err) await mark_agent_uninstalled(sid) return result_item_dict( server_id=sid, server_name=label, success=ok, stdout=r["stdout"][:500] if ok else "", error="" if ok else install_error_msg(r.get("stdout", ""), r.get("stderr", ""), r["exit_code"]), ) except Exception as e: return result_item_dict(server_id=sid, server_name=label, success=False, error=str(e)[:300]) await _run_with_semaphore(live, handler)