TulipeMoutarde.be

About

Table of Contents

Dependency injection in Axum handlers. A quick tour

Written on August 24 2023.

tl;dr; How to share state between handlers in axum. From static dispatch against a struct to dynamic dispatch with trait objects

Dependency injection is an elaborate word to say that we want to introduce flexibility in software systems. Instead of having one big ball of functions, we slice it into components that talk to each other. Sometimes we want to be able to swap a component for another one in certain circumstances. Testing is one of those circumstance: we want to test components in isolation. This is especially true when we have external dependencies (eg., talking to a remote APIs).

Rust is no different than all the other languages when it comes this principle. But Rust being Rust, the user gets to decide how to do it. The language provides the building bricks and the user chooses themself the abstractions they want based on the performance hit they are ready to trade off for convenience.

We will start with the static dispatch approach with no flexibility whatsoever and move to more flexible approach.

Context

The context of our little exploration is an axum app that uses a database of some sort. The database access is provided to all the handlers through the State mechanism provided by Axum. We also have a templating system to generate HTML. We picked handlebars but it could be anything.

We assume that the database access is just one component that is used by the handler. We can imagine that other components are made available through the state (e.g., an AWS service, an external payment API or a redis connection).

All the code is available on github.

Static Dispatch against a struct

This is the straightforward approach. We define a struct that will hold everything and use it in our state.

#[derive(Clone)]
pub struct AppState {
    pub templates: Handlebars<'static>,
    pub db: MemoryDB,
}

#[derive(Clone)]
pub struct MemoryDB {
    items: Arc<HashMap<Uuid, String>>,
}

impl MemoryDB {
    pub fn new() -> Self { /* .. */ }
    pub async fn all_items(&self) -> Vec<(&Uuid, &String)> { /* .. */ }
    pub async fn get_item(&self, item_id: &Uuid) -> Option<&String> { /* .. */ }
}

pub fn build_router() -> Router {
    let state = AppState {
        templates: build_templates(),
        db: MemoryDB::new(),
    };

    Router::new()
        .route("/", get(handlers::index))
        .route("/item/:id", get(handlers::show))
        .with_state(state)
}


mod handlers {
    pub async fn index(State(AppState { db, templates }): State<AppState>) -> Html<String> { 
        /* .. */ 
    }

    pub async fn show(
        Path(id): Path<Uuid>,
        State(AppState { db, templates }): State<AppState>,
    ) -> Html<String> { 
        /* .. */
    }

AppState always holds a MemoryDB and our handlers are tied to it because we've declared that the state holds an AppState (State<AppState> in the handler signatures). This is the simplest way to use a shared state on all the handlers but it is the less flexible as the DB implementation is hardcoded.

Static Dispatch with Generics

To be able to swap the database component with another it can't be a struct anymore but a trait. Coding to interfaces, not implementation would say the purist:

#[async_trait]
pub trait DB: Send + Sync {
    async fn all_items(&self) -> Vec<(&Uuid, &String)>;
    async fn get_item(&self, item_id: &Uuid) -> Option<&String>;
}

#[derive(Clone)]
pub struct MemoryDB {
    items: Arc<HashMap<Uuid, String>>,
}

impl MemoryDB {
    pub fn new() -> Self { /* .. */ }
}

#[async_trait]
impl DB for MemoryDB {
    async fn all_items(&self) -> Vec<(&Uuid, &String)> { /* .. */ }
    async fn get_item(&self, item_id: &Uuid) -> Option<&String> { /* .. */ }
}

Add a generic with a trait bound to AppState and we're done.


#[derive(Clone)]
pub struct AppState<D: DB> {
    pub templates: Handlebars<'static>,
    pub db: D,
}

And that's where the painful part starts. Generics are viral and propagates throughout the code if you don't contain them.

Our handlers are not so simple anymore as we have to define the trait bounds for each of them:

pub async fn index<D: DB>(
    State(AppState { db, templates }): State<AppState<D>>,
) -> Html<String> { 
    /* .. */ 
}

pub async fn show<D: DB>(
    Path(id): Path<Uuid>,
    State(AppState { db, templates }): State<AppState<D>>,
) -> Html<String> { 
    /* .. */ 
}

In this example it is quite simple as the only generic parameter is the database. But the more components we need to swap, the more generics (and their bounds) we have to add for each and every handler. This situation is well described in Julio Merino's article about static dispatch.

Traits to the rescue!

Static Dispatch against a trait

Once again, traits can help. Instead of having AppState as a concrete implementation, we can turn it into a trait which will stop the infamous generics invasion:

pub trait AppState: Clone + Send + Sync + 'static {
    type D: database::DB;

