How to Implement a Nested While Loop in AWS Step Functions

Introduction

About one and a half year has been passed since AWS Step Functions was introduced in November 2016. It improves the possibility of Serverless applications. Previously, when we, the developers, tried to build a complicated business logic in which many AWS Lambda functions interacted with each other, we needed to implement that logic by ourselves. However, by using AWS Step Functions, we can implement it with a simple JSON document (Amazon States language). The service serves the framework for building a complicated logic as simple as possible in the world of Serverless. It's very useful! I use it every day in our project.

In this post, I'd like to give you tips about the service particularly how to loop based on given boolean conditions. In other words, I'd like to introduce how to implement a nested while loop in AWS Step Functions.

Why do you want to talk about a nested loop?

Because it's useful but there is less information about it on the Internet! The Official AWS document explains how to implement a loop in AWS Step Functions, but it's about a flattened loop, not a nested one.

Anyway, let's look at this picture below:

This is the state machine of AWS Step Functions working in our project. We make it continue alerting while the given conditions are evaluating to true.

In the state machine, there are two concepts: notification counts, how many times the state machine should iterate over; and notification user counts, how many users it should notify.

The notification count is evaluated at the IsNotificationCountReached state. If the result of the evaluation is true, then the loop continues, otherwise it ends. The notification user count is evaluated at the IsNotificationUserCountReached state. If the result of the evaluation is true, then the loop transitions to the NotifyAlerts task, otherwise to the IsNotificationCountReached choice.

The point is, we make AWS Step Functions notify alerts to users to maximum notification counts.

How We Implement the Loop in AWS Step Functions

Let's take a closer look at the state machine.

The FetchNotificationCount Task

The function takes the initial values of the loop from AWS DynamoDB and pass them into following states. For example, if there are the values in a variable (e.g, item ), the function should return them like this:

return {
    'notification_count': item['notification_count'],
    'notification_user_accounts': item['user_accounts'],
    'notification_user_count': len(item['user_accounts']),
    'is_first_alert': True
}

Then, following states can receive the data like this:

{
  "input": {
    "notification_count": 3,
    "notification_user_accounts": {
      "1": "foo@example.com",
      "2": "bar@example.com",
      "3": "baz@example.com"
    },
    "notification_user_count": 3,
    "is_first_alert": true
  }
}

AWS Lambda can reference the data with an event variable, or choice states can do that with a $ path.

The IsNotificationCountReached Choice

This is the condition evaluating how many times the loop iterates over. The time is saved in the notification_count variable, so we implement the logic below:

"IsNotificationCountReached": {
  "Type": "Choice",
  "Choices": [
    {
      "Variable": "$.notification_count",
      "NumericGreaterThan": 0,
      "Next": "IsFirstAlert"
    }
  ],
  "Default": "Done"
},

In this case, if the notification_count is greater than 0, the state machine transitions to next states. If not, it finishes its loop.

The IsFirstAlert Choice and the Sleep Wait

We make the state machine sleep for 5 minutes every loop except for the first loop. This prevents "noise alerts". We implement the logic like this:

"IsFirstAlert": {
  "Type": "Choice",
  "Choices": [
    {
      "Variable": "$.is_first_alert",
      "BooleanEquals": true,
      "Next": "DecrementNotificationUserCount"
    }
  ],
  "Default": "Sleep"
},
"Sleep": {
  "Type": "Wait",
  "Seconds": 300,
  "Next": "DecrementNotificationUserCount"
},

The state language expresses that if the is_first_alert variable is false, then it sleeps. If not, it goes to the next state immediately.

The DecrementNotificationUserCount Task

The main features of the task subtract 1 from the notification_user_count variable pass by preceding states and change the is_first_alert value to false. The code can be like this:

return {
    'notification_count': event['notification_count'],
    'notification_user_accounts': event['notification_user_accounts'],
    'notification_user_count': len(event['notification_user_accounts']) - 1,
    'is_first_alert': False
}

Following states can receive the data below:

{
  "input": {
    "notification_count": 3,
    "notification_user_accounts": {
      "1": "foo@example.com",
      "2": "bar@example.com",
      "3": "baz@example.com"
    },
    "notification_user_count": 2,
    "is_first_alert": false
  }
}

In fact, it's not necessary to implement those features in the task (e.g, the NotifyAlerts task can take the role). But, we prefer a simple function to a complicated one. So, for the present, we design the state machine like that.

The NotifyAlerts Task and the IsNotificationUserCountReached Choice

The NotifyAlerts task notifies alerts and the IsNotificationUserCountReached choice manages how many times the task notifies the users.

Let's look at the main features of the task. We make the task alert to Slack, but we will not touch on this matter. The task subtracts 1 from the notification_user_count variable every time when it notifies. If the value is -1, then the task recognizes all notifications have been sent to users and subtracts 1 from the notification_count variable. The main code can be like this:

notification_user_count = event['notification_user_count']
notification_count = event['notification_count']

if notification_user_count == -1:
    notification_count -= 1

output = event
output['notification_count'] = notification_count
output['notification_user_count'] = notification_user_count - 1

return output

Next, the IsNotificationUserCountReached choice. The notification_user_count variable has the information on which how many users the state machine should notify, so we implement the logic using the variable like this:

"IsNotificationUserCountReached": {
  "Type": "Choice",
  "Choices": [
    {
      "Variable": "$.notification_user_count",
      "NumericGreaterThanEquals": 0,
      "Next": "NotifyOverflow"
    }
  ],
  "Default": "IsNotificationCountReached"
},

The Whole of the State Machine

This is the state language described above. ARNs of AWS Lambda are masked.

{
  "Comment": "the state machine notifying alerts",
  "StartAt": "FetchNotificationCount",
  "States": {
    "FetchNotificationCount": {
      "Type": "Task",
      "Resource": "<aws-lambda-arn>",
      "Next": "IsNotificationCountReached"
    },
    "IsNotificationCountReached": {
      "Type": "Choice",
      "Choices": [
        {
          "Variable": "$.notification_count",
          "NumericGreaterThan": 0,
          "Next": "IsFirstAlert"
        }
      ],
      "Default": "Done"
    },
    "IsFirstAlert": {
      "Type": "Choice",
      "Choices": [
        {
          "Variable": "$.is_first_alert",
          "BooleanEquals": true,
          "Next": "DecrementNotificationUserCount"
        }
      ],
      "Default": "Sleep"
    },
    "Sleep": {
      "Type": "Wait",
      "Seconds": 300,
      "Next": "DecrementNotificationUserCount"
    },
    "DecrementNotificationUserCount": {
      "Type": "Task",
      "Resource": "<aws-lambda-arn>",
      "Next": "IsNotificationUserCountReached"
    },
    "IsNotificationUserCountReached": {
      "Type": "Choice",
      "Choices": [
        {
          "Variable": "$.notification_user_count",
          "NumericGreaterThanEquals": 0,
          "Next": "NotifyOverflow"
        }
      ],
      "Default": "IsNotificationCountReached"
    },
    "NotifyAlerts": {
      "Type": "Task",
      "Resource": "<aws-lambda-arn>",
      "Next": "IsNotificationUserCountReached"
    },
    "Done": {
      "Type": "Pass",
      "End": true
    }
  }
}

Conclusion

In this post, the implementation of a nested while loop in AWS Step Functions is discussed. I think it can be adopted in many situations. I hope you find this post helpful.