as-rust/src/main.rs

mod server;
mod control_c;
mod core_libs;
mod clear_core_libs;

use std::str::FromStr;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use tokio::sync::Mutex;
use tracing::{error, info, warn};
use tracing_appender_timezone::non_blocking::WorkerGuard;
use global::cci::CentralControlInfo;
use global::params::Params;

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

// 清仓程序日志级别配置
fn clear_log_level_init(log_str: String, port: u32, account_name: String) -> WorkerGuard {
    let log = global::clear_log_utils::final_init(log_str.as_str(), port, account_name);
    info!("清仓程序日志级别读取成功：{}。", log_str);
    return log;
}

// 获取本地配置
// 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];
//     }
//
//     println!("配置文件路径：{}", path);
//     let params = Params::new(path).unwrap();
//
//     return params;
// }

// 获取本地配置- json 文件解析
fn read_params_json() -> Params {
    let mut path = "config.json";
    let mut call_port = 5555;
    // 运行模式 0.正常策略运行， 1.清理挂单及仓位
    let mut run_mode = 0;
    let mut r_id = "-1".to_string();

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

    for arg in &args {
        // Check for the --config argument and assign its value.
        if arg.starts_with("--config=") {
            let parts: Vec<&str> = arg.split('=').collect();
            if parts.len() == 2 {
                path = parts[1];
            }
        }

        // Check for the --port argument and assign its value.
        if arg.starts_with("--port=") {
            let parts: Vec<&str> = arg.split('=').collect();
            if parts.len() == 2 {
                match u32::from_str(parts[1]) {
                    Ok(num) => call_port = num,
                    Err(_) => eprintln!("Invalid number for port: {}", parts[1]),
                }
            }
        }

        // Check for the --port argument and assign its value.
        if arg.starts_with("--run_mode=") {
            let parts: Vec<&str> = arg.split('=').collect();
            if parts.len() == 2 {
                match u32::from_str(parts[1]) {
                    Ok(num) => run_mode = num,
                    Err(_) => eprintln!("Invalid number for run_mode: {}", parts[1]),
                }
            }
        }

        //上报ID
        if arg.starts_with("--r_id=") {
            let parts: Vec<&str> = arg.split('=').collect();
            if parts.len() == 2 {
                r_id = parts[1].to_string();
            } else {
                error!("启动失败，回执单id参数格式设置错误 --check_id=xxx！");
                panic!("启动失败，回执单id参数格式设置错误 --check_id=xxx！");
            }
        }
    }
    println!("通讯端口：{}, 配置文件路径：{}", call_port, path);
    println!("请求指令携带参数：{:?}",args.clone());
    let mut params = Params::new_json(path, call_port).unwrap();
    if run_mode == 1 {
        params.run_mode = 1;
        if r_id == "-1" {
            error!("启动失败，缺少回执单id参数！");
            panic!("启动失败，缺少回执单id参数！");
        }
    }
    params.r_id = r_id;
    return params;
}

#[tokio::main(flavor = "multi_thread", worker_threads = 2)]
async fn main() {
    // 日志级别配置
    // let params = read_params();
    let params = read_params_json();

    // 日志级别配置
    let _guard;
    if params.run_mode == 1 {
        _guard = clear_log_level_init(params.log_level.clone(), params.port.clone(), params.account_name.clone());
    } else {
        _guard = log_level_init(params.log_level.clone(), params.port.clone());
    }

    info!("--------------------------------程序开始执行-----------------------------");
    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);
        panic_running.store(false, Ordering::Relaxed);
    }));

    // 中央控制器信息
    let cci = CentralControlInfo::new();
    let cci_arc = Arc::new(Mutex::new(cci));

    // ws退出程序
    let ws_running = Arc::new(AtomicBool::new(true));
    if params.run_mode == 1 {
        // core初始化动作
        let mut core_arc = clear_core_libs::init(params.clone(), ws_running.clone(), running.clone(), cci_arc.clone()).await;
        info!("开始执行清仓程序");
        core_arc.exit(params.r_id).await;
        info!("清仓程序执行完毕");
        // 强制退出
        std::process::exit(0);
    } else {
        // core初始化动作
        let core_arc = core_libs::init(params.clone(), ws_running.clone(), running.clone(), cci_arc.clone()).await;
        // 初始化中控服务
        server::run_server(params.port.clone(), running.clone(), cci_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 core = core_arc.lock().await;
        core.exit().await;
        info!("程序已退出！为以防万一，请再次检查仓位和订单！");
        // 等两秒，等中控反应过来
        tokio::time::sleep(Duration::from_secs(2)).await;
        // 强制退出
        std::process::exit(0);
    }
}