Using XState Actors to Model Async Workflows Safely

In my previous post I discussed the challenges of writing async workflows in React that correctly deal with all possible edge cases. Even for a simple case of a client with two dependencies with no error handling, we ended up with this code:

const useClient = (user) => {
  const [client, setClient] = useState(null);

  useEffect(() => {
    let cancelled;
    (async () => {
      const clientAuthToken = await fetchClientToken(user);
      if (cancelled) return;

      const connection = await createWebsocketConnection();
      if (cancelled) {
        connection.close();
        return;
      }

      const client = await createClient(connection, clientAuthToken);
      if (cancelled) {
        client.disconnect();
        return;
      }

      setClient(client);
    })();

    return () => {
      cancelled = true;
    };
  }, [user]);

  useEffect(() => {
    return () => {
      client.disconnect();
    };
  }, [client]);

  return client;
};

This tangle of highly imperative code functions, but it will prove hard to read and change in the future. What we need is a way to express the stateful nature of the various pieces of this workflow and how they interact with each other in a way in which we can easily see if we've missed something or make changes in the future. This is where state machines and the actor model can come in handy.

State machines? Actors?

These are programming patterns that you may or may not have heard of before. I will explain them in a brief and simplified way but you should know that there is a great deal of theoretical and practical background in this area that we will be leveraging even though we won't go over it explicitly.

1. A state machine is an entity consisting of state and a series of rules to be followed to determine the next state from a combination of its previous state and external events it receives. Even though you might rarely think about them, state machines are everywhere. For example, a Promise is a state machine going from pending to resolved state when it receives a value from the asynchronous computation it is wrapping.
2. The actor model is a computing architecture that models asynchronous workflows as the interplay of self-contained units called actors. These units communicate with each other by sending and receiving events, they encapsulate state and they exist in a hierarchical relationship, where parent actors spawn child actors, thus linking their lifecycles.

It's common to combine both patterns so that a single entity is both an actor and a state machine, so that child actors are spawned and messages are sent based on which state the entity is in. I'll be using XState, a Javascript library which allows us to create actors and state machines in an easy declarative style. This won't be a complete introductory tutorial to XState, though. So if you're unfamiliar with the tool and need context for the syntax I'll be using, head to their website to read through the docs.

Setting the stage

The first step is to break down our workflow into the distinct states it can be in. Not every step in a process is a state. Rather, states represent moments in the process where the workflow is waiting for something to happen, whether that is user input or the completion of some external process. In our case we can break our workflow down coarsely into three states:

1. When the workflow is first created, we can immediately start creating the connection, and fetching the auth token. But, we have to wait until those are finished before creating the client. We'll call this state "preparing".
2. Then, we've started the process of creating the client, but we can't use it until the client creation returns it to us. We'll call this state "creatingClient".
3. Finally, everything is ready, and the client can be used. The machine is waiting only for the exit signal so it can release its resources and destroy itself. We'll call this state "clientReady".

This can be represented visually like so (all visualizations produced with Stately):

And in code like so

export const clientFactory = createMachine({
  id: "clientFactory",
  initial: "preparing",
  states: {
    preparing: {
      on: {
        "preparations.complete": {
          target: "creatingClient",
        },
      },
    },
    creatingClient: {
      on: {
        "client.ready": {
          target: "clientReady",
        },
      },
    },
    clientReady: {},
  },
});

However, this is a bit overly simplistic. When we're in our "preparing" state there are actually two separate and independent processes happening, and both of them must complete before we can start creating the client. Fortunately, this is easily represented with parallel child state nodes. Think of parallel state nodes like Promise.all: they advance independently but the parent that invoked them gets notified when they all finish. In XState, "finishing" is defined as reaching a state marked "final", like so

export const clientFactory = createMachine({
  id: "clientFactory",
  initial: "preparing",
  states: {
    preparing: {
      // Parallel state nodes are a good way to model two independent
      // workflows that should happen at the same time
      type: "parallel",
      states: {
        // The token child node
        token: {
          initial: "fetching",
          states: {
            fetching: {
              on: {
                "token.ready": {
                  target: "done",
                },
              },
            },
            done: {
              type: "final",
            },
          },
        },
        // The connection child node
        connection: {
          initial: "connecting",
          states: {
            connecting: {
              on: {
                "connection.ready": {
                  target: "done",
                },
              },
            },
            done: {
              type: "final",
            },
          },
        },
      },
      // The "onDone" transition on parallel state nodes gets called when all child
      // nodes have entered their "final" state. It's a great way to wait until
      // various workflows have completed before moving to the next step!
      onDone: {
        target: "creatingClient",
      },
    },
    creatingClient: {
      on: {
        "client.ready": {
          target: "clientReady",
        },
      },
    },
    clientReady: {},
  },
});

Leaving us with the final shape of our state chart:

Casting call

So far, we only have a single actor: the root actor implicitly created by declaring our state machine. To unlock the real advantages of using actors we need to model all of our disposable resources as actors. We could write them as full state machines using XState but instead let's take advantage of a short and sweet way of defining actors that interact with non-XState code: functions with callbacks. Here is what our connection actor might look like, creating and disposing of a WebSocket:

