asg-graceful-lifecycle-termination.py

import boto3
import json
import logging
import time

logger = logging.getLogger()
logger.setLevel(logging.INFO)

def notify_on_error(message):

    try:

        sns_client = boto3.client('sns')
        
        sns_response = sns_client.publish(
            TopicArn='arn:aws:sns:us-east-1:xxxxxxxxxxxx:Lambda-Errors',
            Message=message,
            Subject='ASG/ECS Lifecycle Termination Error'
        )

        logger.info("SNS Publish HTTP Response: %s" % sns_response[u'ResponseMetadata'][u'HTTPStatusCode'])

    except Exception as e:
        logger.error(e)


def deregister_from_ecs_cluster(instanceid, cluster_name):

    try:
        
        ecs_client = boto3.client('ecs')
        containerarn = None

        ecsclusterinstances = ecs_client.list_container_instances( 
            cluster=cluster_name
        )

        for containerarns in ecsclusterinstances[u'containerInstanceArns']:
     
            response = ecs_client.describe_container_instances(
                cluster=cluster_name,
                containerInstances=[
                    containerarns
                ]
            )
     
            if response[u'containerInstances'][0][u'ec2InstanceId'] == instanceid:
                containerarn = containerarns

        if containerarn:

            logger.info("%s found registered with container ARN (%s) in ECS cluster. Attempting to deregister from cluster." % (instanceid, containerarn))

            ecs_response = ecs_client.deregister_container_instance(
                cluster=cluster_name,
                containerInstance=containerarn,
                force = True
            )

            logger.info("ECS Deregistration HTTP Response: %s" % ecs_response[u'ResponseMetadata'][u'HTTPStatusCode'])

        else:
            logger.info('%s is NOT regesitered with the Cluster.' % instanceid)
            return

    except Exception as e:
        logger.error(e)
        notify_on_error(str(e))


def deregister_from_elb(instanceid, elb):

    try:

        elb_client = boto3.client('elb')

        elb_response = elb_client.describe_load_balancers(
            LoadBalancerNames=[elb]
        )


        for instance in elb_response[u'LoadBalancerDescriptions'][0][u'Instances']:

            if instance[u'InstanceId'] == instanceid:
                logger.info("%s found registered with ELB, attempting to deregister." % instanceid)
                for timer in range(0, 60):

                    elb_response = elb_client.describe_load_balancers(
                        LoadBalancerNames=[elb]
                    )

                    if len(elb_response[u'LoadBalancerDescriptions'][0][u'Instances']) <= 2:
                        logger.info("There are two or fewer instances registered with the ELB, waiting for another instance to register before moving on. (5 min timeout)")
                        logger.info(elb_response[u'LoadBalancerDescriptions'][0][u'Instances'])
                        time.sleep(5)
                    else:
                        logger.info("Sweet! We have more than two instances in the ELB, moving forward with deregistration request.")
                        logger.info(elb_response[u'LoadBalancerDescriptions'][0][u'Instances'])
                        break

                else:
                    logger.info("Whoops! It took more than 5 minutes waiting for more instances to register in the ELB! Giving up on being being graceful!")
                    return

                elb_response = elb_client.deregister_instances_from_load_balancer(
                    LoadBalancerName=elb,
                    Instances=[
                        {
                        'InstanceId': instanceid
                        }
                    ]
                )

                logger.info("ELB Deregistration HTTP Response: %s" % elb_response[u'ResponseMetadata'][u'HTTPStatusCode'])

                for timer in range(0, 24):

                    elb_response = elb_client.describe_instance_health(
                        LoadBalancerName=elb,
                        Instances=[
                            {
                            'InstanceId': instanceid
                            }
                        ]
                    )

                    if elb_response[u'InstanceStates'][0][u'State'] == "InService":
                        logger.info("%s is still InService." % instanceid)
                        logger.info("Description: %s" % elb_response[u'InstanceStates'][0][u'Description'])
                        time.sleep(5)
                    else:
                        logger.info("%s is no longer InService. It has entered State: %s." % (instanceid, elb_response[u'InstanceStates'][0][u'State']))
                        logger.info("%s successfully deregistered from ELB." % instanceid)
                        return

                else:
                    logger.info("Whoops! %s took more than 2 minutes to deregister from the ELB!" % instanceid)
                    return

        logger.info("%s is NOT registered with ELB, moving on." % instanceid)
        return

    except Exception as e:
        logger.error(e)
        notify_on_error(str(e))


def complete_asg_lifecycle(hookname, asg, actiontoken, instanceid):
    
    try:

        asg_client = boto3.client('autoscaling')

        asg_response = asg_client.complete_lifecycle_action(
            LifecycleHookName=hookname,
            AutoScalingGroupName=asg,
            LifecycleActionToken=actiontoken,
            LifecycleActionResult='CONTINUE',
            InstanceId=instanceid
        )

        logger.info("ASG Complete Lifecycle Action Response: %s" % asg_response[u'ResponseMetadata'][u'HTTPStatusCode'])

    except Exception as e:
        logger.error(e)
        notify_on_error(str(e))


def lambda_handler(event, context):

    logger.info(json.dumps(event))
 
    message = json.loads(event[u'Records'][0][u'Sns'][u'Message'])
    
    logger.info(message)
    
    # Parse SNS message for required data.
    ec2instanceid = message['EC2InstanceId']
    stackname = message['NotificationMetadata']
    asgname = message['AutoScalingGroupName']
    lifecycleactiontoken = message['LifecycleActionToken']
    lifecyclehookname = message['LifecycleHookName']

    logger.info("EC2 Instance ID: %s" % ec2instanceid)
    logger.info("CF Stack Name: %s" % stackname)

    # Create Cloudformation connection object.
    cf_client = boto3.client('cloudformation')

    # Get stack details via describe_stacks
    stackdetails = cf_client.describe_stacks(StackName=stackname)
    
    # Initialize ELB var to None.
    elbname = None
    
    # Parse the ELB Name and ECS Cluster Name.
    for output in stackdetails[u'Stacks'][0][u'Outputs']:

        if output[u'OutputKey'] == 'elb':
            logger.info("ELB NAME: %s" % output[u'OutputValue'])
            elbname = output[u'OutputValue']

        if output[u'OutputKey'] == 'ecscluster':
            logger.info("ECS CLUSTER NAME: %s" % output[u'OutputValue'])
            ecsclustername = output[u'OutputValue']

    # Deregister the EC2 instance from the ECS Cluster to
    # prevent new tasks from being launched on it.
    deregister_from_ecs_cluster(ec2instanceid, ecsclustername)

    # Deregister the EC2 instance from the ELB to initiate
    # connection draining before allowing the ASG to terminate.
    if elbname:
        deregister_from_elb(ec2instanceid, elbname)
    else:
        logger.info("No ELB found in stackdetails; skipping ELB de-registration.")
    
    # Notify ASG to complete lifecycle; don't wait for timeout.
    complete_asg_lifecycle(lifecyclehookname, asgname, lifecycleactiontoken, ec2instanceid)