Rust State Machine Part 2

Table of Content (Click to Expand)

tldr

Pass in a Sender to Command
tokio::spawn() takes in a future.
tokio’s rx.recv() returns Option not Result
pub async fn blah(), async is between pub and fn
Two Way Communication code only is here
Full Part 2 Code using tokio is here
There are some unresolved issues but those would be handled in part 3

Overview

In part1 we have something but there were two issues left.

The remote control is essentially fire and forget. In order for the client to get any sorts of responses back the client would need to poll the state.
The main function requires a sleep() to wait for the threads to print statuses.

Add a Sender To The Command

Our current TrafficLightRemote looks like the following


#[derive(Debug)]
enum Command {
    ToGreen,
    ToYellow,
    ToRed,
}

#[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();
    }
}

We would be making two changes here to support two way communication.

Step 1: Update Command and TrafficLightRemote

We first add std::sync::mpsc::Sender<Result<(), String>> to each variant of Command

#[derive(Debug)]
enum Command {
    ToGreen(std::sync::mpsc::Sender<Result<(), String>>),
    ToYellow(std::sync::mpsc::Sender<Result<(), String>>),
    ToRed(std::sync::mpsc::Sender<Result<(), String>>),
}

Result<(), String> means if it is a success we want () and if something is wrong we just have a String.

A few things that could potentially be touched up later

Switch Result<(), String> to Result<(), Box<dyn std::error::Error + Send + Sync>> so the error could have more information than just a String
Instead of passing a Sender<Result<(), String>> in say ToGreen, pass in Option<Sender<Result<(), String>>> for cases where we actually want to fire and forget

With the change above when we call to_green we could block on the receiving side of the sender passed in to the enum.

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

impl TrafficLightRemote {
    pub fn to_green(&self) -> Result<(), String> {
        // tx and rx are a pair
        let (tx, rx) = mpsc::channel();
        self.sender.send(Command::ToGreen(tx)).unwrap();

        // this blocks
        match rx.recv() {
            Ok(inner_result) => inner_result,
            Err(e) => Err(e.to_string()),
        }
    }

    // ...
}

Most important things are commented in code. But let’s go over it again.

When to_green is called, we create a new std::sync::mpsc::channel, this gives you a pair of tx, rx where anything send to tx would come out from rx.

We pass the tx to Command::ToGreen(tx). So later when we pattern match on the ToGreen command the tx channel could be pulled out and have information passed back to it.

We then block by match rx.recv(). If we get an Ok() we just forward the result.

rx.recv() returns Result<Result<(), String>, std::sync::mpsc::RecvError>, we need to pull out the inner_result from it which type is Result<(), String>

If we get an Err we return the to_string of that error.

So far our TrafficLightRemote’s to_green can wait for the results, it’s no longer a fire and forget. But we still need to handle the other part where it actually sends things to the tx in Command::ToGreen(tx).

Step 2: Update Client Code And Send Information To The tx In Command

Let’s have a look at the current code

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

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

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

            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 command {:?} on {:?}", _bad_cmd, original_state);
                    original_state
                }
            });
        }

        println!("worker task done");
    });

    TrafficLightRemote { sender: tx }
}

The part that we are going to change is the match arms of match (state, cmd) {}.

We pattern match out the tx channel that was in the enum and send information back to it then we are all set.

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

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

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

            current_state = Some(match (state, cmd) {
                (TrafficLight::Green(green), Command::ToYellow(reply_channel)) => {
                    // Send Ok(()) back to reply_channel which is the tx passed in Command::ToYellow
                    let _ = reply_channel.send(Ok(()));
                    TrafficLight::Yellow(green.to_yellow())
                }
                (TrafficLight::Yellow(yellow), Command::ToRed(reply_channel)) => {
                    let _ = reply_channel.send(Ok(()));
                    TrafficLight::Red(yellow.to_red())
                }
                (TrafficLight::Red(red), Command::ToGreen(reply_channel)) => {
                    let _ = reply_channel.send(Ok(()));
                    TrafficLight::Green(red.to_green())
                }
                (
                    original_state,
                    Command::ToYellow(reply_channel)
                    | Command::ToRed(reply_channel)
                    | Command::ToGreen(reply_channel),
                ) => {
                    let _ =
                        reply_channel.send(Err("invalid state command combination".to_string()));
                    original_state
                }
            });
        }

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

    TrafficLightRemote { sender: tx }
}

