tina-core 0.0.2

Tina platform
Documentation
//! 基于axum的server mod
#![cfg(feature = "server-axum")]
#![allow(dead_code)]

use crate::tina::server::http::axum::middleware::panic_handler::PanicHandler;
use crate::tina::server::http::axum::middleware::permission_handler::PermissionHandler;
use crate::tina::server::http::axum::middleware::request_init_handler::RequestInitHandler;
use crate::tina::server::http::axum::middleware::request_log_param_handler::RequestLogParamHandler;
use crate::tina::server::http::axum::middleware::token_handler::TokenHandler;
use crate::tina::server::http::axum::middleware::transaction_handler::TransactionHandler;
use crate::tina::server::http::route::RouteBaseConfig;
use crate::tina::server::http::route::RouteRegister;
use crate::tina::server::http::IHttpServer;
use crate::tina::server::Server;
use crate::tina::server::{http::axum::middleware::demo_handler::DemoHandler, random_address_and_port};
use crate::tina::{
    server::application::{AppConfig, Application},
    util::not_empty::INotEmpty,
};
use axum::Router;
use config::Config;
use futures_util::FutureExt;
use regex::Regex;
use std::net::SocketAddr;
use std::ops::Deref;
use std::process::exit;
use std::str::FromStr;
use std::sync::mpsc::{channel, Sender};
use std::sync::Arc;
use std::time::Duration;
use tower::ServiceBuilder;

use self::middleware::response_handler::ResponseHandler;

use super::super::ServerShutdownHook;

pub mod app_error_ext;
pub mod application_ext;
pub(in crate::tina::server) mod delegate;
pub(in crate::tina::server) mod middleware;
pub mod multipart;
pub mod request_conversion;
pub mod request_ext;
pub mod response_data_ext;
pub mod response_ext;
pub mod response_page_ext;
pub mod response_stream_ext;
pub mod route_ext;
pub mod session_ext;

impl IHttpServer for Server {
    fn run_http(self, config: Config) {
        let (application, config) =
            AppConfig::from_config(config).unwrap_or_else(|err| panic!("parse config to application failed: {:?}", err));
        let runtime = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(application.workers)
            .thread_name("axum-worker-thread")
            .enable_all()
            .build()
            .expect("create tokio runtime failed");
        let (shutdown_sender, shutdown_receiver) = channel::<()>();
        runtime.spawn(run_async(application, config, shutdown_sender, self));
        shutdown_receiver.recv().expect("receive shutdown signal failed");
        runtime.shutdown_background();
    }
}

