Rust State Machine Part 3

tldr

• I missed two things when learning from Gemini
• Send the confirmation last
• return JoinHandle for graceful shutdown
• Part 3 code is here
• Every time tokio::sync::mpsc::channel(32) is called a channel is created.
  • This gives you a pair of (tx, rx) which links to each end of the channel
  • When tx is cloned you have multiple tx to send to the same channel
  • In this example, the number of while rx.recv().await {} loops matches the number of channels
  • mpsc’s rx can not be cloned. Though some mpmc implementations allow you to clone rx
• The above means
  • If you call spawn_client twice you created two traffic light states, and they are separate
  • If you call spawn_client and clone the client/sender/remote, you created a remote control that could send commands to change the one and only traffic light state

Overview

In the end of part 2, codex found two issues but my brain was a bit toasted so I decided to leave it and tackle later.

This post is me going through the issues and applying the fixes.

Issue 1: Sending Confirm Before State Change

The first issue is in the code in our tokio worker which tries to find out the next state based on the current state and command.

The reply was sent before the state changed

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

// BAD EXAMPLE
current_state = Some(match (state, cmd) {
    (TrafficLight::Green(green), Command::ToYellow(reply_channel)) => {
        // here
        let _ = reply_channel.send(Ok(()));
        TrafficLight::Yellow(green.to_yellow())
    }
    // ...
});

// ...

The actual state change happens when we put the state back into current_state, but before that the Ok(()) was already sent.

Issue 1 Fix

The solution is to separate deducing the state and setting the state and replying the channel. Make sure the channel reply is the last thing we do.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

while let Some(cmd) = rx.recv().await {
    // take state out from current_state
    // ...

    let (next_state, reply_channel, result) = match (state, cmd) {
        (TrafficLight::Green(green), Command::ToYellow(reply_channel)) => (
            TrafficLight::Yellow(green.to_yellow()),
            reply_channel,
            Ok(()),
        ),
    };
    // ...
    current_state = Some(next_state);
    let _ = reply_channel.send(result);
}

Issue 2: Not Gracefully Shutting Down

This is an issue where the threads/tokio tasks are not coordinating with the main thread.

In part2’s spawn_client there is a println!("all threads done"); after while let Some(cmd) = rx.recv().await {} breaks out.

This line is (1) not being accurate in the tokio version and (2) not guaranteed to be executed.

The while let Some(cmd) = rx.recv().await {} loop breaks out when all the senders are dropped and the existing queued messages are processed.

But in main.rs the drop of remote_control happens right before the main thread exits (i.e. the main program ends)

This means the line after while is racing against main shutting down. So println! could be ran if it won the race, or else the print message is never executed (because main shutting down won).

I didn’t realize how similar this is with goroutines, this is why in go code that needs to coordinate goroutines would use WaitGroup or a done channel to coordinate the goroutines.

Also another issue was the printed text is wrong, in the last part3 version each spawn_client actually spawns tokio workers. So the printed text should be println!("worker done");

Side Note: Difference Between Creating New Remotes vs Cloning Remotes

Client could create two clients by

1. Calling spawn_client twice, or
2. Call spawn_client once and clone() the sender

There exists a big difference between these two.

(1) Is creating two channels and state machines, in other words it creates two traffic light systems.

(2) Is creating one channel and one state machine but creates a duplicated remote control that can control the same state (traffic light)

This matters because when you think about how many "worker done" would be printed.

Since calling spawn_client would create a worker (hence it’s while loop), we could say the number of "worker done" should be printed is supposed to be the number of workers created, which is how many times spawn_client was called.

Issue 2 Fix

To handle the coordinate issue we would return the JoinHandle from tokio::spawn() to main, and (1) make sure the client is dropped before exiting then (2) await on the returned JoinHandle

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

// main.rs
#[tokio::main]
async fn main() {
    let (remote_control, join_handle) = 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);

    drop(remote_control);
    join_handle.await.unwrap();
}

spawn_client also needs to return the JoinHandle

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

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

    let join_handle = tokio::spawn(async move {
        // initial state
        let mut current_state = Some(TrafficLight::Red(Red::new()));
        while let Some(cmd) = rx.recv().await {
            // ...
        }
        println!("worker done");
    });

    (TrafficLightRemote { sender: tx }, join_handle)
}

Do we need to manually drop the client in our example?

The answer to above is yes because otherwise there would be a deadlock.

• Worker is waiting for the while loop to end to signal JoinHandle
• Main is blocked on JoinHandle but won’t drop the client

Another thing important to realize is, if you cloned remotes, all those remotes needs to be dropped or else a deadlock would happen.

Some After Thoughts

Learning from LLM

I am trying to use LLMs to learn things faster, but it is very hard to point out issues while learning a new thing on the fly. After thinking I think one of the solutions is to use it learn then find a way to verify then come back for some notes.

Self Retro, Why Did I Miss It

I didn’t realize that issue 2 was something I am very familiar in Go, it’s either setting up the WaitGroup or block on a channel waiting for the “Done” signal.

Then I was thinking how did I miss it.

I think this happened because I was in the process of learning how to use tokio and how to handle the state by workers. Also the complexity of Rust’s syntax might have contributed a bit. Or it might be I haven’t been doing goroutines for a while so I got Rusty.

Also I think I trusted LLMs a bit too much even in terms of learning.