216 lines
7.5 KiB
Python
216 lines
7.5 KiB
Python
import asyncio
|
|
import os
|
|
import json
|
|
import time
|
|
import secrets
|
|
from typing import Dict, List, Optional, Any, Set
|
|
from fastapi import WebSocket
|
|
from backend.config import JOBS_FILE
|
|
|
|
class JobManager:
|
|
def __init__(self):
|
|
self.active_jobs: Dict[str, Dict[str, Any]] = {}
|
|
self.subscribers: Dict[str, Set[WebSocket]] = {}
|
|
self.global_subscribers: Set[WebSocket] = set()
|
|
self.lock = asyncio.Lock()
|
|
self.type_locks: Dict[str, asyncio.Lock] = {
|
|
"backup": asyncio.Lock(),
|
|
"restore": asyncio.Lock(),
|
|
"update_pangolin": asyncio.Lock(),
|
|
"update_system": asyncio.Lock(),
|
|
}
|
|
|
|
def _load_history(self) -> List[Dict[str, Any]]:
|
|
if not JOBS_FILE.exists():
|
|
return []
|
|
try:
|
|
with open(JOBS_FILE, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except Exception:
|
|
return []
|
|
|
|
def _save_history(self, history: List[Dict[str, Any]]) -> None:
|
|
try:
|
|
with open(JOBS_FILE, "w", encoding="utf-8") as f:
|
|
json.dump(history[:50], f, indent=2)
|
|
except Exception as e:
|
|
print(f"Error saving jobs history: {e}")
|
|
|
|
def list_jobs(self, limit: int = 25) -> List[Dict[str, Any]]:
|
|
history = self._load_history()
|
|
# Merge active jobs at the beginning
|
|
active_list = list(self.active_jobs.values())
|
|
seen_ids = {j["id"] for j in active_list}
|
|
merged = active_list + [j for j in history if j["id"] not in seen_ids]
|
|
merged.sort(key=lambda x: x.get("start_time", 0), reverse=True)
|
|
return merged[:limit]
|
|
|
|
def get_job(self, job_id: str) -> Optional[Dict[str, Any]]:
|
|
if job_id in self.active_jobs:
|
|
return self.active_jobs[job_id]
|
|
history = self._load_history()
|
|
for j in history:
|
|
if j["id"] == job_id:
|
|
return j
|
|
return None
|
|
|
|
async def register_subscriber(self, job_id: str, ws: WebSocket):
|
|
if job_id not in self.subscribers:
|
|
self.subscribers[job_id] = set()
|
|
self.subscribers[job_id].add(ws)
|
|
|
|
def unregister_subscriber(self, job_id: str, ws: WebSocket):
|
|
if job_id in self.subscribers:
|
|
self.subscribers[job_id].discard(ws)
|
|
if not self.subscribers[job_id]:
|
|
del self.subscribers[job_id]
|
|
|
|
async def register_global(self, ws: WebSocket):
|
|
self.global_subscribers.add(ws)
|
|
|
|
def unregister_global(self, ws: WebSocket):
|
|
self.global_subscribers.discard(ws)
|
|
|
|
async def broadcast(self, job_id: str, message: Dict[str, Any]):
|
|
# Broadcast to specific job subscribers
|
|
if job_id in self.subscribers:
|
|
dead = set()
|
|
for ws in list(self.subscribers[job_id]):
|
|
try:
|
|
await ws.send_json(message)
|
|
except Exception:
|
|
dead.add(ws)
|
|
for ws in dead:
|
|
self.subscribers[job_id].discard(ws)
|
|
|
|
# Broadcast to global subscribers
|
|
dead_global = set()
|
|
for ws in list(self.global_subscribers):
|
|
try:
|
|
await ws.send_json(message)
|
|
except Exception:
|
|
dead_global.add(ws)
|
|
for ws in dead_global:
|
|
self.global_subscribers.discard(ws)
|
|
|
|
async def run_command_job(
|
|
self,
|
|
name: str,
|
|
cmd: List[str] | str,
|
|
job_type: str = "general",
|
|
env: Optional[Dict[str, str]] = None,
|
|
cwd: Optional[str] = None,
|
|
description: str = ""
|
|
) -> Dict[str, Any]:
|
|
lock = self.type_locks.get(job_type)
|
|
if lock and lock.locked():
|
|
return {
|
|
"success": False,
|
|
"error": f"A {job_type} task is already running in the background. Please wait for it to complete."
|
|
}
|
|
|
|
job_id = f"job_{int(time.time())}_{secrets.token_hex(4)}"
|
|
job = {
|
|
"id": job_id,
|
|
"name": name,
|
|
"type": job_type,
|
|
"description": description,
|
|
"cmd": cmd if isinstance(cmd, str) else " ".join(cmd),
|
|
"status": "running",
|
|
"start_time": time.time(),
|
|
"end_time": None,
|
|
"duration": None,
|
|
"exit_code": None,
|
|
"logs": []
|
|
}
|
|
|
|
self.active_jobs[job_id] = job
|
|
await self.broadcast(job_id, {"type": "job_started", "job": job})
|
|
|
|
# Run in background asyncio task
|
|
asyncio.create_task(self._execute_process(job_id, cmd, env, cwd, lock))
|
|
return {"success": True, "job_id": job_id, "job": job}
|
|
|
|
async def _execute_process(
|
|
self,
|
|
job_id: str,
|
|
cmd: List[str] | str,
|
|
env: Optional[Dict[str, str]],
|
|
cwd: Optional[str],
|
|
lock: Optional[asyncio.Lock]
|
|
):
|
|
if lock:
|
|
await lock.acquire()
|
|
|
|
job = self.active_jobs[job_id]
|
|
run_env = os.environ.copy()
|
|
run_env["DEBIAN_FRONTEND"] = "noninteractive"
|
|
run_env["TERM"] = "xterm-256color"
|
|
run_env["COLUMNS"] = "120"
|
|
if env:
|
|
run_env.update(env)
|
|
|
|
try:
|
|
if isinstance(cmd, list):
|
|
proc = await asyncio.create_subprocess_exec(
|
|
*cmd,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.STDOUT,
|
|
env=run_env,
|
|
cwd=cwd
|
|
)
|
|
else:
|
|
proc = await asyncio.create_subprocess_shell(
|
|
cmd,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.STDOUT,
|
|
env=run_env,
|
|
cwd=cwd,
|
|
executable="/bin/bash" if os.path.exists("/bin/bash") else None
|
|
)
|
|
|
|
while True:
|
|
line = await proc.stdout.readline()
|
|
if not line:
|
|
break
|
|
text = line.decode('utf-8', errors='replace').rstrip('\r\n')
|
|
timestamp = time.strftime("%H:%M:%S")
|
|
log_entry = f"[{timestamp}] {text}"
|
|
job["logs"].append(log_entry)
|
|
# Keep up to 2000 lines
|
|
if len(job["logs"]) > 2000:
|
|
job["logs"] = job["logs"][-2000:]
|
|
await self.broadcast(job_id, {"type": "log", "job_id": job_id, "data": log_entry})
|
|
|
|
await proc.wait()
|
|
exit_code = proc.returncode
|
|
job["exit_code"] = exit_code
|
|
job["status"] = "completed" if exit_code == 0 else "failed"
|
|
|
|
except Exception as e:
|
|
err_msg = f"Execution error: {str(e)}"
|
|
job["logs"].append(f"[{time.strftime('%H:%M:%S')}] ERROR: {err_msg}")
|
|
job["status"] = "failed"
|
|
job["exit_code"] = -1
|
|
await self.broadcast(job_id, {"type": "log", "job_id": job_id, "data": err_msg})
|
|
|
|
finally:
|
|
job["end_time"] = time.time()
|
|
job["duration"] = round(job["end_time"] - job["start_time"], 1)
|
|
await self.broadcast(job_id, {"type": "job_finished", "job": job})
|
|
|
|
# Save to history file
|
|
history = self._load_history()
|
|
history.insert(0, job.copy())
|
|
self._save_history(history)
|
|
|
|
# Clean from active_jobs after a short delay so clients can finish reading
|
|
await asyncio.sleep(5)
|
|
if job_id in self.active_jobs:
|
|
del self.active_jobs[job_id]
|
|
|
|
if lock:
|
|
lock.release()
|
|
|
|
runner = JobManager()
|