as-rust/src/main.rs

mod server;
mod control_c;
mod quant_libs;

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use tracing::{info, warn};
use tracing_appender_timezone::non_blocking::WorkerGuard;
use global::log_utils::send_remote_err_log;
use global::params::Params;

// 日志级别配置
fn log_level_init(log_str: String, port: u32, account_name: String) -> WorkerGuard {
    info!("日志级别读取成功：{}。", log_str);
    global::log_utils::final_init(log_str.as_str(), port, account_name)
}

// 获取本地配置
fn read_params() -> Params {
    let mut path = "config.toml";

    let args: Vec<String> = std::env::args().collect();

    for arg in &args {
        if !arg.contains("--config") {
            continue;
        }

        let p: Vec<&str> = arg.split("=").collect();
        path = p[1];
    }

    let params = Params::new(path).unwrap();

    return params;
}

#[tokio::main(flavor = "multi_thread", worker_threads = 2)]
async fn main() {
    // 日志级别配置
    let params = read_params();
    // 日志级别配置
    let _guard = log_level_init(params.log_level.clone(), params.port.clone(), params.account_name.clone());
    info!("配置读取成功：{:?}。", params);
    // 主进程控制
    let running = Arc::new(AtomicBool::new(true));

    // panic错误捕获，panic级别的错误直接退出
    let account_name_clone = params.account_name.clone();
    let panic_running = running.clone();
    std::panic::set_hook(Box::new(move |panic_info| {
        let msg = format!("account={}, type=panic, msg={:?}, location={:?}",
                          account_name_clone, panic_info.to_string(), panic_info.location());
        warn!("{}", msg);
        send_remote_err_log(msg);
        panic_running.store(false, Ordering::Relaxed);
    }));

    // ws退出程序
    let ws_running = Arc::new(AtomicBool::new(true));
    // quant初始化动作
    let quant_arc = quant_libs::init(params.clone(), ws_running.clone(), running.clone()).await;
    // 初始化中控服务
    server::run_server(params.port.clone(), running.clone(), quant_arc.clone());
    // ctrl c退出检查程序
    control_c::exit_handler(running.clone());

    // 每一秒检查一次程序是否结束
    while running.load(Ordering::Relaxed) {
        tokio::time::sleep(Duration::from_secs(1)).await;
    }

    info!("检测到退出信号！停止ws订阅……");
    ws_running.store(false, Ordering::Relaxed);
    tokio::time::sleep(Duration::from_secs(1)).await;

    info!("等待清空仓位、订单（再次按control c可以立马结束）……");
    let mut quant = quant_arc.lock().await;
    quant.exit().await;
    info!("程序已退出！为以防万一，请再次检查仓位和订单！");
}