249 lines
8.6 KiB
Python
249 lines
8.6 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
from datetime import datetime, timedelta, timezone
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
def utc_now() -> datetime:
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
def to_iso(value: datetime) -> str:
|
|
return value.astimezone(timezone.utc).replace(microsecond=0).isoformat()
|
|
|
|
|
|
def parse_iso(value: str | None) -> datetime | None:
|
|
raw = (value or "").strip()
|
|
if not raw:
|
|
return None
|
|
try:
|
|
return datetime.fromisoformat(raw.replace("Z", "+00:00")).astimezone(timezone.utc)
|
|
except ValueError:
|
|
return None
|
|
|
|
|
|
def load_state(path: Path) -> dict[str, Any]:
|
|
if not path.exists():
|
|
return {"version": 1, "updated_at": to_iso(utc_now()), "allocations": []}
|
|
try:
|
|
payload = json.loads(path.read_text(encoding="utf-8"))
|
|
except json.JSONDecodeError:
|
|
return {"version": 1, "updated_at": to_iso(utc_now()), "allocations": []}
|
|
if not isinstance(payload, dict):
|
|
return {"version": 1, "updated_at": to_iso(utc_now()), "allocations": []}
|
|
payload.setdefault("version", 1)
|
|
payload.setdefault("updated_at", to_iso(utc_now()))
|
|
allocations = payload.get("allocations")
|
|
if not isinstance(allocations, list):
|
|
payload["allocations"] = []
|
|
return payload
|
|
|
|
|
|
def save_state(path: Path, payload: dict[str, Any]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
payload["updated_at"] = to_iso(utc_now())
|
|
path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
|
|
|
|
def parse_slots(raw: str) -> list[str]:
|
|
slots = [item.strip() for item in raw.split(",") if item.strip()]
|
|
if not slots:
|
|
raise ValueError("at least one slot is required")
|
|
return slots
|
|
|
|
|
|
def render_url(slot: str, template: str | None) -> str | None:
|
|
if not template:
|
|
return None
|
|
return template.replace("{slot}", slot)
|
|
|
|
|
|
def prune_expired(allocations: list[dict[str, Any]], now: datetime) -> tuple[list[dict[str, Any]], int]:
|
|
kept: list[dict[str, Any]] = []
|
|
removed = 0
|
|
for item in allocations:
|
|
expires = parse_iso(str(item.get("expires_at") or ""))
|
|
if expires and expires < now:
|
|
removed += 1
|
|
continue
|
|
kept.append(item)
|
|
return kept, removed
|
|
|
|
|
|
def allocation_sort_key(item: dict[str, Any]) -> datetime:
|
|
last_seen = parse_iso(str(item.get("last_seen_at") or ""))
|
|
allocated = parse_iso(str(item.get("allocated_at") or ""))
|
|
return last_seen or allocated or datetime.fromtimestamp(0, tz=timezone.utc)
|
|
|
|
|
|
def output(payload: dict[str, Any]) -> None:
|
|
print(json.dumps(payload, ensure_ascii=False, indent=2))
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(
|
|
description="Allocate/reuse/release branch preview slots for issue-driven workflows."
|
|
)
|
|
parser.add_argument("--state-file", default=".tmp/preview-slots.json")
|
|
parser.add_argument("--slots", required=True, help="Comma-separated slot names, e.g. preview-a,preview-b")
|
|
parser.add_argument("--repo", default="", help="owner/repo")
|
|
parser.add_argument("--issue", type=int, default=None, help="Issue number")
|
|
parser.add_argument("--branch", default="", help="Branch name")
|
|
parser.add_argument("--slot", default="", help="Slot name (optional release filter)")
|
|
parser.add_argument("--ttl-hours", type=int, default=24, help="Allocation TTL in hours")
|
|
parser.add_argument("--url-template", default="", help="Optional URL template, use {slot} placeholder")
|
|
parser.add_argument("--release", action="store_true", help="Release matching allocation(s)")
|
|
parser.add_argument("--list", action="store_true", help="List active allocations and free slots")
|
|
parser.add_argument("--evict-oldest", action="store_true", help="Evict oldest allocation when slots are full")
|
|
args = parser.parse_args()
|
|
|
|
state_path = Path(args.state_file)
|
|
slots = parse_slots(args.slots)
|
|
state = load_state(state_path)
|
|
|
|
now = utc_now()
|
|
allocations_raw = [item for item in state.get("allocations", []) if isinstance(item, dict)]
|
|
allocations, pruned_count = prune_expired(allocations_raw, now)
|
|
state["allocations"] = allocations
|
|
|
|
if args.list:
|
|
used_slots = {str(item.get("slot") or "") for item in allocations}
|
|
free_slots = [slot for slot in slots if slot not in used_slots]
|
|
output(
|
|
{
|
|
"action": "list",
|
|
"state_file": str(state_path.as_posix()),
|
|
"pruned_expired": pruned_count,
|
|
"total_active": len(allocations),
|
|
"active_allocations": allocations,
|
|
"free_slots": free_slots,
|
|
}
|
|
)
|
|
save_state(state_path, state)
|
|
return
|
|
|
|
if args.release:
|
|
target_repo = args.repo.strip()
|
|
target_branch = args.branch.strip()
|
|
target_slot = args.slot.strip()
|
|
target_issue = args.issue
|
|
|
|
kept: list[dict[str, Any]] = []
|
|
released: list[dict[str, Any]] = []
|
|
for item in allocations:
|
|
match = True
|
|
if target_repo:
|
|
match = match and str(item.get("repo") or "") == target_repo
|
|
if target_branch:
|
|
match = match and str(item.get("branch") or "") == target_branch
|
|
if target_slot:
|
|
match = match and str(item.get("slot") or "") == target_slot
|
|
if target_issue is not None:
|
|
match = match and int(item.get("issue") or -1) == target_issue
|
|
|
|
if match:
|
|
released.append(item)
|
|
else:
|
|
kept.append(item)
|
|
|
|
state["allocations"] = kept
|
|
save_state(state_path, state)
|
|
output(
|
|
{
|
|
"action": "release",
|
|
"state_file": str(state_path.as_posix()),
|
|
"released_count": len(released),
|
|
"released": released,
|
|
"remaining_count": len(kept),
|
|
}
|
|
)
|
|
return
|
|
|
|
repo = args.repo.strip()
|
|
branch = args.branch.strip()
|
|
if not repo or not branch:
|
|
print("allocate requires --repo and --branch", file=sys.stderr)
|
|
sys.exit(2)
|
|
|
|
# Reuse existing allocation for same repo+branch.
|
|
existing = next(
|
|
(item for item in allocations if str(item.get("repo") or "") == repo and str(item.get("branch") or "") == branch),
|
|
None,
|
|
)
|
|
if existing is not None:
|
|
existing["issue"] = args.issue if args.issue is not None else existing.get("issue")
|
|
existing["last_seen_at"] = to_iso(now)
|
|
existing["expires_at"] = to_iso(now + timedelta(hours=max(1, args.ttl_hours)))
|
|
if args.url_template:
|
|
existing["url"] = render_url(str(existing.get("slot") or ""), args.url_template)
|
|
save_state(state_path, state)
|
|
output(
|
|
{
|
|
"action": "allocate",
|
|
"reused": True,
|
|
"state_file": str(state_path.as_posix()),
|
|
"allocation": existing,
|
|
"pruned_expired": pruned_count,
|
|
}
|
|
)
|
|
return
|
|
|
|
used_slots = {str(item.get("slot") or "") for item in allocations}
|
|
free_slots = [slot for slot in slots if slot not in used_slots]
|
|
evicted: dict[str, Any] | None = None
|
|
|
|
if not free_slots:
|
|
if not args.evict_oldest:
|
|
output(
|
|
{
|
|
"action": "allocate",
|
|
"reused": False,
|
|
"allocated": False,
|
|
"reason": "no_free_slots",
|
|
"state_file": str(state_path.as_posix()),
|
|
"active_allocations": allocations,
|
|
}
|
|
)
|
|
sys.exit(3)
|
|
oldest = sorted(allocations, key=allocation_sort_key)[0]
|
|
evicted = oldest
|
|
allocations.remove(oldest)
|
|
free_slots = [str(oldest.get("slot") or "")]
|
|
|
|
slot = free_slots[0]
|
|
allocation = {
|
|
"repo": repo,
|
|
"issue": args.issue,
|
|
"branch": branch,
|
|
"slot": slot,
|
|
"env_id": slot,
|
|
"url": render_url(slot, args.url_template),
|
|
"allocated_at": to_iso(now),
|
|
"last_seen_at": to_iso(now),
|
|
"expires_at": to_iso(now + timedelta(hours=max(1, args.ttl_hours))),
|
|
"status": "active",
|
|
}
|
|
allocations.append(allocation)
|
|
state["allocations"] = allocations
|
|
save_state(state_path, state)
|
|
output(
|
|
{
|
|
"action": "allocate",
|
|
"reused": False,
|
|
"allocated": True,
|
|
"state_file": str(state_path.as_posix()),
|
|
"allocation": allocation,
|
|
"evicted": evicted,
|
|
"pruned_expired": pruned_count,
|
|
}
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|