Skip to main content

mini_apm/
server.rs

1use axum::{
2    Router,
3    extract::DefaultBodyLimit,
4    middleware,
5    routing::{get, post},
6};
7use std::net::SocketAddr;
8use tokio::signal;
9use tower_cookies::CookieManagerLayer;
10use tower_http::trace::TraceLayer;
11
12use crate::{DbPool, api, config::Config, jobs, models, web};
13
14/// Combined state for routes that need both pool and config
15#[derive(Clone)]
16pub struct AppState {
17    pub pool: DbPool,
18    pub config: Config,
19}
20
21/// Maximum request body size (10 MB)
22const MAX_BODY_SIZE: usize = 10 * 1024 * 1024;
23
24pub async fn run(pool: DbPool, config: Config, port: u16) -> anyhow::Result<()> {
25    // Initialize start time for uptime tracking
26    api::health::init_start_time();
27
28    // Always ensure default project and admin exist
29    // When features are disabled, we just skip the UI/auth, not the data
30    let default_project = models::project::ensure_default_project(&pool)?;
31    models::user::ensure_default_admin(&pool)?;
32
33    if !config.enable_projects {
34        tracing::info!("Single-project mode - API key: {}", default_project.api_key);
35    }
36
37    // Start background jobs
38    jobs::start(pool.clone(), config.clone());
39
40    // Build router
41    let app = Router::new()
42        // Health check (no auth)
43        .route("/health", get(api::health_handler))
44        // Ingestion API (with API key auth)
45        .nest(
46            "/ingest",
47            Router::new()
48                .route("/deploys", post(api::ingest_deploys))
49                .route("/v1/traces", post(api::ingest_spans))
50                .route("/errors", post(api::ingest_errors))
51                .route("/errors/batch", post(api::ingest_errors_batch))
52                .layer(middleware::from_fn_with_state(
53                    pool.clone(),
54                    api::auth_middleware,
55                )),
56        )
57        // Auth routes (always available)
58        .merge(web::auth_routes())
59        // Web UI (protected when user accounts enabled)
60        .merge(web::routes(pool.clone()))
61        // Static files
62        .nest_service("/static", tower_http::services::ServeDir::new("static"))
63        // State and middleware
64        .with_state(pool)
65        .layer(DefaultBodyLimit::max(MAX_BODY_SIZE))
66        .layer(CookieManagerLayer::new())
67        .layer(TraceLayer::new_for_http());
68
69    let addr = SocketAddr::from(([0, 0, 0, 0], port));
70    tracing::info!("MiniAPM server listening on http://{}", addr);
71
72    if config.enable_user_accounts {
73        tracing::info!("User accounts ENABLED - login required");
74    }
75
76    let listener = tokio::net::TcpListener::bind(addr).await?;
77
78    // Graceful shutdown
79    axum::serve(listener, app)
80        .with_graceful_shutdown(shutdown_signal())
81        .await?;
82
83    tracing::info!("Server shutdown complete");
84    Ok(())
85}
86
87async fn shutdown_signal() {
88    let ctrl_c = async {
89        signal::ctrl_c()
90            .await
91            .expect("Failed to install Ctrl+C handler");
92    };
93
94    #[cfg(unix)]
95    let terminate = async {
96        signal::unix::signal(signal::unix::SignalKind::terminate())
97            .expect("Failed to install SIGTERM handler")
98            .recv()
99            .await;
100    };
101
102    #[cfg(not(unix))]
103    let terminate = std::future::pending::<()>();
104
105    tokio::select! {
106        _ = ctrl_c => {
107            tracing::info!("Received Ctrl+C, starting graceful shutdown...");
108        }
109        _ = terminate => {
110            tracing::info!("Received SIGTERM, starting graceful shutdown...");
111        }
112    }
113}