AWS Step Function to start a server and set DNS

In my previous post here I shared information on how to start an EC2 server and then set its DNS. This needed a host to run the python code with the appropriate IAM role. But what if you do not have a spare server to run python code. In this post, I will show you how to accomplish the task using AWS Lambda and Step Functions.

To accomplish this, I will define a workflow that will utilize two Lambda functions.

Note: The code described here is sampl code and not the complete listing. Please add error/exception handling as needed.

• First, let us create an IAM role for Step Functions and call it StepFunctionsLambdaRole. The role should allow Lambda Invocation.

• We will create two Lambda function, one to start our server and a second one to set its DNS. Refer here to create lambda functions.
• Let us call the first function start-server and the second one server-dns.
• Lambda Function start-server

import json, boto3

def start_server():
    id = ''
    ec2client = boto3.client('ec2')
    response = ec2client.describe_instances(
       Filters=[
          {
             'Name' : 'tag:Name', 'Values':[ '<name of your server>']
          },
          {
             'Name' : 'instance-state-name', 'Values':['stopped']
          }
         ]
    )
    for reservation in response["Reservations"]:
      for instance in reservation["Instances"]:
          id = instance["InstanceId"]
          ec2client.start_instances(InstanceIds=[id])          

    return id
    
def lambda_handler(event, context):
    id = start_server()
    return {
        'statusCode': 200,
        'body': id
    }

• Lambda function server-dns

import json, boto3

HOSTED_ZONE_ID = 'XXXXXXXXXX'
route53 = boto3.client('route53')
ec2client = boto3.client('ec2')

def server_dns(id):
   iresponse = ec2client.describe_instances(InstanceIds=[id])
   for ireservation in iresponse['Reservations']:
       for i in ireservation['Instances']:           
           print(i['PrivateIpAddress'])
           print(i['PublicIpAddress'])

   dns_changes = {
        'Changes': [
            {
                'Action': 'UPSERT',
                'ResourceRecordSet': {
                    'Name': '<your DNS record>',
                    'Type': 'A',
                    'ResourceRecords': [
                        {
                            'Value': i['PublicIpAddress']
                        }
                    ],
                    'TTL': 300
                }
            }
        ]
   }
   response = route53.change_resource_record_sets(
        HostedZoneId=HOSTED_ZONE_ID,
        ChangeBatch=dns_changes
   )
   return response['ChangeInfo']['Status']
  
def lambda_handler(event, context):
    if 'body' in event:
       id = event['body']
       if id:
          print(id)
          out = server_dns(id)
          return {
             'statusCode': 200,
             'body': json.dumps(out, indent=2)
          }

• Select Step Function from AWS Services and then select create State Machine

• Select ‘Author from scratch’ option and use the IAM role created above

• Modify the code as shown below

{
  "Comment": "Step function to start up server and set DNS",
  "StartAt": "StartServer",
  "States": {
    "StartServer": {
      "Type": "Task",
      "Resource": "arn:aws:lambda:aws region:account number:function:server-start",
      "Next": "ChoiceState"
    },
    "ChoiceState": {
      "Type" : "Choice",
      "Choices": [
        {
          "Variable": "$.body",
          "StringEquals": "",
          "Next": "EndState"
        }
      ],
      "Default": "WaitState"
    },
      "WaitState": {
      "Type": "Wait",
      "Seconds": 60,
      "Next": "FinalState"
    },
    "FinalState": {
      "Type": "Task",
      "Resource": "arn:aws:lambda:aws region:account number:function:server-dns",
      "End": true
    },
    "EndState": {
      "Type": "Pass",
      "End": true
    }
  }
}

A closer look at the state machine definition:

• The first Task is to invoke our lambda function to start our server.
• It is followed by a ChoiceState which decides what the next action will be performed next.
• In our lambda function, I am looking for a specific server in the stopped state. If none is found the id returned is blank in which case our state is to End.
• If our previous lambda function returned an id, the next step is a WaitState to allow our server to startup. It may take some time for the server to start and get a public IP. We must wait a few seconds before we try to get public ip for it.
• After the WaitState, we move on to the FinalState where we invoke our lambda function to set the DNS record.

You might ask why was there a need to split the task. In the previous post, the task was accomplished in one py execution. The answer is that AWS charges for Lambda based on function execution time. So if we were to add a sleep in the function we would be paying for it. Using Step function, we split the task and took the Sleep out and added it as a  WaitState. Step transitions are cheaper and every account gets 4000 transitions free.

Further reading and improvements:

• Instead of adding a WaitState, we could have our second function check if the server is running and if not exit and go back into a loop until we can get a public ip for it. If the output of start was not a success add a few seconds WaitState and call the function to try setting DNS again.
• Add error/exception handling in code.
• In this scenario, I am dealing with a fixed server name and DNS record. However, this can be easily modified so that the server name and DNS record to be set is passed as a parameter to the first function.
• https://docs.aws.amazon.com/lambda/latest/dg/intro-permission-model.html
• https://docs.aws.amazon.com/step-functions/latest/dg/tutorial-creating-lambda-state-machine.html
• https://docs.aws.amazon.com/step-functions/latest/dg/how-step-functions-works.html