axum.rs

use autometrics::{autometrics, encode_global_metrics, global_metrics_exporter};
use axum::response::{IntoResponse, Response};
use axum::routing::{get, post};
use axum::{http::StatusCode, Json, Router};
use rand::{random, thread_rng, Rng};
use serde::{Deserialize, Serialize};
use std::{net::SocketAddr, time::Duration};
use strum::IntoStaticStr;
use thiserror::Error;
use tokio::time::sleep;
use util::{run_prometheus, sleep_random_duration};

mod util;

// This is a full example of how to use Autometrics with Axum and thiserror
// and how to collect the generated metrics with Prometheus.

// Starting simple, hover over the function name to see the Autometrics graph links in the Rust Docs!
/// This is a simple endpoint that never errors
#[autometrics]
async fn get_index() -> &'static str {
    "Hello, World!"
}

/// This is a function that returns an error ~25% of the time
/// The call counter metric generated by autometrics will have a label
/// `result` = `ok` or `error`, depending on what the function returned
#[autometrics]
async fn get_random_error() -> Result<(), ApiError> {
    let error = thread_rng().gen_range(0..4);

    sleep_random_duration().await;

    match error {
        0 => Err(ApiError::NotFound),
        1 => Err(ApiError::BadRequest),
        2 => Err(ApiError::Internal),
        _ => Ok(()),
    }
}

// Now, to look at the actual error.
//
// We're using `thiserror` to define our error type, and we're using `strum` to
// enable the error variants to be turned into &'static str's, which
// will actually become another label on the call counter metric.
//
// In this case, the label will be `error` = `not_found`, `bad_request`, or `internal`.
//
// Instead of looking at high-level HTTP status codes in our metrics,
// we'll instead see the actual variant name of the error.
#[derive(Debug, Error, IntoStaticStr)]
#[strum(serialize_all = "snake_case")]
enum ApiError {
    #[error("User not found")]
    NotFound,
    #[error("Bad request")]
    BadRequest,
    #[error("Internal server error")]
    Internal,
}

impl IntoResponse for ApiError {
    fn into_response(self) -> Response {
        let status_code = match self {
            ApiError::NotFound => StatusCode::NOT_FOUND,
            ApiError::BadRequest => StatusCode::BAD_REQUEST,
            ApiError::Internal => StatusCode::INTERNAL_SERVER_ERROR,
        };
        (status_code, format!("{:?}", self)).into_response()
    }
}

// This handler calls another internal function that is also instrumented with autometrics.
//
// Unlike other instrumentation libraries, autometrics is designed to give you more
// granular metrics that allow you to dig into the internals of your application
// before even reaching for logs or traces.
//
// Try hovering over the function name to see the graph links and pay special attention
// to the links for the functions _called by this function_.
// You can also hover over the load_details_from_database function to see the graph links for that function.
#[autometrics]
async fn create_user(Json(payload): Json<CreateUser>) -> Result<Json<User>, ApiError> {
    let user = User {
        id: 1337,
        username: payload.username,
    };

    load_details_from_database().await?;

    sleep_random_duration().await;

    Ok(Json(user))
}

// The input to our `create_user` handler
#[derive(Serialize, Deserialize)]
struct CreateUser {
    username: String,
}

// The output to our `create_user` handler
#[derive(Serialize)]
struct User {
    id: u64,
    username: String,
}

/// An internal function that is also instrumented with autometrics
#[autometrics]
async fn load_details_from_database() -> Result<(), ApiError> {
    let should_error = random::<bool>();

    sleep_random_duration().await;

    if should_error {
        Err(ApiError::Internal)
    } else {
        Ok(())
    }
}

/// This handler serializes the metrics into a string for Prometheus to scrape
async fn get_metrics() -> (StatusCode, String) {
    match encode_global_metrics() {
        Ok(metrics) => (StatusCode::OK, metrics),
        Err(err) => (StatusCode::INTERNAL_SERVER_ERROR, format!("{:?}", err)),
    }
}

/// The main function runs the server, spawns a Prometheus instance, and generates random traffic
#[tokio::main]
async fn main() {
    // Run Prometheus and generate random traffic for the app
    // (You would not actually do this in production, but it makes it easier to see the example in action)
    run_prometheus();
    tokio::spawn(generate_random_traffic());

    // Set up the exporter to collect metrics
    let _exporter = global_metrics_exporter();

    let app = Router::new()
        .route("/", get(get_index))
        .route("/users", post(create_user))
        .route("/random-error", get(get_random_error))
        // Expose the metrics for Prometheus to scrape
        .route("/metrics", get(get_metrics));

    let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
    let server = axum::Server::bind(&addr);

    println!(
        "The example API server is now running on: {addr}

Wait a few seconds for the traffic generator to create some fake traffic, \
then hover over one of the HTTP handler functions (in your editor) \
to bring up the Rust Docs.

Click on one of the Autometrics links to see the graph for that handler's metrics in Prometheus."
    );

    server
        .serve(app.into_make_service())
        .await
        .expect("Error starting example API server");
}

/// Make some random API calls to generate data that we can see in the graphs
pub async fn generate_random_traffic() {
    let client = reqwest::Client::new();
    loop {
        let request_type = thread_rng().gen_range(0..3);
        let sleep_duration = Duration::from_millis(thread_rng().gen_range(10..50));
        match request_type {
            0 => {
                let _ = client.get("http://localhost:3000").send().await;
            }
            1 => {
                let _ = client
                    .post("http://localhost:3000/users")
                    .json(&CreateUser {
                        username: "test".to_string(),
                    })
                    .send()
                    .await;
            }
            2 => {
                let _ = reqwest::get("http://localhost:3000/random-error").await;
            }
            _ => unreachable!(),
        }
        sleep(sleep_duration).await
    }
}