loco-rs 1.0.1

//! # Application Bootstrapping and Logic
//! This module contains functions and structures for bootstrapping and running
//! your application.
use std::{
    env,
    path::{Path, PathBuf},
    sync::Arc,
};

use axum::Router;
#[cfg(feature = "with-db")]
use sea_orm_migration::MigratorTrait;
use tokio::{select, signal, task::JoinHandle};
use tracing::{debug, error, info, warn};

#[cfg(feature = "with-db")]
use crate::db;
use crate::{
    app::{AppContext, Hooks, Initializer},
    banner::print_banner,
    bgworker, cache,
    config::{self, Config, WorkerMode},
    controller::ListRoutes,
    env_vars,
    environment::Environment,
    errors::Error,
    mailer::{EmailSender, MailerWorker},
    prelude::BackgroundWorker,
    scheduler::{self, Scheduler},
    storage::{self, Storage},
    task::{self, Tasks},
    Result,
};

/// Represents the application startup mode.
#[derive(Debug)]
pub enum StartMode {
    /// Run the application as a server only. when running web server only,
    /// workers job will not handle.
    ServerOnly,
    /// Run the application web server and the worker in the same process.
    ServerAndWorker,
    /// Run the server and scheduler without workers.
    ServerAndScheduler,
    /// Pulling job worker and execute them
    WorkerOnly {
        /// Specifies that the worker should only handle jobs associated with one of these tags.
        /// If empty, the worker handles all jobs.
        tags: Vec<String>,
    },
    /// Run workers and scheduler without the HTTP server.
    WorkerAndScheduler {
        /// Specifies that the worker should only handle jobs associated with one of these tags.
        /// If empty, the worker handles all jobs.
        tags: Vec<String>,
    },
    /// Run the app with all available components in the same process.
    All,
}

pub struct BootResult {
    /// Application Context
    pub app_context: AppContext,
    /// Web server routes
    pub router: Option<Router>,
    /// worker processor
    pub worker: Option<Vec<String>>,
    /// scheduler processor
    pub run_scheduler: bool,
}

/// Configuration structure for serving an application.
#[derive(Debug)]
pub struct ServeParams {
    /// The port number on which the server will listen for incoming
    /// connections.
    pub port: i32,
    /// The network address to which the server will bind. It specifies the
    /// interface to listen on.
    pub binding: String,
}

/// Runs the application based on the provided `BootResult`.
///
/// This function is responsible for starting the application, including the
/// server and/or workers.
///
/// # Errors
///
/// When could not initialize the application.
pub async fn start<H: Hooks>(
    boot: BootResult,
    server_config: ServeParams,
    no_banner: bool,
) -> Result<()> {
    if boot.run_scheduler {
        let scheduler = scheduler::<H>(&boot.app_context, None, None, None)?;
        tokio::spawn(async move {
            if let Err(err) = scheduler.run().await {
                error!(err = err.to_string(), "error while running scheduler");
            }
        });
    }

    if !no_banner {
        print_banner(&boot, &server_config);
    }

    let BootResult {
        router,
        worker,
        run_scheduler: _,
        app_context,
    } = boot;

    match (router, worker) {
        (Some(router), None) => {
            H::serve(router, &app_context, &server_config).await?;
        }
        (Some(router), Some(tags)) => {
            let handle = if app_context.config.workers.mode == WorkerMode::BackgroundQueue {
                Some(start_queue_worker(&app_context, tags)?)
            } else {
                None
            };

            H::serve(router, &app_context, &server_config).await?;

            if let Some(handle) = handle {
                shutdown_and_await_queue_worker(&app_context, handle).await?;
            }
        }
        (None, Some(tags)) => {
            let handle = if app_context.config.workers.mode == WorkerMode::BackgroundQueue {
                Some(start_queue_worker(&app_context, tags)?)
            } else {
                None
            };

            await_shutdown_then_run_hook::<H>(&app_context, shutdown_signal()).await;

            if let Some(handle) = handle {
                shutdown_and_await_queue_worker(&app_context, handle).await?;
            }
        }
        _ => {}
    }
    Ok(())
}

