Scheduling managed image updates for AppStream 2.0

Amazon AppStream 2.0 has released the Managed Image updates feature. This feature applies the AppStream 2.0 component and Windows updates to your existing AppStream 2.0 images with a single operation. Our previous blog covers the feature details. Customers ask, how can we schedule this?

In this blog, we describe how to schedule your AppStream 2.0 image updates, and apply them in batches.

Overview

We can use EventBridge rules to start a Step Functions workflow, which launches a Lambda function. The Lambda evaluates if the AppStream 2.0 image is tagged for update. Lambda calls the Create_Updated_Image API for each image up to the BatchSize defined in the EventBridge rule. The names of the images pending, and the images beyond the BatchSize limit are passed back to the Step Functions workflow. After a wait period, the Step Functions workflow will launch the Lambda function again to check the status of pending updates, and start the update on any waiting images. Amazon Simple Email Service is used to message when the process is completed, or if an error occurs.

Walkthrough

In this walk-through, you complete the following tasks:

Create an IAM Policy
Create a Role for the AWS Lambda function.
Create an AWS Lambda function.
Create an AWS Step Functions state machine.
Create Amazon EventBridge rule to start the Lambda function.

Prerequisites

Make sure you meet the following requirements before getting started:

Step 1. Create the IAM Policy

In this step, you create an IAM Policy, and attach it to an IAM Role that the Lambda function can assume.

Navigate to the IAM console.
In the navigation pane, choose Policies.
Choose Create policy.
Choose the JSON tab.
Copy and paste the following JSON policy.
When you’re done, choose Review policy.
Enter a name of your choosing.
Choose Create policy.

IAM Policy document example:

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "ses:SendEmail",
                "appstream:UntagResource",
                "appstream:DescribeImages",
                "appstream:CreateUpdatedImage",
                "appstream:TagResource",
                "appstream:ListTagsForResource",
                "logs:CreateLogGroup",
                "logs:CreateLogStream",
                "logs:PutLogEvents"
            ],
            "Resource": "*"
        }
    ]
}

Step 2. Create the IAM Role

Now that the IAM Policy has been created, create the IAM Role for Lambda to assume and attach the policy you created in Step 1.

Open the IAM console.
In the navigation pane, choose Roles.
Choose Create role.
For Select type of trusted entity, keep AWS service selected.
Choose Lambda, and then choose Next: Permissions.
In the filter policies search box, type name of the policy created in the previous step. When the policy appears in the list, select the check box next to the policy name.
Choose Next: Tags. Although you can specify a tag for the policy, a tag is not required.
Choose Next: Review.
Enter a name for your Role to help you identify it.
Choose Create role.

Step 3. Create a Lambda function

In this step, create your Lambda function that lists the images, creates image update requests, and checks the status of existing update requests. Modify the functions timeout.

Open the Lambda console
Choose Create function.
Enter a meaningful name in Function name.
Select Python 3.8 as the Runtime.
Expand the permissions section, select Use an existing role, and from the list choose the role created in step 1.
Choose Create function.
Under the Function code section, choose lambda_function.py from the file tree.
Replace the placeholder code with the following code.
Choose Deploy.
Choose Configuration.
Choose General configuration, and then choose Edit.
Set the Timeout to 1 minute.
Choose Save.
Make a note of the Lambda Arn for the step 4.

Lambda function example:

import sys
from pip._internal import main
main(['install', '-I', '-q', 'boto3', '--target', '/tmp/', '--no-cache-dir', '--disable-pip-version-check'])
sys.path.insert(0,'/tmp/') 
import logging
from boto3 import client
from botocore import config
from botocore import exceptions
from datetime import datetime

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

config = config.Config(
   retries = {
      'max_attempts': 10,
      'mode': 'standard'
   })

as2 = client('appstream', config=config)
ses = client('ses')

def checkStatus(ImageName):
    logger.info("Checking status of image: " + ImageName)
    try:
        status = as2.describe_images(Names=[ImageName]).get('Images')
        if status[0].get('State') == 'AVAILABLE':
            logger.info("Image Creation succeed for image: " + ImageName)
            return 'Successful'
        elif status[0].get('State') == 'CREATING':
            logger.info("Image Creation is pending image: " + ImageName)
            return 'Pending'
        else:
            logger.error("Image Creation has failed for image: " + ImageName)
            return 'Failed'
    except:
        logger.error("Image Creation has failed for image: " + ImageName)
        return 'Failed'

