ECS Rolling Deployment with Automatic Rollback on Health Check Failure

import boto3
import time
import sys
import logging

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")


class ECSRollingDeployer:
    def __init__(self, cluster: str, service: str, region: str = "us-east-1"):
        """
        cluster: ECS cluster name or ARN (e.g. "prod-cluster")
        service: ECS service name (e.g. "payment-service")
        """
        self.ecs     = boto3.client("ecs", region_name=region)
        self.cluster = cluster
        self.service = service

    # ── Step 1: Get current task definition ──────────────────────
    def get_current_task_definition(self) -> str:
        """
        describe_services() returns the full service configuration.
        services[0]["taskDefinition"] is the ARN of the currently
        deployed task definition, e.g.:
          arn:aws:ecs:us-east-1:123456789012:task-definition/payment:8
        """
        response = self.ecs.describe_services(
            cluster=self.cluster,
            services=[self.service],
        )
        if not response["services"]:
            raise ValueError(f"Service not found: {self.service} in {self.cluster}")
        return response["services"][0]["taskDefinition"]

    # ── Step 2: Register new task definition ─────────────────────
    def create_new_task_definition(self, current_td_arn: str, new_image: str) -> str:
        """
        ECS task definitions are immutable — you can't edit them.
        To deploy a new image, you must register a NEW task definition revision
        with the updated container image.

        describe_task_definition() returns the full spec of an existing revision.
        We then strip the read-only fields AWS adds (taskDefinitionArn, revision,
        status, etc.) and register a new revision with the updated image.

        containerDefinitions[0]["image"] is the Docker image URI.
        We update only the first container — adjust the index if your task
        has multiple containers and you want to update a specific one.
        """
        td = self.ecs.describe_task_definition(
            taskDefinition=current_td_arn
        )["taskDefinition"]

        # Update the image in container 0 (main application container)
        containers = td["containerDefinitions"]
        old_image = containers[0]["image"]
        containers[0]["image"] = new_image
        logger.info(f"Image: {old_image} → {new_image}")

        # Strip fields that are set by AWS and cannot be provided on register
        readonly_fields = [
            "taskDefinitionArn", "revision", "status",
            "requiresAttributes", "compatibilities",
            "registeredAt", "registeredBy",
        ]
        for field in readonly_fields:
            td.pop(field, None)

        response = self.ecs.register_task_definition(**td)
        new_td_arn = response["taskDefinition"]["taskDefinitionArn"]
        logger.info(f"Registered new task definition: {new_td_arn}")
        return new_td_arn

    # ── Step 3: Update service ────────────────────────────────────
    def trigger_deployment(self, new_td_arn: str) -> None:
        """
        update_service() tells ECS to replace running tasks with new ones
        using the specified task definition.

        deploymentConfiguration controls the rolling update strategy:
        - minimumHealthyPercent=100: Never go below 100% capacity
          (ensures zero downtime — requires extra capacity)
        - maximumPercent=200: Allow up to 2× the desired count during deploy
          (starts new tasks before stopping old ones — blue/green-like)

        forceNewDeployment=True is needed if the task definition ARN
        hasn't changed (same tag like "latest") — forces ECS to pull
        the latest image even if the ARN is the same.
        """
        self.ecs.update_service(
            cluster=self.cluster,
            service=self.service,
            taskDefinition=new_td_arn,
            deploymentConfiguration={
                "minimumHealthyPercent": 100,
                "maximumPercent":        200,
            },
            forceNewDeployment=True,
        )

    # ── Step 4: Poll deployment health ───────────────────────────
    def get_deployment_status(self) -> str:
        """
        describe_services() returns a list of deployments for the service.
        There is always a PRIMARY deployment (the latest) and optionally
        ACTIVE deployments (older tasks being drained).

        Primary deployment fields:
        - desiredCount:  how many tasks ECS wants to run
        - runningCount:  how many are actually running
        - failedTasks:   how many task launches have failed

        Once runningCount == desiredCount and failedTasks == 0,
        the deployment is complete.
        """
        response = self.ecs.describe_services(
            cluster=self.cluster,
            services=[self.service],
        )
        deployments = response["services"][0]["deployments"]

        # Find the PRIMARY (most recent) deployment
        primary = next(
            (d for d in deployments if d["status"] == "PRIMARY"),
            None,
        )
        if not primary:
            return "UNKNOWN"

        desired = primary["desiredCount"]
        running = primary["runningCount"]
        failed  = primary["failedTasks"]

        if failed > 0:
            logger.error(f"Deployment failure: {failed} failed task(s)")
            return "FAILED"

        if running == desired and desired > 0:
            return "HEALTHY"

        return f"IN_PROGRESS ({running}/{desired} tasks running)"

    # ── Step 5: Roll back ─────────────────────────────────────────
    def rollback(self, previous_td_arn: str) -> None:
        """
        Rollback by calling update_service() with the PREVIOUS task definition.
        ECS will drain the failing new tasks and start tasks from the old definition.
        This is the same API call as deployment — ECS handles the swap gracefully.
        """
        logger.warning(f"Rolling back to: {previous_td_arn}")
        self.ecs.update_service(
            cluster=self.cluster,
            service=self.service,
            taskDefinition=previous_td_arn,
        )
        logger.info("Rollback triggered. Previous version will restore.")

    # ── Orchestrator ──────────────────────────────────────────────
    def deploy(
        self,
        new_image: str,
        health_check_retries: int = 15,
        health_check_delay: int   = 30,
    ) -> bool:
        """
        Full deployment lifecycle:
        1. Capture current state for rollback
        2. Register new task definition
        3. Trigger rolling update
        4. Poll health status
        5. Rollback if failed or timed out

        Returns True if deployment succeeded, False if rolled back.
        """
        logger.info(f"Starting deployment of: {new_image}")

        # Save the current (stable) task definition for rollback
        current_td = self.get_current_task_definition()
        logger.info(f"Current task definition (rollback target): {current_td}")

        # Create and deploy new task definition
        new_td = self.create_new_task_definition(current_td, new_image)
        self.trigger_deployment(new_td)
        logger.info(f"Deployment triggered. Monitoring health ({health_check_retries} checks × {health_check_delay}s)...")

        # Monitor deployment
        for attempt in range(1, health_check_retries + 1):
            time.sleep(health_check_delay)
            status = self.get_deployment_status()
            logger.info(f"[{attempt}/{health_check_retries}] Status: {status}")

            if status == "HEALTHY":
                logger.info(f"✅ Deployment succeeded!")
                return True

            if status == "FAILED":
                logger.error("❌ Deployment failed! Rolling back...")
                self.rollback(current_td)
                return False

        # Timeout — roll back
        logger.error("⏰ Health check timed out. Rolling back...")
        self.rollback(current_td)
        return False


