Cleanly scaling in instances that are running Sidekiq

instance-iam-policy.json

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Action": [
                "autoscaling:CompleteLifecycleAction",
                "autoscaling:RecordLifecycleActionHeartbeat"
            ],
            "Effect": "Allow",
            "Resource": "*"
        }
    ]
}

lambda-iam-policy.json

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "Stmt1509124372000",
            "Effect": "Allow",
            "Action": [
                "ssm:SendCommand"
            ],
            "Resource": [
                "arn:aws:ec2:*:*:instance/*",
                "arn:aws:ssm:*:*:document/PrepInstanceForTermination"
            ]
        }
    ]
}

lambda.py

import boto3
import json
import logging

logger = logging.getLogger()
logger.setLevel(logging.INFO)
ssm_client = boto3.client("ssm")

def handler(event, context):
    logger.info(json.dumps(event))
    response = ssm_client.send_command(
        InstanceIds = [ event['detail']['EC2InstanceId'] ], 
        DocumentName = 'PrepInstanceForTermination',
        TimeoutSeconds = 120
        )
    logger.info("Command sent: %s", response)

prep_for_termination.sh

#!/bin/bash

INSTANCE_ID=$(ec2metadata --instance-id)
REGION=us-east-1
ASG_NAME=$(aws ec2 describe-tags --filters "Name=resource-id,Values=${INSTANCE_ID}" "Name=key,Values=aws:autoscaling:groupName" --region ${REGION} | jq '.Tags[0].Value' -r)
CMD_OPTS="--region ${REGION} --auto-scaling-group-name ${ASG_NAME} --lifecycle-hook-name prep-for-termination --instance-id ${INSTANCE_ID}"

if aws autoscaling record-lifecycle-action-heartbeat $CMD_OPTS; then
  su -c "cd ~/current && RAILS_ENV=production ./bin/rails runner ./script/stop_sidekiq_workers.rb" honeybadger
  aws autoscaling complete-lifecycle-action $CMD_OPTS --lifecycle-action-result CONTINUE
fi

ssm-doc.json

{
    "schemaVersion":"1.2",
    "description": "Do cleanup tasks before shutting down an instance",
    "parameters": {},
    "runtimeConfig":{
        "aws:runShellScript":{
            "properties":[
                {
                    "id":"0.aws:runShellScript",
                    "runCommand":["/sbin/prep_for_termination.sh"],
                    "workingDirectory":"/",
                    "timeoutSeconds":"600"
                }
            ]
        }
    }
}

stop_sidekiq_workers.rb

#!/usr/bin/env ruby

hostname = Socket.gethostname
ps = Sidekiq::ProcessSet.new

puts "Stopping workers"
ps.select {|p| p['hostname'] == hostname }.map(&:quiet!)

while ps.select {|p| p['hostname'] == hostname }.map {|p| p['busy'] }.any?(&:nonzero?) do
  putc "."
  sleep 1
end

puts ""