Skip to content

MkUnion and state machines in Go

This document will show how to use mkunion to manage application state using the example of an Order Service. You will learn:

  • how to model state machines in Go, and find similarities to "clean architecture"
  • How to test state machines (with fuzzing), and as a bonus, you will get mermaid diagrams for free
  • How to persist state in a database and how optimistic concurrency helps resolve concurrency conflicts
  • How to handle errors in state machines and build foundations for self-healing systems

Working example

As a driving example, we will use an e-commerce inspired Order Service that can be in one of the following states:

  • Pending - The order is created and is waiting for someone to process it.
  • Processing - The order is being processed; a human is going to pick up items from the warehouse and pack them.
  • Cancelled - The order was cancelled. There can be many reasons, one of them being that the warehouse is out of stock.
  • Completed - The order is completed and can be shipped to the customer.

Such states have rules that govern transitions; for example, an order cannot be cancelled if it's already completed, and so on.

We need to have a way to trigger changes in state, like creating an order that is pending for processing, or cancelling an order. We will call these triggers commands.

Some of these rules could change in the future, and we want to be able to change them without rewriting the whole application. This also informs us that our design should be open for extension.

Side note: if you want to go straight to the final code product, then go into the example/state/ directory and have fun exploring.

Modeling commands and states

Our example can be represented as a state machine that looks like this: simple_machine_test.go.state_diagram.mmd 

stateDiagram
    OrderCancelled: *state.OrderCancelled
    OrderCompleted: *state.OrderCompleted
    OrderError: *state.OrderError
    OrderPending: *state.OrderPending
    OrderProcessing: *state.OrderProcessing

    OrderProcessing --> OrderCancelled: *state.CancelOrderCMD
    [*] --> OrderPending: *state.CreateOrderCMD
    OrderPending --> OrderProcessing: *state.MarkAsProcessingCMD
    OrderProcessing --> OrderCompleted: *state.MarkOrderCompleteCMD
    OrderProcessing --> OrderError: *state.MarkOrderCompleteCMD
    OrderError --> OrderCompleted: *state.TryRecoverErrorCMD

In this diagram, we can see that we have 5 states and 6 commands that can trigger transitions between states, shown as arrows.

Because this diagram is generated from code, it has names that represent types in Go that we use in the implementation.

For example, *state.CreateOrderCMD:

  • state is the package name.
  • CreateOrderCMD is a struct name in that package.
  • The CMD suffix is a naming convention that is optional, but I find it makes the code more readable and easier to distinguish commands from states.

Below is a code snippet that demonstrates a complete model of the state and commands of the Order Service that we discussed.

