Kohei Nozaki's blog 

The State Pattern


Posted on Thursday Aug 11, 2022 at 01:21PM in Technology


What is the State Pattern?

It’s one of the GoF Design Patterns which helps to organize code that handles multiple states. If you have a code base where there are many boolean fields or an enum field used to determine the state, similar if/switch branches and unclear state transitions, a giant class where there is all of the stuff in messy code causing a maintainability issue, applying the State Pattern might be worth considering.

The State Pattern consists of the following participants:

fa37dd6b 3a52 4a9a afa4 ccc5079d9f10

The Context class provides the clients of your code with the public API. All of the other participants are implementation details which the clients of your class don’t have to know about.

The Context class holds a reference to an object which implements the State interface. The Context class forwards some of the API calls to the state object. The State interface defines part of the API whose behavior has to be changed depending on the state.

The ConcreteState classes implement the State interface. Each implementation exists for each possible state of your application. It handles the API calls forwarded from the Context class appropriately according to the requirement of the state it corresponds to. A ConcreteState class can trigger a state transition.

Case study: the SMTP protocol

As an example application for the State Pattern, let’s think about an SMTP server. SMTP is a protocol used for email transport. Usually email client software like Mozilla Thunderbird supports this protocol. When such software sends an email, it connects to an SMTP server and talks to the server using this protocol to send an email.

This is a typical flow of an SMTP communication session (taken from Wikipedia):

S: 220 smtp.example.com ESMTP Postfix
C: HELO relay.example.com
S: 250 smtp.example.com, I am glad to meet you
C: MAIL FROM:<bob@example.com>
S: 250 Ok
C: RCPT TO:<alice@example.com>
S: 250 Ok
C: RCPT TO:<theboss@example.com>
S: 250 Ok
C: DATA
S: 354 End data with <CR><LF>.<CR><LF>
C: From: "Bob Example" <bob@example.com>
C: To: Alice Example <alice@example.com>
C: Cc: theboss@example.com
C: Date: Tue, 15 Jan 2008 16:02:43 -0500
C: Subject: Test message
C:
C: Hello Alice.
C: This is a test message with 5 header fields and 4 lines in the message body.
C: Your friend,
C: Bob
C: .
S: 250 Ok: queued as 12345
C: QUIT
S: 221 Bye
{The server closes the connection}

The lines beginning with "S:" are sent by the server and the ones beginning with "C: " are sent by the client.

Typically, after establishing a TCP connection, the server sends a short message which indicates it is ready to accept a command from the client. The client first sends an HELO message with its hostname, then begins an email transaction with specifying the email address of the sender of the email with a MAIL FROM command. After that, the client sends the email addresses of the recipients of the email with RCPT TO commands. Then finally, the client sends a DATA command, sends the content of the email and finishes it with a line which contains only a period. If everything goes fine, the server responds with an OK message which means the email is accepted by the server and the email will be delivered to the recipients specified.

What we can see from here is that in this protocol a client has to send necessary information to the server in a specific order. The server reacts differently for each state of the communication process. For example, the client must provide a sender identification first. Otherwise, any other command from the client doesn’t get processed and the server returns an error. It means that the server remembers the current communication state at any given moment.

So, what are those states? A good way to find that out is drawing a state diagram. It will look like the following:

0334ad5d 100e 4b11 98a3 8705c9cb832b

There are 4 states in the communication process. The first one is the idle state where the client has to start a session with an HELO command. Then it proceeds to the initial state, where the client has to provide the email address of the sender of the email with a MAIL FROM command. Then we proceed to the transaction started state where the client provides the recipients of the email with an RCPT TO command. The client might stay in this state until it finishes sending all of the recipients and that is why there is this transition which goes back to the same state. After that, finally it proceeds to the data transfer state with a DATA command where the client sends the body of the email and finishes the transaction with a line which contains only a period. Then it gets back to the initial state and if the client wants to send another email, it can start over from there.

Implementing an SMTP server with the State Pattern

In order to implement such a communication process which consists of multiple stages or states, applying the State Pattern can be a good way to keep the implementation clean and organized. Otherwise, we might end up having a giant class where there are a lot of mutable states and if/switch conditionals which are highly likely to be a maintenance problem.

Let’s look at one possible design based on the State Pattern which can handle the SMTP protocol:

53b553bd 9d02 4c68 965b 2adad582e346

The SMTPSessionHandler class corresponds to the Context class in the previous class diagram. An instance of this class exists for each SMTP conversation session. There is a public method called handleCommand() which receives a command from the client and returns the response for the command. This method has to handle the commands from the client appropriately depending on the current state. In order to achieve that, we have those 4 concrete state classes that correspond to the states in the state diagram given earlier, and what this method does is basically just forwarding the method calls to the currentState object.

In this design, state transitions are done by each concrete state class. For that purpose, the SMTPSessionHandler class has the setCurrentState() method. And when the SMTPSessionHandler class forwards the method calls to the currentState object, it also passes its own reference as an additional parameter. It allows each concrete state class to call the setCurrentState() method.

Let’s look at the source code of each participant. The SMTPSessionHandler class just holds the currentState object and forwards any handleCommand() calls to the currentState object with a reference to this object. It also has the setCurrentState() method for the concrete state classes.

