AWS Step Functions: Creating a ‘Busy Waiting’ flow to wait for successful lambda executions.

Hi everyone, welcome to this article. In the past few days, I had to create a logic in AWS Step functions which involved waiting for a lambda that queries athena before going forward in the flow.

When I was discussing the logic, it felt like an easy thing to accomplish in minutes but I was surprised upon knowing that Step Functions does not have its native technique/solve to wait for a state to succeed before starting another. This made me search for hours before coming up with a solution.

Here, I will share with you a technique which will help you in the future.

First of all, I will create a simple lambda function. All we need to know from the lambda is whether it did its job. So, this simple function will make sure to return its status.

Lambda function:

import json

def lambda_handler(event, context):
    response = {
        'status': 'success',
        'message': 'Hello, World!'
    }

    return response

You can copy both the function and the Step function definition which I will share below to replicate all of this in your own environment.

Now, lets go ahead and create a state machine in AWS Step Functions.

Click on ‘Create state machine’.

Choose the option to write your own definition and choose ‘standard’.

Now, in the definition section, copy and paste the following definition.

Step Functions definition:

{
  "Comment": "A Hello World example of the Amazon States Language using Pass states",
  "StartAt": "Main",
  "States": {
    "Main": {
      "Type": "Task",
      "Resource": "arn:aws:states:::lambda:invoke",
      "OutputPath": "$.Payload",
      "Parameters": {
        "Payload.$": "$",
        "FunctionName": "arn:aws:lambda:us-east-1:031342435657:function:stepfunctionlambda_test:$LATEST"
      },
      "Retry": [
        {
          "ErrorEquals": [
            "Lambda.ServiceException",
            "Lambda.AWSLambdaException",
            "Lambda.SdkClientException",
            "Lambda.TooManyRequestsException"
          ],
          "IntervalSeconds": 2,
          "MaxAttempts": 6,
          "BackoffRate": 2
        }
      ],
      "Next": "WaitforInvocation"
    },
    "WaitforInvocation": {
      "Type": "Wait",
      "Seconds": 2,
      "Next": "CheckInvocationsStatus"
    },
    "CheckInvocationsStatus": {
      "Type": "Choice",
      "Choices": [
        {
          "Variable": "$.status",
          "StringEquals": "success",
          "Next": "Final"
        }
      ],
      "Default": "WaitforInvocation"
    },
    "Final": {
      "Type": "Succeed"
    }
  }
}

Save the definition and click next.

Specify the name of your state machine and then leave everything else as default. A new execution role will be created if you don’t specify any.

Save it and then you will have created a state machine which looks like this:

Visual flow:

To explain the definition above, I have pointed out how this state machine works:

1. The state machine starts at the "Main" state.

2. The "Main" state is a Task state that invokes a Lambda function specified by the "FunctionName" parameter. The input to the Lambda function is the entire input passed to the state machine.

3. The "OutputPath" parameter is set to "$.Payload" to capture the output of the Lambda function.

4. If the Lambda function invocation fails with certain exceptions specified in the "Retry" section (Lambda.ServiceException, Lambda.AWSLambdaException, Lambda.SdkClientException, Lambda.TooManyRequestsException), the state machine will retry the invocation up to 6 times with a backoff interval of 2 seconds.

5. After invoking the Lambda function, the state machine proceeds to the "WaitforInvocation" state.

6. The "WaitforInvocation" state is a Wait state that pauses the execution for 2 seconds.

7. After the wait period, the state machine goes to the "CheckInvocationsStatus" state.

8. The "CheckInvocationsStatus" state is a Choice state that checks the value of the "$.status" field in the output of the Lambda function. If the status is "success", the state machine proceeds to the "Final" state. Otherwise, it goes back to the "WaitforInvocation" state.

9. The "Final" state is a Succeed state that indicates the successful completion of the state machine.

   Now, lets go ahead and execute it!

Click on ‘Start execution’ and leave the input field empty.

Incase of a lambda that requires input, you can insert your input here as a Json

After you start your execution, each state will turn into different colours which indicates the following.

• Green: Indicates successful completion of a state or step.
• Red: Indicates a failure or error occurred during the execution of a state or step.
• Yellow: Indicates a state is in a waiting or paused state.
• Blue: Represents a decision or choice point where the state machine evaluates conditions to determine the next state.
• Gray: Represents a passive state that simply passes input to output without performing significant work.
• White: Represents transitions or connections between states, indicating the flow of execution from one state to another.

Finally this is how the execution will look like:

Thank you. Hope you can now go ahead and apply the “Busy Waiting” logic in your step functions.