#!/usr/bin/env python3 """Asyncio timeout handling examples for reliable network tasks. This script provides a small, practical toolkit for applying asyncio timeouts in network automation and async service workflows. Features: - Single-operation timeout wrapper - Phase-based timeout budgeting - Fan-out task execution with partial results - Timeout-aware cleanup patterns - CLI for quick demonstrations and validation No credentials, secrets, or environment-specific endpoints are included. """ from __future__ import annotations import argparse import asyncio import json import sys from dataclasses import dataclass, asdict from typing import Any, Awaitable, Callable, Iterable, List, Optional @dataclass class TaskResult: name: str status: str elapsed_seconds: float result: Optional[Any] = None error: Optional[str] = None async def with_timeout(coro: Awaitable[Any], timeout_seconds: float) -> Any: """Run an awaitable within a timeout boundary. Raises: TimeoutError: if the timeout expires. """ if timeout_seconds <= 0: raise ValueError("timeout_seconds must be greater than 0") async with asyncio.timeout(timeout_seconds): return await coro async def run_step(name: str, duration: float) -> str: """Simulate a network step or remote call. This is a safe placeholder for a real network coroutine. """ if duration < 0: raise ValueError(f"duration for {name!r} must be non-negative") try: await asyncio.sleep(duration) return f"{name} completed" finally: # Put cleanup here for sockets, temp buffers, tasks, or sessions. # The finally block ensures cancellation-aware cleanup paths remain safe. pass async def bounded_operation( name: str, duration: float, timeout_seconds: float, ) -> TaskResult: """Execute one operation with a timeout and return a structured result.""" start = asyncio.get_running_loop().time() try: result = await with_timeout(run_step(name, duration), timeout_seconds) elapsed = asyncio.get_running_loop().time() - start return TaskResult(name=name, status="ok", elapsed_seconds=elapsed, result=result) except TimeoutError: elapsed = asyncio.get_running_loop().time() - start return TaskResult( name=name, status="timeout", elapsed_seconds=elapsed, error=f"Operation exceeded {timeout_seconds:.2f} seconds", ) except Exception as exc: elapsed = asyncio.get_running_loop().time() - start return TaskResult(name=name, status="error", elapsed_seconds=elapsed, error=str(exc)) async def fan_out_operations( operations: Iterable[tuple[str, float]], timeout_seconds: float, ) -> list[TaskResult]: """Run multiple operations concurrently and keep partial results. Pending tasks are cancelled when the overall timeout expires. """ if timeout_seconds <= 0: raise ValueError("timeout_seconds must be greater than 0") loop = asyncio.get_running_loop() start = loop.time() tasks = [asyncio.create_task(run_step(name, duration), name=name) for name, duration in operations] results: list[TaskResult] = [] try: async with asyncio.timeout(timeout_seconds): completed = await asyncio.gather(*tasks, return_exceptions=True) elapsed = loop.time() - start for (name, _), item in zip(operations, completed): if isinstance(item, Exception): results.append(TaskResult(name=name, status="error", elapsed_seconds=elapsed, error=str(item))) else: results.append(TaskResult(name=name, status="ok", elapsed_seconds=elapsed, result=item)) return results except TimeoutError: elapsed = loop.time() - start for task in tasks: if not task.done(): task.cancel() await asyncio.gather(*tasks, return_exceptions=True) for task in tasks: name = task.get_name() if task.cancelled(): results.append( TaskResult( name=name, status="cancelled", elapsed_seconds=elapsed, error="Cancelled due to overall timeout", ) ) elif task.exception() is not None: results.append( TaskResult( name=name, status="error", elapsed_seconds=elapsed, error=str(task.exception()), ) ) else: results.append( TaskResult( name=name, status="ok", elapsed_seconds=elapsed, result=task.result(), ) ) return results async def phase_demo(connect_seconds: float, read_seconds: float, overall_seconds: float) -> list[TaskResult]: """Demonstrate phase-oriented timeouts. This pattern is useful when connection setup should fail fast while the read phase may reasonably take longer. """ results: list[TaskResult] = [] results.append(await bounded_operation("connect", connect_seconds, timeout_seconds=2.0)) if results[-1].status != "ok": return results results.append(await bounded_operation("read", read_seconds, timeout_seconds=overall_seconds)) return results def _parse_operation(value: str) -> tuple[str, float]: if ":" not in value: raise argparse.ArgumentTypeError("operation must use NAME:DURATION format") name, raw_duration = value.split(":", 1) name = name.strip() if not name: raise argparse.ArgumentTypeError("operation name cannot be empty") try: duration = float(raw_duration) except ValueError as exc: raise argparse.ArgumentTypeError(f"invalid duration for {name!r}") from exc if duration < 0: raise argparse.ArgumentTypeError("duration must be non-negative") return name, duration async def main_async(args: argparse.Namespace) -> int: if args.mode == "single": result = await bounded_operation(args.name, args.duration, args.timeout) print(json.dumps(asdict(result), indent=2)) return 0 if args.mode == "fanout": operations = args.operation results = await fan_out_operations(operations, args.timeout) print(json.dumps([asdict(r) for r in results], indent=2)) return 0 if args.mode == "phase": results = await phase_demo(args.connect_duration, args.read_duration, args.read_timeout) print(json.dumps([asdict(r) for r in results], indent=2)) return 0 raise ValueError(f"Unsupported mode: {args.mode}") def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description="Demonstrate safe asyncio timeout handling patterns for network tasks.") subparsers = parser.add_subparsers(dest="mode", required=True) single = subparsers.add_parser("single", help="Run one operation with a timeout") single.add_argument("--name", default="request", help="Operation name") single.add_argument("--duration", type=float, default=1.0, help="Simulated duration in seconds") single.add_argument("--timeout", type=float, default=2.0, help="Timeout in seconds") fanout = subparsers.add_parser("fanout", help="Run multiple operations concurrently") fanout.add_argument( "--operation", action="append", type=_parse_operation, default=[("fast", 0.5), ("medium", 1.5), ("slow", 3.0)], help="Operation in NAME:DURATION format; can be repeated", ) fanout.add_argument("--timeout", type=float, default=2.0, help="Overall timeout in seconds") phase = subparsers.add_parser("phase", help="Run connect and read phases with separate budgets") phase.add_argument("--connect-duration", type=float, default=0.5, help="Simulated connect duration") phase.add_argument("--read-duration", type=float, default=1.5, help="Simulated read duration") phase.add_argument("--read-timeout", type=float, default=3.0, help="Read-phase timeout in seconds") return parser def validate_args(args: argparse.Namespace) -> None: if getattr(args, "timeout", 1.0) <= 0: raise ValueError("timeout must be greater than 0") if getattr(args, "duration", 0.0) < 0: raise ValueError("duration must be non-negative") if getattr(args, "connect_duration", 0.0) < 0: raise ValueError("connect-duration must be non-negative") if getattr(args, "read_duration", 0.0) < 0: raise ValueError("read-duration must be non-negative") if getattr(args, "read_timeout", 1.0) <= 0: raise ValueError("read-timeout must be greater than 0") def main(argv: Optional[list[str]] = None) -> int: parser = build_parser() args = parser.parse_args(argv) try: validate_args(args) return asyncio.run(main_async(args)) except KeyboardInterrupt: print("Interrupted", file=sys.stderr) return 130 except Exception as exc: print(f"Error: {exc}", file=sys.stderr) return 1 if __name__ == "__main__": raise SystemExit(main())