Production-ready microservice in Rust: 5. Application state

Adding shared state to our application

I demonstrated earlier how to load application configuration from environment variables or files, but our axum handler methods cannot use this information yet. To solve this problem we have to introduce application state and distribute this state to all the handler methods.

One small change first: we have to make our Settings struct cloneable. We need this because we pass it to the serve function as a reference, but the application state has to own its own dedicated copy, otherwise we cannot satisfy the Rust borrow checker. All we have to do is add the Clone trait to the derive macros in shelter_main/src/settings.rs:

#[derive(Debug, Deserialize, Default, Clone)]
#[allow(unused)]
pub struct Database {
    pub url: Option<String>,
}

#[derive(Debug, Deserialize, Default, Clone)]
#[allow(unused)]
pub struct Logging {
    pub log_level: Option<String>,
}

#[derive(Debug, Deserialize, Default, Clone)]
#[allow(unused)]
pub struct ConfigInfo {
    pub location: Option<String>,
    pub env_prefix: Option<String>,
}

#[derive(Debug, Deserialize, Default, Clone)]
#[allow(unused)]
pub struct Settings {
    #[serde(default)]
    pub config: ConfigInfo,
    #[serde(default)]
    pub database: Database,
    #[serde(default)]
    pub logging: Logging,
}

Now we can introduce the ApplicationState struct. Let's create the shelter_main/src/state/mod.rs file, and include mod state in lib.rs:

use crate::settings::Settings;
use std::sync::Arc;
use arc_swap::ArcSwap;

pub struct ApplicationState {
    pub settings: ArcSwap<Settings>,
}

impl ApplicationState {
    pub fn new(settings: &Settings) -> anyhow::Result<Self> {
        Ok(Self {
            settings: ArcSwap::new(Arc::new((*settings).clone())),
        })
    }
}

What is Arc and ArcSwap? Quote from rust-lang.org: "A thread-safe reference-counting pointer. Arc stands for Atomically Reference Counted." So Arc is a thread-safe alternative to Rc, the basic reference-counting pointer.

When you clone a struct in Rust you create a complete deep copy, replicating its whole in-memory representation. This is a costly operation, so reference counting provides a cheap alternative: after cloning an Rc or Arc container we hold multiple references to the same in-memory location.

The ArcSwap construct is a container holding an Arc that can be replaced in a thread-safe way atomically, without locking. It allows lock-free internal mutability in a multi-threaded environment. We need this because the tokio runtime may execute the axum async handler methods on multiple threads concurrently.

In our ApplicationState struct we store our own copy: (*settings).clone() of the the settings loaded earlier. We encapsulate the settings in an ArcSwap so we can update the settings using internal mutability later, without replacing the whole ApplicationState structure.

Now in shelter_main/src/commands/serve.rs we can instantiate ApplicationState and pass it to the axum route configuration:

fn start_tokio(port: u16, settings: &Settings) -> anyhow::Result<()> {
    tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .unwrap()
        .block_on(async move {
            let state = Arc::new(ApplicationState::new(settings)?);

            let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), port);

            let routes = crate::api::configure(state).layer(TraceLayer::new_for_http());

            tracing::info!("starting axum on port {}", port);

            axum::Server::bind(&addr)
                .serve(routes.into_make_service())
                .await?;

            Ok::<(), anyhow::Error>(())
        })?;

    std::process::exit(0);
}

We have to change the signature of the configure methods too. First in api/mod.rs:

use crate::state::ApplicationState;
use axum::Router;
use std::sync::Arc;

mod handlers;
mod v1;

pub fn configure(state: Arc<ApplicationState>) -> Router {
    Router::new().nest("/v1", v1::configure(state))
}

then in api/v1.rs:

use super::handlers;
use crate::state::ApplicationState;
use axum::routing::get;
use axum::Router;
use std::sync::Arc;

pub fn configure(state: Arc<ApplicationState>) -> Router {
    Router::new().route("/hello", get(handlers::hello::hello).with_state(state))
}

Here comes a new axum extractor, the with_state state extractor. It allows us to receive the state in handlers/hello.rs:

use crate::state::ApplicationState;
use axum::extract::State;
use axum::http::StatusCode;
use std::sync::Arc;

pub async fn hello(State(state): State<Arc<ApplicationState>>) -> Result<String, StatusCode> {
    Ok(format!(
        "\nHello world! Using configuration from {}\n\n",
        state
            .settings
            .load()
            .config
            .location
            .clone()
            .unwrap_or("-".to_string())
    ))
}

And that's it, now our handler methods can use application state and the loaded settings with it. Build the application, run it with the serve subcommand, and test the endpoint with curl:

$ curl -v http://127.0.0.1:8080/v1/hello
*   Trying 127.0.0.1:8080...
* Connected to 127.0.0.1 (127.0.0.1) port 8080 (#0)
> GET /v1/hello HTTP/1.1
> Host: 127.0.0.1:8080
> User-Agent: curl/7.81.0
> Accept: */*
> 
* Mark bundle as not supporting multiuse
< HTTP/1.1 200 OK
< content-type: text/plain; charset=utf-8
< content-length: 52
< date: Mon, 29 May 2023 12:08:00 GMT
< 

Hello world! Using configuration from config.json

* Connection #0 to host 127.0.0.1 left intact

You can find the sample code on GitHub

Next article »