feat: add prometheus metrics and tracing instrumentation

auth-api/src/main.rs

@@ -10,6 +10,7 @@ use axum::{
 };
 use chrono::{Duration, Utc};
 use jsonwebtoken::{Algorithm, DecodingKey, EncodingKey, Header, Validation, decode, encode};
+use metrics_exporter_prometheus::{PrometheusBuilder, PrometheusHandle};
 use redis::AsyncCommands;
 use serde::{Deserialize, Serialize};
 use tracing::{info, warn};
@@ -18,6 +19,7 @@ use tracing::{info, warn};
 struct AppState {
     jwt_secret: Arc<String>,
     redis: Option<redis::aio::ConnectionManager>,
+    metrics: PrometheusHandle,
 }
 
 #[derive(Debug, Serialize, Deserialize, Clone)]
@@ -78,6 +80,9 @@ async fn main() -> Result<()> {
 
     let bind = std::env::var("AUTH_BIND").unwrap_or_else(|_| "0.0.0.0:8080".into());
     let jwt_secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| "dev-secret-change-me".into());
+    let metrics = PrometheusBuilder::new()
+        .install_recorder()
+        .context("install prometheus recorder")?;
     let redis = if let Ok(url) = std::env::var("REDIS_URL") {
         let client = redis::Client::open(url.clone()).context("open redis client")?;
         match redis::aio::ConnectionManager::new(client).await {
@@ -97,10 +102,12 @@ async fn main() -> Result<()> {
     let state = AppState {
         jwt_secret: Arc::new(jwt_secret),
         redis,
+        metrics,
     };
 
     let app = Router::new()
         .route("/healthz", get(healthz))
+        .route("/metrics", get(metrics_endpoint))
         .route("/v1/token/dev", post(issue_dev_token))
         .route("/v1/token/validate", post(validate_token))
         .route("/v1/stripe/webhook", post(stripe_webhook))
@@ -119,10 +126,16 @@ async fn healthz() -> &'static str {
     "ok"
 }
 
+async fn metrics_endpoint(State(state): State<AppState>) -> impl IntoResponse {
+    state.metrics.render()
+}
+
+#[tracing::instrument(skip(state))]
 async fn issue_dev_token(
     State(state): State<AppState>,
     Json(req): Json<DevTokenRequest>,
 ) -> Result<Json<TokenResponse>, ApiError> {
+    let started = std::time::Instant::now();
     let now = Utc::now();
     let exp = now + Duration::hours(24);
     let claims = Claims {
@@ -155,6 +168,9 @@ async fn issue_dev_token(
             .await
             .map_err(ApiError::internal)?;
     }
+    metrics::counter!("auth_jwt_issued_total").increment(1);
+    metrics::histogram!("auth_issue_token_latency_ms")
+        .record(started.elapsed().as_secs_f64() * 1000.0);
 
     Ok(Json(TokenResponse {
         token,
@@ -162,10 +178,12 @@ async fn issue_dev_token(
     }))
 }
 
+#[tracing::instrument(skip(state, req))]
 async fn validate_token(
     State(state): State<AppState>,
     Json(req): Json<ValidateRequest>,
 ) -> Result<Json<ValidateResponse>, ApiError> {
+    let started = std::time::Instant::now();
     let decoded = decode::<Claims>(
         &req.token,
         &DecodingKey::from_secret(state.jwt_secret.as_bytes()),
@@ -185,6 +203,9 @@ async fn validate_token(
                 .to_string();
                 let _: () = redis.set_ex(key, payload, 300).await.map_err(ApiError::internal)?;
             }
+            metrics::counter!("auth_token_validate_total", "result" => "valid").increment(1);
+            metrics::histogram!("auth_validate_token_latency_ms")
+                .record(started.elapsed().as_secs_f64() * 1000.0);
             Ok(Json(ValidateResponse {
                 valid: true,
                 user_id: Some(c.sub),
@@ -197,14 +218,21 @@ async fn validate_token(
             user_id: None,
             tier: None,
             max_tunnels: None,
-        })),
+        })).map(|resp| {
+            metrics::counter!("auth_token_validate_total", "result" => "invalid").increment(1);
+            metrics::histogram!("auth_validate_token_latency_ms")
+                .record(started.elapsed().as_secs_f64() * 1000.0);
+            resp
+        }),
     }
 }
 
+#[tracing::instrument(skip(state))]
 async fn stripe_webhook(
     State(state): State<AppState>,
     Json(event): Json<StripeWebhookEvent>,
 ) -> Result<impl IntoResponse, ApiError> {
+    let started = std::time::Instant::now();
     if let Some(mut redis) = state.redis.clone() {
         let key = format!("plan:user:{}", event.user_id);
         let payload = serde_json::json!({
@@ -217,6 +245,9 @@ async fn stripe_webhook(
         .to_string();
         let _: () = redis.set_ex(key, payload, 300).await.map_err(ApiError::internal)?;
     }
+    metrics::counter!("stripe_webhook_events_total", "event_type" => event.event_type.clone()).increment(1);
+    metrics::histogram!("stripe_webhook_latency_ms")
+        .record(started.elapsed().as_secs_f64() * 1000.0);
     Ok(StatusCode::NO_CONTENT)
 }