def CreateUpdate(imageName, dateStr):
    logger.info("Creating Update Image Request for Image: " + imageName)
    try:
        image = as2.describe_images(Names=[imageName]).get('Images')[0]
        tags = as2.list_tags_for_resource(ResourceArn=image.get('Arn')).get('Tags')
        if image.get('Description'):
            descrip = image.get('Description')
        else:
            descrip = ''
        if tags.get('BaseName'):
            baseName = tags.get('BaseName')
        else:
            baseName = image.get('Name')
        tags['BaseName'] = baseName
        newImage = as2.create_updated_image(
            existingImageName=image.get('Name'),
            newImageName=baseName + '-Updated-' + dateStr,
            newImageDescription=descrip,
            newImageDisplayName=baseName + ' Updated: ' + dateStr,
            newImageTags=tags,
            dryRun=False
        )
        tags['AutoUpdate'] = 'Completed'
        tags['NextVersion'] = newImage.get('image').get('Name')
        as2.tag_resource(
            ResourceArn=image.get('Arn'),
            Tags=tags
        )
        logger.info("Updated Image: " + newImage.get('image').get('Name'))
        return newImage.get('image').get('Name')
            
    except exceptions.ClientError as err:
        logger.error(err)
        return err

    except exceptions.ParamValidationError as err:
        logger.error(err)
        return err

def send_notification(message, toAddress, fromAddress, messageType):
    if messageType == 'Error':
        logger.error('Notifying: ' + toAddress + ' of failure')
        SubjectStr = 'Attention Required: AppStream 2.0 Image Update Failed'
        MessageTxt = "The following error has occurred while updating an AppStream 2.0 Image: " + str(message)
        MessageHTML = "<html><body>The following error has occurred while updating an AppStream 2.0 Image:<p>" + str(message) + "</body></html>"
    else:
        logger.info('Notifying: ' + toAddress + ' of success')
        SubjectStr = 'AppStream 2.0 Image Update Successful'
        MessageTxt = "The following AppStream 2.0 Images have been updated successfully: " + str(message)
        MessageHTML = "<html><body>The following AppStream 2.0 Images have been updated successfully:<p>" + str(message) + "</body></html>"
    sentEmail = ses.send_email(
        Source=fromAddress,
        Destination={
            'ToAddresses': [
                toAddress
            ]
        },
        Message={
            'Subject': {
                'Data': SubjectStr,
                'Charset': 'UTF-8'
            },
            'Body': {
                'Text': {
                    'Data': MessageTxt,
                    'Charset': 'UTF-8'
                },
                'Html': {
                    'Data': MessageHTML,
                    'Charset': 'UTF-8'
                }
            }
        }
    )

def lambda_handler(event, context):
    dateObj = datetime.today()
    dateStr = str(dateObj.day) + str(dateObj.month) + str(dateObj.year)
    counter = 0
    batchSize = event.get('BatchSize', 5)
    toAddress = event.get('ToAddress')
    fromAddress = event.get('FromAddress')
    if event.get('Results'):
        logger.info("Checking Image Update Status.")
        pending = event.get('Results').get('Pending', [])
        waiting = event.get('Results').get('Waiting', [])
        successful = event.get('Results').get('Success', [])
        cycleCount = event.get('Results').get('CycleCount', 0)
        if pending:
            for imageName in pending:
                status = checkStatus(imageName)
                if status == 'Pending':
                    stillpending = True
                else:
                    if status == 'Successful':
                        successful.append(imageName)
                    else:
                        send_notification(imageName, toAddress, fromAddress, 'Error')
                    pending.remove(imageName)
            cycleCount += 1

        if not pending:
            cycleCount = 0
            if waiting:
                for imageName in waiting:
                    if counter < batchSize:
                        newImageName = CreateUpdate(imageName, dateStr)
                        if type(newImageName) == str:
                            pending.append(newImageName)
                            counter += 1
                        else:
                            send_notification(newImageName, toAddress, fromAddress, 'Error')
                        waiting.remove(imageName)
            else:
                send_notification(successful, toAddress, fromAddress, 'Success')

    else:
        logger.info("Starting Image Update process.")
        pending = []
        waiting = []
        successful = []
        cycleCount = 0
        paginator = as2.get_paginator('describe_images')
        pages = paginator.paginate(Type='PRIVATE')
        images = as2.describe_images(Type='PRIVATE')
        for page in pages:
            for image in page.get('Images'):
                tags = as2.list_tags_for_resource(ResourceArn=image.get('Arn')).get('Tags')
                if tags.get('AutoUpdate') == 'Enabled' and image.get('State') == 'AVAILABLE':
                    if counter < batchSize:
                        newImageName = CreateUpdate(image.get('Name'), dateStr)
                        if type(newImageName) == str:
                            pending.append(newImageName)
                            counter += 1
                        else:
                            send_notification(newImageName, toAddress, fromAddress, 'Error')
                    else:
                        waiting.append(image.get('Name'))

    if not pending and not waiting:
        logger.info("No updates pending.")
        return 0
    else:
        logger.info("Returning pending updates.")
        return {
            "Pending" : pending,
            "Waiting" : waiting,
            "Success": successful,
            "CycleCount": cycleCount
        }

