State Management with axum::extract::State

1. Welcome and Icebreaker:

If you could store any piece of data globally in your app, what would it be? (e.g., a user count, a config value, or a fun fact!)

2. Outline the Session:

Agenda:

  1. Concept Explanation: Sharing state in Axum using State.
  2. Live Coding: Add shared state to your app and access it in handlers.
  3. Best Practices & Common Pitfalls.
  4. Tasks: Practice with your own stateful endpoints.

Learning Outcomes:

  • Understand how to share application state across handlers.
  • Use axum::extract::State and Arc for shared data.
  • Know when and how to use Mutex/RwLock for mutable state.

Expectations: Ask questions, experiment, and try adding your own state!

3. Concept Explanation:

In real applications, you often need to share data (like configuration, a database pool, or counters) between handlers. Axum provides the State extractor for this purpose.

How it works:

  • You define an AppState struct to hold your shared data.
  • Wrap it in an Arc when clones would otherwise be expensive; state must be Clone + Send + Sync and handlers may run concurrently. Database pools are usually already cheap to clone.
  • Pass the state to the router using .with_state().
  • Handlers can access the state using the State extractor.

Why use Arc? Rust’s safety guarantees mean you can’t share data between threads without synchronization. Arc allows multiple threads to read the same data. If you need to mutate state, wrap the field in a Mutex or RwLock.

Example use cases:

  • App configuration (read-only)
  • Database connection pools
  • Caches
  • Counters (with Mutex)

4. Live Coding Demonstration:

Let’s build a simple Axum app that shares state!

use axum::{extract::State, Router, routing::get};
use std::sync::Arc;

// 1. Define your shared state
struct AppState {
    app_name: String,
}

// 2. Handler that reads from state
async fn app_info(State(state): State<Arc<AppState>>) -> String {
    format!("App: {}", state.app_name)
}

#[tokio::main]
async fn main() {
    // 3. Initialize state and wrap in Arc
    let state = Arc::new(AppState {
        app_name: "My Axum App".to_string(),
    });
    // 4. Pass state to router
    let app = Router::new()
        .route("/info", get(app_info))
        .with_state(state);
    // 5. Run the server
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

Step-by-step:

  • Define AppState for your shared data.
  • Use Arc to allow safe sharing between threads.
  • Use the State extractor in your handler to access the state.
  • Attach state to your router with .with_state().
  • Start the server with tokio::net::TcpListener and axum::serve.

Mutable State Example: If you want to mutate state (like a counter), use Mutex:

use std::sync::{Arc, Mutex};

struct AppState {
    counter: Mutex<u64>,
}

async fn get_counter(State(state): State<Arc<AppState>>) -> String {
    let count = state.counter.lock().unwrap();
    format!("Counter: {}", *count)
}

Best Practices:

  • Use Arc if cloning your state is expensive; it does not make interior mutation safe on its own.
  • Use Mutex/RwLock for fields that need mutation.
  • Keep your AppState struct focused and clear.
  • Avoid putting large or complex objects in state unless necessary.

Common Pitfalls:

  • Holding a std::sync::Mutex guard across .await (release the guard first).
  • Forgetting to add .with_state() to the router (handlers won’t receive state).
  • Deadlocks or panics from improper Mutex usage (always lock/unlock carefully).
  • Trying to mutate state without a lock (Rust will prevent this, but it’s a common confusion).

5. Tasks

  • Add a Counter: Add a counter: Mutex<u64> to AppState and create an endpoint to display its value.
  • Share a Config Value: Add a version string or other config value to your state and expose it via an endpoint.
  • Experiment with Mutation: Try incrementing the counter in a handler (hint: lock the mutex, then update).
  • Bonus: What happens if you try to mutate state without a Mutex? Try it and see what error you get!

Feel free to ask for help or share your solutions!