There you go, the successful state transitions would send back Ok(()) to tx and the errors would be sent back as Err("invalid state command combination".to_string())

Let’s test this out in main without sleep(). There’s still a timing issue with the final output, which I’ll fix in Part 3.

fn main() {
    let remote_control = state::spawn_client();
    let res = remote_control.to_green();
    println!("{:?}", res);

    let res = remote_control.to_yellow();
    println!("{:?}", res);

    let res = remote_control.to_red();
    println!("{:?}", res);

    let res = remote_control.to_green();
    println!("{:?}", res);
}

shengy@Mac ~/repos/ruststuff/practice-statemachine-types ((part2-two-way-communication)|clean) [12:54:58]
> cargo run
   Compiling practice-statemachine-types v0.1.0 (/Users/shengy/repos/ruststuff/practice-statemachine-types)
    Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.48s
     Running `target/debug/practice-statemachine-types`
start as red
red to green
Ok(())
green to yellow
Ok(())
yellow to red
Ok(())
red to green
Ok(())
all threads done

Switch To tokio

The end goal of this state machine is to learn to write a BlazingMQ client. tokio is the defacto library for networking and is the defacto async executor. To plug the state machine in there it’s best to move from std::sync::mpsc to tokio::sync::mpsc.

into_split might be a separate post

At high level the steps should be

Add tokio as a dependency
Switch all std::sync::mpsc to tokio::sync::mpsc. For the reply channel we just created, tokio::sync::oneshot is actually a better fit
Update TrafficLightRemote’s methods
Update spawn_client code from threads to async
Update main to be a tokio main

Step 1: Install tokio

This is simple, cargo add tokio -F full and voilà.

Technically speaking you might want to go over the actual features you need but I’m a bit lazy for that now.

Step 2: Switch Channels To tokio

Replace std::sync::mpsc to tokio::sync::mpsc, but for the channel in Command we use tokio::sync::oneshot

Relevant snippet below

enum Command {
    ToGreen(tokio::sync::oneshot::Sender<Result<(), String>>),
    ToYellow(tokio::sync::oneshot::Sender<Result<(), String>>),
    ToRed(tokio::sync::oneshot::Sender<Result<(), String>>),
}

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

Step 3: Update TrafficLightRemote methods

This would be touching the same code we just did, I’ll use the to_green function as an example.

impl TrafficLightRemote {
    pub async fn to_green(&self) -> Result<(), String> {
        let (tx, rx) = tokio::sync::oneshot::channel();
        self.sender.send(Command::ToGreen(tx)).await.unwrap();
        match rx.await {
            Ok(inner_result) => inner_result,
            Err(e) => Err(e.to_string()),
        }
    }
}

Now sender.send is returning a future, so we would need to .await on it
Because we are .awaiting in the function, to_green needs to be async
The way we wait for the channel rx is rx.await, note that this is special for the oneshot channel, for tokio::sync::mpsc channels to await for rx is rx.recv().await (We would see this in the spawn_client function later).

Step 4: Update spawn_client

The notable changes are

let (tx, mut rx) = tokio::sync::mpsc::channel(32);
- rx needs to be mut
- channel takes in 32 which is the buffer size, note that 32 here means 32 messages
std::thread::spawn() to tokio::spawn()
- The biggest difference here is thread::spawn(move || {}) takes in FnOnce where tokio::spawn(async move {}) takes in a future created by an async move block.
while let Some(cmd) = rx.recv().await {}
- rx.recv().await returns Option not Result

See code below

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

    tokio::spawn(async move {
        // initial state
        let mut current_state = Some(TrafficLight::Red(Red::new()));

        while let Some(cmd) = rx.recv().await {
            let Some(state) = current_state.take() else {
                println!("fatal error, invalid state");
                break;
            };

            current_state = Some(match (state, cmd) {
                //
                (TrafficLight::Green(green), Command::ToYellow(reply_channel)) => {
                    let _ = reply_channel.send(Ok(()));
                    TrafficLight::Yellow(green.to_yellow())
                }
                (TrafficLight::Yellow(yellow), Command::ToRed(reply_channel)) => {
                    let _ = reply_channel.send(Ok(()));
                    TrafficLight::Red(yellow.to_red())
                }
                (TrafficLight::Red(red), Command::ToGreen(reply_channel)) => {
                    let _ = reply_channel.send(Ok(()));
                    TrafficLight::Green(red.to_green())
                }
                (
                    original_state,
                    Command::ToYellow(reply_channel)
                    | Command::ToRed(reply_channel)
                    | Command::ToGreen(reply_channel),
                ) => {
                    let _ =
                        reply_channel.send(Err("invalid state command combination".to_string()));
                    original_state
                }
            });
        }

        // this should be "worker task done"
        // this "all thread done" inaccuracy would be addressed in part3
        println!("all threads done");
    });

    TrafficLightRemote { sender: tx }
}

