An AWS step function to register and ping multiple virtual Buildkite agents.

buildkite_agent_step_function.json

{
	"StartAt": "RegisterBuildkiteAgents",
	"TimeoutSeconds": 1800,
	"States": {
		"RegisterBuildkiteAgents": {
			"Type": "Pass",
			"Result": {
				"TaskDefinitions": {
					"Definitions": [],
					"Count": 0
				},
				"Agents": {
					"Connected": [],
					"Count": 0
				},
				"Poll": {
					"PingRemaining": "15",
					"Status": "Connecting"
				}
			},
			"Next": "DiscoverBuildkiteAgentTaskDefinitions"
		},
		"DiscoverBuildkiteAgentTaskDefinitions": {
			"Comment": "Find ECR images that can act as buildkite agents.",
			"Type": "Task",
			"Resource": "arn:aws:lambda:ap-southeast-2:358740724786:function:DiscoverBuildkiteAgentTaskDefinitions",
			"Retry": [
				{
					"ErrorEquals": [ "States.TaskFailed" ],
					"IntervalSeconds": 1,
					"MaxAttempts": 5,
					"BackoffRate": 2
				}
			],
			"ResultPath": "$.TaskDefinitions",
			"Next": "CheckTaskDefinitionsFinished"	
		},
		"CheckTaskDefinitionsFinished": {
			"Comment": "Check length of task definitions.",
			"Type": "Choice",
			"Choices": [
				{
					"Variable": "$.TaskDefinitions.Count",
					"NumericGreaterThan": 0,
					"Next": "RegisterBuildkiteAgent"
				}
			],
			"Default": "StartPollingBuildkiteAgents"
		},
		"RegisterBuildkiteAgent": {
			"Comment": "Register an agent using the Buildkite agent API.",
			"Type": "Task",
			"Resource": "arn:aws:lambda:ap-southeast-2:358740724786:function:RegisterBuildkiteAgent",
			"InputPath": "$.TaskDefinitions.Definitions[0]",
			"Parameters": {
				"TaskDefinitionFamily.$": "$.TaskDefinitionFamily",
				"AgentTags.$": "$.AgentTags"
			},
			"Retry": [
				{
					"ErrorEquals": [ "States.TaskFailed" ],
					"IntervalSeconds": 1,
					"MaxAttempts": 5,
					"BackoffRate": 2
				}
			],
            "Catch": [
              {
                "ErrorEquals": [ "States.ALL" ],
                "ResultPath": null,
                "Next": "PopTaskDefinition"
              }
            ],
			"ResultPath": "$.Agents.Current",
            "Next": "PutBuildkiteAgentAccessToken"
		},
		"PutBuildkiteAgentAccessToken": {
			"Comment": "Store the input agent's access token in SSM ParameterStore /buildkite/agent/${Id}/access-token.",
			"Type": "Task",
			"Resource": "arn:aws:lambda:ap-southeast-2:358740724786:function:PutBuildkiteAgentAccessToken",
			"InputPath": "$.Agents.Current.Config",
			"Retry": [
				{
					"ErrorEquals": [ "States.TaskFailed" ],
					"IntervalSeconds": 1,
					"MaxAttempts": 5,
					"BackoffRate": 1.2
				}
			],
			"ResultPath": null,
			"Catch": [
				{
					"ErrorEquals": [ "States.ALL" ],
					"ResultPath": null,
					"Next": "PopTaskDefinition"
				}
			],
			"Next": "ConnectBuildkiteAgent"
		},
		"ConnectBuildkiteAgent": {
			"Type": "Task",
			"Comment": "Connect the given agent using the Buildkite agent API.",
			"Resource": "arn:aws:lambda:ap-southeast-2:358740724786:function:ConnectBuildkiteAgent",
			"InputPath": "$.Agents.Current.Config",
			"ResultPath": null,
			"Retry": [
				{
					"ErrorEquals": [ "States.TaskFailed" ],
					"IntervalSeconds": 1,
					"MaxAttempts": 3,
					"BackoffRate": 1.2
				}
			],
			"Catch": [
				{
					"ErrorEquals": [ "States.ALL" ],
					"ResultPath": null,
					"Next": "PopTaskDefinition"
				}
			],
			"Next": "PushBuildkiteAgent"
		},
		"PushBuildkiteAgent": {
			"Comment": "Move and remove 'Current' to 'Connected' array, update 'Count'.",
			"Type": "Task",
			"Resource": "arn:aws:lambda:ap-southeast-2:358740724786:function:PushBuildkiteAgent",
			"InputPath": "$.Agents",
			"ResultPath": "$.Agents",
			"Retry": [
				{
					"ErrorEquals": [ "States.TaskFailed" ],
					"IntervalSeconds": 1,
					"MaxAttempts": 5,
					"BackoffRate": 2
				}
			],
			"Catch": [
				{
					"ErrorEquals": [ "States.ALL" ],
					"ResultPath": null,
					"Next": "PopTaskDefinition"
				}
			],
			"Next": "PopTaskDefinition"
		},
		"PopTaskDefinition": {
			"Comment": "Decrement count of top task definition if >0, pop from array if 1. Set $.TaskDefinitions.Count to the remaining count.",
			"Type": "Task",
			"Resource": "arn:aws:lambda:ap-southeast-2:358740724786:function:PopTaskDefinition",
			"Retry": [
				{
					"ErrorEquals": [ "States.TaskFailed" ],
					"IntervalSeconds": 1,
					"MaxAttempts": 5,
					"BackoffRate": 2
				}
			],
			"Catch": [
				{
					"ErrorEquals": [ "States.ALL" ],
					"ResultPath": null,
					"Next": "StopPollingBuildkiteAgents"
				}
			],
			"InputPath": "$.TaskDefinitions",
			"ResultPath": "$.TaskDefinitions",
			"Next": "CheckTaskDefinitionsFinished"
		},
		"StartPollingBuildkiteAgents": {
			"Comment": "Ping the agents, sleep for 60s between pings, loop until all agents are gone.",
			"Type": "Pass",
			"Result": "Connected",
			"ResultPath": "$.Poll.Status",
			"Next": "CheckAgentExit"
		},
		"CheckAgentExit": {
			"Type": "Choice",
			"Choices": [
				{
					"Variable": "$.Agents.ConnectedCount",
					"NumericEquals": 0,
					"Next": "Succeed"
				},
				{
					"And": [
						{
							"Variable": "$.Poll.PingRemaining",
							"NumericEquals": 0
						},
						{
							"Variable": "$.Poll.Status",
							"StringEquals": "Connected"
						}
					],
					"Next": "StopPollingBuildkiteAgents"
				}
			],
			"Default": "DecrementPingRemaining"
		},
		"DecrementPingRemaining": {
			"Type": "Task",
			"Resource": "arn:aws:lambda:ap-southeast-2:358740724786:function:DecrementPingRemaining",
			"InputPath": "$.Poll.PingRemaining",
			"ResultPath": "$.Poll.PingRemaining",
			"Retry": [
				{
					"ErrorEquals": [ "States.TaskFailed" ],
					"IntervalSeconds": 1,
					"MaxAttempts": 5,
					"BackoffRate": 2
				}
			],
			"Next": "Sleep"
		},
		"PingBuildkiteAgents": {
			"Type": "Pass",
			"Parameters": {
				"Pending.$": "$.Agents.Connected",
				"PendingCount.$": "$.Agents.ConnectedCount"
			},
			"ResultPath": "$.Agents",
			"Next": "CheckPingBuildkiteAgentsFinished"
		},
		"CheckPingBuildkiteAgentsFinished": {
			"Type": "Choice",
			"Choices": [
				{
					"Variable": "$.Agents.PendingCount",
					"NumericGreaterThan": 0,
					"Next": "PingBuildkiteAgent"
				}
			],
			"Default": "CheckAgentExit"
		},
		"PingBuildkiteAgent": {
			"Comment": "Ping the given agent using the Buildkite agent API.",
			"Type": "Task",
			"Resource": "arn:aws:lambda:ap-southeast-2:358740724786:function:PingBuildkiteAgent",
			"InputPath": "$.Agents.Pending[0].Config",
			"ResultPath": "$.Agents.Pending[0].Action",
			"Retry": [
				{
					"ErrorEquals": [ "States.TaskFailed" ],
					"IntervalSeconds": 1,
					"MaxAttempts": 5,
					"BackoffRate": 1
				}
			],
			"Catch": [
				{
					"ErrorEquals": [ "States.ALL" ],
					"ResultPath": null,
					"Next": "DisconnectPendingBuildkiteAgent"
				}
			],
			"Next": "PingAction"
		},
		"PingAction": {
			"Type": "Choice",
			"InputPath": "$.Agents.Pending[0].Action",
			"Choices": [
				{
					"Variable": "$",
					"StringEquals": "Disconnect",
					"Next": "DisconnectPendingBuildkiteAgent"
				}
			],
			"Default": "SavePendingBuildkiteAgent"
		},
		"DisconnectPendingBuildkiteAgent": {
			"Type": "Task",
			"Resource": "arn:aws:lambda:ap-southeast-2:358740724786:function:DisconnectBuildkiteAgent",
			"InputPath": "$.Agents.Pending[0].Config",
			"ResultPath": null,
			"Retry": [
				{
					"ErrorEquals": [ "States.TaskFailed" ],
					"IntervalSeconds": 1,
					"MaxAttempts": 5,
					"BackoffRate": 1
				}
			],
			"Catch": [
				{
					"ErrorEquals": [ "States.ALL" ],
					"ResultPath": null,
					"Next": "DeletePendingBuildkiteAgentAccessToken"
				}
			],
			"Next": "DeletePendingBuildkiteAgentAccessToken"
		},
		"DeletePendingBuildkiteAgentAccessToken": {
			"Type": "Task",
			"Resource": "arn:aws:lambda:ap-southeast-2:358740724786:function:DeleteBuildkiteAgentAccessToken",
			"InputPath": "$.Agents.Pending[0].Config.Id",
			"ResultPath": null,
			"Retry": [
				{
					"ErrorEquals": [ "States.TaskFailed" ],
					"IntervalSeconds": 1,
					"MaxAttempts": 5,
					"BackoffRate": 2
				}
			],
			"Catch": [
				{
					"ErrorEquals": [ "States.ALL" ],
					"ResultPath": null,
					"Next": "PopPendingBuildkiteAgent"
				}
			],
			"Next": "PopPendingBuildkiteAgent"
		},
		"PopPendingBuildkiteAgent": {
			"Comment": "Drop the head of Pending",
			"Type": "Task",
			"Resource": "arn:aws:lambda:ap-southeast-2:358740724786:function:PopPendingBuildkiteAgent",
			"InputPath": "$.Agents",
			"ResultPath": "$.Agents",
			"Retry": [
				{
					"ErrorEquals": [ "States.TaskFailed" ],
					"IntervalSeconds": 1,
					"MaxAttempts": 5,
					"BackoffRate": 1
				}
			],
			"Next": "CheckPingBuildkiteAgentsFinished"
		},
		"SavePendingBuildkiteAgent": {
			"Comment": "Move the head of Pending into Connected",
			"Type": "Task",
			"Resource": "arn:aws:lambda:ap-southeast-2:358740724786:function:SavePendingBuildkiteAgent",
			"InputPath": "$.Agents",
			"ResultPath": "$.Agents",
			"Retry": [
				{
					"ErrorEquals": [ "States.TaskFailed" ],
					"IntervalSeconds": 1,
					"MaxAttempts": 5,
					"BackoffRate": 1
				}
			],
			"Next": "CheckPingBuildkiteAgentsFinished"
		},
		"StopPollingBuildkiteAgents": {
			"Type": "Pass",
			"Result": "Stopping",
			"ResultPath": "$.Poll.Status",
			"Next": "StopBuildkiteAgents"
		},
		"StopBuildkiteAgents": {
			"Comment": "Send the stop command to the agents using the Buildkite API.",
			"Type": "Task",
			"Resource": "arn:aws:lambda:ap-southeast-2:358740724786:function:StopBuildkiteAgents",
			"Retry": [
                {
                  "ErrorEquals": [ "States.TaskFailed" ],
				  "IntervalSeconds": 1,
    			  "MaxAttempts": 5,
	  			  "BackoffRate": 1.5
  			    }
            ],
			"InputPath": "$.Agents.Connected[*].Config.Id",
			"ResultPath": null,
			"Next": "CheckAgentExit"
		},
		"Sleep": {
			"Type": "Wait",
			"Seconds": 60,
			"Next": "PingBuildkiteAgents"
		},		
		"Succeed": {
			"Type": "Succeed"
		}
	}
}