shoutdotdev/src/main.rs

use app_state::App;
use axum::middleware::map_request;
use chrono::{TimeDelta, Utc};
use clap::{Parser, Subcommand};
use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness};
use dotenvy::dotenv;
use middleware::lowercase_uri_path;
use tokio::time::sleep;
use tower::ServiceBuilder;
use tower_http::{
    compression::CompressionLayer, normalize_path::NormalizePathLayer, trace::TraceLayer,
};
use tracing_subscriber::EnvFilter;

use crate::{app_state::AppState, router::new_router, settings::Settings, worker::run_worker};

pub mod api_keys;
pub mod app_error;
pub mod app_state;
pub mod auth;
pub mod channel_selections;
pub mod channels;
mod channels_router;
pub mod csrf;
pub mod email;
pub mod governors;
pub mod guards;
pub mod messages;
pub mod middleware;
mod nav_state;
pub mod projects;
mod projects_router;
pub mod router;
pub mod schema;
pub mod sessions;
pub mod settings;
pub mod team_memberships;
pub mod teams;
mod teams_router;
pub mod users;
mod v0_router;
pub mod worker;

#[derive(Parser)]
#[command(version, about, long_about = None)]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Run web server
    Serve,
    /// Run background worker
    Worker {
        /// Loop the every n seconds instead of exiting after execution
        #[arg(long)]
        auto_loop_seconds: Option<u32>,
    },
    // TODO: add a low-frequency worker task exclusively for self-healing
    // mechanisms like Governor::reset_all()
}

/// Run CLI
#[tokio::main]
async fn main() {
    // Attempt to pre-load .env in case it contains a RUST_LOG variable
    dotenv().ok();
    tracing_subscriber::fmt()
        .with_env_filter(EnvFilter::from_default_env())
        .init();

    let settings = Settings::load().unwrap();

    let state: AppState = App::from_settings(settings.clone()).await.unwrap().into();

    if settings.run_database_migrations == Some(1) {
        const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations/");
        // Run migrations on server startup
        let conn = state.db_pool.get().await.unwrap();
        conn.interact(|conn| conn.run_pending_migrations(MIGRATIONS).map(|_| ()))
            .await
            .unwrap()
            .unwrap();
    }

    let cli = Cli::parse();
    match &cli.command {
        Commands::Serve => {
            let router = new_router(state.clone()).layer(
                ServiceBuilder::new()
                    .layer(map_request(lowercase_uri_path))
                    .layer(TraceLayer::new_for_http())
                    .layer(CompressionLayer::new())
                    .layer(NormalizePathLayer::trim_trailing_slash()),
            );

            let listener = tokio::net::TcpListener::bind((settings.host.clone(), settings.port))
                .await
                .unwrap();
            tracing::info!(
                "App running at http://{}:{}{}",
                settings.host,
                settings.port,
                settings.base_path
            );

            axum::serve(listener, router).await.unwrap();
        }
        Commands::Worker { auto_loop_seconds } => {
            if let Some(loop_seconds) = auto_loop_seconds {
                let loop_delta = TimeDelta::seconds(i64::from(*loop_seconds));
                loop {
                    let t_next_loop = Utc::now() + loop_delta;

                    if let Err(err) = run_worker(state.clone()).await {
                        tracing::error!("{}", err)
                    }

                    let sleep_delta = t_next_loop - Utc::now();
                    match sleep_delta.to_std() {
                        Ok(duration) => {
                            sleep(duration).await;
                        }
                        Err(_) => { /* sleep_delta was < 0, so don't sleep */ }
                    }
                }
            } else {
                run_worker(state).await.unwrap();
            }
        }
    }
}