Step 5: Update main To Be a tokio main

This is just slapping on the #[tokio::main] procedural macro to main and call .await on the remote control

mod state;

#[tokio::main]
async fn main() {
    let remote_control = state::spawn_client();
    let res = remote_control.to_green().await;
    println!("{:?}", res);

    let res = remote_control.to_yellow().await;
    println!("{:?}", res);

    let res = remote_control.to_red().await;
    println!("{:?}", res);

    let res = remote_control.to_green().await;
    println!("{:?}", res);
}

shengy@Mac ~/repos/ruststuff/practice-statemachine-types ((part2)|clean) [13:34:05]
> cargo run
   Compiling practice-statemachine-types v0.1.0 (/Users/shengy/repos/ruststuff/practice-statemachine-types)
    Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.26s
     Running `target/debug/practice-statemachine-types`
start as red
red to green
Ok(())
green to yellow
Ok(())
yellow to red
Ok(())
red to green
Ok(())

Further things to think about

How does async/await work in Rust?
async-await-in-asyncio was something I wrote before on python’s asyncio
A few issues remaining with the code that would be addressed in part3
- The reply is sent before the state change happened.
- How to actually determine “all threads done”

Review Notes Written by Codex

This section, including the diagrams, was written by Codex, the AI assistant that reviewed this post. The explanations and suggested fixes below are Codex’s review notes.

There are two separate timing issues here, and both apply to the threaded and Tokio versions.

First, the worker sends Ok(()) before calling the state transition method. The caller can receive that reply and continue while the transition is still pending. For the last command, main could exit before the transition prints its output. So the reply currently means “command accepted,” not “transition completed.” The suggested fix for Part 3 is to perform the transition and store the new state before sending the reply.

One possible ordering in the threaded version looks like this. Time moves downward:

Worker                              Main
------                              ----
Receive the final command
Send Ok(()) ----------------------> Receive Ok(())
                                    Print Ok(())
                                    Exit the program
-----------------------------------------------
Transition and its print may never run

In the Tokio version, the caller can also receive the reply before the transition finishes. Returning from main starts runtime shutdown; receiving that reply is not a guarantee of task completion.

Second, receiving the last reply does not mean the worker has finished. The println!("all threads done") runs after the receive loop exits. In this example, that happens once all remote handles are dropped and the queued commands have been processed. Dropping the handles lets the worker exit, but does not wait for it to do so, which means the final message is not guaranteed to appear.

Even after fixing the reply ordering, these are separate events:

Worker                              Main
------                              ----
Perform transition and store state
Send Ok(()) ----------------------> Receive Ok(())
Wait for another command            Drop all remote handles
                                    Return from main

May still be pending without waiting for the worker:
  Receive loop ends
  Print "worker done"
  Worker returns

The suggested shutdown fix is to keep the JoinHandle returned by std::thread::spawn or tokio::spawn. After dropping all remote handles, including any clones, call .join() for the thread or .await the Tokio task’s handle. Dropping the remotes needs to happen first; otherwise the worker can keep waiting for another command while main waits for the worker. Also, “worker done” would be a more accurate message: finishing this one worker does not mean all threads have finished.

With both fixes, the intended sequence is:

Worker                              Main
------                              ----
Perform transition and store state
Send Ok(()) ----------------------> Receive Ok(())
                                    Drop all remote handles
                                    Join / await worker handle
Drain any remaining commands        |
Receive loop ends                   | Wait for worker completion
Print "worker done"                 |
Worker returns -------------------> Join / await completes
                                    Return from main

Update: 2026-09-26

Tags:

Rust

Learning

<< Rust State Machine Part 1 Rust State Machine Part 3 >>