Skip to main content

Finite State Machine with Annotations

Posted by carcassi on February 2, 2007 at 6:16 PM PST

This is an exercise that I did more or less an year ago to learn how to use annotations.
The idea is to use the then new Java 5.0 language features to create finite state
machines with no separate configuration files. By this time, probably there are
already a million implementation around, but the excercise did help me quite a bit
to understand how annotations can be used. Hope you like it!

This is an exercise that I did more or less an year ago to learn how to use annotations.
The idea is to use the then new Java 5.0 language features to create finite state
machines with no separate configuration files. By this time, probably there are
already a million implementation around, but the excercise did help me quite a bit
to understand how annotations can be used. So here it goes.

The idea is simply to:

  • Have one class for each state
  • Have one method for each transition
  • Have annotations for each method describing what the next state should be
  • Have an interface that defines the overall machine, that is the list of
    all possible states, the initial state and the final state
  • Have a generic Builder that takes the interface and provides an object instance
    that implements all the plumbing

The first four items are what the user would write, and the annotations will
provide the information to the Builder on how to implement the state changes. The
Builder needs to trap all the methods calls, route them to the correct object-state,
check if there is an annotation that prescribes a state change and change the object-state
for the next call. We can introspect to get all the information we need about the
annotation, but how do we "trap" the method calls? Since Java 1.4 there is a class
called Proxy: it allows you to create a handler that it's called every time a method
is called. The handler is passed the method name and the parameters, which allows
us to route the call to whatever object-state and add some logic before and after.
So, we have a plan, let's see how it works in practice.

The On/Off switch test case

We should start with a simple example, and what could possibly be simpler than
an on/off switch?

We said we would have one class for each state, so here is the Off class:

package jafsm.examples.onoff;

import jafsm.Transition;

public class Off {

@Transition(newState = On.class)
public void turnOn() {
        System.out.println("Lights are on!");
    }

}

The Off class has only one method, turnOn which will change the state to On.
The On class:

package jafsm.examples.onoff;

import jafsm.Transition;

public class On {

    @Transition(newState = Off.class)
    public void turnOff() {
        System.out.println("Lights are off!");
    }

}

has only the turnOff method which will change the state to Off. Now we need a
third class, an interface is better since all the implementation is actually in
the state classes, that puts the two state classes together: the Switch class.

package jafsm.examples.onoff;

import jafsm.FiniteStateMachine;

@FiniteStateMachine(states = {On.class, Off.class}, initialState = Off.class)
public interface Switch {

    void turnOn();
    void turnOff();

}

This says that a Switch is a finite state machine, with On and Off as the set
of states and the initial state as Off. Yeah, well... it's not that difficult to
figure out, right? :-)

Our JUnit test for the Switch would look like this:

    public void testSwitch() {
// We expect to just give the Switch class to a builder
        Switch fsm = FSMBuilder.buildFSM(Switch.class);

// We expect to get some sort of control to
// know in which state the FSM is
        FSMControl fsmControl = (FSMControl) fsm;

// The initial state should be Off
        assertEquals("Off", fsmControl.getCurrentState());

// If we turn the switch on, the state should be On
        fsm.turnOn();
        assertEquals("On", fsmControl.getCurrentState());

// If we turn it back off, the state should be Off
        fsm.turnOff();
        assertEquals("Off", fsmControl.getCurrentState());

// If we turn it off again, it should complain as the
// switch is already Off. The state should remain Off
        try {
            fsm.turnOff();
            this.fail("Shouldn't allow to turn off twice");
        } catch (Exception e) {
            // Failed to turn off twice: good
        }
        assertEquals("Off", fsmControl.getCurrentState());
    }

An Annotated Finite State Machine framework

So, the "only" thing to do is implement the annotations and the builder!. Let's
start with the annotations first. The Transition annotation will be:

package jafsm;

import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPoliy;

@Retention(value=RetentionPolicy.RUNTIME)
public @interface Transition {

    Class<?> newState();

}