Step 4. Create state machine

In this step, create a state machine that coordinates the multiple executions of our Lambda function.

Open the AWS Step Functions console.
Do one of the following:
- If you haven’t created any state machines functions, a Getting Started page displays. Choose Getting Started, and then choose State Machines.
- If you have created a state machines function, in the upper right corner of the State machines page, choose Create state machine.
On the Define State Machine page, keep Author from code snippet selected.
In the Type section, keep Standard selected.
In the State machine definition section, delete the placeholder code and paste the code that follows in the edit window.
Replace <Lambda_Arn> in two places with the Arn recorded in step 3.
Choose Next.
For Name, enter a meaningful name for identification later.
Keep the rest of the options as default.
Choose Create state machine.

Step Function state machine example:

{
    "Comment": "State machine to wait for AppStream 2.0 Managed Image Update to complete and verify they all completed successfully.",
    "StartAt": "Start",
    "States": {
      "Start": {
        "Type": "Task",
        "ResultPath": "$.Results",
        "Resource": "<Lambda_Arn>",
        "Next": "StatusCheck"
      },
      "StatusCheck": {
        "Type": "Choice",
        "Choices": [
          {
            "Variable": "$.Results.Pending",
            "IsPresent": true,
            "Next": "wait"
          },
          {
            "Variable": "$.Results",
            "NumericEquals": 0,
            "Next": "Succeed"
          }
        ],
        "Default": "Failed"
      },
      "wait": {
        "Type": "Wait",
        "Seconds": 21600,
        "Next": "StatusCheckTask"
      },
      "StatusCheckTask": {
        "Type": "Task",
        "ResultPath": "$.Results",
        "Resource": "<Lambda_Arn>",
        "Next": "CountCheck"
      },
      "CountCheck": {
        "Type": "Choice",
        "Choices": [
          {
            "Variable": "$.Results.CycleCount",
            "NumericGreaterThan": 2,
            "Next": "Failed"
          }
        ],
        "Default": "StatusCheck"
      },
      "Succeed": {
        "Type": "Succeed"
      },
      "Failed": {
        "Type": "Fail"
      }
    }
  }

Step 5. Create Amazon EventBridge rule

In this step, create an EventBridge rule to execute the Lambda function we created. The rule triggers once a month, on the first of the month.

Open the EventBridge console.
Choose Create rule.
Enter a Name, and optionally a Description.
Select Schedule, and then Cron expression.
Enter the cron expression 0 12 1 * ? * into the Cron expression box.
Under Select targets, select Step Functions State Machine, then select state machine created in step 4.
Expand the Configure input, and select Constant (JSON text).
Enter JSON that follows into the Constant (JSON text) box.
Replace the following values:
- <Email_Address_for_notification> with the address notifications should be sent to.
- < Verified_SES_Sender> with a verified SES sender address.
Choose Create.

EventBridge input example:

{
    "BatchSize" : 5,
    "ToAddress" : "<Email_Address_for_notification>",
    "FromAddress" : "<Verified_SES_Sender>"
}

Clean up

To avoid incurring future charges, remove the resources that were created. Delete the EventBridge rule, Lambda function, state machine, IAM Policy and Role.

Conclusion

In this blog you created a scheduled workflow to create updated versions of AppStream 2.0 images on a monthly basis. The updated images are tagged with the AutoUpdate:Enabled key-value pair to ensure they are updated on the next EventBridge scheduled run. This ensures each month updated images will be ready for testing, and assignment to your AppStream 2.0 Fleets.

Desktop and Application Streaming