/// Waits for `shutdown` to resolve, then runs [`Hooks::on_shutdown`].
///
/// This is used by the worker-only / worker-and-scheduler start modes, which
/// (unlike the server modes) have no Axum graceful-shutdown hook of their own
/// to invoke `on_shutdown` from. Extracted out of [`start`] purely so the
/// "wait, then hook" sequence can be exercised in tests with an
/// already-resolved `shutdown` future, instead of the real
/// [`shutdown_signal`] (which blocks on OS signals).
async fn await_shutdown_then_run_hook<H: Hooks>(
    app_context: &AppContext,
    shutdown: impl std::future::Future<Output = ()>,
) {
    shutdown.await;
    H::on_shutdown(app_context).await;
}

fn start_queue_worker(app_context: &AppContext, tags: Vec<String>) -> Result<JoinHandle<()>> {
    debug!("note: worker is run in-process (tokio spawn)");

    if let Some(queue) = &app_context.queue_provider {
        let cloned_queue = queue.clone();
        let handle = tokio::spawn(async move {
            if let Err(err) = cloned_queue.run(tags).await {
                error!(err = err.to_string(), "error while running worker");
            }
        });
        return Ok(handle);
    }

    Err(Error::QueueProviderMissing)
}

async fn shutdown_and_await_queue_worker(
    app_context: &AppContext,
    handle: JoinHandle<()>,
) -> Result<(), Error> {
    if let Some(queue) = &app_context.queue_provider {
        queue.shutdown()?;
    }

    println!("press ctrl-c again to force quit");
    select! {
        _ = handle => {}
        () = shutdown_signal() => {}
    }
    Ok(())
}

/// Run task
///
/// # Errors
///
/// When running could not run the task.
pub async fn run_task<H: Hooks>(
    app_context: &AppContext,
    task: Option<&String>,
    vars: &task::Vars,
) -> Result<()> {
    let mut tasks = Tasks::default();
    H::register_tasks(&mut tasks);

    if let Some(task) = task {
        let task_span = tracing::span!(tracing::Level::DEBUG, "task", task,);
        let _guard = task_span.enter();
        tasks.run(app_context, task, vars).await?;
    } else {
        let list = tasks.list();
        for item in &list {
            println!("{:<30}[{}]", item.name, item.detail);
        }
    }
    Ok(())
}

/// Initializes a new scheduler instance based on the provided configuration and context.
fn scheduler<H: Hooks>(
    app_context: &AppContext,
    config: Option<&PathBuf>,
    name: Option<String>,
    tag: Option<String>,
) -> Result<Scheduler> {
    let env_config_path = env::var(env_vars::SCHEDULER_CONFIG).ok();

    let config_path: Option<&Path> = config.map_or_else(
        || env_config_path.as_deref().map(Path::new),
        |path| Some(path.as_path()),
    );

    let scheduler = match config_path {
        Some(path) => Scheduler::from_config::<H>(path, &app_context.environment)?,
        None => {
            if let Some(config) = &app_context.config.scheduler {
                Scheduler::new::<H>(config, &app_context.environment)?
            } else {
                return Err(Error::Scheduler(scheduler::Error::Empty));
            }
        }
    };

    Ok(scheduler.by_spec(&scheduler::Spec { name, tag }))
}

/// Runs the scheduler with the given configuration and context. in case if list
/// args is true prints scheduler job configuration
///
/// This function initializes the scheduler, registers tasks through the
/// provided [`Hooks`], and executes the scheduler based on the specified
/// configuration or context. The scheduler continuously runs, managing and
/// executing scheduled tasks until a signal is received to shut down.
/// Upon receiving this signal, the function gracefully shuts down all running
/// tasks and exits safely.
///
/// # Errors
///
/// When running could not run the scheduler.
pub async fn run_scheduler<H: Hooks>(
    app_context: &AppContext,
    config: Option<&PathBuf>,
    name: Option<String>,
    tag: Option<String>,
    list: bool,
) -> Result<()> {
    let task_span = tracing::span!(tracing::Level::DEBUG, "scheduler_jobs");
    let _guard = task_span.enter();

    let scheduler = scheduler::<H>(app_context, config, name, tag)?;
    if list {
        println!("{scheduler}");
        Ok(())
    } else {
        Ok(scheduler.run().await?)
    }
}

/// Represents commands for handling database-related operations.
#[derive(Debug)]
pub enum RunDbCommand {
    /// Apply pending migrations.
    Migrate,
    /// Run one or more down migrations.
    Down(u32),
    /// Drop all tables, then reapply all migrations.
    Reset,
    /// Check the status of all migrations.
    Status,
    /// Generate entity.
    Entities,
    /// Truncate tables, by executing the implementation in [`Hooks::seed`]
    /// (without dropping).
    Truncate,
    /// Seed database.
    Seed {
        reset: bool,
        from: PathBuf,
        dump: bool,
        dump_tables: Option<Vec<String>>,
    },
    /// Dump database schema
    Schema,
}