// Demonstrated here is the simplest and most versatile form of actor: a function that
// takes a callback that sends events to the parent actor, and returns a function that
// will be called when it is stopped.
const createConnectionActor = () => (send) => {
  const connection = new WebSocket("wss://example.com");

  connection.onopen = () =>
    // We send an event to the parent that contains the ready connection
    send({ type: "connection.ready", data: connection });

  // Actors are automatically stopped when their parent stops so simple actors are a great
  // way to manage resources that need to be disposed of. The function returned by an
  // actor will be called when it receives the stop signal.
  return () => {
    connection.close();
  };
};

And here is one for the client, which demonstrates the use of promises inside a callback actor. You can spawn promises as actors directly but they provide no mechanism for responding to events, cleaning up after themselves, or sending any events other than "done" and "error", so they are a poor choice in most cases. It's better to invoke your promise-creating function inside a callback actor, and use the Promise methods like .then() to control async responses.

// We can have the actor creation function take arguments, which we will populate
// when we spawn it
const createClientActor => (token, connection) => (send) => {
  const clientPromise = createClient(token, connection);
  clientPromise.then((client) =>
    send({ type: "client.ready", data: client })
  );

  return () => {
    // A good way to make sure the result of an async function is
    // always cleaned up is by invoking cleanup through .then()
    // If this executes before the promise is resolved, it will cleanup
    // on resolution, whereas if it executes after it's resolved, it will
    // clean up immediately
    clientPromise.then((client) => {
      client.disconnect();
    });
  };
};

Actors are spawned with the spawn action creator from XState, but we also need to save the reference to the running actor somewhere, so spawn is usually combined with assign to create an actor, and save it into the parent's context.

// We put this as the machine options. Machine options can be customised when the
// machine is interpreted, which gives us a way to use values from e.g. React context
// to define our actions, although this is not demonstrated here
const clientFactoryOptions = {
  spawnConnection: assign({
    connectionRef: () => spawn(createConnectionActor()),
  }),
  spawnClient: assign({
    // The assign action creator lets us use the machine context when defining the
    // state to be assigned, this way actors can inherit parent state
    clientRef: (context) =>
      spawn(createClientActor(context.token, context.connection)),
  }),
};

And then it becomes an easy task to trigger these actions when certain states are entered:

export const clientFactory = createMachine({
  id: "clientFactory",
  initial: "preparing",
  states: {
    preparing: {
      type: "parallel",
      states: {
        token: {
          initial: "fetching",
          states: {
            fetching: {
              // Because there's no resource to manage once it's done, we
              // can simply invoke a promise here. Invoked services are like
              // actors, but they're automatically spawned when the state node
              // is entered, and destroyed when it is exited,
              invoke: {
                src: "fetchToken",
                // Invoking a promise provides us with a handy "onDone" transition
                // that triggers when the promise resolves. To handle rejections,
                // we would similarly implement "onError"
                onDone: {
                  // These "save" actions will save the result to the machine
                  // context. They're simple assigners, but you can see them in
                  // the full code example linked at the end.
                  actions: "saveToken",
                  target: "done",
                },
              },
            },
            done: {
              type: "final",
            },
          },
        },
        // The connection child node
        connection: {
          initial: "connecting",
          states: {
            connecting: {
              // We want our connection actor to stick around, because by design,
              // the actor destroys the connection when it exits, so we store
              // it in state by using a "spawn" action
              entry: "spawnConnection",
              on: {
                // Since we're dealing with a persistent actor, we don't get an
                // automatic "onDone" transition. Instead, we rely on the actor
                // to send us an event.
                "connection.ready": {
                  actions: "saveConnection",
                  target: "done",
                },
              },
            },
            done: {
              type: "final",
            },
          },
        },
      },
      onDone: {
        target: "creatingClient",
      },
    },
    creatingClient: {
      // The same pattern as the connection actor. We spawn a  persistent actor
      // that takes care of creating and destroying the client.
      // entry: "spawnClient",
      on: {
        "client.ready": {
          actions: "saveClient",
          target: "clientReady",
        },
      },
    },
    // Even though this node can't be exited, it is not "final". A final node would
    // cause the machine to stop operating, which would stop the child actors!
    clientReady: {},
  },
});

Putting on the performance

XState provides hooks that simplify the process of using state machines in React, making this the equivalent to our async hook at the start:

const useClient = (user) => {
  const [state] = useMachine(clientFactory, clientFactoryOptions);

  if (state.matches("clientReady")) return state.context.client;
  return null;
};

Of course, combined with the machine definition, the action definitions and the actor code are hardly less code, or even simpler code. The advantage of breaking a workflow down like this include:

1. Each part can be tested independently. You can verify that the machine follows the logic set out without invoking the actors, and you can verify that the actors clean up after themselves without running the whole machine.
2. The parts can be shuffled around, and added to without having to rewrite them. We could easily add an extra step between connecting and creating the client, or introduce error handling and error states.
3. We can read and visualize every state and transition of the workflow to make sure we've accounted for all of them. This is a particular improvement over long async/await chains where every await implicitly creates a new state and two transitions — success and error — and the precise placement of catch blocks can drastically change the shape of the state chart.

You won't need to break out these patterns very often in an application. Maybe once or twice, or maybe never. After all, many applications never have to worry about complex workflows and disposable resources. However, having these ideas in your back pocket can get you out of some jams, particularly if you're already using state machines to model UI behaviour — something you should definitely consider doing if you're not already.

A complete code example with everything discussed above using Typescript, and with mock actors and services, that actually run in the visualizer, can be found here.