Notice that we use mkunion to group commands and states. (Look for //go:tag mkunion:"Command")

This is one example of how union types can be used in Go. Historically in Go, it would be very hard to achieve such a thing, and it would require a lot of boilerplate code. Here, the interface that groups these types is generated automatically. You can focus on modeling your domain.

example/state/model.go
package state

import "time"

//go:tag mkunion:"Command"
type (
    CreateOrderCMD struct {
        OrderID OrderID
        Attr    OrderAttr
    }
    MarkAsProcessingCMD struct {
        OrderID  OrderID
        WorkerID WorkerID
    }
    CancelOrderCMD struct {
        OrderID OrderID
        Reason  string
    }
    MarkOrderCompleteCMD struct {
        OrderID OrderID
    }
    // TryRecoverErrorCMD is a special command that can be used to recover from error state
    // you can have different "self-healing" rules based on the error code or even return to previous healthy state
    TryRecoverErrorCMD struct {
        OrderID OrderID
    }
)

//go:tag mkunion:"State"
type (
    OrderPending struct {
        Order Order
    }
    OrderProcessing struct {
        Order Order
    }
    OrderCompleted struct {
        Order Order
    }
    OrderCancelled struct {
        Order Order
    }
    // OrderError is a special state that represent an error
    // during order processing, you can have different "self-healing jobs" based on the error code
    // like retrying the order, cancel the order, etc.
    // treating error as state is a good practice in state machine, it allow you to centralise the error handling
    OrderError struct {
        // error information
        Retried   int
        RetriedAt *time.Time

        ProblemCode ProblemCode

        ProblemCommand Command
        ProblemState   State
    }
)

type (
    // OrderID Price, Quantity are placeholders for value objects, to ensure better data semantic and type safety
    OrderID  = string
    Price    = float64
    Quantity = int

    OrderAttr struct {
        // placeholder for order attributes
        // like customer name, address, etc.
        // like product name, price, etc.
        // for simplicity we only have Price and Quantity
        Price    Price
        Quantity Quantity
    }

    // WorkerID represent human that process the order
    WorkerID = string

    // Order everything we know about order
    Order struct {
        ID               OrderID
        OrderAttr        OrderAttr
        WorkerID         WorkerID
        StockRemovedAt   *time.Time
        PaymentChargedAt *time.Time
        DeliveredAt      *time.Time
        CancelledAt      *time.Time
        CancelledReason  string
    }
)

type ProblemCode int

const (
    ProblemWarehouseAPIUnreachable ProblemCode = iota
    ProblemPaymentAPIUnreachable
)

Modeling transitions

One thing that is missing is the implementation of transitions between states. There are a few ways to do it. I will show you how to do it using a functional approach (think reduce or map function).

Let's name the function that we will build Transition and define it as:

func Transition(ctx context.Context, dep Dependencies, cmd Command, state State) (State, error)

Our function has a few arguments; let's break them down:

  • ctx is the standard Go context, which is used to pass deadlines, cancellation signals, etc.
  • dep encapsulates dependencies like API clients, database connections, configuration, context, etc. — everything that is needed for a complete production implementation.
  • cmd is a command that we want to apply to the state, and it has the Command interface, which was generated by mkunion when it was used to group commands.
  • state is a state that we want to apply our command to and change, and it has the State interface, which was generated similarly to the Command interface.

Our function must return either a new state or an error if something went wrong during the transition, like a network error or a validation error.

Below is a snippet of the implementation of the Transition function for our Order Service:

example/state/machine.go
        func(x *CreateOrderCMD) (State, error) {
            if x.OrderID == "" {
                return nil, ErrOrderIDRequired
            }

            switch state.(type) {
            case nil:
                o := Order{
                    ID:        x.OrderID,
                    OrderAttr: x.Attr,
                }
                return &OrderPending{
                    Order: o,
                }, nil
            }

            return nil, ErrOrderAlreadyExist
        },
// ...
// rest removed for brevity 
// ...

You can notice a few patterns in this snippet:

  • The Dependency interface helps us to keep dependencies well-defined, which greatly helps in testability and readability of the code.
  • The use of the generated function MatchCommandR2 to exhaustively match all commands. This is powerful; when a new command is added, you can be sure that you will get a compile-time error if you don't handle it.
  • Validation of commands is done in the transition function. The current implementation is simple, but you can use go-validate to make it more robust, or refactor the code and introduce domain helper functions or methods to the types.
  • Each command checks the state to which it is being applied using a switch statement; it ignores states that it doesn't care about. This means as an implementer, you have to focus only on a small part of the picture and not worry about the rest of the states. This is also an example where non-exhaustive use of the switch statement is welcome.

Simple, isn't it? Simplicity also comes from the fact that we don't have to worry about marshalling/unmarshalling data or working with the database; those are things that will be done in other parts of the application, keeping this part clean and focused on business logic.

Note: The implementation for educational purposes is kept in one big function, but for large projects, it may be better to split it into smaller functions, or define an OrderService struct that conforms to the visitor pattern interface, which was also generated for you:

example/state/model_union_gen.go
type CommandVisitor interface {
    VisitCreateOrderCMD(v *CreateOrderCMD) any
    VisitMarkAsProcessingCMD(v *MarkAsProcessingCMD) any
    VisitCancelOrderCMD(v *CancelOrderCMD) any
    VisitMarkOrderCompleteCMD(v *MarkOrderCompleteCMD) any
    VisitTryRecoverErrorCMD(v *TryRecoverErrorCMD) any
}

Testing state machines & self-documenting

Before we go further, let's talk about testing our implementation.

Testing will not only help us ensure that our implementation is correct but also help us document our state machine and discover transitions that we didn't think about, that should or shouldn't be possible.

Here is how you can test a state machine in a declarative way, using the mkunion/x/machine package:

example/state/machine_test.go
    var di Dependency = &DependencyMock{
        TimeNowFunc: func() *time.Time {
            return &now
        },
    }

    order := OrderAttr{
        Price:    100,
        Quantity: 3,
    }

    suite := machine.NewTestSuite(di, NewMachine)
    suite.Case(t, "happy path of order state transition",
        func(t *testing.T, c *machine.Case[Dependency, Command, State]) {
            c.
                GivenCommand(&CreateOrderCMD{OrderID: "123", Attr: order}).
                ThenState(t, &OrderPending{
                    Order: Order{
                        ID:        "123",
                        OrderAttr: order,
                    },
                }).
                ForkCase(t, "start processing order", func(t *testing.T, c *machine.Case[Dependency, Command, State]) {
                    c.
                        GivenCommand(&MarkAsProcessingCMD{
                            OrderID:  "123",
                            WorkerID: "worker-1",
                        }).
                        ThenState(t, &OrderProcessing{
                            Order: Order{
                                ID:        "123",
                                OrderAttr: order,
                                WorkerID:  "worker-1",
                            },
                        }).
                        ForkCase(t, "mark order as completed", func(t *testing.T, c *machine.Case[Dependency, Command, State]) {
                            c.
                                GivenCommand(&MarkOrderCompleteCMD{
                                    OrderID: "123",
                                }).
                                ThenState(t, &OrderCompleted{
                                    Order: Order{
                                        ID:               "123",
                                        OrderAttr:        order,
                                        WorkerID:         "worker-1",
                                        DeliveredAt:      &now,
                                        StockRemovedAt:   &now,
                                        PaymentChargedAt: &now,
                                    },
                                })
                        }).
                        ForkCase(t, "cancel order", func(t *testing.T, c *machine.Case[Dependency, Command, State]) {
                            c.
                                GivenCommand(&CancelOrderCMD{
                                    OrderID: "123",
                                    Reason:  "out of stock",
                                }).
                                ThenState(t, &OrderCancelled{
                                    Order: Order{
                                        ID:              "123",
                                        OrderAttr:       order,
                                        WorkerID:        "worker-1",
                                        CancelledAt:     &now,
                                        CancelledReason: "out of stock",
                                    },
                                })
                        }).
                        ForkCase(t, "try complete order but removing products from stock fails", func(t *testing.T, c *machine.Case[Dependency, Command, State]) {
                            c.
                                GivenCommand(&MarkOrderCompleteCMD{
                                    OrderID: "123",
                                }).
                                BeforeCommand(func(t testing.TB, di Dependency) {
                                    di.(*DependencyMock).ResetCalls()
                                    di.(*DependencyMock).WarehouseRemoveStockFunc = func(ctx context.Context, quantity int) error {
                                        return fmt.Errorf("warehouse api unreachable")
                                    }
                                }).
                                AfterCommand(func(t testing.TB, di Dependency) {
                                    dep := di.(*DependencyMock)
                                    dep.WarehouseRemoveStockFunc = nil
                                    if assert.Len(t, dep.WarehouseRemoveStockCalls(), 1) {
                                        assert.Equal(t, order.Quantity, dep.WarehouseRemoveStockCalls()[0].Quantity)
                                    }

                                    assert.Len(t, dep.PaymentChargeCalls(), 0)
                                }).
                                ThenState(t, &OrderError{
                                    Retried:        0,
                                    RetriedAt:      nil,
                                    ProblemCode:    ProblemWarehouseAPIUnreachable,
                                    ProblemCommand: &MarkOrderCompleteCMD{OrderID: "123"},
                                    ProblemState: &OrderProcessing{
                                        Order: Order{
                                            ID:        "123",
                                            OrderAttr: order,
                                            WorkerID:  "worker-1",
                                        },
                                    },
                                }).
                                ForkCase(t, "successfully recover", func(t *testing.T, c *machine.Case[Dependency, Command, State]) {
                                    c.
                                        GivenCommand(&TryRecoverErrorCMD{OrderID: "123"}).
                                        BeforeCommand(func(t testing.TB, di Dependency) {
                                            di.(*DependencyMock).ResetCalls()
                                        }).
                                        AfterCommand(func(t testing.TB, di Dependency) {
                                            dep := di.(*DependencyMock)
                                            if assert.Len(t, dep.WarehouseRemoveStockCalls(), 1) {
                                                assert.Equal(t, order.Quantity, dep.WarehouseRemoveStockCalls()[0].Quantity)
                                            }
                                            if assert.Len(t, dep.PaymentChargeCalls(), 1) {
                                                assert.Equal(t, order.Price, dep.PaymentChargeCalls()[0].Price)
                                            }
                                        }).
                                        ThenState(t, &OrderCompleted{
                                            Order: Order{
                                                ID:               "123",
                                                OrderAttr:        order,
                                                WorkerID:         "worker-1",
                                                DeliveredAt:      &now,
                                                StockRemovedAt:   &now,
                                                PaymentChargedAt: &now,
                                            },
                                        })
                                })
                        })
                })
        },
    )

    if suite.AssertSelfDocumentStateDiagram(t, "machine_test.go") {
        suite.SelfDocumentStateDiagram(t, "machine_test.go")
    }
A few things to notice in this test:

  • We use standard Go testing.
  • We use machine.NewTestSuite as a standard way to test state machines.
  • We start with describing the happy path and use suite.Case to define a test case.
  • But most importantly, we define test cases using GivenCommand and ThenState functions, which help make the test more readable and hopefully self-documenting.
  • You can see the use of the ForkCase command, which allows you to take the definition of a state declared in the ThenState command, apply a new command to it, and expect a new state.
  • Less visible is the use of moq to generate DependencyMock for dependencies, but it's still important for writing more concise code.

I know it's subjective, but I find it very readable and easy to understand, even for non-programmers.

Generating state diagram from tests

The last bit is this line at the bottom of the test file:

example/state/machine_test.go
if suite.AssertSelfDocumentStateDiagram(t, "machine_test.go") {
   suite.SelfDocumentStateDiagram(t, "machine_test.go")
}

This code takes all inputs provided in the test suite and fuzzes them, applies commands to random states, and records the results of those transitions.

  • SelfDocumentStateDiagram - produces two mermaid diagrams that show all possible transitions that are possible in our state machine.
  • AssertSelfDocumentStateDiagram can be used to compare newly generated diagrams to diagrams committed to the repository and fail the test if they are different. You don't have to use it, but it's good practice to ensure that your state machine is well-tested and doesn't regress without you noticing.

There are two diagrams that are generated.

One is a diagram of ONLY successful transitions that you saw at the beginning of this post.

stateDiagram
    OrderCancelled: *state.OrderCancelled
    OrderCompleted: *state.OrderCompleted
    OrderError: *state.OrderError
    OrderPending: *state.OrderPending
    OrderProcessing: *state.OrderProcessing

    OrderProcessing --> OrderCancelled: *state.CancelOrderCMD
    [*] --> OrderPending: *state.CreateOrderCMD
    OrderPending --> OrderProcessing: *state.MarkAsProcessingCMD
    OrderProcessing --> OrderCompleted: *state.MarkOrderCompleteCMD
    OrderProcessing --> OrderError: *state.MarkOrderCompleteCMD
    OrderError --> OrderCompleted: *state.TryRecoverErrorCMD

Second is a diagram that includes commands that resulted in errors: 

stateDiagram
    OrderCancelled: *state.OrderCancelled
    OrderCompleted: *state.OrderCompleted
    OrderError: *state.OrderError
    OrderPending: *state.OrderPending
    OrderProcessing: *state.OrderProcessing

    %% error=cannot cancel order, order must be processing to cancel it; invalid transition 
    OrderCancelled --> OrderCancelled: ❌*state.CancelOrderCMD
    %% error=cannot cancel order, order must be processing to cancel it; invalid transition 
    OrderCompleted --> OrderCompleted: ❌*state.CancelOrderCMD
    %% error=cannot cancel order, order must be processing to cancel it; invalid transition 
    OrderError --> OrderError: ❌*state.CancelOrderCMD
    %% error=cannot cancel order, order must be processing to cancel it; invalid transition 
    OrderPending --> OrderPending: ❌*state.CancelOrderCMD
    OrderProcessing --> OrderCancelled: *state.CancelOrderCMD
    %% error=cannot cancel order, order must be processing to cancel it; invalid transition 
    [*] --> [*]: ❌*state.CancelOrderCMD
    %% error=cannot attemp order creation, order exists: invalid transition 
    OrderCancelled --> OrderCancelled: ❌*state.CreateOrderCMD
    %% error=cannot attemp order creation, order exists: invalid transition 
    OrderCompleted --> OrderCompleted: ❌*state.CreateOrderCMD
    %% error=cannot attemp order creation, order exists: invalid transition 
    OrderError --> OrderError: ❌*state.CreateOrderCMD
    %% error=cannot attemp order creation, order exists: invalid transition 
    OrderPending --> OrderPending: ❌*state.CreateOrderCMD
    %% error=cannot attemp order creation, order exists: invalid transition 
    OrderProcessing --> OrderProcessing: ❌*state.CreateOrderCMD
    [*] --> OrderPending: *state.CreateOrderCMD
    %% error=invalid transition 
    OrderCancelled --> OrderCancelled: ❌*state.MarkAsProcessingCMD
    %% error=invalid transition 
    OrderCompleted --> OrderCompleted: ❌*state.MarkAsProcessingCMD
    %% error=invalid transition 
    OrderError --> OrderError: ❌*state.MarkAsProcessingCMD
    OrderPending --> OrderProcessing: *state.MarkAsProcessingCMD
    %% error=invalid transition 
    OrderProcessing --> OrderProcessing: ❌*state.MarkAsProcessingCMD
    %% error=invalid transition 
    [*] --> [*]: ❌*state.MarkAsProcessingCMD
    %% error=cannot mark order as complete, order is not being process; invalid transition 
    OrderCancelled --> OrderCancelled: ❌*state.MarkOrderCompleteCMD
    %% error=cannot mark order as complete, order is not being process; invalid transition 
    OrderCompleted --> OrderCompleted: ❌*state.MarkOrderCompleteCMD
    %% error=cannot mark order as complete, order is not being process; invalid transition 
    OrderError --> OrderError: ❌*state.MarkOrderCompleteCMD
    %% error=cannot mark order as complete, order is not being process; invalid transition 
    OrderPending --> OrderPending: ❌*state.MarkOrderCompleteCMD
    OrderProcessing --> OrderCompleted: *state.MarkOrderCompleteCMD
    OrderProcessing --> OrderError: *state.MarkOrderCompleteCMD
    %% error=cannot mark order as complete, order is not being process; invalid transition 
    [*] --> [*]: ❌*state.MarkOrderCompleteCMD
    %% error=cannot recover from non error state; invalid transition 
    OrderCancelled --> OrderCancelled: ❌*state.TryRecoverErrorCMD
    %% error=cannot recover from non error state; invalid transition 
    OrderCompleted --> OrderCompleted: ❌*state.TryRecoverErrorCMD
    OrderError --> OrderCompleted: *state.TryRecoverErrorCMD
    %% error=cannot recover from non error state; invalid transition 
    OrderPending --> OrderPending: ❌*state.TryRecoverErrorCMD
    %% error=cannot recover from non error state; invalid transition 
    OrderProcessing --> OrderProcessing: ❌*state.TryRecoverErrorCMD
    %% error=cannot recover from non error state; invalid transition 
    [*] --> [*]: ❌*state.TryRecoverErrorCMD

Those diagrams are stored in the same directory as the test file and are prefixed with the name used in the AssertSelfDocumentStateDiagram function. 

machine_test.go.state_diagram.mmd
machine_test.go.state_diagram_with_errors.mmd

State machines builder

MkUnion provides a *machine.Machine[Dependency, Command, State] struct that wires the Transition, dependencies, and state together. It provides methods like:

  • Handle(ctx context.Context, cmd C) error that applies a command to the state and returns an error if something went wrong during the transition.
  • State() S that returns the current state of the machine.
  • Dep() D that returns the dependencies the machine was built with.

This standard helps build on top of it; for example, the testing library we use in Testing state machines & self-documenting leverages it.

Another good practice is that every package that defines a state machine in the way described here should provide a NewMachine function that will return a bootstrapped machine with package types, like so:

example/state/machine.go
func NewMachine(di Dependency, init State) *machine.Machine[Dependency, Command, State] {
    return machine.NewMachine(di, Transition, init)
}

Conclusion

Now we have all the pieces in place, and we can start building our application.

  • We have a NewMachine constructor that will give us an object to use in our application.
  • We have tests that will ensure that our state machine is correct; fuzzy tests help discover edge cases, and lastly, we get diagrams showing which paths we tested and covered.
  • We saw how this approach focuses on business logic and keeps it separate from other concerns like the database or API clients. This is one of the principles of clean architecture.

For a comprehensive guide on best practices, patterns, and advanced techniques, see our State Machine Best Practices guide.

Next steps