#[cfg(feature = "with-db")]
/// Handles database commands.
///
/// # Errors
///
/// Return an error when the given command fails. mostly return
/// [`sea_orm::DbErr`]
#[allow(clippy::cognitive_complexity)]
pub async fn run_db<H: Hooks, M: MigratorTrait>(
    app_context: &AppContext,
    cmd: RunDbCommand,
) -> Result<()> {
    match cmd {
        RunDbCommand::Migrate => {
            tracing::warn!("migrate:");
            db::migrate::<M>(&app_context.db).await?;
        }
        RunDbCommand::Down(steps) => {
            tracing::warn!("down:");
            db::down::<M>(&app_context.db, steps).await?;
        }
        RunDbCommand::Reset => {
            tracing::warn!("reset:");
            db::reset::<M>(&app_context.db).await?;
        }
        RunDbCommand::Status => {
            tracing::warn!("status:");
            db::status::<M>(&app_context.db).await?;
        }
        RunDbCommand::Entities => {
            tracing::warn!("entities:");

            tracing::warn!("{}", db::entities::<M>(app_context).await?);
        }
        RunDbCommand::Truncate => {
            tracing::warn!("truncate:");
            H::truncate(app_context).await?;
        }
        RunDbCommand::Seed {
            reset,
            from,
            dump,
            dump_tables,
        } => {
            tracing::warn!(reset = reset, from = %from.display(), "seed:");

            if let Some(tables) = dump_tables {
                // Explicit table list: schema-introspection dump of just those tables.
                db::dump_tables(&app_context.db, from.as_path(), Some(tables)).await?;
            } else if dump {
                // Plain `--dump`: route through Hooks::dump so apps can override
                // it with typed, streaming db::dump per entity.
                db::run_app_dump::<H>(app_context, &from).await?;
            } else {
                if reset {
                    db::reset::<M>(&app_context.db).await?;
                }
                db::run_app_seed::<H>(app_context, &from).await?;
            }
        }
        RunDbCommand::Schema => {
            db::dump_schema(app_context, "schema_dump.json").await?;
            println!("Database schema dumped to 'schema_dump.json'");
        }
    }
    Ok(())
}

/// Initializes the application context by loading configuration and
/// establishing connections.
///
/// # Errors
/// When has an error to create DB connection.
pub async fn create_context<H: Hooks>(
    environment: &Environment,
    config: Config,
) -> Result<AppContext> {
    if config.logger.pretty_backtrace {
        // SAFETY: `create_context` runs during boot, before the server or any
        // background-worker threads are spawned, so no other thread is reading
        // or writing the environment concurrently.
        unsafe { std::env::set_var("RUST_BACKTRACE", "1") };
        warn!(
            "pretty backtraces are enabled (this is great for development but has a runtime cost \
             for production. disable with `logger.pretty_backtrace` in your config yaml)"
        );
    }
    #[cfg(feature = "with-db")]
    let db = db::connect(&config.database).await?;

    let mailer = if let Some(cfg) = config.mailer.as_ref() {
        create_mailer(cfg)?
    } else {
        None
    };

    let queue_provider = bgworker::create_queue_provider(&config).await?;
    let ctx = AppContext {
        environment: environment.clone(),
        #[cfg(feature = "with-db")]
        db,
        queue_provider,
        storage: Storage::single(storage::drivers::null::new()).into(),
        cache: cache::create_cache_provider(&config).await?,
        config,
        mailer,
        shared_store: Arc::new(crate::app::SharedStore::default()),
    };

    H::after_context(ctx).await
}

#[cfg(feature = "with-db")]
/// Creates an application based on the specified mode and environment.
///
/// # Errors
///
/// When could not create the application
pub async fn create_app<H: Hooks, M: MigratorTrait>(
    mode: StartMode,
    environment: &Environment,
    config: Config,
) -> Result<BootResult> {
    let app_context = create_context::<H>(environment, config).await?;
    db::converge::<H, M>(&app_context, &app_context.config.database).await?;

    if let (Some(queue), Some(config)) = (&app_context.queue_provider, &app_context.config.queue) {
        bgworker::converge(queue, config).await?;
    }

    run_app::<H>(&mode, app_context).await
}