# ── Entry point ───────────────────────────────────────────────────
if __name__ == "__main__":
    deployer = ECSRollingDeployer(
        cluster="prod-cluster",
        service="payment-service",
        region="ap-south-1",
    )

    success = deployer.deploy(
        new_image="123456789012.dkr.ecr.ap-south-1.amazonaws.com/payment:v2.4.1",
        health_check_retries=15,   # Wait up to 7.5 min (15 × 30s)
        health_check_delay=30,
    )

    # Exit with non-zero code on failure — triggers CI/CD pipeline failure
    sys.exit(0 if success else 1)

self.ecs     = boto3.client("ecs", region_name=region)
self.cluster = cluster
self.service = service

response = self.ecs.describe_services(
    cluster=self.cluster,
    services=[self.service],
)
return response["services"][0]["taskDefinition"]

td = self.ecs.describe_task_definition(
    taskDefinition=current_td_arn
)["taskDefinition"]

containers = td["containerDefinitions"]
old_image = containers[0]["image"]
containers[0]["image"] = new_image

readonly_fields = [
    "taskDefinitionArn", "revision", "status",
    "requiresAttributes", "compatibilities",
    "registeredAt", "registeredBy",
]
for field in readonly_fields:
    td.pop(field, None)

response = self.ecs.register_task_definition(**td)
new_td_arn = response["taskDefinition"]["taskDefinitionArn"]

self.ecs.update_service(
    cluster=self.cluster,
    service=self.service,
    taskDefinition=new_td_arn,
    deploymentConfiguration={
        "minimumHealthyPercent": 100,
        "maximumPercent":        200,
    },
    forceNewDeployment=True,
)

deployments = response["services"][0]["deployments"]
primary = next(
    (d for d in deployments if d["status"] == "PRIMARY"),
    None,
)

desired = primary["desiredCount"]
running = primary["runningCount"]
failed  = primary["failedTasks"]

if failed > 0:   return "FAILED"
if running == desired and desired > 0:   return "HEALTHY"
return f"IN_PROGRESS ({running}/{desired} tasks running)"

self.ecs.update_service(
    cluster=self.cluster,
    service=self.service,
    taskDefinition=previous_td_arn,
)

current_td = self.get_current_task_definition()   # Save for rollback
new_td = self.create_new_task_definition(current_td, new_image)
self.trigger_deployment(new_td)

for attempt in range(1, health_check_retries + 1):
    time.sleep(health_check_delay)
    status = self.get_deployment_status()
    if status == "HEALTHY":   return True
    if status == "FAILED":
        self.rollback(current_td)
        return False

self.rollback(current_td)   # Timeout
return False