The Typecast Pattern: Harnessing State Transitions

Whether you’re dealing with network connections, user sessions, or any state-dependent operations, ensuring that transitions are valid and safe is crucial. The complexity arises because states often dictate what actions can be performed, and performing an invalid action can lead to unpredictable behavior or even system failures.

state transition

The Problem with Invalid State Transitions

Consider a simple TCP connection example where we have three states: Closed, Open, and Listening. Let’s see what could go wrong if we don’t manage these states properly.

Here’s a code snippet illustrating an issue with invalid state transitions in a more loosely managed system:

enum State {
    Closed,
    Open,
    Listening,
}

struct TcpConnection {
    state: State,
}

impl TcpConnection {
    fn new() -> Self {
        TcpConnection { state: State::Closed }
    }

    fn open(&mut self) {
        match self.state {
            State::Closed => self.state = State::Open,
            _ => println!("Invalid state transition!"),
        }
    }

    fn close(&mut self) {
        match self.state {
            State::Open | State::Listening => self.state = State::Closed,
            _ => println!("Invalid state transition!"),
        }
    }

    fn listen(&mut self) {
        match self.state {
            State::Open => self.state = State::Listening,
            _ => println!("Invalid state transition!"),
        }
    }

    fn send(&self, data: &str) {
        if let State::Open = self.state {
            println!("Sending data: {}", data);
        } else {
            println!("Cannot send data in current state!");
        }
    }
}

fn main() {
    let mut connection = TcpConnection::new();

    // Trying to send data while the connection is closed
    connection.send("Hello, world!"); // Outputs: Cannot send data in current state!

    // Opening the connection and then sending data
    connection.open();
    connection.send("Hello, world!"); // Outputs: Sending data: Hello, world!

    // Trying to listen while in Open state
    connection.listen();

    // Invalid transition: trying to send data while listening
    connection.send("Hello again!"); // Outputs: Cannot send data in current state!
}

What’s Wrong Here?

  • Invalid State Transitions: The system allows state transitions that don’t make sense, like trying to send data while the connection is Listening.
  • Manual State Checks: Each method has to manually check the current state, which can lead to duplicated and error-prone code.
  • Lack of Compile-Time Safety: These errors are only caught at runtime, making the system more prone to bugs that could have been caught earlier.

Type-Safe State Management

To address these issues, we can use the Typecast Pattern to ensure that state transitions are valid and enforced at compile time. This approach leverages the type system to manage state transitions more safely and efficiently.

Let’s revisit the example, applying the Typecast Pattern:

use states::*;

pub mod states {
    #[derive(Debug, Clone)]
    pub struct Open {
        pub data: Option<String>
    }

    #[derive(Debug, Clone)]
    pub struct Closed {}

    #[derive(Debug, Clone)]
    pub struct Listening {}


    pub trait StateMarker {}
    impl StateMarker for () {}
    impl StateMarker for Open {}
    impl StateMarker for Closed {}
    impl StateMarker for Listening {}
}

struct TcpConnection<S: StateMarker> {
    _state: S,
}

impl TcpConnection<()> {
    fn new() -> TcpConnection<Closed> {
        TcpConnection {
            _state: states::Closed {},
        }
    }

    fn close(self) -> TcpConnection<Closed> {
        TcpConnection {
            _state: states::Closed {}
        }
    }
}

impl TcpConnection<Closed> {
    fn open(self) -> TcpConnection<Open> {
        TcpConnection {
            _state: states::Open {
                data: None
            }
        }
    }

    fn is_open(&self) -> bool {
        false
    }
}

impl TcpConnection<Open> {
    fn listen(self) -> TcpConnection<Listening> {
        TcpConnection {
            _state: states::Listening {}
        }
    }

    fn close(self) -> TcpConnection<Closed> {
        TcpConnection {
            _state: states::Closed {}
        }
    }

    fn is_open(&self) -> bool {
        true
    }

    fn is_listening(&self) -> bool {
        false
    }

    fn send(&mut self, data: &str) {
        self._state.data = Some(String::from(data));
    }

    fn receive(self) -> String {
        match self._state.data.clone() {
            None => String::from(""),
            Some(data) => data
        }
    }
}

impl TcpConnection<Listening> {
    fn is_open(&self) -> bool {
        true
    }

    fn is_listening(&self) -> bool {
        true
    }

    fn accept(&self) {
        println!("I am accepting data")
    }
}

Here every action is returning a new state, we cannot call e.g. listen() on a closed connection. Also every action is consuming the state, so the following is not possible.

    let connection = TcpConnection::new();
    let mut open_connection = connection.open();
    open_connection.close();
    open_connection.send("Hello, world!");

We cannot call send() on a closed connection, the compiler will complain with

error[E0382]: borrow of moved value: `open_connection`
   --> tcp_typestate_pattern/src/main.rs:143:9
    |
141 |         let mut open_connection = connection.open();
    |             ------------------- move occurs because `open_connection` has type `TcpConnection<states::Open>`, which does not implement the `Copy` trait
142 |         open_connection.close();
    |                         ------- `open_connection` moved due to this method call
143 |         open_connection.send("Hello, world!");
    |         ^^^^^^^^^^^^^^^ value borrowed here after move
  • Type Safety: Ensures that only valid transitions can occur, caught at compile time.
  • Reduced Boilerplate: Eliminates the need for repetitive state checks within methods.
  • Clarity and Maintainability: Makes the code easier to understand and maintain by encapsulating state-specific logic within state-specific implementations.
Published 10 Jun 2024

Tüftler (someone who enjoys working on and solving technical problems, often in a meticulous and innovative manner). Opinions are my own and not necessarily the views of my employer.
Avraam Mavridis on Twitter