#[cfg(not(feature = "with-db"))]
pub async fn create_app<H: Hooks>(
    mode: StartMode,
    environment: &Environment,
    config: Config,
) -> Result<BootResult> {
    let app_context = create_context::<H>(environment, config).await?;

    if let (Some(queue), Some(config)) = (&app_context.queue_provider, &app_context.config.queue) {
        bgworker::converge(queue, config).await?;
    }

    run_app::<H>(&mode, app_context).await
}

/// Run the application with the  given mode
/// # Errors
///
/// When could not create the application
pub async fn run_app<H: Hooks>(mode: &StartMode, app_context: AppContext) -> Result<BootResult> {
    H::before_run(&app_context).await?;
    let initializers = H::initializers(&app_context).await?;

    info!(
        initializers = ?initializers.iter().map(|init| init.name()).collect::<Vec<_>>().join(","),
        "initializers loaded"
    );

    for initializer in &initializers {
        initializer.before_run(&app_context).await?;
    }

    match mode {
        StartMode::ServerOnly => {
            let router = setup_routes::<H>(&app_context, &initializers).await?;
            Ok(BootResult {
                app_context,
                router: Some(router),
                worker: None,
                run_scheduler: false,
            })
        }
        StartMode::ServerAndWorker => {
            register_workers::<H>(&app_context).await?;
            let router = setup_routes::<H>(&app_context, &initializers).await?;
            Ok(BootResult {
                app_context,
                router: Some(router),
                worker: Some(vec![]),
                run_scheduler: false,
            })
        }
        StartMode::ServerAndScheduler => {
            let router = setup_routes::<H>(&app_context, &initializers).await?;
            Ok(BootResult {
                app_context,
                router: Some(router),
                worker: None,
                run_scheduler: true,
            })
        }
        StartMode::All => {
            register_workers::<H>(&app_context).await?;
            let router = setup_routes::<H>(&app_context, &initializers).await?;
            Ok(BootResult {
                app_context,
                router: Some(router),
                worker: Some(vec![]),
                run_scheduler: true,
            })
        }
        StartMode::WorkerOnly { tags } => {
            register_workers::<H>(&app_context).await?;
            Ok(BootResult {
                app_context,
                router: None,
                worker: Some(tags.clone()),
                run_scheduler: false,
            })
        }
        StartMode::WorkerAndScheduler { tags } => {
            register_workers::<H>(&app_context).await?;
            Ok(BootResult {
                app_context,
                router: None,
                worker: Some(tags.clone()),
                run_scheduler: true,
            })
        }
    }
}

/// Sets up the application's routes based on the provided initializers and hooks.
async fn setup_routes<H: Hooks>(
    app_context: &AppContext,
    initializers: &[Box<dyn Initializer>],
) -> Result<Router> {
    let app = H::before_routes(app_context).await?;
    let app = H::routes(app_context).to_router::<H>(app_context.clone(), app)?;
    let mut router = H::after_routes(app, app_context).await?;

    for initializer in initializers {
        router = initializer.after_routes(router, app_context).await?;
    }

    Ok(router)
}

async fn register_workers<H: Hooks>(app_context: &AppContext) -> Result<()> {
    if app_context.config.workers.mode == WorkerMode::BackgroundQueue {
        if let Some(queue) = &app_context.queue_provider {
            queue.register(MailerWorker::build(app_context)).await?;
            H::connect_workers(app_context, queue).await?;
        } else {
            return Err(Error::QueueProviderMissing);
        }

        debug!("done registering workers and queues");
    }
    Ok(())
}

#[must_use]
pub fn list_endpoints<H: Hooks>(ctx: &AppContext) -> Vec<ListRoutes> {
    H::routes(ctx).collect()
}