This simple annotation defines a transition to a new state, independent of the
arguments or return value of the method. It's retained at runtime, which means that
we will be able to use introspection to find it. The other annotation we used was
the FiniteStateMachine:

package jafsm;

import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;

@Retention(value=RetentionPolicy.RUNTIME)
public @interface FiniteStateMachine {

    Class<?>[] states();
    Class<?> initialState();

    // TODO Object as final state flags the machine has no final state
    Class<?> finalState() default Object.class;

}

This annotation is also retained at runtime. It defines the list of state as
the set of possible states, an initialState and a possible finalState. If one doesn't
specify the finalState it means the machine goes on forever. In this case, the default
value (Object) is used to discriminate that case. Probably a "NoEnd" class would
be more elegant. I'll pretend this is left as an exercise to cover for the fact
that I was lazy.

So, we have our annotation, now we "only" need our builder. The first part of
our builder class should be:

public class FSMBuilder {

    private FSMBuilder() {
    }

    public static <T> T buildFSM(Class<T> classToken) {
        return buildFSM(classToken, extractStateClasses(classToken),
                        extractInitialState(classToken),
                        extractFinalState(classToken));
    }
    ...

The buildFSM is the only public method. It gets an instance of a class and returns
an object of type of that class. We are using bits of generics there to avoid a
type cast. We'll have three methods that will take the list of state classes, the
initial and the final state and call another method. Let's first see how the states
are extracted:

    private static List<Class<?>> extractStateClasses(Class<?> classToken) {
        FiniteStateMachine fsm = classToken.getAnnotation(FiniteStateMachine.class);
        return Arrays.asList(fsm.states());
    }

    private static Class<?> extractInitialState(Class<?> classToken) {
        FiniteStateMachine fsm = classToken.getAnnotation(FiniteStateMachine.class);
        return fsm.initialState();
    }

    private static Class<?> extractFinalState(Class<?> classToken) {
        FiniteStateMachine fsm = classToken.getAnnotation(FiniteStateMachine.class);
        return fsm.finalState();
    }

As you see, getting the values from the annotation is straight forward: you get
the annotation from the classToken object, and you get the annotation properties.
Now let's see what we do with those values:

    private static <T> T buildFSM(Class<T> fsmClassToken, List<Class<?>> stateClasses,
                         Class<?> initialStateClass, Class<?> finalStateClass) {
        return buildFSM(fsmClassToken, new FSMInvocationHandler(stateClasses,
                        initialStateClass, finalStateClass));
    }

    @SuppressWarnings("unchecked")
    static <T> T buildFSM(Class<T> classToken, InvocationHandler handler) {
        // Suppressed unchecked warning
        // The proxy will be of class T as the classToken is passed as one of the
        // interfaces to be implemented by the proxy.
        return (T) Proxy.newProxyInstance(classToken.getClassLoader(),
                        new Class<?>[] {classToken, FSMControl.class}, handler);
    }

}

The first method, which is the one called by the public one, simply creates an
FSMInvocationHandler and calls yet another method. The last method uses the Proxy
class to create a proxy object that implements 2 interfaces, the classToken (which
would be the Switch class in our testcase) and the FSMControl (which our test uses
to query the machine state). The handler is actually where the magic happens: it's
the object that will trap all the calls and implement the logic.

Notice the use of @SuppressWarnings: if it wasn't there the compiler would complain
that it cannot guarantee that the output of Proxy is of type classToken. We know
that is always going to be the case, as classToken is one of the interfaces that
the Proxy will implement. So we pass on that guarantee to the compiler by turning
off the warning. Such warning should be on the deepest possible method as you do
not want to accidentally suppress other warnings.

As we said, the whole magic happens in the FSMInvocationHandler. So let's have
a look at that.

class FSMInvocationHandler implements InvocationHandler {

    List<Class<?>> stateClasses;
    Class<?> initialStateClass;
    Class<?> finalStateClass;

    Map<Class<?>, Object> stateObjects = new Hashtable<Class<?>, Object>();

    Object currentState;

