#!/usr/bin/env python3 """ Async Batch Processor with Checkpointing, Idempotency & Graceful Interruption. Runs unattended, re-runnable automation glue. Uses only Python 3.11+ stdlib. """ from __future__ import annotations import asyncio import json import os import random import signal import sys import tempfile import time from typing import Any, Dict, List, Optional, Tuple class ProcessingError(Exception): """Raised when the external call fails.""" pass async def process(item: str) -> str: """ Simulates a flaky external API call. - Latency: 10-80ms random uniform - Failure rate: ~20% (raises ProcessingError) """ await asyncio.sleep(random.uniform(0.01, 0.08)) if random.random() < 0.2: raise ProcessingError(f"Flaky timeout/network error for {item}") return f"result-{item}" # --------------------------------------------------------------------------- # # Checkpoint I/O # --------------------------------------------------------------------------- # def load_checkpoint(path: str) -> Dict[str, Any]: """Load checkpoint or return empty structure if missing.""" if os.path.exists(path): with open(path, "r", encoding="utf-8") as f: return json.load(f) return {"items": {}, "summary": {"succeeded": 0, "failed": 0}} def save_checkpoint( items_dict: Dict[str, str], succeeded: int, failed: int, path: str ) -> None: """Atomically write checkpoint using temp file + rename.""" data = { "items": items_dict, "summary": {"succeeded": succeeded, "failed": failed}, } dir_path = os.path.dirname(os.path.abspath(path)) or "." fd, tmp_path = tempfile.mkstemp(dir=dir_path, suffix=".tmp") try: with os.fdopen(fd, "w", encoding="utf-8") as f: json.dump(data, f) os.replace(tmp_path, path) # Atomic on POSIX & Windows (Vista+) except Exception: if os.path.exists(tmp_path): os.unlink(tmp_path) raise # --------------------------------------------------------------------------- # # Core Processor # --------------------------------------------------------------------------- # async def run_batch( items: List[str], checkpoint_path: str = "checkpoint.json", concurrency_limit: int = 8, stop_event: Optional[asyncio.Event] = None, ) -> Tuple[Dict[str, Any], int]: """ Process a list of items with bounded concurrency, retries, and checkpointing. Returns: (summary_dict, max_concurrent_tasks_observed) """ if stop_event is None: stop_event = asyncio.Event() checkpoint = load_checkpoint(checkpoint_path) completed_items: Dict[str, str] = checkpoint.get("items", {}) semaphore = asyncio.Semaphore(concurrency_limit) running_tasks: List[asyncio.Task[Any]] = [] succeeded_count = checkpoint["summary"]["succeeded"] failed_count = checkpoint["summary"]["failed"] skipped_count = 0 max_concurrent = 0 current_concurrent = 0 start_time = time.monotonic() async def process_item(item: str) -> None: nonlocal succeeded_count, failed_count, max_concurrent, current_concurrent async with semaphore: current_concurrent += 1 if current_concurrent > max_concurrent: max_concurrent = current_concurrent try: # Retry up to 3 times on ProcessingError (4 attempts total) for attempt in range(4): try: await process(item) completed_items[item] = "success" succeeded_count += 1 save_checkpoint(completed_items, succeeded_count, failed_count, checkpoint_path) break # Success, exit retry loop except ProcessingError: if attempt == 3: # Last attempt exhausted completed_items[item] = "failed" failed_count += 1 save_checkpoint(completed_items, succeeded_count, failed_count, checkpoint_path) else: backoff = 0.1 * (2 ** attempt) # 0.1, 0.2, 0.4 await asyncio.sleep(backoff) finally: current_concurrent -= 1 for item in items: if stop_event.is_set(): break status = completed_items.get(item) if status in ("success", "failed"): skipped_count += 1 continue task = asyncio.create_task(process_item(item)) running_tasks.append(task) # Let in-flight tasks finish before exiting if stop_event.is_set(): await asyncio.gather(*running_tasks, return_exceptions=True) elapsed_ms = int((time.monotonic() - start_time) * 1000) summary = { "succeeded": succeeded_count, "failed": failed_count, "skipped": skipped_count, "total": len(items), "elapsed_ms": elapsed_ms, } print(json.dumps(summary)) return summary, max_concurrent # --------------------------------------------------------------------------- # # Test Suite # --------------------------------------------------------------------------- # async def test_resumability() -> None: """a) Run batch, kill mid-run, restart. Assert no reprocessing & terminal state.""" print("Running Test A: Resumability & Idempotency...") items = [f"resum-{i}" for i in range(20)] cp_path = "test_checkpoint_a.json" # First run: cancel after 150ms to simulate SIGINT/crash mid-flight async def run_and_timeout(): await asyncio.wait_for(run_batch(items, checkpoint_path=cp_path), timeout=0.15) task = asyncio.create_task(run_and_timeout()) try: await task except asyncio.TimeoutError: pass # Allow cancellation to propagate and in-flight tasks to flush checkpoints await asyncio.sleep(0.2) # Second run (resume) summary, _ = await run_batch(items, checkpoint_path=cp_path) assert summary["total"] == 20 assert summary["succeeded"] + summary["failed"] + summary["skipped"] == 20 # Verify all items are terminal in checkpoint with open(cp_path) as f: cp = json.load(f) for item in items: assert cp["items"][item] in ("success", "failed"), \ f"{item} is not terminal: {cp['items'][item]}" # Verify no completed item was reprocessed (skipped count == previously finished) with open(cp_path) as f: final_cp = json.load(f) expected_skipped = sum(1 for v in final_cp["items"].values() if v in ("success", "failed")) assert summary["skipped"] == expected_skipped, \ f"Reprocessing detected. Expected skipped={expected_skipped}, got {summary['skipped']}" print(" PASSED") async def test_checkpoint_integrity() -> None: """b) Watcher reads checkpoint repeatedly during run, asserts valid JSON always.""" print("Running Test B: Checkpoint Integrity...") items = [f"integ-{i}" for i in range(10)] cp_path = "test_checkpoint_b.json" integrity_ok = True async def watcher(): nonlocal integrity_ok while True: if os.path.exists(cp_path): try: with open(cp_path, "r", encoding="utf-8") as f: json.load(f) # Raises on invalid JSON except (json.JSONDecodeError, FileNotFoundError): integrity_ok = False return await asyncio.sleep(0.05) watcher_task = asyncio.create_task(watcher()) await run_batch(items, checkpoint_path=cp_path) try: await asyncio.wait_for(watcher_task, timeout=1.0) except asyncio.TimeoutError: pass assert integrity_ok, "Checkpoint contained invalid JSON at some point during execution" print(" PASSED") async def test_concurrency_cap() -> None: """c) Assert concurrency cap of 8 holds.""" print("Running Test C: Concurrency Cap...") items = [f"conc-{i}" for i in range(50)] cp_path = "test_checkpoint_c.json" _, max_conc = await run_batch(items, checkpoint_path=cp_path) assert max_conc <= 8, f"Concurrency cap violated: observed {max_conc} > 8" print(" PASSED") async def main_tests() -> None: """Run all tests sequentially.""" try: await test_resumability() await test_checkpoint_integrity() await test_concurrency_cap() print("\n✅ All tests passed.") except AssertionError as e: print(f"\n❌ Test failed: {e}") sys.exit(1) # --------------------------------------------------------------------------- # # Entry Point # --------------------------------------------------------------------------- # if __name__ == "__main__": # If run with --batch, execute the processor normally. Otherwise run tests. if "--batch" in sys.argv: items = [f"job-{i}" for i in range(50)] # Example workload loop = asyncio.get_running_loop() stop_event = asyncio.Event() def handle_sigint(): stop_event.set() old_handler = signal.signal(signal.SIGINT, lambda s, f: loop.call_soon_threadsafe(handle_sigint)) try: asyncio.run(run_batch(items, stop_event=stop_event)) except KeyboardInterrupt: pass # Handled by stop_event + gather finally: signal.signal(signal.SIGINT, old_handler) else: asyncio.run(main_tests())