Rust State Machine Part 1

Table of Content (Click to Expand)

tldr

Define the states.
- struct A and optionallyState<S> for type safe calls
- enum to hold the structs
Define the commands so it could be “send” by the client
Define the control panel which has a sender tx
Spawn the control panel
- initial current_state in Option
- while let Ok(cmd) = rx.recv()
- current_state = Some(match(state, cmd){});
Full example code here

Overview

I’ve been learning how to write a state machine in Rust and I am at a point where I could stop and write something down.

I would be starting from a very simple one then go to one which allows multiple threads to try access the same state machine.

Note that the code is not production ready, it’s more of a study note post.

There would be a part 2 once I learned how to handle two-way communications and work with tokio.

The Scenario

Say we want to model states of traffic lights. It only has 3 states. Red, Green and Yellow.

  stateDiagram-v2
    [*] --> Red
    Red --> Green
    Green --> Yellow
    Yellow --> Red

Simple Traffic Light

Here let’s write some code to model the state machine above.

// state.rs

pub struct Red;
pub struct Green;
pub struct Yellow;

impl Red {
    pub fn new() -> Self {
        println!("start with red");
        Red {}
    }
    pub fn to_green(self) -> Green {
        println!("red to green");
        Green {}
    }
}

impl Green {
    pub fn to_yellow(self) -> Yellow {
        println!("green to yellow");
        Yellow {}
    }
}

impl Yellow {
    pub fn to_red(self) -> Red {
        println!("yellow to red");
        Red {}
    }
}

The above defines three empty structs and implements their transitions.

A few things worth noting.

Only Red has pub fn new() -> Self, this is because our state machine can only start from Red (The code in this example actually allows constructing directly, let’s ignore that for the sake of this post)
Each of the to_xxx takes in self, meaning it is actually consuming self instead of borrowing it. For instance when we call to_green(self) -> Green, the existing state is consumed and a new Green state is created.

The benefit of this structure is that the state machine is enforced by the rust compiler.

mod state;
use state::Red;

fn main() {
    let red = Red::new();
    let green = red.to_green();
    let yellow = green.to_yellow();
    let _red = yellow.to_red();
}

The output of the above is

start with red
red to green
green to yellow
yellow to red

Honestly this isn’t really that impressive. Also this might be fine to be used within a single thread but what if it needs to be shared by multiple threads? This might be a nightmare since all the state transitions require consuming the state.

There is an elegant way to handle this. It’s called the actor pattern. Let’s have a look at it.

Traffic Light Remote Control

Assume that I want to create a traffic light remote control, where on the remote control I can press buttons to change the light’s color. And not only I have the control, my friends could have their own control but we all control the same traffic light.

There are 4 steps, and I’ll go over them one by one

Step 1: Define The Physical States

We would still be using the empty structs but on top of that we need to create an enum that wraps over them. The reason an enum is needed is because we needed a “box” to hold the states. You’ll see this later.

// state.rs

// the original empty structs
pub struct Red;
pub struct Green;
pub struct Yellow;

impl Red {
    pub fn new() -> Self {
        println!("start with red");
        Red {}
    }
    pub fn to_green(self) -> Green {
        println!("red to green");
        Green {}
    }
}

impl Green {
    pub fn to_yellow(self) -> Yellow {
        println!("green to yellow");
        Yellow {}
    }
}

impl Yellow {
    pub fn to_red(self) -> Red {
        println!("yellow to red");
        Red {}
    }
}

// the enum wrapper
pub enum TrafficLight {
    Red(Red),
    Green(Green),
    Yellow(Yellow),
}

Note that in pub enum TrafficLight, Red(Red) reads as “Defining an enum called Red which wraps around state::Red”.

Step 2: Define The Commands

We need to define the commands our remote control would send, these would be defined as an enum.

// state.rs

pub enum Command {
    ToRed,
    ToGreen,
    ToYellow,
}

Step 3: Define The Control Panel (The Remote Control)

Let’s define the struct that clients would use to send the commands.

state.rs

use std::sync::mpsc;

#[derive(Clone)]
pub struct TrafficLightRemote {
    sender: mpsc::Sender<Command>,
}

impl TrafficLightRemote {
    pub fn to_green(&self) {
        self.sender.send(Command::ToGreen).unwrap();
    }

    pub fn to_yellow(&self) {
        self.sender.send(Command::ToYellow).unwrap();
    }

    pub fn to_red(&self) {
        self.sender.send(Command::ToRed).unwrap();
    }
}

Ignore the unwrap() here, the idea is that TrafficLightRemote holds a mpsc::Sender<Command> and can send commands to attempt to modify the state.

Step 4: Spawn The Control Panel (The Remote Control)

This is where we match the current state and command then decide what to do. Let’s have a look at the code and I’ll point out a few important things.

// state.rs
use std::sync::mpsc;

pub fn spawn_client() -> TrafficLightRemote {
    let (tx, rx) = mpsc::channel();

    std::thread::spawn(move || {
        // initial state for the state machine
        let mut current_state = Some(TrafficLight::Red(Red::new()));

        while let Ok(cmd) = rx.recv() {
            // take the state out from Option
            let Some(state) = current_state.take() else {
                println!("fatal error, invalid state");
                break;
            };

            // based on current state and command, move to the next state
            current_state = Some(match (state, cmd) {
                (TrafficLight::Green(green), Command::ToYellow) => {
                    TrafficLight::Yellow(green.to_yellow())
                }
                (TrafficLight::Yellow(yellow), Command::ToRed) => {
                    TrafficLight::Red(yellow.to_red())
                }
                (TrafficLight::Red(red), Command::ToGreen) => {
                    TrafficLight::Green(red.to_green())
                }
                (original_state, _bad_cmd) => {
                    println!("invalid state command combination");
                    original_state
                }
            });
        }

        println!("all threds done");
    });

    // returns the remote instance with tx wired to sender
    TrafficLightRemote { sender: tx }
}

while let Ok(cmd) = rx.recv() would exit when all senders are dropped.
The last (original_state, _bad_cmd) match is needed to satisfy exhaustiveness
If (_, _) was used instead of (original_state, _bad_cmd), then state would actually be dropped, so current_state would either be () which triggers a compile error or I force None in it which would break the next command.
green in (TrafficLight::Green(green), Command::ToYellow) is of type struct Green. This is where the type safe calls is enforced.
current_state has to be mut because it’s being reassigned at current_state = Some(match (state, cmd) {...});
move moves rx in the thread
When calling client.clone() only the sender is cloned, the rx lives in one and only one thread

What is missing?

See the following main.rs, it sometimes does not print because the main thread exits so the client thread never got the chance to print out the messages. A sleep was added just to demo the print behavior.

use std::{thread::sleep, time::Duration};

mod state;

fn main() {
    let remote_control = state::spawn_client();
    remote_control.to_green();
    remote_control.to_yellow();
    remote_control.to_red();
    remote_control.to_green();

    sleep(Duration::from_secs(1));
}