public class SMTPSessionHandler {
    private State currentState = new IdleState();

    public String handleCommand(String command) {
        return currentState.handleCommand(command, this);
    }

    void setCurrentState(State newState) {
        this.currentState = newState;
    }
}

The SMTPSessionHandler class calls the underlying current state object through this State interface, which has 4 concrete implementations.

interface State {
    String handleCommand(String command, SMTPSessionHandler context);
}

The IdleState class corresponds to the very first state when a connection is established with the client. The only command it accepts is the HELO command. When it receives the HELO command, it makes a state transition by calling the setCurrentState() method of the context object with a new instance of the InitialState class and returns a greeting to the client. Otherwise, it returns an error.

class IdleState implements State {
    @Override
    public String handleCommand(String command, SMTPSessionHandler context) {
        if (command.startsWith("HELO")) {
            context.setCurrentState(new InitialState());
            return "250 smtp.example.com, I am glad to meet you";
        } else {
            return "500 5.5.1 Invalid command";
        }
    }
}

When the state transition in the IdleState class happens, the InitialState class takes over. The only command it accepts is the MAIL FROM command. When it happens, it extracts the email address which the client has sent and makes another state transition with a new instance of the TransactionStartedState class. When it creates the instance, it passes the email address it has extracted in order to let the subsequent process use it. Otherwise it returns an error.

class InitialState implements State {

    private static final Pattern PATTERN_FOR_EXTRACTING_EMAIL =
            Pattern.compile("MAIL FROM:<([^>]+)>");

    @Override
    public String handleCommand(String command, SMTPSessionHandler context) {
        Matcher matcher = PATTERN_FOR_EXTRACTING_EMAIL.matcher(command);
        if (matcher.find()) {
            String from = matcher.group(1);
            context.setCurrentState(new TransactionStartedState(from));
            return "250 Ok";
        } else {
            return "500 5.5.1 Invalid command";
        }
    }
}

The TransactionStartedState class is where we receive the destinations of the email. After specifying at least one destination, we can proceed to the next state but if there is none, it returns an error. The client has to send the destinations with the RCPT TO command. When it receives the RCPT TO command, it extracts the email address and keeps it in the List object. After sending at least one destination, the client can proceed to the next state with the DATA command. At this point, we have the address of the sender and the destinations and those are passed as the parameters of the constructor of the DataTransferState class.

class TransactionStartedState implements State {
    private static final Pattern PATTERN_FOR_EXTRACTING_EMAIL =
            Pattern.compile("RCPT TO:<([^>]+)>");
    private final String from;
    private final List<String> destinations = new ArrayList<>();

    TransactionStartedState(String from) {
        this.from = from;
    }

    @Override
    public String handleCommand(String command, SMTPSessionHandler context) {
        if (command.equals("DATA")) {
            if (destinations.isEmpty()) {
                return "500 5.5.1 Invalid command";
            } else {
                context.setCurrentState(new DataTransferState(from, destinations));
                return "354 End data with <CR><LF>.<CR><LF>";
            }
        }

        Matcher matcher = PATTERN_FOR_EXTRACTING_EMAIL.matcher(command);
        if (matcher.find()) {
            String to = matcher.group(1);
            destinations.add(to);
            return "250 Ok";
        } else {
            return "500 5.5.1 Invalid command";
        }
    }
}

The DataTransferState class corresponds to the final state of this communication process. What it does is just accumulating the body of the email in a StringBuilder object until it receives one single period which means the end of the body. After that, it will trigger the email delivery process which involves a DNS lookup, relaying the message to another SMTP server using the pieces of the data we received from the client during the process and making a state transition to the initial state. If the client wants to send another email, the client can start all over from there.

class DataTransferState implements State {
    private final String from;
    private final List<String> destinations;
    private final StringBuilder body = new StringBuilder();

    static DeliverySystem deliverySystem = (from, destinations, body) -> {
        // looks up MX records, connects to external SMTP servers and relays the message
    };

    DataTransferState(String from, List<String> destinations) {
        this.from = from;
        this.destinations = destinations;
    }

    @Override
    public String handleCommand(String command, SMTPSessionHandler context) {
        if (command.equals(".")) {
            deliverySystem.deliver(from, destinations, body.toString());
            context.setCurrentState(new InitialState());
            return "250 Ok: the email has been delivered";
        } else {
            body.append(command);
            body.append('\n');
            return null;
        }
    }
}

Conclusion: benefits of the State Pattern

That’s all of the implementation. The responsibility of each concrete state class is clear and it helps to make each class concise and short. The state transitions are also clear because when they happen, they are done in a clear way, which is the setCurrentState() method in our example. There are no cryptic boolean flags that maintain states in an unclear way. It makes the code readable, easy to maintain and extend.

For example, when you need to add a new state, what you need to do is write a new class which corresponds to the new state and add state transitions to the relevant existing state classes that are likely to be much simpler than a giant single class maintaining all of the states. It will reduce the risk of degradation.

In my experience, there are not so many use cases where this pattern is a great match, but it greatly improves the maintainability if it matches the use case and is applied appropriately.