"""Minimal ECS task-protection experiment worker. Pure stdlib. Exercises the UpdateTaskProtection agent endpoint in a tight enable/simulated-work/disable loop. Emits one structured JSON log line per event so the experiment runbook can reconstruct a precise timeline. Configuration via env vars: - WORK_DURATION_MS: how long to "work" between enable() and disable(). Short values simulate a saturated queue where the disable→enable window is tiny; longer values simulate idle or per-message work. - IDLE_MS: how long to wait between disable() and the next enable(). Zero means "simulate a saturated queue, re-enable immediately." - EXPIRES_MINUTES: value passed to the agent endpoint; caps the protection expiry window. - WORKER_LABEL: optional string stamped on every log line. Useful for differentiating workers in a multi-task deploy when parsing logs. """ import json import os import random import signal import sys import time import urllib.error import urllib.request _AGENT_URI = os.environ.get("ECS_AGENT_URI", "") _METADATA_URI = os.environ.get("ECS_CONTAINER_METADATA_URI_V4", "") _WORK_DURATION_MS = int(os.environ.get("WORK_DURATION_MS", "200")) _IDLE_MS = int(os.environ.get("IDLE_MS", "50")) _EXPIRES_MINUTES = int(os.environ.get("EXPIRES_MINUTES", "5")) _WORKER_LABEL = os.environ.get("WORKER_LABEL", "") # Phase-jitter spreads workers across the cycle rather than aligning them. # STARTUP_JITTER_MS: random sleep at worker startup (0 to this value). # IDLE_JITTER_MS: random addition to each idle period (0 to this value). # Set both to 0 for the synchronized baseline (default). _STARTUP_JITTER_MS = int(os.environ.get("STARTUP_JITTER_MS", "0")) _IDLE_JITTER_MS = int(os.environ.get("IDLE_JITTER_MS", "0")) # PHASE_MODE controls how startup phase is chosen: # "random" (default when STARTUP_JITTER_MS > 0): uniform 0..STARTUP_JITTER_MS # "bimodal": 50/50 split between 0 and half of cycle length. Adversarial # case where protected_count oscillates around desired and DEPLOYMENT_BLOCKED # might never cleanly trip. _PHASE_MODE = os.environ.get("PHASE_MODE", "random") _shutdown = False _metadata = {} def log(event: str, **fields: object) -> None: record = { "ts": time.time(), "event": event, "label": _WORKER_LABEL, "family": _metadata.get("family"), "revision": _metadata.get("revision"), "task_arn": _metadata.get("task_arn"), **fields, } sys.stdout.write(json.dumps(record) + "\n") sys.stdout.flush() def handle_sigterm(signum: int, frame: object) -> None: global _shutdown _shutdown = True log("sigterm_received", signum=signum) def fetch_metadata() -> None: if not _METADATA_URI: return try: with urllib.request.urlopen(f"{_METADATA_URI}/task", timeout=5) as resp: data = json.loads(resp.read()) except urllib.error.URLError as exc: log("metadata_fetch_failed", error=str(exc)) return _metadata["task_arn"] = data.get("TaskARN", "") _metadata["family"] = data.get("Family", "") _metadata["revision"] = data.get("Revision", "") _metadata["cluster"] = data.get("Cluster", "") def _single_put(body: dict) -> tuple[int, dict]: url = f"{_AGENT_URI}/task-protection/v1/state" req = urllib.request.Request( url, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"}, method="PUT", ) start = time.monotonic() try: with urllib.request.urlopen(req, timeout=10) as resp: elapsed_ms = int((time.monotonic() - start) * 1000) return elapsed_ms, json.loads(resp.read()) except urllib.error.HTTPError as exc: elapsed_ms = int((time.monotonic() - start) * 1000) try: body_text = exc.read().decode() except Exception: # noqa: BLE001 body_text = "" return elapsed_ms, {"_http_error": exc.code, "_body": body_text} except urllib.error.URLError as exc: elapsed_ms = int((time.monotonic() - start) * 1000) return elapsed_ms, {"_url_error": str(exc)} # Retry 3 times on ThrottlingException, 5xx, or network errors. Matches the # production `task_protection.py` retry policy (3 retries, exponential # backoff). A sustained rate-limit beyond the retry budget still surfaces # as an error — production can't mask that either. _RETRY_BACKOFFS_S = (0.5, 1.0, 2.0) def put_protection(body: dict) -> tuple[int, dict]: """Return (total_elapsed_ms_including_retries, final_response_or_error).""" total_start = time.monotonic() _, resp = _single_put(body) for backoff_s in _RETRY_BACKOFFS_S: if "_url_error" in resp: retryable = True elif "_http_error" in resp: retryable = resp["_http_error"] >= 500 or "ThrottlingException" in resp.get("_body", "") else: retryable = False if not retryable: break time.sleep(backoff_s) _, resp = _single_put(body) total_elapsed_ms = int((time.monotonic() - total_start) * 1000) return total_elapsed_ms, resp def _log_disable(iteration: int, elapsed_ms: int, resp: dict, *, after_shutdown: bool) -> None: event_prefix = "disable_after_shutdown" if after_shutdown else "disable" failure = resp.get("failure") if failure: log( f"{event_prefix}_failed", iteration=iteration, elapsed_ms=elapsed_ms, reason=failure.get("Reason"), detail=failure.get("Detail"), ) return if "_http_error" in resp or "_url_error" in resp: log(f"{event_prefix}_error", iteration=iteration, elapsed_ms=elapsed_ms, response=resp) return log(f"{event_prefix}_success", iteration=iteration, elapsed_ms=elapsed_ms) def main() -> int: signal.signal(signal.SIGTERM, handle_sigterm) signal.signal(signal.SIGINT, handle_sigterm) fetch_metadata() log( "worker_start", agent_uri_set=bool(_AGENT_URI), metadata_uri_set=bool(_METADATA_URI), work_duration_ms=_WORK_DURATION_MS, idle_ms=_IDLE_MS, expires_minutes=_EXPIRES_MINUTES, startup_jitter_ms=_STARTUP_JITTER_MS, idle_jitter_ms=_IDLE_JITTER_MS, ) if not _AGENT_URI: log("no_agent_uri_exiting") return 1 if _PHASE_MODE == "bimodal": # Half the workers start at phase 0, half at phase (WORK+IDLE)/2. # Produces the adversarial case where ~50% of tasks are always # protected and ~50% are always unprotected — so protected_count # never cleanly exceeds desired_count. cycle_ms = _WORK_DURATION_MS + _IDLE_MS half_cycle_seconds = (cycle_ms / 2) / 1000 group = random.choice(["A", "B"]) # noqa: S311 if group == "B": log("startup_jitter", mode="bimodal", group=group, jitter_seconds=half_cycle_seconds) time.sleep(half_cycle_seconds) else: log("startup_jitter", mode="bimodal", group=group, jitter_seconds=0) elif _STARTUP_JITTER_MS > 0: jitter_seconds = random.uniform(0, _STARTUP_JITTER_MS / 1000) # noqa: S311 log("startup_jitter", mode="random", jitter_seconds=jitter_seconds) time.sleep(jitter_seconds) iteration = 0 while not _shutdown: iteration += 1 enable_ms, enable_resp = put_protection( {"ProtectionEnabled": True, "ExpiresInMinutes": _EXPIRES_MINUTES}, ) failure = enable_resp.get("failure") if isinstance(enable_resp, dict) else None if failure: log( "enable_failed", iteration=iteration, elapsed_ms=enable_ms, reason=failure.get("Reason"), detail=failure.get("Detail"), ) # Mirror what the real worker would do: exit on rejection. break if "_http_error" in enable_resp or "_url_error" in enable_resp: log("enable_error", iteration=iteration, elapsed_ms=enable_ms, response=enable_resp) break expiration = None if isinstance(enable_resp, dict): expiration = enable_resp.get("protection", {}).get("ExpirationDate") log("enable_success", iteration=iteration, elapsed_ms=enable_ms, expiration=expiration) # Simulate work log("work_start", iteration=iteration) time.sleep(_WORK_DURATION_MS / 1000) log("work_end", iteration=iteration) if _shutdown: # SIGTERM landed during work; still disable cleanly so the task # becomes eligible for termination without waiting for expiry. disable_ms, disable_resp = put_protection({"ProtectionEnabled": False}) _log_disable(iteration, disable_ms, disable_resp, after_shutdown=True) break disable_ms, disable_resp = put_protection({"ProtectionEnabled": False}) _log_disable(iteration, disable_ms, disable_resp, after_shutdown=False) if _IDLE_MS > 0 or _IDLE_JITTER_MS > 0: jitter_ms = random.randint(0, _IDLE_JITTER_MS) if _IDLE_JITTER_MS > 0 else 0 # noqa: S311 total_idle_ms = _IDLE_MS + jitter_ms log("idle_start", iteration=iteration, idle_ms=total_idle_ms, jitter_ms=jitter_ms) time.sleep(total_idle_ms / 1000) log("worker_exit", reason="shutdown" if _shutdown else "loop_exited", iterations=iteration) return 0 if __name__ == "__main__": sys.exit(main())