584 lines
21 KiB
Python
584 lines
21 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import subprocess
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from uuid import uuid4
|
||
|
||
import redis
|
||
|
||
from app.core.config import settings
|
||
from app.core.redis_client import get_redis
|
||
from app.services.runtime_settings_service import get_runtime_settings
|
||
|
||
|
||
WORKER_CONTROL_CHANNEL = "domain_tool:worker_control"
|
||
WORKER_PENDING_COMMAND_KEY = "domain_tool:worker_pending_command"
|
||
|
||
|
||
def _normalize_target_node_codes(payload: dict | None) -> list[str]:
|
||
if not isinstance(payload, dict):
|
||
return []
|
||
|
||
normalized_targets: list[str] = []
|
||
|
||
def append_target(raw_value: object) -> None:
|
||
normalized_value = str(raw_value or "").strip()
|
||
if normalized_value and normalized_value not in normalized_targets:
|
||
normalized_targets.append(normalized_value)
|
||
|
||
for key in ("target_node_codes", "node_codes"):
|
||
raw_value = payload.get(key)
|
||
if isinstance(raw_value, (list, tuple, set)):
|
||
for item in raw_value:
|
||
append_target(item)
|
||
elif isinstance(raw_value, str) and raw_value.strip():
|
||
for item in raw_value.split(","):
|
||
append_target(item)
|
||
if normalized_targets:
|
||
return normalized_targets
|
||
|
||
for key in ("target_node_code", "node_code"):
|
||
raw_value = payload.get(key)
|
||
if raw_value not in (None, ""):
|
||
append_target(raw_value)
|
||
if normalized_targets:
|
||
return normalized_targets
|
||
|
||
return normalized_targets
|
||
|
||
|
||
def _pending_command_keys(command_payload: dict) -> list[str]:
|
||
target_node_codes = _normalize_target_node_codes(command_payload)
|
||
if not target_node_codes:
|
||
return [WORKER_PENDING_COMMAND_KEY]
|
||
return [f"{WORKER_PENDING_COMMAND_KEY}:{node_code}" for node_code in target_node_codes]
|
||
|
||
|
||
def _dedupe_target_node_codes(node_codes: list[str]) -> list[str]:
|
||
deduped: list[str] = []
|
||
seen: set[str] = set()
|
||
for raw_value in list(node_codes or []):
|
||
normalized_value = str(raw_value or "").strip()
|
||
if not normalized_value or normalized_value in seen:
|
||
continue
|
||
seen.add(normalized_value)
|
||
deduped.append(normalized_value)
|
||
return deduped
|
||
|
||
|
||
def _expand_local_linux_worker_target_node_codes(service_name: str) -> list[str]:
|
||
base_node_code = str(settings.node_code or "").strip()
|
||
normalized_service_name = str(service_name or "").strip()
|
||
if not base_node_code or not normalized_service_name:
|
||
return []
|
||
|
||
target_node_codes = [base_node_code]
|
||
for unit in _expand_linux_worker_control_units(normalized_service_name):
|
||
normalized_unit = str(unit or "").strip()
|
||
if not normalized_unit:
|
||
continue
|
||
if normalized_unit.endswith(".service"):
|
||
normalized_unit = normalized_unit[:-8]
|
||
if normalized_unit == normalized_service_name:
|
||
continue
|
||
template_prefix = f"{normalized_service_name}@"
|
||
if not normalized_unit.startswith(template_prefix):
|
||
continue
|
||
instance_suffix = str(normalized_unit.split("@", 1)[1] or "").strip()
|
||
if instance_suffix:
|
||
target_node_codes.append(f"{base_node_code}-{instance_suffix}")
|
||
return _dedupe_target_node_codes(target_node_codes)
|
||
|
||
|
||
def _publish_worker_command(redis_client, *, serialized: str, pending_keys: list[str]) -> None:
|
||
for key in pending_keys:
|
||
redis_client.set(key, serialized, ex=120)
|
||
redis_client.publish(WORKER_CONTROL_CHANNEL, serialized)
|
||
|
||
|
||
def _build_direct_redis_client() -> redis.Redis:
|
||
return redis.Redis(
|
||
host=settings.redis_host,
|
||
port=settings.redis_port,
|
||
password=settings.redis_password or None,
|
||
db=settings.redis_db,
|
||
decode_responses=True,
|
||
socket_connect_timeout=5,
|
||
socket_timeout=5,
|
||
retry_on_timeout=True,
|
||
client_name=f"domain-api-workerctl:{settings.node_code}:{os.getpid()}",
|
||
)
|
||
|
||
|
||
def _domain_root() -> Path:
|
||
return Path(settings.domain_root)
|
||
|
||
|
||
def _runtime_config() -> dict:
|
||
return get_runtime_settings()
|
||
|
||
|
||
def _run_powershell(command: str, timeout: int = 20) -> subprocess.CompletedProcess[str]:
|
||
return subprocess.run(
|
||
["powershell", "-NoProfile", "-ExecutionPolicy", "Bypass", "-Command", command],
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=timeout,
|
||
)
|
||
|
||
|
||
def _run_shell(command: list[str], timeout: int = 20) -> subprocess.CompletedProcess[str]:
|
||
return subprocess.run(command, capture_output=True, text=True, timeout=timeout)
|
||
|
||
|
||
def _parse_process_count_output(result: subprocess.CompletedProcess[str]) -> int | None:
|
||
raw_output = (result.stdout or "").strip()
|
||
if raw_output.isdigit():
|
||
return max(0, int(raw_output or 0))
|
||
return None
|
||
|
||
|
||
def _probe_linux_worker_process_count() -> int:
|
||
probe_commands = (
|
||
(["bash", "-lc", "pgrep -fc '[d]etect_worker.py' || true"], 8),
|
||
(["bash", "-lc", "ps -eo args= | grep '[d]etect_worker.py' | wc -l"], 12),
|
||
)
|
||
for command, timeout in probe_commands:
|
||
try:
|
||
result = _run_shell(command, timeout=timeout)
|
||
except Exception:
|
||
continue
|
||
parsed_count = _parse_process_count_output(result)
|
||
if parsed_count is not None:
|
||
return parsed_count
|
||
return 0
|
||
|
||
|
||
def _probe_linux_worker_instance_count(service_name: str) -> int:
|
||
normalized_service_name = str(service_name or "").strip()
|
||
if not normalized_service_name:
|
||
return 0
|
||
template_prefix = normalized_service_name[:-8] if normalized_service_name.endswith(".service") else normalized_service_name
|
||
try:
|
||
result = _run_systemctl(
|
||
[
|
||
"list-units",
|
||
f"{template_prefix}@*",
|
||
"--type=service",
|
||
"--all",
|
||
"--no-legend",
|
||
"--plain",
|
||
],
|
||
timeout=12,
|
||
require_sudo=False,
|
||
)
|
||
except Exception:
|
||
return 0
|
||
if result.returncode != 0:
|
||
return 0
|
||
count = 0
|
||
for raw_line in (result.stdout or "").splitlines():
|
||
line = str(raw_line or "").strip()
|
||
if not line:
|
||
continue
|
||
parts = line.split()
|
||
if len(parts) < 4:
|
||
continue
|
||
if parts[2] != "active" or parts[3] != "running":
|
||
continue
|
||
count += 1
|
||
return count
|
||
|
||
|
||
def _dedupe_units(units: list[str]) -> list[str]:
|
||
deduped: list[str] = []
|
||
seen: set[str] = set()
|
||
for item in units:
|
||
normalized = str(item or "").strip()
|
||
if not normalized or normalized in seen:
|
||
continue
|
||
seen.add(normalized)
|
||
deduped.append(normalized)
|
||
return deduped
|
||
|
||
|
||
def _list_linux_worker_instance_units(service_name: str) -> list[str]:
|
||
normalized_service_name = str(service_name or "").strip()
|
||
if not normalized_service_name:
|
||
return []
|
||
template_prefix = normalized_service_name[:-8] if normalized_service_name.endswith(".service") else normalized_service_name
|
||
units: list[str] = []
|
||
try:
|
||
result = _run_systemctl(
|
||
[
|
||
"list-units",
|
||
f"{template_prefix}@*",
|
||
"--type=service",
|
||
"--all",
|
||
"--no-legend",
|
||
"--plain",
|
||
],
|
||
timeout=12,
|
||
require_sudo=False,
|
||
)
|
||
for raw_line in (result.stdout or "").splitlines():
|
||
parts = str(raw_line or "").strip().split()
|
||
if parts:
|
||
units.append(str(parts[0] or "").strip())
|
||
except Exception:
|
||
pass
|
||
|
||
managed_prefix = f"{normalized_service_name}-"
|
||
try:
|
||
for candidate in Path("/etc/default").iterdir():
|
||
if not candidate.is_file():
|
||
continue
|
||
if not candidate.name.startswith(managed_prefix):
|
||
continue
|
||
suffix = str(candidate.name[len(managed_prefix):] or "").strip()
|
||
if suffix:
|
||
units.append(f"{normalized_service_name}@{suffix}")
|
||
except Exception:
|
||
pass
|
||
return _dedupe_units(units)
|
||
|
||
|
||
def _expand_linux_worker_control_units(service_name: str) -> list[str]:
|
||
normalized_service_name = str(service_name or "").strip()
|
||
if not normalized_service_name:
|
||
return []
|
||
return _dedupe_units([normalized_service_name, *_list_linux_worker_instance_units(normalized_service_name)])
|
||
|
||
|
||
def _run_systemctl(
|
||
command: list[str],
|
||
timeout: int = 20,
|
||
require_sudo: bool = True,
|
||
no_block: bool = False,
|
||
) -> subprocess.CompletedProcess[str]:
|
||
systemctl_command = ["systemctl"]
|
||
if no_block and "--no-block" not in command:
|
||
systemctl_command.append("--no-block")
|
||
systemctl_command.extend(command)
|
||
if os.geteuid() == 0 or not require_sudo:
|
||
return _run_shell(systemctl_command, timeout=timeout)
|
||
return _run_shell(["sudo", "-n", *systemctl_command], timeout=timeout)
|
||
|
||
|
||
def normalize_systemctl_error(raw_message: str, *, service_name: str = "") -> str:
|
||
message = str(raw_message or "").strip()
|
||
normalized_service_name = str(service_name or "").strip()
|
||
if not message:
|
||
target = normalized_service_name or "目标服务"
|
||
return f"{target} 控制失败，未返回可用错误信息"
|
||
|
||
lowered = message.lower()
|
||
if (
|
||
"sudo: a password is required" in lowered
|
||
or "authentication is required" in lowered
|
||
or "interactive authentication required" in lowered
|
||
or "authorization not available" in lowered
|
||
):
|
||
target = normalized_service_name or "systemd 服务"
|
||
return f"{target} 控制失败：当前运行用户没有免密 systemctl 权限，请为 API 进程授予对应 sudo/systemd 权限"
|
||
if "unit " in lowered and " could not be found" in lowered:
|
||
target = normalized_service_name or "目标服务"
|
||
return f"{target} 控制失败：systemd 中未找到该服务，请检查服务名配置是否正确"
|
||
if "access denied" in lowered or "permission denied" in lowered:
|
||
target = normalized_service_name or "systemd 服务"
|
||
return f"{target} 控制失败：权限不足，请检查当前 API 进程用户是否具备 systemctl 控制权限"
|
||
if "host is down" in lowered or "failed to connect to bus" in lowered:
|
||
return "systemctl 调用失败：当前环境未接入可用的 systemd/dbus，请确认部署方式与 worker_mode 是否匹配"
|
||
return message
|
||
|
||
|
||
def _parse_systemd_timestamp(raw_timestamp: str) -> str:
|
||
raw_timestamp = (raw_timestamp or "").strip()
|
||
if not raw_timestamp:
|
||
return ""
|
||
try:
|
||
parsed = datetime.strptime(raw_timestamp, "%a %Y-%m-%d %H:%M:%S %Z")
|
||
return parsed.strftime("%Y-%m-%d %H:%M:%S")
|
||
except ValueError:
|
||
pass
|
||
try:
|
||
parsed = datetime.strptime(raw_timestamp.rsplit(" ", 1)[0], "%a %Y-%m-%d %H:%M:%S")
|
||
return parsed.strftime("%Y-%m-%d %H:%M:%S")
|
||
except ValueError:
|
||
return ""
|
||
|
||
|
||
def probe_systemd_service(service_name: str, *, mode: str = "linux-systemd") -> dict:
|
||
result = _run_systemctl(
|
||
[
|
||
"show",
|
||
service_name,
|
||
"--no-page",
|
||
"--property=ActiveState,SubState,MainPID,ExecMainStartTimestamp,ActiveEnterTimestamp",
|
||
],
|
||
require_sudo=False,
|
||
)
|
||
output = (result.stdout or result.stderr or "").strip()
|
||
if result.returncode != 0:
|
||
return {
|
||
"mode": mode,
|
||
"service_name": service_name,
|
||
"running": False,
|
||
"process_count": 0,
|
||
"latest_start_time": "",
|
||
"message": output or f"systemd service {service_name} not available",
|
||
}
|
||
|
||
data: dict[str, str] = {}
|
||
for line in output.splitlines():
|
||
if "=" in line:
|
||
key, value = line.split("=", 1)
|
||
data[key] = value
|
||
main_pid = int(data.get("MainPID", "0") or 0)
|
||
active_state = data.get("ActiveState", "")
|
||
sub_state = data.get("SubState", "")
|
||
latest_start_time = ""
|
||
for raw_timestamp in (data.get("ExecMainStartTimestamp", ""), data.get("ActiveEnterTimestamp", "")):
|
||
latest_start_time = _parse_systemd_timestamp(raw_timestamp)
|
||
if latest_start_time:
|
||
break
|
||
return {
|
||
"mode": mode,
|
||
"service_name": service_name,
|
||
"running": active_state == "active",
|
||
"process_count": 1 if main_pid > 0 else 0,
|
||
"latest_start_time": latest_start_time,
|
||
"message": f"{active_state}/{sub_state}" if active_state else "",
|
||
}
|
||
|
||
|
||
def _windows_runtime() -> dict:
|
||
command = """
|
||
$targets = Get-CimInstance Win32_Process -Filter "name='python.exe'" |
|
||
Where-Object { $_.CommandLine -like '*detect_worker.py*' } |
|
||
Select-Object ProcessId, CommandLine
|
||
if (-not $targets) {
|
||
Write-Output '{"running":false,"process_count":0,"latest_start_time":"","mode":"windows-local"}'
|
||
exit 0
|
||
}
|
||
$latest = $null
|
||
foreach ($item in $targets) {
|
||
try {
|
||
$proc = Get-Process -Id $item.ProcessId -ErrorAction Stop
|
||
if (-not $latest -or $proc.StartTime -gt $latest.StartTime) {
|
||
$latest = $proc
|
||
}
|
||
} catch {}
|
||
}
|
||
$payload = @{
|
||
running = $true
|
||
process_count = @($targets).Count
|
||
latest_start_time = if ($latest) { $latest.StartTime.ToString('yyyy-MM-dd HH:mm:ss') } else { '' }
|
||
mode = 'windows-local'
|
||
} | ConvertTo-Json -Compress
|
||
Write-Output $payload
|
||
"""
|
||
result = _run_powershell(command)
|
||
output = (result.stdout or "").strip()
|
||
if result.returncode != 0 or not output:
|
||
return {
|
||
"mode": "windows-local",
|
||
"running": False,
|
||
"process_count": 0,
|
||
"latest_start_time": "",
|
||
"message": (result.stderr or result.stdout or "worker runtime probe failed").strip(),
|
||
}
|
||
try:
|
||
payload = json.loads(output)
|
||
except json.JSONDecodeError:
|
||
return {
|
||
"mode": "windows-local",
|
||
"running": False,
|
||
"process_count": 0,
|
||
"latest_start_time": "",
|
||
"message": output,
|
||
}
|
||
payload.setdefault("message", "")
|
||
return payload
|
||
|
||
|
||
def _linux_runtime() -> dict:
|
||
runtime = _runtime_config()
|
||
service_name = runtime["worker_service_name"]
|
||
service_probe = probe_systemd_service(service_name, mode="linux-systemd")
|
||
instance_count = _probe_linux_worker_instance_count(service_name)
|
||
process_count = _probe_linux_worker_process_count()
|
||
|
||
if process_count <= 0 and instance_count <= 0:
|
||
return service_probe
|
||
|
||
latest_start_time = str(service_probe.get("latest_start_time") or "").strip()
|
||
if service_probe.get("running", False) or instance_count > 0:
|
||
message = service_probe.get("message") or f"multi-instance active ({process_count})"
|
||
if instance_count > 0 and not service_probe.get("running", False):
|
||
message = f"template instances active ({instance_count})"
|
||
return {
|
||
**service_probe,
|
||
"running": True,
|
||
"process_count": process_count,
|
||
"latest_start_time": latest_start_time,
|
||
"message": message,
|
||
}
|
||
message = str(service_probe.get("message") or "").strip()
|
||
if message:
|
||
message = f"{message}; detected {process_count} unmanaged worker processes"
|
||
else:
|
||
message = f"detected {process_count} unmanaged worker processes"
|
||
return {
|
||
**service_probe,
|
||
"process_count": process_count,
|
||
"latest_start_time": latest_start_time,
|
||
"message": message,
|
||
}
|
||
|
||
|
||
def detect_sync_agent_runtime() -> dict:
|
||
runtime = _runtime_config()
|
||
worker_mode = runtime.get("worker_mode", settings.worker_mode)
|
||
service_name = runtime.get("sync_agent_service_name", settings.sync_agent_service_name)
|
||
if worker_mode != "linux-systemd":
|
||
return {
|
||
"mode": worker_mode,
|
||
"service_name": service_name,
|
||
"running": False,
|
||
"process_count": 0,
|
||
"latest_start_time": "",
|
||
"message": "sync-agent 仅在 Linux systemd 多机部署中使用",
|
||
}
|
||
return probe_systemd_service(service_name, mode="linux-systemd")
|
||
|
||
|
||
def detect_worker_runtime() -> dict:
|
||
runtime = _runtime_config()
|
||
worker_mode = runtime["worker_mode"]
|
||
if worker_mode == "linux-systemd":
|
||
return _linux_runtime()
|
||
if os.name == "nt":
|
||
return _windows_runtime()
|
||
return {
|
||
"mode": worker_mode,
|
||
"running": False,
|
||
"process_count": 0,
|
||
"latest_start_time": "",
|
||
"message": f"unsupported worker_mode: {worker_mode}",
|
||
}
|
||
|
||
|
||
def start_worker() -> tuple[bool, str]:
|
||
runtime = _runtime_config()
|
||
worker_mode = runtime["worker_mode"]
|
||
service_name = runtime["worker_service_name"]
|
||
if worker_mode == "linux-systemd":
|
||
units = _expand_linux_worker_control_units(service_name)
|
||
result = _run_systemctl(["start", *units], timeout=max(30, 15 * max(1, len(units))))
|
||
if result.returncode != 0:
|
||
return False, normalize_systemctl_error(result.stderr or result.stdout or "启动 Linux Worker 失败", service_name=service_name)
|
||
extra_units = max(0, len(units) - 1)
|
||
if extra_units > 0:
|
||
return True, f"Linux Worker 启动命令已发送: {service_name}，附带 {extra_units} 个实例"
|
||
return True, f"Linux Worker 启动命令已发送: {service_name}"
|
||
|
||
if os.name != "nt":
|
||
return False, "当前仅实现 Windows 本地 Worker 启动，Linux 请将 worker_mode 设为 linux-systemd。"
|
||
|
||
script_path = _domain_root() / "start_worker.ps1"
|
||
if not script_path.exists():
|
||
return False, f"未找到启动脚本: {script_path}"
|
||
|
||
command = (
|
||
"Start-Process powershell "
|
||
f"-ArgumentList '-ExecutionPolicy Bypass -File \"{script_path}\"' "
|
||
f"-WorkingDirectory '{_domain_root()}'"
|
||
)
|
||
result = _run_powershell(command)
|
||
if result.returncode != 0:
|
||
return False, (result.stderr or result.stdout or "启动检测端失败").strip()
|
||
return True, "检测端启动命令已发送"
|
||
|
||
|
||
def stop_worker() -> tuple[bool, str]:
|
||
runtime = _runtime_config()
|
||
worker_mode = runtime["worker_mode"]
|
||
service_name = runtime["worker_service_name"]
|
||
if worker_mode == "linux-systemd":
|
||
units = _expand_linux_worker_control_units(service_name)
|
||
result = _run_systemctl(["stop", *units], timeout=max(30, 15 * max(1, len(units))))
|
||
if result.returncode != 0:
|
||
return False, normalize_systemctl_error(result.stderr or result.stdout or "停止 Linux Worker 失败", service_name=service_name)
|
||
extra_units = max(0, len(units) - 1)
|
||
if extra_units > 0:
|
||
return True, f"Linux Worker 停止命令已发送: {service_name}，附带 {extra_units} 个实例"
|
||
return True, f"Linux Worker 停止命令已发送: {service_name}"
|
||
|
||
if os.name != "nt":
|
||
return False, "当前仅实现 Windows 本地 Worker 停止，Linux 请将 worker_mode 设为 linux-systemd。"
|
||
|
||
command = """
|
||
$targets = Get-CimInstance Win32_Process -Filter "name='python.exe'" |
|
||
Where-Object { $_.CommandLine -like '*detect_worker.py*' } |
|
||
Select-Object -ExpandProperty ProcessId
|
||
if (-not $targets) {
|
||
Write-Output 'NO_PROCESS'
|
||
exit 0
|
||
}
|
||
$targets | ForEach-Object { Stop-Process -Id $_ -Force }
|
||
Write-Output ('STOPPED:' + (($targets | Measure-Object).Count))
|
||
"""
|
||
result = _run_powershell(command)
|
||
output = (result.stdout or result.stderr or "").strip()
|
||
if result.returncode != 0:
|
||
return False, output or "停止检测端失败"
|
||
if "NO_PROCESS" in output:
|
||
return True, "当前没有运行中的检测端进程"
|
||
return True, output or "检测端已停止"
|
||
|
||
|
||
def send_worker_command(action: str, payload: dict | None = None) -> tuple[bool, str]:
|
||
command_payload = {"action": action}
|
||
if payload:
|
||
command_payload.update(payload)
|
||
runtime = _runtime_config()
|
||
explicit_targets = _normalize_target_node_codes(command_payload)
|
||
if not explicit_targets and str(runtime.get("worker_mode") or "").strip() == "linux-systemd":
|
||
expanded_targets = _expand_local_linux_worker_target_node_codes(
|
||
str(runtime.get("worker_service_name") or "").strip()
|
||
)
|
||
if expanded_targets:
|
||
command_payload["target_node_codes"] = expanded_targets
|
||
command_payload["request_id"] = str(command_payload.get("request_id") or f"workerctl-{uuid4().hex[:12]}")
|
||
serialized = json.dumps(command_payload, ensure_ascii=False)
|
||
pending_keys = _pending_command_keys(command_payload)
|
||
target_node_codes = _normalize_target_node_codes(command_payload)
|
||
|
||
direct_client = None
|
||
try:
|
||
try:
|
||
_publish_worker_command(get_redis(), serialized=serialized, pending_keys=pending_keys)
|
||
except Exception:
|
||
direct_client = _build_direct_redis_client()
|
||
_publish_worker_command(direct_client, serialized=serialized, pending_keys=pending_keys)
|
||
|
||
if pending_keys == [WORKER_PENDING_COMMAND_KEY]:
|
||
return True, f"已发送 Worker 控制指令: {action}"
|
||
if len(target_node_codes) <= 4:
|
||
target_summary = ",".join(target_node_codes)
|
||
else:
|
||
target_summary = f"{len(target_node_codes)} targets"
|
||
return True, f"已发送 Worker 控制指令: {action} -> {target_summary}"
|
||
except Exception as exc:
|
||
return False, f"发送 Worker 控制指令失败: {exc}"
|
||
finally:
|
||
if direct_client is not None:
|
||
try:
|
||
direct_client.close()
|
||
except Exception:
|
||
pass
|