import asyncio import json import os import random import signal import time from dataclasses import dataclass, asdict, field from pathlib import Path from typing import Set, List, Dict, Any # --- Domain Exceptions --- class ProcessingError(Exception): """Raised when the flaky external call fails.""" pass # --- Models --- @dataclass class Checkpoint: completed_ids: Set[str] = field(default_factory=set) failed_ids: Set[str] = field(default_factory=set) def to_dict(self) -> Dict[str, Any]: return { "completed_ids": list(self.completed_ids), "failed_ids": list(self.failed_ids) } @classmethod def from_dict(cls, data: Dict[str, Any]) -> "Checkpoint": return cls( completed_ids=set(data.get("completed_ids", [])), failed_ids=set(data.get("failed_ids", [])) ) # --- Core Logic --- class BatchProcessor: def __init__( self, items: List[str], checkpoint_path: str = "checkpoint.json", max_concurrency: int = 8, max_retries: int = 3 ): self.items = items self.checkpoint_path = Path(checkpoint_path) self.semaphore = asyncio.Semaphore(max_concurrency) self.max_retries = max_retries self.checkpoint = Checkpoint() self.shutdown_event = asyncio.Event() self.start_time = 0.0 self._pending_tasks: Set[asyncio.Task] = set() def _load_checkpoint(self): if self.checkpoint_path.exists(): try: with open(self.checkpoint_path, "r") as f: data = json.load(f) self.checkpoint = Checkpoint.from_dict(data) except (json.JSONDecodeError, KeyError): # If corrupt, we start fresh or could raise error. # Requirement says "crash must leave valid checkpoint", so we assume it's valid. pass def _save_checkpoint(self): """Atomic write using a temporary file and os.replace.""" temp_path = self.checkpoint_path.with_suffix(".tmp") try: with open(temp_path, "w") as f: json.dump(self.checkpoint.to_dict(), f) # os.replace is atomic on both POSIX and Windows (for existing files) os.replace(temp_path, self.checkpoint_path) except Exception as e: print(f"Failed to save checkpoint: {e}") async def _process_item_with_retry(self, item: str, process_func): """Handles the retry logic and backoff for a single item.""" async with self.semaphore: # Check if we are shutting down before starting a new item if self.shutdown_event.is_set(): return for attempt in range(self.max_retries + 1): try: await process_func(item) # Success path self.checkpoint.completed_ids.add(item) self._save_checkpoint() return "succeeded" except ProcessingError: if attempt < self.max_retries: # Exponential backoff: 0.1, 0.2, 0.4... wait_time = 0.1 * (2 ** attempt) await asyncio.sleep(wait_time) else: # Terminal failure after retries self.checkpoint.failed_ids.add(item) self._save_checkpoint() return "failed" except asyncio.CancelledError: # If the task is cancelled (e.g. during shutdown), # we don't record it as completed or failed. raise async def run(self, process_func) -> Dict[str, Any]: self._load_checkpoint() self.start_time = time.perf_counter() # Identify items to skip to_process = [ i for i in self.items if i not in self.checkpoint.completed_ids and i not in self.checkpoint.failed_ids ] skipped_count = len(self.items) - len(to_process) - len(self.checkpoint.failed_ids) # Note: skipped is items already in checkpoint.completed_ids skipped_count = len([i for i in self.items if i in self.checkpoint.completed_ids]) # We use a set to track tasks so we can await them during graceful shutdown tasks = [] for item in to_process: if self.shutdown_event.is_set(): break task = asyncio.create_task(self._process_item_with_retry(item, process_func)) tasks.append(task) self._pending_tasks.add(task) task.add_done_callback(self._pending_tasks.discard) # Wait for all spawned tasks to finish if tasks: await asyncio.gather(*tasks, return_exceptions=True) elapsed_ms = int((time.perf_counter() - self.start_time) * 1000) # Final Summary Calculation # We must re-calculate based on the final state of checkpoint + what was skipped succeeded = len(self.checkpoint.completed_ids) failed = len(self.checkpoint.failed_ids) # Items that were already done before we started already_done = len([i for i in self.items if i in self.checkpoint.completed_ids and i not in to_process]) # Actually, the requirement asks for a summary of this run's context. # Let's provide the absolute state: return { "succeeded": succeeded, "failed": failed, "skipped": skipped_count, "total": len(self.items), "elapsed_ms": elapsed_ms } # --- Mock External Service --- async def flaky_process(item: str) -> str: """The provided flaky function.""" # Simulate random latency 10-80ms await asyncio.sleep(random.uniform(0.01, 0.08)) # 80% success rate if random.random() < 0.2: raise ProcessingError(f"Failed to process {item}") return f"Result for {item}" # --- Tests --- async def run_tests(): print("\n--- Starting Tests ---") checkpoint_file = "test_checkpoint.json" if os.path.exists(checkpoint_file): os.remove(checkpoint_file) # 1. Concurrency Test print("Test: Bounded Concurrency (Max 8)...", end="", flush=True) active_count = 0 max_observed = 0 lock = asyncio.Lock() async def concurrency_mock(item: str): nonlocal active_count, max_observed async with lock: active_count += 1 max_observed = max(max_observed, active_count) await asyncio.sleep(0.05) async with lock: active_count -= 1 proc = BatchProcessor(["item-" + str(i) for i in range(20)], checkpoint_path=checkpoint_file, max_concurrency=8) await proc.run(concurrency_mock) assert max_observed <= 8, f"Concurrency exceeded! Observed {max_observed}" print("PASSED") # 2. Integrity Test (Watcher) print("Test: Checkpoint Integrity (JSON Validity)...", end="", flush=True) stop_watcher = asyncio.Event() async def integrity_watcher(): while not stop_watcher.is_set(): if os.path.exists(checkpoint_file): try: with open(checkpoint_file, "r") as f: json.load(f) except Exception as e: raise AssertionError(f"Corrupt JSON detected: {e}") await asyncio.sleep(0.005) watcher_task = asyncio.create_task(integrity_watcher()) proc = BatchProcessor(["item-" + str(i) for i in range(50)], checkpoint_path=checkpoint_file) await proc.run(flaky_process) stop_watcher.set() await watcher_task print("PASSED") # 3. Resumability / Idempotency Test print("Test: Resumability & Idempotency...", end="", flush=True) items = [f"resumable-{i}" for i in range(100)] checkpoint_file = "resumable.json" if os.path.exists(checkpoint_file): os.remove(checkpoint_file) # Run 1: Start and interrupt mid-way proc = BatchProcessor(items, checkpoint_path=checkpoint_file) # We simulate an interruption by cancelling the main task after a short delay main_task = asyncio.create_task(proc.run(flaky_process)) await asyncio.sleep(0.2) proc.shutdown_event.set() # Signal shutdown # We don't cancel the task immediately to allow in-flight items to finish (graceful) await main_task # Run 2: Restart proc_retry = BatchProcessor(items, checkpoint_path=checkpoint_file) final_summary = await proc_retry.run(flaky_process) # Assertions # 1. All items must eventually be in completed or failed sets all_processed = proc_retry.checkpoint.completed_ids | proc_retry.checkpoint.failed_ids assert len(all_processed) == len(items), f"Not all items processed. Got {len(all_processed)}/{len(items)}" # 2. Checkpoint must be valid and contain all items assert final_summary["total"] == 100 print("PASSED") # Cleanup for f in [checkpoint_file, "test_checkpoint.json", "resumable.json"]: if os.path.exists(f): os.remove(f) print("--- All Tests Passed Successfully ---\n") # --- Main Entry Point --- async def main(): # Setup signal handling for SIGINT (Ctrl+C) loop = asyncio.get_running_loop() # We'll use a global processor instance to allow the signal handler access processor = None def handle_sigint(): if processor: print("\n[!] SIGINT received. Shutting down gracefully (finishing in-flight tasks)...") processor.shutdown_event.set() # Register signal handler try: for sig in (signal.SIGINT, signal.SIGTERM): loop.add_signal_handler(sig, handle_sigint) except NotImplementedError: # signal.add_signal_handler is not implemented on Windows in some Python versions pass # Input Data work_items = [f"job-{i:03d}" for i in range(100)] global processor processor = BatchProcessor(work_items, max_concurrency=8) try: summary = await processor.run(flaky_process) print(json.dumps(summary)) except KeyboardInterrupt: # Fallback for environments where signal handler doesn't catch it pass except Exception as e: print(f"Fatal error: {e}") finally: # If we are running tests, don't exit yet import sys if len(sys.argv) > 1 and sys.argv[1] == "--test": pass if __name__ == "__main__": import sys if len(sys.argv) > 1 and sys.argv[1] == "--test": asyncio.run(run_tests()) else: try: asyncio.run(main()) except KeyboardInterrupt: # Final catch for the top-level loop sys.exit(0)