/// Waits for a shutdown signal, either via Ctrl+C or termination signal.
///
/// # Panics
///
/// This function will panic if it fails to install the signal handlers for
/// Ctrl+C or the terminate signal on Unix-based systems.
pub async fn shutdown_signal() {
    let ctrl_c = async {
        signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("failed to install signal handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        () = ctrl_c => {},
        () = terminate => {},
    }
}

pub struct MiddlewareInfo {
    pub id: String,
    pub enabled: bool,
    pub detail: String,
}

#[must_use]
pub fn list_middlewares<H: Hooks>(ctx: &AppContext) -> Vec<MiddlewareInfo> {
    H::middlewares(ctx)
        .iter()
        .map(|m| MiddlewareInfo {
            id: m.name().to_string(),
            enabled: m.is_enabled(),
            detail: m.config().unwrap_or_default().to_string(),
        })
        .collect::<Vec<_>>()
}

/// Initializes an [`EmailSender`] based on the mailer configuration settings
/// ([`config::Mailer`]).
fn create_mailer(config: &config::Mailer) -> Result<Option<EmailSender>> {
    if config.stub {
        return Ok(Some(EmailSender::stub()));
    }
    if let Some(smtp) = config.smtp.as_ref()
        && smtp.enable
    {
        return Ok(Some(EmailSender::smtp(smtp)?));
    }
    Ok(None)
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::{AtomicBool, Ordering};

    use async_trait::async_trait;

    use super::*;
    use crate::{bgworker::Queue, controller::AppRoutes, tests_cfg::app::get_app_context};

    /// A `Hooks` impl whose only interesting behavior is `on_shutdown`: it
    /// flips a flag stashed in the `AppContext`'s `shared_store` so tests can
    /// observe whether the hook actually ran. Every other method is an
    /// unreachable stub because `await_shutdown_then_run_hook` never calls
    /// them.
    struct ShutdownFlagHooks;

    #[async_trait]
    impl Hooks for ShutdownFlagHooks {
        fn app_name() -> &'static str {
            "shutdown_flag_hooks_test"
        }

        async fn boot(
            _mode: StartMode,
            _environment: &Environment,
            _config: Config,
        ) -> Result<BootResult> {
            unreachable!("not exercised by this test")
        }

        async fn connect_workers(_ctx: &AppContext, _queue: &Queue) -> Result<()> {
            unreachable!("not exercised by this test")
        }

        fn register_tasks(_tasks: &mut task::Tasks) {
            unreachable!("not exercised by this test")
        }

        async fn truncate(_ctx: &AppContext) -> Result<()> {
            unreachable!("not exercised by this test")
        }

        async fn seed(_ctx: &AppContext, _path: &Path) -> Result<()> {
            unreachable!("not exercised by this test")
        }

        fn routes(_ctx: &AppContext) -> AppRoutes {
            unreachable!("not exercised by this test")
        }

        async fn on_shutdown(ctx: &AppContext) {
            if let Some(flag) = ctx.shared_store.get::<Arc<AtomicBool>>() {
                flag.store(true, Ordering::SeqCst);
            }
        }
    }

    #[tokio::test]
    async fn test_await_shutdown_then_run_hook_invokes_on_shutdown() {
        let app_context = get_app_context().await;
        let shutdown_hook_fired = Arc::new(AtomicBool::new(false));
        app_context.shared_store.insert(shutdown_hook_fired.clone());

        // An already-resolved future stands in for `shutdown_signal()`, so
        // the test doesn't have to wait on (or fake) an OS signal.
        await_shutdown_then_run_hook::<ShutdownFlagHooks>(&app_context, std::future::ready(()))
            .await;

        assert!(
            shutdown_hook_fired.load(Ordering::SeqCst),
            "on_shutdown should fire once the shutdown future resolves"
        );
    }

    #[tokio::test]
    async fn test_await_shutdown_then_run_hook_waits_for_shutdown_before_running_hook() {
        let app_context = get_app_context().await;
        let shutdown_hook_fired = Arc::new(AtomicBool::new(false));
        app_context.shared_store.insert(shutdown_hook_fired.clone());

        let (tx, rx) = tokio::sync::oneshot::channel::<()>();
        let shutdown = async move {
            let _ = rx.await;
        };

        let fut = await_shutdown_then_run_hook::<ShutdownFlagHooks>(&app_context, shutdown);
        futures_util::pin_mut!(fut);

        // Poll once: the shutdown future hasn't resolved yet, so the hook
        // must not have run.
        let still_pending = futures_util::poll!(&mut fut);
        assert!(matches!(still_pending, std::task::Poll::Pending));
        assert!(
            !shutdown_hook_fired.load(Ordering::SeqCst),
            "on_shutdown must not fire before shutdown resolves"
        );

        tx.send(()).expect("receiver dropped");
        fut.await;

        assert!(
            shutdown_hook_fired.load(Ordering::SeqCst),
            "on_shutdown should fire after shutdown resolves"
        );
    }
}