boto3 Retry Decorator with Exponential Backoff for ThrottlingException

import boto3
import time
import random
import functools
import logging
from botocore.exceptions import ClientError
from botocore.config import Config

logger = logging.getLogger(__name__)


# ── Decorator: add retry to any function ─────────────────────────
def aws_retry(
    max_retries: int   = 5,
    base_delay: float  = 0.5,
    max_delay: float   = 30.0,
    jitter: bool       = True,
    retryable_errors: set = None,
):
    """
    Decorator that wraps any function with automatic retry logic.

    max_retries:      total number of retry attempts (not counting the first try)
    base_delay:       initial wait before first retry (seconds)
    max_delay:        cap the wait time at this many seconds
    jitter:           add random variation to prevent thundering herd
    retryable_errors: set of AWS error codes to retry on

    Exponential backoff formula:
      delay = min(base_delay × 2^attempt, max_delay)
      with jitter: delay × uniform(0.75, 1.25)

    functools.wraps(func) copies the original function's __name__,
    __doc__, __module__, __qualname__ and __annotations__ to the
    wrapper — essential for debugging and introspection.
    """
    if retryable_errors is None:
        retryable_errors = {
            "ThrottlingException",
            "RequestLimitExceeded",
            "TooManyRequestsException",
            "ServiceUnavailable",
            "InternalServerError",
            "RequestTimeout",
            "ProvisionedThroughputExceededException",
            "LimitExceededException",
            "RequestExpired",
            "Throttling",               # Some services use just "Throttling"
            "SlowDown",                 # S3 uses this for throttling
        }

    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            last_exception = None

            for attempt in range(max_retries + 1):   # +1: attempt 0 is the initial try
                try:
                    return func(*args, **kwargs)

                except ClientError as e:
                    error_code = e.response["Error"]["Code"]
                    error_msg  = e.response["Error"]["Message"]

                    # Non-retryable: re-raise immediately
                    if error_code not in retryable_errors:
                        raise

                    last_exception = e

                    # We've exhausted all retries
                    if attempt == max_retries:
                        logger.error(
                            f"{func.__name__} failed after {max_retries} retries. "
                            f"Last error: {error_code}: {error_msg}"
                        )
                        raise

                    # Calculate wait time
                    delay = min(base_delay * (2 ** attempt), max_delay)
                    if jitter:
                        # Multiply by random value between 0.75 and 1.25
                        # This spreads retries across time instead of bunching them
                        delay *= (0.75 + random.random() * 0.5)

                    logger.warning(
                        f"{func.__name__} attempt {attempt + 1}/{max_retries} failed "
                        f"({error_code}). Retrying in {delay:.2f}s..."
                    )
                    time.sleep(delay)

            raise last_exception   # Should never reach here, but satisfies type checkers

        return wrapper
    return decorator


# ── Usage: function-level decorator ──────────────────────────────
@aws_retry(max_retries=5, base_delay=1.0)
def list_all_instances(region: str) -> list:
    """List all EC2 instances in a region with auto-retry on throttle."""
    ec2 = boto3.client("ec2", region_name=region)
    instances = []
    paginator = ec2.get_paginator("describe_instances")
    for page in paginator.paginate():
        for reservation in page["Reservations"]:
            instances.extend(reservation["Instances"])
    return instances


@aws_retry(max_retries=3, base_delay=0.5)
def get_secret(secret_name: str) -> dict:
    """Retrieve secret from Secrets Manager with retry."""
    import json
    sm = boto3.client("secretsmanager")
    response = sm.get_secret_value(SecretId=secret_name)
    return json.loads(response["SecretString"])


# ── Usage: class-based wrapper (all methods auto-retry) ───────────
class AWSClientWithRetry:
    """
    Wraps a boto3 client so that EVERY method call is automatically
    retried on throttling. Useful when you use a single client extensively.

    __getattr__ is called when Python can't find an attribute on the object.
    We intercept it to wrap any callable (boto3 method) with the retry decorator.
    """

    def __init__(
        self,
        service: str,
        region: str = "us-east-1",
        max_retries: int = 5,
        **boto_kwargs,
    ):
        self._client      = boto3.client(service, region_name=region, **boto_kwargs)
        self._max_retries = max_retries

    def __getattr__(self, name: str):
        """
        Called when accessing any attribute not found on this object.
        Returns the boto3 method wrapped with the retry decorator.
        """
        attr = getattr(self._client, name)
        if callable(attr):
            return aws_retry(max_retries=self._max_retries)(attr)
        return attr


# ── Method 3: botocore built-in retry config (simpler) ───────────
def get_client_with_builtin_retry(service: str, region: str = "us-east-1"):
    """
    botocore has built-in retry logic via Config.
    retry.mode options:
      "legacy"  — default, 3 retries with fixed delays
      "standard" — 3 retries with exponential backoff
      "adaptive" — dynamic retry with token bucket algorithm (best for throttling)

    max_attempts includes the initial attempt + retries.
    So max_attempts=5 means 1 initial + 4 retries.
    """
    config = Config(
        retries={
            "mode":         "adaptive",   # Adaptive token bucket algorithm
            "max_attempts": 10,           # Up to 9 retries
        },
        connect_timeout=5,
        read_timeout=30,
    )
    return boto3.client(service, region_name=region, config=config)


# ── Combining approaches ──────────────────────────────────────────
if __name__ == "__main__":
    # Approach 1: Function decorator (best for specific functions)
    instances = list_all_instances("ap-south-1")
    print(f"Found {len(instances)} instances")

    # Approach 2: Class wrapper (best when reusing a client heavily)
    ec2 = AWSClientWithRetry("ec2", region="ap-south-1", max_retries=5)

    # All these calls will auto-retry on ThrottlingException:
    response = ec2.describe_vpcs()
    sgs = ec2.describe_security_groups()
    subnets = ec2.describe_subnets()
    print(f"VPCs: {len(response['Vpcs'])}")

    # Approach 3: botocore adaptive mode (simplest — built-in)
    s3 = get_client_with_builtin_retry("s3")
    buckets = s3.list_buckets()["Buckets"]
    print(f"S3 buckets: {len(buckets)}")

def aws_retry(
    max_retries: int   = 5,
    base_delay: float  = 0.5,
    max_delay: float   = 30.0,
    jitter: bool       = True,
    retryable_errors: set = None,
):

if retryable_errors is None:
    retryable_errors = {
        "ThrottlingException",
        "RequestLimitExceeded",
        "SlowDown",
        ...
    }

def decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):

for attempt in range(max_retries + 1):   # +1: attempt 0 is the initial try
    try:
        return func(*args, **kwargs)
    except ClientError as e:
        error_code = e.response["Error"]["Code"]
        error_msg  = e.response["Error"]["Message"]
        if error_code not in retryable_errors:
            raise

delay = min(base_delay * (2 ** attempt), max_delay)
if jitter:
    delay *= (0.75 + random.random() * 0.5)
time.sleep(delay)

def __getattr__(self, name: str):
    attr = getattr(self._client, name)
    if callable(attr):
        return aws_retry(max_retries=self._max_retries)(attr)
    return attr

config = Config(
    retries={"mode": "adaptive", "max_attempts": 10},
    connect_timeout=5,
    read_timeout=30,
)