async fn run_async(mut application: AppConfig, config: Config, shutdown_sender: Sender<()>, server: Server) {
    let Server {
        init,
        post_init,
        after_start,
        shutdown_hook,
    } = server;
    if let Some(init_fn) = init {
        let config2 = config.clone();
        let app_config = match init_fn(application, config2).await {
            Ok(v) => v,
            Err(err) => {
                println!("init app config failed: {:?}", err);
                tracing::error!("init app config failed: {:?}", err);
                exit(-1)
            }
        };
        application = app_config;
    }
    let application = match application.init(&config).await {
        Ok(v) => v,
        Err(err) => {
            println!("application init failed: {:?}", err);
            exit(-1);
        }
    };

    if let Some(post_init_fn) = post_init {
        let application2 = application.clone();
        match async move { post_init_fn(application2).await }.await {
            Ok(_) => {}
            Err(err) => {
                println!("post init app config failed: {:?}", err);
                tracing::error!("post init app config failed: {:?}", err);
                exit(-1)
            }
        }
    }

    let bind_address = match application.address.not_empty() && application.port.is_some() {
        true => {
            let address = application.address.as_str();
            let port = application.port.unwrap_or_default();
            SocketAddr::from_str(format!("{}:{}", address, port).as_str())
                .unwrap_or_else(|err| panic!("绑定端口失败: {}:{}, reason: {:?}", address, port, err))
        }
        false => random_address_and_port(application.address.as_str()),
    };

    let application_for_after_start: Application = application.to_owned();
    let application_for_signal: Application = application.to_owned();
    let application = Arc::new(application);

    let app = {
        let mut route_configs = RouteRegister::get_registed_routes();
        let mut app = Router::new();

        let param_regex = Regex::new("(\\{(.*?)\\}(\\*?))").expect("build PARAM_REGEX failed");
        for config in route_configs.iter_mut() {
            let base = &mut config.config;
            let path = RouteBaseConfig::get_route_path(application.deref(), base);
            let mut found_wildcard = false;
            if let Some(cs) = param_regex.captures(&path) {
                if let Some(g) = cs.get(3) {
                    if !g.as_str().is_empty() {
                        found_wildcard = true;
                    }
                }
            }
            let replace_path = match found_wildcard {
                true => param_regex.replace_all(&path, "$3$2"),
                false => param_regex.replace_all(&path, ":$2"),
            };
            // println!("{}  ->  {}", path, replace_path);
            base.path = replace_path.into_owned();
        }
        route_configs.sort_by(|o1, o2| {
            let path1 = &o1.config.path;
            let path2 = &o2.config.path;
            path1.cmp(path2)
        });

        for config in route_configs.into_iter() {
            let base = config.config;
            let path = &base.path;
            if let Some(mut route) = config.route {
                let base_config = Arc::new(base.clone());
                // 注意这里的中间件的执行顺序是: 执行接口函数之前的部分(从上往下执行), 执行接口函数之后的部分(从下往上执行)
                let layer = ServiceBuilder::new()
                    .layer(RequestInitHandler::new(&application, &base_config)) // 请求信息初始化处理
                    .layer(ResponseHandler) // 响应处理
                    .layer(PanicHandler) // panic捕获处理
                    .layer(DemoHandler) // 演示模式处理
                    .layer(RequestLogParamHandler) // 请求参数日志处理
                    .layer(PanicHandler) // panic捕获处理
                    .layer(TransactionHandler) // 自动事务处理
                    .layer(TokenHandler) // 用户登录检查
                    .layer(PermissionHandler) // 基于RBAC的访问控制
                    .into_inner();
                route = route.route_layer(layer);
                app = app.route(path, route);
            }
        }
        app
    };
    let app = app.into_make_service_with_connect_info::<SocketAddr>();
    axum::Server::bind(&bind_address)
        .serve(app)
        .with_graceful_shutdown(async move {
            tracing::info!("启动Axum Server成功 {}", bind_address);
            match after_start {
                Some(after_start) => {
                    match after_start(application_for_after_start).await {
                        Ok(_) => {
                            process_system_signals(application_for_signal, shutdown_hook).await;
                        }
                        Err(err) => {
                            tracing::error!("Failed to process after server started: {err:?}");
                        }
                    }
                    shutdown_sender.send(()).expect("send shutdown signal failed");
                }
                None => {
                    process_system_signals(application_for_signal, shutdown_hook).await;
                    shutdown_sender.send(()).expect("send shutdown signal failed");
                }
            }
        })
        .await
        .expect("启动Axum Server失败!");
}

async fn process_system_signals(application: Application, shutdown_hook: Option<ServerShutdownHook>) {
    let hook_application = application.clone();
    let shutdown_hook = match shutdown_hook {
        None => async move {}.boxed(),
        Some(hook) => async move {
            if let Err(err) = hook(hook_application).await {
                tracing::error!("{:?}", err);
            }
            if let Err(err) = crate::tina::log::wait_for_closeable_appender_shutdown() {
                println!("{:?}", err);
            }
        }
        .boxed(),
    };
    if application.shutdown_after_start.unwrap_or_default() {
        shutdown_hook.await;
        tokio::time::sleep(Duration::from_millis(1000)).await;
        application.terminate();
        return;
    }
    let ctrl_c = async {
        tokio::signal::ctrl_c().await.expect("failed to install Ctrl+C handler");
    };

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

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

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

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

    tokio::select! {
        _ = ctrl_c => {
            tracing::info!("处理系统{}信号...", "Ctrl+C中断");
            shutdown_hook.await;
            tokio::time::sleep(Duration::from_millis(1000)).await;
            application.terminate();
        },
        _ = terminate => {
            tracing::info!("处理系统{}信号...", "Kill中断");
            shutdown_hook.await;
            tokio::time::sleep(Duration::from_millis(1000)).await;
            application.terminate();
        },
        _ = interrupt => {
            tracing::info!("处理系统{}信号...", "Ctrl+C中断");
            shutdown_hook.await;
            tokio::time::sleep(Duration::from_millis(1000)).await;
            application.terminate();
        },
        _ = quit => {
            tracing::info!("处理系统{}信号...", "Ctrl+C中断");
            shutdown_hook.await;
            tokio::time::sleep(Duration::from_millis(1000)).await;
            application.terminate();
        },
    }
}