325 lines
10 KiB
Python
325 lines
10 KiB
Python
"""autoscalerctl CLI entry point."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import http.client
|
|
import json
|
|
import socket
|
|
from collections.abc import Sequence
|
|
from typing import Any
|
|
|
|
|
|
class UnixHTTPConnection(http.client.HTTPConnection):
|
|
"""HTTPConnection that dials a Unix domain socket."""
|
|
|
|
def __init__(self, socket_path: str, timeout: float = 5.0) -> None:
|
|
super().__init__("localhost", timeout=timeout)
|
|
self._socket_path = socket_path
|
|
|
|
def connect(self) -> None:
|
|
self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
|
self.sock.connect(self._socket_path)
|
|
|
|
|
|
def _uds_request(
|
|
socket_path: str,
|
|
method: str,
|
|
path: str,
|
|
body: dict[str, Any] | None = None,
|
|
) -> tuple[int, dict[str, Any] | list[dict[str, Any]] | str]:
|
|
conn = UnixHTTPConnection(socket_path)
|
|
headers = {"Host": "localhost", "Accept": "application/json"}
|
|
payload: str | None = None
|
|
if body is not None:
|
|
payload = json.dumps(body)
|
|
headers["Content-Type"] = "application/json"
|
|
|
|
try:
|
|
conn.request(method, path, body=payload, headers=headers)
|
|
resp = conn.getresponse()
|
|
raw = resp.read()
|
|
text = raw.decode() if raw else ""
|
|
content_type = resp.getheader("Content-Type", "")
|
|
if text and "application/json" in content_type:
|
|
parsed = json.loads(text)
|
|
if isinstance(parsed, dict | list):
|
|
return resp.status, parsed
|
|
return resp.status, text
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def _print_table(headers: Sequence[str], rows: Sequence[Sequence[str]]) -> None:
|
|
widths = [len(h) for h in headers]
|
|
for row in rows:
|
|
for idx, cell in enumerate(row):
|
|
widths[idx] = max(widths[idx], len(cell))
|
|
|
|
header_line = " ".join(h.ljust(widths[idx]) for idx, h in enumerate(headers))
|
|
separator = " ".join("-" * widths[idx] for idx in range(len(headers)))
|
|
print(header_line)
|
|
print(separator)
|
|
for row in rows:
|
|
print(" ".join(cell.ljust(widths[idx]) for idx, cell in enumerate(row)))
|
|
|
|
|
|
def _print_slots(data: list[dict[str, Any]]) -> None:
|
|
rows: list[list[str]] = []
|
|
for slot in data:
|
|
rows.append(
|
|
[
|
|
str(slot.get("slot_id", "")),
|
|
str(slot.get("state", "")),
|
|
str(slot.get("instance_id") or "-"),
|
|
str(slot.get("instance_ip") or "-"),
|
|
str(slot.get("lease_count", 0)),
|
|
]
|
|
)
|
|
_print_table(["slot_id", "state", "instance_id", "ip", "leases"], rows)
|
|
|
|
|
|
def _print_reservations(data: list[dict[str, Any]]) -> None:
|
|
rows: list[list[str]] = []
|
|
for resv in data:
|
|
rows.append(
|
|
[
|
|
str(resv.get("reservation_id", "")),
|
|
str(resv.get("phase", "")),
|
|
str(resv.get("system", "")),
|
|
str(resv.get("slot") or "-"),
|
|
str(resv.get("instance_id") or "-"),
|
|
]
|
|
)
|
|
_print_table(["reservation_id", "phase", "system", "slot", "instance_id"], rows)
|
|
|
|
|
|
def _print_status_summary(data: dict[str, Any]) -> None:
|
|
slots = data.get("slots", {})
|
|
reservations = data.get("reservations", {})
|
|
ec2 = data.get("ec2", {})
|
|
haproxy = data.get("haproxy", {})
|
|
rows = [
|
|
["slots.total", str(slots.get("total", 0))],
|
|
["slots.ready", str(slots.get("ready", 0))],
|
|
["slots.launching", str(slots.get("launching", 0))],
|
|
["slots.booting", str(slots.get("booting", 0))],
|
|
["slots.binding", str(slots.get("binding", 0))],
|
|
["slots.terminating", str(slots.get("terminating", 0))],
|
|
["slots.empty", str(slots.get("empty", 0))],
|
|
["slots.error", str(slots.get("error", 0))],
|
|
["reservations.pending", str(reservations.get("pending", 0))],
|
|
["reservations.ready", str(reservations.get("ready", 0))],
|
|
["reservations.failed", str(reservations.get("failed", 0))],
|
|
["ec2.api_ok", str(ec2.get("api_ok", False))],
|
|
["haproxy.socket_ok", str(haproxy.get("socket_ok", False))],
|
|
]
|
|
_print_table(["metric", "value"], rows)
|
|
|
|
|
|
def _bulk_slot_action(socket_path: str, action: str) -> dict[str, Any]:
|
|
if action == "drain":
|
|
eligible_states = {"ready"}
|
|
action_path = "/v1/admin/drain"
|
|
elif action == "unquarantine":
|
|
eligible_states = {"error"}
|
|
action_path = "/v1/admin/unquarantine"
|
|
else:
|
|
msg = f"unknown bulk action: {action}"
|
|
raise ValueError(msg)
|
|
|
|
status, data = _uds_request(socket_path, "GET", "/v1/slots")
|
|
if status < 200 or status >= 300 or not isinstance(data, list):
|
|
msg = "failed to list slots for bulk action"
|
|
raise RuntimeError(msg)
|
|
|
|
results: list[dict[str, Any]] = []
|
|
summary: dict[str, Any] = {
|
|
"action": action,
|
|
"matched": 0,
|
|
"attempted": 0,
|
|
"succeeded": 0,
|
|
"failed": 0,
|
|
"skipped": 0,
|
|
"results": results,
|
|
}
|
|
|
|
for slot in data:
|
|
slot_id = str(slot.get("slot_id", ""))
|
|
state = str(slot.get("state", ""))
|
|
if not slot_id:
|
|
continue
|
|
|
|
if state not in eligible_states:
|
|
summary["skipped"] += 1
|
|
results.append(
|
|
{
|
|
"slot_id": slot_id,
|
|
"state": state,
|
|
"result": "skipped",
|
|
"reason": "ineligible_state",
|
|
}
|
|
)
|
|
continue
|
|
|
|
summary["matched"] += 1
|
|
summary["attempted"] += 1
|
|
try:
|
|
action_status, action_data = _uds_request(
|
|
socket_path,
|
|
"POST",
|
|
action_path,
|
|
body={"slot_id": slot_id},
|
|
)
|
|
except OSError as err:
|
|
summary["failed"] += 1
|
|
results.append(
|
|
{
|
|
"slot_id": slot_id,
|
|
"state": state,
|
|
"result": "failed",
|
|
"error": str(err),
|
|
}
|
|
)
|
|
continue
|
|
|
|
if 200 <= action_status < 300:
|
|
summary["succeeded"] += 1
|
|
results.append(
|
|
{
|
|
"slot_id": slot_id,
|
|
"state": state,
|
|
"result": "ok",
|
|
}
|
|
)
|
|
else:
|
|
summary["failed"] += 1
|
|
results.append(
|
|
{
|
|
"slot_id": slot_id,
|
|
"state": state,
|
|
"result": "failed",
|
|
"status": action_status,
|
|
"response": action_data,
|
|
}
|
|
)
|
|
|
|
return summary
|
|
|
|
|
|
def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(prog="autoscalerctl", description="Autoscaler CLI")
|
|
parser.add_argument(
|
|
"--socket",
|
|
default="/run/nix-builder-autoscaler/daemon.sock",
|
|
help="Daemon Unix socket path",
|
|
)
|
|
parser.add_argument(
|
|
"--json",
|
|
action="store_true",
|
|
help="Output JSON for status command.",
|
|
)
|
|
subparsers = parser.add_subparsers(dest="command")
|
|
subparsers.add_parser("status", help="Show state summary")
|
|
subparsers.add_parser("slots", help="List slots")
|
|
subparsers.add_parser("reservations", help="List reservations")
|
|
|
|
parser_drain = subparsers.add_parser("drain", help="Drain a slot")
|
|
parser_drain.add_argument("slot_id")
|
|
parser_unq = subparsers.add_parser("unquarantine", help="Unquarantine a slot")
|
|
parser_unq.add_argument("slot_id")
|
|
subparsers.add_parser("drain-all", help="Drain all eligible slots (state=ready)")
|
|
subparsers.add_parser("unquarantine-all", help="Unquarantine all error slots")
|
|
subparsers.add_parser("reconcile-now", help="Trigger immediate reconcile tick")
|
|
args = parser.parse_args(argv)
|
|
if not args.command:
|
|
parser.print_help()
|
|
raise SystemExit(0)
|
|
return args
|
|
|
|
|
|
def _print_error(data: object) -> None:
|
|
if isinstance(data, dict | list):
|
|
print(json.dumps(data, indent=2))
|
|
else:
|
|
print(str(data))
|
|
|
|
|
|
def main() -> None:
|
|
"""Entry point for the autoscalerctl CLI."""
|
|
args = _parse_args()
|
|
|
|
if args.command in {"drain-all", "unquarantine-all"}:
|
|
action = "drain" if args.command == "drain-all" else "unquarantine"
|
|
try:
|
|
summary = _bulk_slot_action(args.socket, action)
|
|
except OSError as err:
|
|
print(f"Error: cannot connect to daemon at {args.socket}")
|
|
raise SystemExit(1) from err
|
|
except RuntimeError as err:
|
|
print(str(err))
|
|
raise SystemExit(1) from err
|
|
print(json.dumps(summary, indent=2))
|
|
raise SystemExit(0 if summary["failed"] == 0 else 1)
|
|
|
|
method = "GET"
|
|
path = ""
|
|
body: dict[str, Any] | None = None
|
|
if args.command == "status":
|
|
path = "/v1/state/summary"
|
|
elif args.command == "slots":
|
|
path = "/v1/slots"
|
|
elif args.command == "reservations":
|
|
path = "/v1/reservations"
|
|
elif args.command == "drain":
|
|
method = "POST"
|
|
path = "/v1/admin/drain"
|
|
body = {"slot_id": args.slot_id}
|
|
elif args.command == "unquarantine":
|
|
method = "POST"
|
|
path = "/v1/admin/unquarantine"
|
|
body = {"slot_id": args.slot_id}
|
|
elif args.command == "reconcile-now":
|
|
method = "POST"
|
|
path = "/v1/admin/reconcile-now"
|
|
else:
|
|
raise SystemExit(1)
|
|
|
|
try:
|
|
status, data = _uds_request(args.socket, method, path, body=body)
|
|
except OSError as err:
|
|
print(f"Error: cannot connect to daemon at {args.socket}")
|
|
raise SystemExit(1) from err
|
|
|
|
if status < 200 or status >= 300:
|
|
_print_error(data)
|
|
raise SystemExit(1)
|
|
|
|
if args.command == "status":
|
|
if not isinstance(data, dict):
|
|
_print_error(data)
|
|
raise SystemExit(1)
|
|
if args.json:
|
|
print(json.dumps(data, indent=2))
|
|
else:
|
|
_print_status_summary(data)
|
|
elif args.command in {"drain", "unquarantine", "reconcile-now"}:
|
|
print(json.dumps(data, indent=2))
|
|
elif args.command == "slots":
|
|
if isinstance(data, list):
|
|
_print_slots(data)
|
|
else:
|
|
_print_error(data)
|
|
raise SystemExit(1)
|
|
elif args.command == "reservations":
|
|
if isinstance(data, list):
|
|
_print_reservations(data)
|
|
else:
|
|
_print_error(data)
|
|
raise SystemExit(1)
|
|
|
|
raise SystemExit(0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|