    fn db(&self) -> &Self::D;
    fn templates(&self) -> &Handlebars<'static>;
}

#[derive(Clone)]
pub struct RegularAppState {
    pub templates: Handlebars<'static>,
    pub db: MemoryDB,
}

impl AppState for RegularAppState {
    type D = MemoryDB;

    fn db(&self) -> &Self::D {
        &self.db
    }

    fn templates(&self) -> &Handlebars<'static> {
        &self.templates
    }
}

The users of the AppState trait don't have to know the actual type of D, only types implementing the trait have to provide the actual type. For example, the RegularState implementation og AppState says that it uses MemoryDB.

We can now build a router without knowing the explicit type of the state:

pub fn build_router() -> Router {
    let state = RegularAppState {
        templates: build_templates(),
        db: MemoryDB::new(),
    };
    build(state)
}

fn build<A: AppState>(state: A) -> Router {
    Router::new()
        .route("/", get(handlers::index::<A>))
        .route("/item/:id", get(handlers::show::<A>))
        .with_state(state)
}

and the handlers are back to a sane definition:

pub async fn index<S: AppState>(State(state): State<S>) -> Html<String> { 
    /* .. */ 
}

pub async fn show<S: AppState>(Path(id): Path<Uuid>, State(state): State<S>) -> Html<String> { 
    /* .. */ 
}

A bit of trait dance and no dynamic dispatch. It is quite handy when performance is important.

There are two drawbacks to this approach though:

  1. For every combination of dependencies, we have to define a new struct that implements the AppState trait. Again, we only have the DB here but it might become crowded quite quickly in a real application with different test scenarios.
  2. As AppState is a trait, there is no way to directly match agains the actual fields in the handler definition. Components are exposed through methods on the trait. Some might argue that it is a good thing.

Dynamic Dispatch against a trait object

This approach will be familiar for anyone coming from an object oriented language. In Rust parlance, the trick is to use a trait object via the dyn keyword.

This incurs a small performance hit because the application will determine at runtime the actual implementation to use when a function is called. But in the context of a web application it does not matter that much (until it does) as the overhead of all the other things (eg., database query, templating) is much bigger.

#[derive(Clone)]
pub struct AppState {
    pub templates: Handlebars<'static>,
    pub db: Arc<dyn DB + Send + Sync>,
}

pub fn build_router() -> Router {
    let state = AppState {
        templates: build_templates(),
        db: Arc::new(MemoryDB::new()),
    };

    Router::new()
        .route("/", get(handlers::index))
        .route("/item/:id", get(handlers::show))
        .with_state(state)
}

Handlers are back to their initial implementation:


pub async fn index(State(AppState { db, templates }): State<AppState>) -> Html<String> {
    /* .. */
}

pub async fn show(
        Path(id): Path<Uuid>,
        State(AppState { db, templates }): State<AppState>,
    ) -> Html<String> {
    /* .. */
}

Conclusion

The full running code is available on github.

As usual, Rust lets you choose the abstraction level that fits the need for the task at hand. It is sometimes daunting to know which approach is the right one for a specific use case.

When in doubt, I would go with the simplest solution that provides enough performance. In this case the dynamic dispatch route is the most flexible while being more than fast enough for a web application that is driven by SQL calls and external APIs.

If you are interested in the subject, Julio Merino's blog is worth a read as he explores this space. Luca Palmieri's blog and book Zero To Production in Rust are wonderful resources to get started in web development with rust.