Instrumenting Axum projects

mod instrumentation;
mod logger;

use clap::Parser;
use std::net::{IpAddr, Ipv6Addr, SocketAddr};

#[derive(Parser)]
pub(crate) struct Cli {
    #[clap(long, env = "DEMO_BIND", default_value_t = SocketAddr::new(
        IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0)), 8080)
    )]
    pub(crate) bind: SocketAddr,

    #[clap(flatten)]
    pub(crate) instrumentation: instrumentation::Instrumentation,
}

#[derive(Clone, Default, Debug, clap::ValueEnum)]
pub(crate) enum Logger {
    #[default]
    Compact,
    Full,
    Pretty,
    Json,
}

impl std::fmt::Display for Logger {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let logger = match self {
            Logger::Compact => "compact",
            Logger::Full => "full",
            Logger::Pretty => "pretty",
            Logger::Json => "json",
        };
        write!(f, "{}", logger)
    }
}

use color_eyre::eyre::WrapErr;
use std::{error::Error, io::IsTerminal};
use tracing::Subscriber;
use tracing_subscriber::{
    filter::Directive,
    layer::{Layer, SubscriberExt},
    registry::LookupSpan,
    util::SubscriberInitExt,
    EnvFilter,
};

use super::logger::Logger;

#[derive(clap::Args, Debug, Default)]
pub(crate) struct Instrumentation {
    /// Enable debug logs, -vv for trace
    #[clap(
        short = 'v',
        env = "DEMO_VERBOSITY",
        long, action = clap::ArgAction::Count,
        global = true
    )]
    pub verbose: u8,
    /// Which logger to use
    #[clap(
        long,
        env = "DEMO_LOGGER",
        default_value_t = Default::default(),
        global = true
    )]
    pub(crate) logger: Logger,
    /// Tracing directives
    ///
    /// See https://docs.rs/tracing-subscriber/latest/tracing_subscriber/filter/struct.EnvFilter.html#directives
    #[clap(long = "log-directive", global = true, env = "DEMO_LOG_DIRECTIVES", value_delimiter = ',', num_args = 0..)]
    pub(crate) log_directives: Vec<Directive>,
}

impl Instrumentation {
    pub(crate) fn log_level(&self) -> String {
        match self.verbose {
            0 => "info",
            1 => "debug",
            _ => "trace",
        }
        .to_string()
    }
    // (continued below)
}

use axum::{
    http::StatusCode,
    response::{IntoResponse, Response},
};

pub type Result<T, E = Report> = color_eyre::Result<T, E>;
// A generic error report
// Produced via `Err(some_err).wrap_err("Some context")`
// or `Err(color_eyre::eyre::Report::new(SomeError))`
pub struct Report(color_eyre::Report);

impl std::fmt::Debug for Report {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.0.fmt(f)
    }
}

impl<E> From<E> for Report
where
    E: Into<color_eyre::Report>,
{
    fn from(err: E) -> Self {
        Self(err.into())
    }
}

// Tell axum how to convert `Report` into a response.
impl IntoResponse for Report {
    fn into_response(self) -> Response {
        let err = self.0;
        let err_string = format!("{err:?}");

        tracing::error!("{err_string}");

        if let Some(err) = err.downcast_ref::<DemoError>() {
            return err.response()
        }

        // Fallback
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            "Something went wrong".to_string(),
        )
            .into_response()
    }
}

#[derive(thiserror::Error, Debug)]
pub(crate) enum DemoError {
    #[error("A spooky thing happened")]
    Spooky,
}

// Tell axum how to convert `DemoError` into a response.
impl DemoError {
    fn response(&self) -> Response {
        match self {
            Self::Spooky => (
                StatusCode::IM_A_TEAPOT,
                "A user-facing message about a Spooky".to_string(),
            )
                .into_response(),
        }
    }
}

mod cli;
mod error;
mod routes;
mod trace_layer;

use crate::{cli::Cli, error::Result};
use axum::routing::get;
use clap::Parser;
use std::{io::IsTerminal, net::SocketAddr, process::ExitCode};
use tower_http::trace::TraceLayer;

