Axum routes with a Handler

As your application grows, it is probably not a good idea to group the logic for all your routes within the web server. This is where the Handler concept comes in. The Handler is a struct that contains the business logic to drive the web application. The Handler is responsible for processing the request and returning a response. This is a common pattern in web development.

Let's refactor our greeting application to use a Handler struct. We will create a WebHandler struct that will contain the logic for greeting a visitor and the logic for saying goodbye. We'll also take the opportunity to separate the data model into a separate module.

main.rs

use std::sync::Arc;

use axum::{
    extract::{Path, State},
    Json,
    Router, routing::{delete, get},
};

use crate::handler::WebHandler;
use crate::model::Greeting;

mod handler;
mod model;

#[tokio::main]
async fn main() {
    // Create a shared state for our application. We use an Arc so that we clone the pointer to the state and
    // not the state itself.
    let app_state = Arc::new(WebHandler::default());

    // set up our application with "hello world" route at "/
    let app = Router::new()
        .route("/hello/:visitor", get(greet_visitor))
        .route("/bye", delete(say_goodbye))
        .with_state(app_state);

    // start the server on port 3000
    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

/// Extract the `visitor` path parameter and use it to greet the visitor.
/// We also use the `State` extractor to access the shared `Handler` and call the `greet` method.
/// We use `Json` to automatically serialize the `Greeting` struct to JSON.
async fn greet_visitor(
    State(handler): State<Arc<WebHandler>>,
    Path(visitor): Path<String>,
) -> Json<Greeting> {
    Json(handler.greet(visitor))
}

/// Say goodbye to the visitor.
async fn say_goodbye(State(handler): State<Arc<WebHandler>>) -> String {
    handler.say_goodbye()
}

handler.rs

use std::{
    sync::atomic::AtomicU16,
    sync::atomic::Ordering::Relaxed,
};

use crate::model::Greeting;

/// A handler for our web application.
pub struct WebHandler {
    number_of_visits: AtomicU16,
}

impl WebHandler {
    /// Greet the visitor and increment the number of visits.
    pub fn greet(&self, visitor: String) -> Greeting {
        let visits = self
            .number_of_visits
            .fetch_add(1, Relaxed);
        Greeting::new("Hello", visitor, visits)
    }

    /// Say goodbye to the visitor.
    pub fn say_goodbye(&self) -> String {
        "Goodbye".to_string()
    }
}

impl Default for WebHandler {
    fn default() -> Self {
        WebHandler {
            number_of_visits: AtomicU16::new(0),
        }
    }
}

model.rs

use serde::{Deserialize, Serialize};

#[derive(Serialize, Deserialize)]
pub struct Greeting {
    greeting: String,
    visitor: String,
    visits: u16,
}

impl Greeting {
    pub(crate) fn new(greeting: &str, visitor: String, visits: u16) -> Self {
        Greeting {
            greeting: greeting.to_string(),
            visitor,
            visits,
        }
    }
}

By splitting the logic into a Handler struct, we have made our code more modular and easier to maintain. The Handler struct contains the business logic for our application, while the model module contains the data model. This separation of concerns makes our code more organized and easier to understand.

Keeping the business logic separate from the web I/O is a good practice. It helps with testing and also allows you to add other interfaces to your application, such as a command-line interface or a gRPC API.

Handler Trait

In the above example, we used a Handler struct to contain the business logic for our application. However, we can also use a trait to define the interface for our handler. This allows us to define multiple handlers that implement the same interface. This is extremely useful when you want to create a pluggable architecture for your application.

Using a trait to define the interface for your handler allows you to 'mock' the handler in your tests. This makes it easier to test your web endpoints in isolation.

Let's refactor our example to use a Handler trait instead of a concrete WebHandler struct.

main.rs

use std::sync::Arc;

use axum::{
    extract::{Path, State},
    Json,
    Router, routing::{delete, get},
};

use crate::handler::{GreetingHandler, WebHandler};
use crate::model::Greeting;

mod handler;
mod model;

type AppState = Arc<dyn GreetingHandler>;

#[tokio::main]
async fn main() {
    // Create a shared state for our application. We use an Arc so that we clone the pointer to the state and
    // not the state itself.
    let app_state: AppState = Arc::new(WebHandler::default());

    // set up our application with "hello world" route at "/
    let app = Router::new()
        .route("/hello/:visitor", get(greet_visitor))
        .route("/bye", delete(say_goodbye))
        .with_state(app_state);

    // start the server on port 3000
    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

/// Extract the `visitor` path parameter and use it to greet the visitor.
/// We also use the `State` extractor to access the shared `Handler` and call the `greet` method.
/// We use `Json` to automatically serialize the `Greeting` struct to JSON.
async fn greet_visitor(
    State(handler): State<AppState>,
    Path(visitor): Path<String>,
) -> Json<Greeting> {
    Json(handler.greet(visitor))
}

/// Say goodbye to the visitor.
async fn say_goodbye(State(handler): State<AppState>) -> String {
    handler.say_goodbye()
}

handler.rs

use std::{
    sync::atomic::AtomicU16,
    sync::atomic::Ordering::Relaxed,
};

use crate::model::Greeting;

/// A trait for handling greetings.
pub trait GreetingHandler: Send + Sync {
    fn greet(&self, visitor: String) -> Greeting;
    fn say_goodbye(&self) -> String;
}

/// A greeting handler implementation for our web application.
pub struct WebHandler {
    number_of_visits: AtomicU16,
}

impl GreetingHandler for WebHandler {
    /// Greet the visitor and increment the number of visits.
    fn greet(&self, visitor: String) -> Greeting {
        let visits = self
            .number_of_visits
            .fetch_add(1, Relaxed);
        Greeting::new("Hello", visitor, visits)
    }

    /// Say goodbye to the visitor.
    fn say_goodbye(&self) -> String {
        "Goodbye".to_string()
    }
}

impl Default for WebHandler {
    fn default() -> Self {
        WebHandler {
            number_of_visits: AtomicU16::new(0),
        }
    }
}

Our model.rs remains the same.

By using a Handler trait, we have made our code more modular and easier to extend. In this example we use dynamic dispatch to allow different implementations of the GreetingHandler trait to be used at runtime. This allows us to create different handlers for different environments, such as a test environment or a production environment.

Dynamic dispatching comes with a small performance penalty, as the compiler cannot optimize the code as much as with static dispatch. We can make a final update to our code to use static dispatch.

Static Dispatch

In the previous example, we used dynamic dispatch to allow different implementations of the GreetingHandler trait to be used at runtime. However, we can also use static dispatch to allow the compiler to optimize the code more efficiently. This is done by specifying the concrete type of the GreetingHandler trait at compile time.

main.rs

use crate::{
    handler::WebHandler,
    web::AxumWebServer,
};

mod handler;
mod model;
mod web;

#[tokio::main]
async fn main() {
    let handler = WebHandler::default();
    let web_server = AxumWebServer::new(handler);
    web_server.start().await;
}

handler.rs

use std::{
    sync::atomic::AtomicU16,
    sync::atomic::Ordering::Relaxed,
};


use crate::model::Greeting;

/// A trait for handling greetings.
pub trait GreetingHandler: Send + Sync + 'static {
    fn greet(&self, visitor: String) -> Greeting;
    fn say_goodbye(&self) -> String;
}

/// A greeting handler implementation for our web application.
pub struct WebHandler {
    number_of_visits: AtomicU16,
}

impl GreetingHandler for WebHandler {
    /// Greet the visitor and increment the number of visits.
    fn greet(&self, visitor: String) -> Greeting {
        let visits = self
            .number_of_visits
            .fetch_add(1, Relaxed);
        Greeting::new("Hello", visitor, visits)
    }

    /// Say goodbye to the visitor.
    fn say_goodbye(&self) -> String {
        "Goodbye".to_string()
    }
}

impl Default for WebHandler {
    fn default() -> Self {
        WebHandler {
            number_of_visits: AtomicU16::new(0),
        }
    }
}

web.rs

use std::sync::Arc;

use axum::{
    extract::{Path, State},
    Json,
    Router,
    routing::{delete, get},
};

use crate::{
    handler::GreetingHandler,
    model::Greeting,
};

type AppState<G> = Arc<G>;

pub struct AxumWebServer<G: GreetingHandler> {
    app_state: AppState<G>,
}

impl<G: GreetingHandler> AxumWebServer<G> {
    pub fn new(handler: G) -> Self {
        // Create a shared state for our application. We use an Arc so that we clone the pointer to the state and
        // not the state itself.
        let app_state: AppState<G> = Arc::new(handler);
        AxumWebServer { app_state }
    }

    pub async fn start(&self) {
        let app_state = self.app_state.clone();

        // set up our application with "hello world" route at "/
        let app = Router::new()
            .route("/hello/:visitor", get(Self::greet_visitor))
            .route("/bye", delete(Self::say_goodbye))
            .with_state(app_state);

        // start the server on port 3000
        let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
        axum::serve(listener, app).await.unwrap();
    }


    /// Extract the `visitor` path parameter and use it to greet the visitor.
    /// We also use the `State` extractor to access the shared `Handler` and call the `greet` method.
    /// We use `Json` to automatically serialize the `Greeting` struct to JSON.
    async fn greet_visitor(
        State(handler): State<AppState<G>>,
        Path(visitor): Path<String>,
    ) -> Json<Greeting> {
        Json(handler.greet(visitor))
    }

    /// Say goodbye to the visitor.
    async fn say_goodbye(State(handler): State<AppState<G>>) -> String {
        handler.say_goodbye()
    }
}

The model.rs stays the same.

By using static dispatch, we have eliminated the performance penalty of dynamic dispatch. This makes our code more efficient and easier to optimize. As you can see the statically dispatched version of our code comes with some more cognitive overhead. This is where you have to decide what is more important to you: performance or ease of development.

Finally, we have isolated the web server logic in the web.rs file. Our main.rs file is now very clean and only contains the minimum code to start the application.

With these latest changes, our code is getting a little bit closer to a production-ready state.

Benchmarking is really the only way to know if the performance improvements are worth the added complexity. The overhead that the indirection of dynamic dispatch introduces is usually negligible. Remember that the execution time of the business logic will typically dwarf the overhead of the dynamic dispatch of the REST API calls.