A Simple State Machine Library in Rust

Rust, a systems programming language celebrated for its performance and safety guarantees, is perfect for building high-performance applications. In this article, we’ll explore how to create a simple state machine library using Rust, suitable for a wide range of designs.

Code is provided at https://github.com/nano-e/ne-fsm

Introduction to State Machines

State machines, or Finite State Machines (FSMs), effectively manage system behavior and control flow in applications. They simplify complex systems using a finite number of states, events, and transitions, making them easier to maintain and debug.

FSMs are essential for:

1. Simplifying complex systems: FSMs break down systems into manageable parts, streamlining understanding and issue identification.
2. Ensuring predictable behavior: FSMs enforce transition rules, reducing unexpected behavior.
3. Facilitating maintenance and updates: Discrete states and transitions allow for easier updates without affecting the overall system.
4. Improving debugging and testing: Clear representations and modularity enable better issue identification and isolated testing.

Finite State Machines (FSMs) are often represented using state machine diagrams, which provide a visual representation of the system’s states, transitions, and events, enabling a clearer understanding of the system’s behavior and logic.

Here’s an example of a state machine diagram for a Call processing:

State Diagram for call processing

This diagram represents a simple call state machine with five states: Idle, Dialing, Ringing, Connected, and Disconnected. Transitions between these states are triggered by events such as Dial, Reject, IncomingCall , Answer, Hangup and Reset.

Implementation

The Rust state machine implementation has three main components: FsmEnum, Stateful traits, and StateMachine struct.

1. FsmEnum: A trait that defines creating new state machine states based on enums, simplifying state object creation.
2. Stateful: A trait outlining state transition event handling with on_enter, on_event, and on_exit methods for custom behavior.
3. StateMachine: A generic struct representing a state machine instance, managing current state and providing initialization, event processing, and state retrieval methods.

The library has synchronous and asynchronous versions, both using FsmEnum, Stateful, and StateMachine components. The async version uses the async_trait crate for async traits and methods, enabling smooth integration with Rust’s async/await syntax.

States are cached in a HashMap for performance and memory efficiency. A global event handler can be defined for non-state-specific events, increasing flexibility.

Leveraging Generics

The Rust state machine library takes advantage of Rust’s powerful generics feature to make the library portable and adaptable to various state machine designs. Generics enable the creation of reusable code that works with different types, allowing developers to build flexible and efficient state machines that cater to their specific needs.

This means that developers can use the same library to implement different state machines with various state and event types, without having to rewrite or modify the library code.

Here’s an example of how the StateMachine struct uses generics:

pub struct StateMachine<S: Hash + PartialEq + Eq + Clone + FsmEnum<S, CTX, E>, CTX, E: Debug> {
        states: HashMap<S, Box<dyn Stateful<S, CTX, E> + Send>>,
        current_state: Option<S>,
        context: CTX,
        global_event_handler: Option<Box<dyn EventHandler<S, CTX, E> + Send>>,
    }

In this example, S and E are generic types representing the states and events, respectively. The C generic type represents the context object (see below). This design allows the StateMachine struct to be instantiated with any specific set of enums for states and events.

Context Objects for Shared Data and Functions

Context objects play a crucial role in the Rust state machine library by providing a way to share data and functions between states. They can be used to store state-specific data, global data, or references to other parts of the system that the state machine interacts with.

The context object is passed as a mutable reference to the on_enter, on_event, and on_exit methods of the Stateful trait. This allows each state to access and manipulate shared data and functions when handling state transitions and events.

Here’s an example of how the context object is used in the Stateful trait:

// Define the Stateful trait, which contains the event handling methods for each state
    pub trait Stateful<S: Hash + PartialEq + Eq + Clone, CTX, E: Debug>
    {
        fn on_enter(&mut self, context: &mut CTX) -> Response<S>;
        fn on_event(&mut self, event: &E, context: &mut CTX) -> Response<S>;
        fn on_exit(&mut self, context: &mut CTX);
    }

The context object enables state machine states to interact with shared resources, communicate with other parts of the system, or store data that persists across state transitions. This feature enhances the flexibility and power of the Rust state machine library, making it suitable for a wide range of applications and scenarios.

Example (Implementing FSM for call state diagram)

We will implement the state diagram for a simple call (telecom), as described in the above state diagram. We will use our implementation of a Finite State Machine (FSM) in Rust with asynchronous support. This implementation will demonstrate how to define states and events, and how to handle state transitions based on events.

Import the required modules and dependencies:

use nefsm::Async::{self, FsmEnum, Stateful, Response};
use tokio::sync::mpsc::{Receiver, channel, Sender};
use std::fmt::Debug;
use async_trait::async_trait;

Define the CallState enum for the different states:

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum CallState {
    Idle,
    Dialing,
    Ringing,
    Connected,
    Disconnected,
}

Define the CallEvent enum for the different events that can occur:

#[derive(Debug)]
pub enum CallEvent {
    Dial,
    IncomingCall,
    Answer,
    Reject,
    HangUp,
    Reset,
}

Implement the FsmEnum trait for the CallState enum. This trait defines how to create a new state machine state based on a given enum value:

impl FsmEnum<CallState, CallContext, CallEvent> for CallState {
    fn create(enum_value: &CallState) -> Box<dyn Stateful<CallState, CallContext, CallEvent> + Send> {
        match enum_value {
            CallState::Idle => Box::new(IdleState {}),
            CallState::Dialing => Box::new(DialingState {}),
            CallState::Ringing => Box::new(RingingState {}),
            CallState::Connected => Box::new(ConnectedState {}),
            CallState::Disconnected => Box::new(DisconnectedState {}),
        }
    }
}

Define the CallContext struct to store the number of retries for dialing:

pub struct CallContext {
    pub retries: u32,
}
impl CallContext {
    pub fn new() -> Self {
        Self { retries: 0 }
    }
    pub fn increment_retries(&mut self) {
        self.retries += 1;
    }
    pub fn reset_retries(&mut self) {
        self.retries = 0;
    }
}

Implement the Idle, Dialing, Ringing, Connected, and Disconnected states using the Stateful trait. The Stateful trait outlines how a state should handle state transition events, including methods for handling events when entering a state (on_enter), receiving an event (on_event), and exiting a state (on_exit). This trait allows for custom behavior to be implemented for each state:

pub struct IdleState;
#[async_trait]
impl Stateful<CallState, CallContext, CallEvent> for IdleState {
    async fn on_enter(&mut self, _context: &mut CallContext) -> Response<CallState> {
        println!("Entering Idle state");
        Response::Handled        
}
    async fn on_event(&mut self, event: &CallEvent, _context: &mut CallContext) -> Response<CallState> {
        match event {
            CallEvent::Dial => Response::Transition(CallState::Dialing),
            CallEvent::IncomingCall => Response::Transition(CallState::Ringing),
            _ => {
                println!("Invalid event for Idle state");
                Response::Handled
            }
        }
    }
async fn on_exit(&mut self, _context: &mut CallContext) {
        println!("Exiting Idle state");
    }
}
//.
//.
//.
//The complete listing can be found in the github repository.

Conclusion

In conclusion, we demonstrated implementing a finite state machine using Rust and the nefsm library through a telecom call example. By defining states, events, and context, we showcased a modular and flexible approach for managing complex state transitions in various applications. Leveraging Rust's type system and asynchronous programming, developers can efficiently create maintainable state machines for real-world scenarios.