#[tokio::main]
async fn main() -> Result<ExitCode> {
    color_eyre::config::HookBuilder::default()
        .theme(if !std::io::stderr().is_terminal() {
            // Don't attempt color
            color_eyre::config::Theme::new()
        } else {
            color_eyre::config::Theme::dark()
        })
        .install()?;

    let cli = Cli::parse();
    cli.instrumentation.setup()?;

    let trace_layer = TraceLayer::new_for_http()
        .make_span_with(trace_layer::trace_layer_make_span_with)
        .on_request(trace_layer::trace_layer_on_request)
        .on_response(trace_layer::trace_layer_on_response);

    let app = axum::Router::new()
        .route("/", get(routes::get_home))
        .route("/errors", get(routes::get_error))
        .route("/errors/demo", get(routes::get_demo_error))
        .fallback(routes::fallback_404)
        .layer(trace_layer);

    tracing::info!("Listening on {}", cli.bind);
    axum::Server::bind(&cli.bind)
        .serve(app.into_make_service_with_connect_info::<SocketAddr>())
        .await?;

    Ok(ExitCode::SUCCESS)
}

use axum::{body::BoxBody, extract::ConnectInfo, response::Response};
use hyper::{Body, Request};
use std::{net::SocketAddr, time::Duration};
use tracing::Span;

pub(crate) fn trace_layer_make_span_with(request: &Request<Body>) -> Span {
    tracing::error_span!("request",
        uri = %request.uri(),
        method = %request.method(),
        // This is not particularly robust, but suitable for a demo
        // You'll need to change this if you deploy behind a proxy
        // (eg the `X-forwarded-for` header)
        source = request.extensions()
            .get::<ConnectInfo<SocketAddr>>()
            .map(|connect_info|
                tracing::field::display(connect_info.ip().to_string()),
            ).unwrap_or_else(||
                tracing::field::display(String::from("<unknown>"))
            ),
        // Fields must be defined to be used, define them as empty if they populate later
        status = tracing::field::Empty,
        latency = tracing::field::Empty,
    )
}

pub(crate) fn trace_layer_on_request(_request: &Request<Body>, _span: &Span) {
    tracing::trace!("Got request")
}

pub(crate) fn trace_layer_on_response(
    response: &Response<BoxBody>,
    latency: Duration,
    span: &Span,
) {
    span.record(
        "latency",
        tracing::field::display(format!("{}μs", latency.as_micros())),
    );
    span.record("status", tracing::field::display(response.status()));
    tracing::trace!("Responded");
}

use crate::error::{Result, DemoError};
use axum::response::IntoResponse;
use color_eyre::eyre::eyre;

#[tracing::instrument]
pub(crate) async fn get_home() -> Result<impl IntoResponse> {
    tracing::info!("Got a homepage request");
    Ok("Welcome to my super cute home page")
}

#[tracing::instrument]
pub(crate) async fn get_error() -> Result<impl IntoResponse> {
    tracing::info!("Got an error request");

    Err(eyre!("Whoopsies"))?;

    Ok(())
}

#[tracing::instrument]
pub(crate) async fn get_demo_error() -> Result<impl IntoResponse> {
    tracing::info!("Got an error request");

    Err(DemoError::Spooky)?;

    Ok(())
}

#[tracing::instrument(skip_all)]
pub(crate) async fn fallback_404() -> Result<impl IntoResponse> {
    tracing::info!("Got request, but no route was found.");
    Ok("You failed to find my super cute home page")
}

[package]
name = "demo"
version = "0.1.0"
edition = "2021"

[dependencies]
axum = "0.6"
clap = { version = "4.3", features = ["derive", "env"] }
color-eyre = { version = "0.6", default-features = false, features = [ "issue-url", "tracing-error", "capture-spantrace", "color-spantrace" ] }
hyper = "0.14"
thiserror = "1"
tokio = { version = "1", features = ["full"] }
tower = "0.4"
tower-http = { version = "0.4", features = ["trace"] }
tracing = "0.1"
tracing-error = "0.2"
tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }