Designing a CDK State Machine Builder - Part 2
Great on paper, terrible on grass?
In Part 1 of this series, I went through my process of designing an alternative API for defining state machines using CDK. In this part, I document my trials and tribulations of implementing that API.
All the code for this post can be found in this GitHub repo.
Edit: The resulting code is now available as an npm package.
One question I always ask myself before writing any code is how am I going to test it. That is, how can I have any confidence that the code is running as I expect? I might not take a full-blown Test-driven development (TDD) approach, but I need to have some sort of plan of how I am going to execute the code and verify the results. Ideally, this plan involves a straightforward way to both create and repeat those tests.
In Part 1, I created a set of examples covering different aspects of defining state machines using CDK. In Visualising a CDK State Machine using a custom Construct, I created a custom construct that outputs the graph JSON for such definitions. It seemed logical to me to combine the two; use the examples as test cases and compare the graph objects from the two implementations to verify. The result is shown below.
it('renders simple chain', async () => {
//
const cdkStateMachine = new StateMachineWithGraph(new cdk.Stack(), 'SimpleChain-CDK', {
getDefinition: (definitionScope): sfn.IChainable => {
//
const state1 = new sfn.Pass(definitionScope, 'State1');
// Define other states...
const definition = sfn.Chain.start(
state1.next(state2.next(state3.next(state4.next(state5.next(state6)))))
);
return definition;
},
});
const builderStateMachine = new StateMachineWithGraph(new cdk.Stack(), 'SimpleChain-Builder', {
getDefinition: (definitionScope): sfn.IChainable => {
//
const state1 = new sfn.Pass(definitionScope, 'State1');
// Define other states...
const definition = new StateMachineBuilder()
.perform(state1)
// Perform other states...
.build(definitionScope);
return definition;
},
});
const cdkGraph = JSON.parse(cdkStateMachine.graphJson);
const builderGraph = JSON.parse(builderStateMachine.graphJson);
expect(builderGraph).to.deep.equal(cdkGraph);
});
Running this test resulted in the following expected, but informative, failure.
AssertionError: expected { Object (StartAt, States) } to deeply equal { Object (StartAt, States) }
at ... <snip>
+ expected - actual
{
- "StartAt": "TODO"
+ "StartAt": "State1"
"States": {
- "TODO": {
+ "State1": {
+ "Next": "State2"
+ "Type": "Pass"
+ }
With my testing strategy in place, I turned my attention to getting the test to pass.
The first method to implement was perform
, where we supply a state to be added to the definition when we call the build
method. To do this, we need to capture the details for the build
method to use. I knew we would need to capture details for other methods, such as choice
, so I created a an interface and class to capture these. As TypeScript doesn't have a true type system, I included a type
enumeration to make introspection easy at runtime.
enum StepType {
Perform = 'Perform',
}
interface BuilderStep {
type: StepType;
id: string;
}
class PerformStep implements BuilderStep {
//
constructor(public state: sfn.State) {
this.type = StepType.Perform;
this.id = state.id;
}
type: StepType;
id: string;
}
With this in place, I added a class-level array to StateMachineBuilder
to hold the steps and updated the perform
method to capture the details to perform.
private readonly steps = new Array<BuilderStep>();
perform(state: sfn.State): StateMachineBuilder {
this.steps.push(new PerformStep(state));
return this;
}
With this in place, I started to look at the build
method and how we could use these details to build a CDK state machine definition. The build
method takes a scope parameter and returns an instance that implements IChainable
. In my mind, I could see that we would need to recurse through the steps to replicate the CDK approach. It seemed logical to me that the build
method should just initiate the recursion by returning the IChainable
for the first step. The getStepChain
method would then recursively call itself to build the structure.
build(scope: cdk.Construct): sfn.IChainable {
return this.getStepChain(scope, 0);
}
private getStepChain(scope: cdk.Construct, stepIndex: number): sfn.IChainable {
// TODO: Recursively call getStepChain
}
I could see that the getStepChain
method would need to handle the various step types, so I added a switch
and deferred the processing to a specific handler method.
private getStepChain(scope: cdk.Construct, stepIndex: number): sfn.IChainable {
//
const step = this.steps[stepIndex];
let stepChain: sfn.IChainable;
switch (step.type) {
//
case StepType.Perform:
stepChain = this.getPerformStepChain(scope, stepIndex);
break;
default:
throw new Error(`Unhandled step type: ${JSON.stringify(step)}`);
}
return stepChain;
}
The getPerformStepChain
method is the place where the real work was to be done. I.e., the place where the states would be wired together to build the state machine. The logic I had in mind was as follows.
- Get the state for the current step
- If there is a next step:
- Invoke the
next
method on the current step state, passing in the chain for the next step
- Invoke the
- Else
- Return the current step state
This was implemented as below.
private getPerformStepChain(scope: cdk.Construct, stepIndex: number): sfn.IChainable {
//
const step = this.steps[stepIndex] as PerformStep;
const stepState = (step as PerformStep).state;
const stepChain = stepIndex < this.steps.length - 1
? stepState.next(this.getStepChain(scope, stepIndex + 1))
: stepState;
return stepChain;
}
All looked fine, but there was a problem. I could see the following error.
Property 'next' does not exist on type 'State'. Did you mean '_next'?
I had assumed that the State
class had a next
method. However, by looking at the definition for the Pass
state, I could see the following.
export declare class Pass extends State implements INextable
It turned out that the next
method lives on a separate interface. What I wanted was something that encapsulated a State
with a next
method, so I created my own INextableState
interface.
interface INextableState extends sfn.State, sfn.INextable {}
With this, I could replace the references to State
and the problem with next
went away. When I re-ran the unit test, all was good. We now had an alternative way of defining state machines in CDK. The only caveat being, they can only consist of a sequence of states. Good, but not that useful, so the next thing to look at was implementing choices.
In Part 1, we designed the API to define a choice as follows.
.choice('Choice1', {
choices: [{ when: sfn.Condition.booleanEquals('$.var1', true), next: 'Choice2' }],
otherwise: 'Choice3',
})
As with perform
, we need to capture these details in the choices
method. To do this, I extended the StepType
enumeration, created a ChoiceStep
class, and amended the choice
method to store a ChoiceStep
instance containing the captured values.
enum StepType {
// Snip
Choice = 'Choice',
}
class ChoiceStep implements BuilderStep {
//
constructor(public id: string, public props: BuilderChoiceProps) {
this.type = StepType.Choice;
}
type: StepType;
}
export default class StateMachineBuilder {
// Snip
choice(id: string, props: BuilderChoiceProps): StateMachineBuilder {
this.steps.push(new ChoiceStep(id, props));
return this;
}
}
With this in place, I could extend the getStepChain
method to handle the Choice
step type and call a new getChoiceStepChain
method.
switch (step.type) {
// Snip
case StepType.Choice:
stepChain = this.getChoiceStepChain(scope, stepIndex);
break;
The implementation of the getChoiceStepChain
required a slightly different approach, as it needed to instantiate the State
as well as invoking the appropriate methods on it. It was for this reason that we added the scope
parameter to the build
method.
To build the resulting Choice
state, I needed to invoke the when
and otherwise
methods with IChainable
values. However, the choices
method only captures the string id
values. The solution was straightforward and was to create a getStepIndexById
method to covert one to the other. I went with a simple linear lookup for now, but if performance was paramount, then a indexed lookup could be implemented.
private getStepIndexById(id: string): number {
//
const stepIndex = this.steps.findIndex((s) => s.id === id);
if (stepIndex === -1) {
throw new Error(`Could not find index for id: ${id}`);
}
return stepIndex;
}
private getChoiceStepChain(scope: cdk.Construct, stepIndex: number): sfn.IChainable {
//
const step = this.steps[stepIndex] as ChoiceStep;
const stepChain = new sfn.Choice(scope, step.id, step.props);
step.props.choices.forEach((choice) => {
const nextStepIndex = this.getStepIndexById(choice.next);
const nextStepChain = this.getStepChain(scope, nextStepIndex);
stepChain.when(choice.when, nextStepChain);
});
const otherwiseStepIndex = this.getStepIndexById(step.props.otherwise);
const otherwiseStepChain = this.getStepChain(scope, otherwiseStepIndex);
stepChain.otherwise(otherwiseStepChain);
return stepChain;
}
In the getStepIndexById
method, I made sure to shout loudly and clearly when the the id
could not be found. In my experience, you will thank yourself if you throw informative errors when an unhandled value is encountered.
We were nearly there, but there was still one more piece of the choices
puzzle. To separate the various end states of the state machine, we have calls to the end
method as follows.
.perform(state1)
.end()
.perform(state2)
.end()
Our intention here was to tell the build
method to stop recursing and so make the previous state an 'end' state. To do this, I needed to create a new BuilderStep
and amend the end
method to add an instance to steps captured.
enum StepType {
// Snip
End = 'End',
}
class EndStep implements BuilderStep {
//
constructor() {
this.type = StepType.End;
}
id: string;
type: StepType;
}
export default class StateMachineBuilder {
// Snip
end(): StateMachineBuilder {
this.steps.push(new EndStep());
return this;
}
}
In getPerformStepChain
we had a test for whether we should continue and recursively add a 'next' state. This test relied on the last state being the last state in the steps
array. With the existence of the end
states, this assumption was no longer true. To cater for this, I extended the test to check ahead for an 'end' state created and encapsulated the test in a hasNextStep
method.
export default class StateMachineBuilder {
// Snip
private hasNextStep(stepIndex: number): boolean {
//
const isLastStep = stepIndex === this.steps.length - 1;
const isNextStepEnd = !isLastStep && this.steps[stepIndex + 1].type === StepType.End;
const hasNextStep = !(isLastStep || isNextStepEnd);
return hasNextStep;
}
private getPerformStepChain(scope: cdk.Construct, stepIndex: number): sfn.IChainable {
// Snip
const stepChain = this.hasNextStep(stepIndex)
? stepState.next(this.getStepChain(scope, stepIndex + 1))
: stepState;
return stepChain;
}
}
With this in place, I re-ran the unit tests and was met with unalloyed success. In the next part, I look to implement the map
and parallel
methods, and to implement functionality to add error handlers too.
Edit: The final StateMachineBuilder
component is now available on npm.