398 lines
13 KiB
Python
398 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""Exercise GovOPlaN worker delivery guarantees against a real Redis broker."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from datetime import UTC, datetime
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from typing import Any
|
|
from urllib.parse import urlsplit, urlunsplit
|
|
import uuid
|
|
|
|
|
|
TERMINAL_STATES = {"FAILURE", "REVOKED", "SUCCESS"}
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(
|
|
description=(
|
|
"Prove Celery publish/consume, retry, warm shutdown, and worker-loss "
|
|
"redelivery against an isolated Redis database."
|
|
)
|
|
)
|
|
parser.add_argument(
|
|
"--redis-url",
|
|
default=os.environ.get("GOVOPLAN_WORKER_DRILL_REDIS_URL", ""),
|
|
)
|
|
parser.add_argument("--python", default=sys.executable)
|
|
parser.add_argument("--timeout-seconds", type=float, default=120.0)
|
|
parser.add_argument(
|
|
"--visibility-timeout-seconds",
|
|
type=int,
|
|
default=30,
|
|
)
|
|
parser.add_argument("--output", type=Path)
|
|
return parser
|
|
|
|
|
|
def _redacted_redis_url(value: str) -> str:
|
|
parsed = urlsplit(value)
|
|
hostname = parsed.hostname or ""
|
|
port = f":{parsed.port}" if parsed.port else ""
|
|
return urlunsplit((parsed.scheme, f"{hostname}{port}", parsed.path, "", ""))
|
|
|
|
|
|
def _wait_until(predicate, *, timeout_seconds: float, detail: str):
|
|
deadline = time.monotonic() + timeout_seconds
|
|
last_value: Any = None
|
|
while time.monotonic() < deadline:
|
|
last_value = predicate()
|
|
if last_value:
|
|
return last_value
|
|
time.sleep(0.2)
|
|
raise TimeoutError(f"Timed out waiting for {detail}; last value: {last_value!r}")
|
|
|
|
|
|
def _worker_command(python: str, hostname: str) -> list[str]:
|
|
return [
|
|
python,
|
|
"-m",
|
|
"celery",
|
|
"-A",
|
|
"govoplan_core.celery_app:celery",
|
|
"worker",
|
|
"--pool",
|
|
"solo",
|
|
"--queues",
|
|
"default",
|
|
"--hostname",
|
|
hostname,
|
|
"--loglevel",
|
|
"WARNING",
|
|
"--without-mingle",
|
|
]
|
|
|
|
|
|
def _start_worker(
|
|
*,
|
|
python: str,
|
|
hostname: str,
|
|
environment: dict[str, str],
|
|
log_path: Path,
|
|
) -> tuple[subprocess.Popen[bytes], Any]:
|
|
log = log_path.open("ab", buffering=0)
|
|
process = subprocess.Popen(
|
|
_worker_command(python, hostname),
|
|
env=environment,
|
|
stdin=subprocess.DEVNULL,
|
|
stdout=log,
|
|
stderr=subprocess.STDOUT,
|
|
start_new_session=True,
|
|
)
|
|
return process, log
|
|
|
|
|
|
def _stop_worker(
|
|
process: subprocess.Popen[bytes] | None,
|
|
log: Any,
|
|
*,
|
|
timeout_seconds: float = 30.0,
|
|
) -> None:
|
|
if process is not None and process.poll() is None:
|
|
process.send_signal(signal.SIGTERM)
|
|
try:
|
|
process.wait(timeout=timeout_seconds)
|
|
except subprocess.TimeoutExpired:
|
|
process.kill()
|
|
process.wait(timeout=10)
|
|
if log is not None:
|
|
log.close()
|
|
|
|
|
|
def _wait_for_worker(app, process: subprocess.Popen[bytes], prefix: str, timeout: float) -> str:
|
|
def ping() -> str | None:
|
|
if process.poll() is not None:
|
|
raise RuntimeError(f"Worker exited before readiness with code {process.returncode}")
|
|
replies = app.control.ping(timeout=1.0) or []
|
|
for reply in replies:
|
|
for hostname, payload in reply.items():
|
|
if hostname.startswith(prefix) and payload.get("ok") == "pong":
|
|
return hostname
|
|
return None
|
|
|
|
return _wait_until(ping, timeout_seconds=timeout, detail=f"worker {prefix}")
|
|
|
|
|
|
def _wait_for_started(result, *, timeout_seconds: float) -> dict[str, Any]:
|
|
def started() -> dict[str, Any] | None:
|
|
state = result.state
|
|
if state in TERMINAL_STATES and state != "SUCCESS":
|
|
raise RuntimeError(f"Probe {result.id} became {state}: {result.result!r}")
|
|
info = result.info
|
|
if state in {"PROGRESS", "STARTED"} and isinstance(info, dict):
|
|
return dict(info)
|
|
return None
|
|
|
|
return _wait_until(
|
|
started,
|
|
timeout_seconds=timeout_seconds,
|
|
detail=f"probe {result.id} to start",
|
|
)
|
|
|
|
|
|
def _wait_for_result(result, *, timeout_seconds: float) -> dict[str, Any]:
|
|
value = result.get(timeout=timeout_seconds, propagate=True, disable_sync_subtasks=False)
|
|
if not isinstance(value, dict):
|
|
raise RuntimeError(f"Probe {result.id} returned an invalid result: {value!r}")
|
|
return dict(value)
|
|
|
|
|
|
def _publish(probe_task, probe_id: str, *, mode: str, delay_seconds: float):
|
|
return probe_task.apply_async(
|
|
args=(probe_id,),
|
|
kwargs={
|
|
"mode": mode,
|
|
"delay_seconds": delay_seconds,
|
|
"track_delivery": True,
|
|
},
|
|
queue="default",
|
|
)
|
|
|
|
|
|
def _assert_probe(value: dict[str, Any], *, probe_id: str) -> None:
|
|
if value.get("probe_id") != probe_id:
|
|
raise RuntimeError(f"Probe identity mismatch: {value!r}")
|
|
|
|
|
|
def run_drill(args: argparse.Namespace) -> dict[str, Any]:
|
|
redis_url = str(args.redis_url or "").strip()
|
|
parsed_redis = urlsplit(redis_url)
|
|
if parsed_redis.scheme not in {"redis", "rediss"} or not parsed_redis.hostname:
|
|
raise ValueError("--redis-url must be an explicit redis:// or rediss:// URL")
|
|
if args.visibility_timeout_seconds < 30:
|
|
raise ValueError("--visibility-timeout-seconds must be at least 30")
|
|
|
|
run_id = uuid.uuid4().hex
|
|
started_at = datetime.now(UTC)
|
|
with tempfile.TemporaryDirectory(prefix="govoplan-worker-drill-") as directory:
|
|
root = Path(directory)
|
|
database_url = f"sqlite:///{root / 'runtime.db'}"
|
|
environment = dict(os.environ)
|
|
environment.update(
|
|
{
|
|
"REDIS_URL": redis_url,
|
|
"CELERY_ENABLED": "true",
|
|
"CELERY_QUEUES": "default",
|
|
"CELERY_VISIBILITY_TIMEOUT_SECONDS": str(
|
|
args.visibility_timeout_seconds
|
|
),
|
|
"DATABASE_URL": database_url,
|
|
"ENABLED_MODULES": "access",
|
|
"APP_ENV": "development",
|
|
"GOVOPLAN_EXPECTED_WORKER_REPLICAS": "1",
|
|
"GOVOPLAN_WORKER_POOL": "acceptance-drill",
|
|
}
|
|
)
|
|
os.environ.update(environment)
|
|
|
|
from govoplan_core.db.migrations import migrate_database
|
|
|
|
migrate_database(
|
|
database_url=database_url,
|
|
enabled_modules=("access",),
|
|
)
|
|
from govoplan_core.celery_app import celery, worker_acceptance_probe
|
|
|
|
celery.backend.client.ping()
|
|
celery.control.purge()
|
|
worker: subprocess.Popen[bytes] | None = None
|
|
worker_log: Any = None
|
|
evidence: dict[str, Any] = {
|
|
"schema_version": "1.0",
|
|
"run_id": run_id,
|
|
"started_at": started_at.isoformat(),
|
|
"redis": _redacted_redis_url(redis_url),
|
|
"visibility_timeout_seconds": args.visibility_timeout_seconds,
|
|
"checks": [],
|
|
}
|
|
try:
|
|
prefix = f"govoplan-drill-{run_id[:8]}-a@"
|
|
worker, worker_log = _start_worker(
|
|
python=args.python,
|
|
hostname=prefix + "%h",
|
|
environment=environment,
|
|
log_path=root / "worker-a.log",
|
|
)
|
|
hostname = _wait_for_worker(
|
|
celery,
|
|
worker,
|
|
prefix,
|
|
args.timeout_seconds,
|
|
)
|
|
evidence["checks"].append(
|
|
{"id": "worker_startup", "state": "passed", "worker": hostname}
|
|
)
|
|
|
|
probe_id = f"{run_id}-publish"
|
|
result = _publish(
|
|
worker_acceptance_probe,
|
|
probe_id,
|
|
mode="complete",
|
|
delay_seconds=0,
|
|
)
|
|
value = _wait_for_result(result, timeout_seconds=args.timeout_seconds)
|
|
_assert_probe(value, probe_id=probe_id)
|
|
if value.get("delivery_count") != 1:
|
|
raise RuntimeError(f"Publish probe was not delivered exactly once: {value!r}")
|
|
evidence["checks"].append(
|
|
{
|
|
"id": "publish_consume",
|
|
"state": "passed",
|
|
"task_id": result.id,
|
|
"delivery_count": value["delivery_count"],
|
|
}
|
|
)
|
|
|
|
probe_id = f"{run_id}-retry"
|
|
result = _publish(
|
|
worker_acceptance_probe,
|
|
probe_id,
|
|
mode="retry_once",
|
|
delay_seconds=0,
|
|
)
|
|
value = _wait_for_result(result, timeout_seconds=args.timeout_seconds)
|
|
_assert_probe(value, probe_id=probe_id)
|
|
if value.get("retries") != 1 or value.get("delivery_count") != 2:
|
|
raise RuntimeError(f"Retry probe did not execute twice: {value!r}")
|
|
evidence["checks"].append(
|
|
{
|
|
"id": "application_retry",
|
|
"state": "passed",
|
|
"task_id": result.id,
|
|
"delivery_count": value["delivery_count"],
|
|
}
|
|
)
|
|
|
|
probe_id = f"{run_id}-warm"
|
|
result = _publish(
|
|
worker_acceptance_probe,
|
|
probe_id,
|
|
mode="complete",
|
|
delay_seconds=2,
|
|
)
|
|
_wait_for_started(result, timeout_seconds=args.timeout_seconds)
|
|
worker.send_signal(signal.SIGTERM)
|
|
value = _wait_for_result(result, timeout_seconds=args.timeout_seconds)
|
|
_assert_probe(value, probe_id=probe_id)
|
|
worker.wait(timeout=args.timeout_seconds)
|
|
if worker.returncode != 0 or value.get("delivery_count") != 1:
|
|
raise RuntimeError(
|
|
f"Warm worker shutdown did not finish in-flight work: {value!r}"
|
|
)
|
|
worker_log.close()
|
|
worker = None
|
|
worker_log = None
|
|
evidence["checks"].append(
|
|
{
|
|
"id": "graceful_shutdown",
|
|
"state": "passed",
|
|
"task_id": result.id,
|
|
"delivery_count": value["delivery_count"],
|
|
}
|
|
)
|
|
|
|
prefix = f"govoplan-drill-{run_id[:8]}-b@"
|
|
worker, worker_log = _start_worker(
|
|
python=args.python,
|
|
hostname=prefix + "%h",
|
|
environment=environment,
|
|
log_path=root / "worker-b.log",
|
|
)
|
|
_wait_for_worker(celery, worker, prefix, args.timeout_seconds)
|
|
probe_id = f"{run_id}-loss"
|
|
result = _publish(
|
|
worker_acceptance_probe,
|
|
probe_id,
|
|
mode="complete",
|
|
delay_seconds=min(120.0, args.visibility_timeout_seconds + 20.0),
|
|
)
|
|
first_started = _wait_for_started(
|
|
result,
|
|
timeout_seconds=args.timeout_seconds,
|
|
)
|
|
if first_started.get("delivery_count") != 1:
|
|
raise RuntimeError(f"Worker-loss probe did not start once: {first_started!r}")
|
|
worker.kill()
|
|
worker.wait(timeout=10)
|
|
worker_log.close()
|
|
worker = None
|
|
worker_log = None
|
|
|
|
prefix = f"govoplan-drill-{run_id[:8]}-c@"
|
|
worker, worker_log = _start_worker(
|
|
python=args.python,
|
|
hostname=prefix + "%h",
|
|
environment=environment,
|
|
log_path=root / "worker-c.log",
|
|
)
|
|
_wait_for_worker(celery, worker, prefix, args.timeout_seconds)
|
|
value = _wait_for_result(
|
|
result,
|
|
timeout_seconds=args.timeout_seconds
|
|
+ args.visibility_timeout_seconds
|
|
+ 30,
|
|
)
|
|
_assert_probe(value, probe_id=probe_id)
|
|
if value.get("delivery_count") != 2:
|
|
raise RuntimeError(f"Worker-loss probe was not redelivered: {value!r}")
|
|
evidence["checks"].append(
|
|
{
|
|
"id": "worker_loss_redelivery",
|
|
"state": "passed",
|
|
"task_id": result.id,
|
|
"delivery_count": value["delivery_count"],
|
|
"broker_redelivered": bool(value.get("redelivered")),
|
|
}
|
|
)
|
|
except BaseException as exc:
|
|
evidence["error"] = f"{type(exc).__name__}: {exc}"
|
|
evidence["result"] = {"state": "failed"}
|
|
raise
|
|
finally:
|
|
_stop_worker(worker, worker_log)
|
|
celery.control.purge()
|
|
evidence["completed_at"] = datetime.now(UTC).isoformat()
|
|
evidence["result"] = {"state": "passed"}
|
|
return evidence
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
args = build_parser().parse_args(argv)
|
|
try:
|
|
evidence = run_drill(args)
|
|
except (OSError, RuntimeError, TimeoutError, ValueError) as exc:
|
|
print(f"worker runtime drill failed: {exc}", file=sys.stderr)
|
|
return 1
|
|
payload = json.dumps(evidence, indent=2, sort_keys=True) + "\n"
|
|
if args.output:
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = args.output.with_suffix(args.output.suffix + ".tmp")
|
|
temporary.write_text(payload, encoding="utf-8")
|
|
temporary.replace(args.output)
|
|
print(f"Worker runtime evidence written to {args.output}")
|
|
else:
|
|
print(payload, end="")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|