    FSMInvocationHandler(List<Class<?>> stateClasses, Class<?> initialStateClass,
                         Class<?> finalStateClass) {

        this.stateClasses = new ArrayList<Class<?>>(stateClasses);
        this.initialStateClass = initialStateClass;
        this.finalStateClass = finalStateClass;

        initializeStateObjects();
        setInitialState();
    }

The FSMInvocationHandler will have the list of stateClasses, the initalState
and the finalState, which will be fed in by the constructor. It will also have a
map stateObject that will contain one instance for each object that represent a
state. It will also contain a pointer to the object that represent the current state.
As part of the constructor, the stateObject map will be initialized and the initial
state will be set. Here are the two method responsible:

    private void initializeStateObjects() {
        for (Cl2ass<?> stateClass : stateClasses) {
            try {
                stateObjects.put(stateClass, stateClass.newInstance());
            } catch (InstantiationException ex) {
                throw new RuntimeException("State " + stateClass.getName() +
                  " can't be instanciated: " + ex.getMessage(), ex);
            } catch (IllegalAccessException ex) {
                throw new RuntimeException("State " + stateClass.getName() +
                  " has no accessible default constructor: " + ex.getMessage(), ex);
            }
        }
    }   

    private void setInitialState() {
        currentState = stateObjects.get(initialStateClass);
    }

To initialize the state object we simply iterate over the state classes and create
an object for each one, and to set the initial state, we simply retrieve the object
for the initial state class. And now the core of the plumbing:

    public Object invoke(Object proxy, Method method, Object[] args)
                         throws Throwable {
        if ("getCurrentState".equals(method.getName()) &&
            (method.getParameterTypes().length == 0)) {
            return currentState.getClass().getSimpleName();
        }
        Method stateMethod = lookupCurrentStateMethod(method, args);
        Object methodReturn = stateMethod.invoke(currentState, args);
        Transition transition = stateMethod.getAnnotation(Transition.class);
        if (transition != null) {
            Class<?> nextStateClass = transition.newState();
            currentState = stateObjects.get(nextStateClass);
        }
        return methodReturn;
    }   

    private Method lookupCurrentStateMethod(Method method, Object[] args) {
        try {
            return currentState.getClass().getMethod(method.getName(),
                                             method.getParameterTypes());
        } catch (NoSuchMethodException ex) {
            throw new RuntimeException("State '" + currentState.getClass().getName()
              + "' doesn't allow '" + method.getName() + "'");
        }
    }

The invoke method is going to be called for any method called on the switch object.
First thing we check is if the getCurrentState method was called, which should return
the name of the current state. In all the other cases, we try to find the method
in the current object state; we don't it means that state doesn't support it so
we throw an exception. If the method exists, then we forward the call with the approriate
arguments to the current state object. After that we check if there is an annotation
with the state we should transition to, and we change the current state object if
that's the case.

Missing pieces and considerations

As you can see, there are a lot of details missing:

  • the finalState is not actually used
  • we should implement equals, hashcode, clone, etc.. more consistently
  • we should provide annotation for those transition that depends on the arguments
    or return value of the method
  • the control could be extended to have bean-like event notification for state
    change
  • ...

But it should be clear how to go about to solve them: it's just a matter of making
the FSMInvocationHandler a bit smarter and possibly creating another couple of Annotations.

Annotations and the Proxy class combined are pretty powerful. The amount of code
for this simple example framework is really not so much: it took a Sunday afternoon
to write. But it does quite a lot of things already. For example, say you where
modeling some GUI that needed to change the panel displayed depending on the state
of a complicated process. The annotated FSM could have a getPanel, which each state
reimplements by returning the appropriate panel. Or say you are implementing an
old text adventure, making a state-class out of each location, with a getDescription
for what can be seen and the classic east/west/north/south for navigation.

And the same approach can be used whenever we need a very complex object to be
partitioned in several ones to make things more manageble, or we need to inject
at runtime some behavior before or after the method implementation of some interface.

One of the things I took out of this exercise on year ago was how powerful the
Java language has become with 5.0. I wasn't really sure how these new language functionalities
could be used, and this helped quite a bit.

Related Topics >>