as-rust/exchanges/src/binance_swap_ws.rs

use std::sync::Arc;
use std::sync::atomic::AtomicBool;

use futures_channel::mpsc::{UnboundedReceiver, UnboundedSender};
use serde_json::json;
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::{Error, Message};
use tracing::{info, trace};

use crate::response_base::ResponseData;
use crate::socket_tool::{AbstractWsMode, HeartbeatType};
use crate::utils::get_time_microsecond;

//类型
pub enum BinanceSwapWsType {
    PublicAndPrivate,
}

//订阅频道
#[derive(Clone)]
pub enum BinanceSwapSubscribeType {
    PuBookTicker,
    PuAggTrade,
    PuDepth20levels100ms,
}

//账号信息
#[derive(Clone)]
#[allow(dead_code)]
pub struct BinanceSwapLogin {
    api_key: String,
    api_secret: String,
}

#[derive(Clone)]
#[allow(dead_code)]
pub struct BinanceSwapWs {
    //类型
    label: String,
    //地址
    address_url: String,
    //账号
    login_param: Option<BinanceSwapLogin>,
    //币对
    symbol_s: Vec<String>,
    //订阅
    subscribe_types: Vec<BinanceSwapSubscribeType>,
    //心跳间隔
    heartbeat_time: u64,
}

impl BinanceSwapWs {
    /*******************************************************************************************************/
    /*****************************************获取一个对象****************************************************/
    /*******************************************************************************************************/
    pub fn new(is_colo: bool, login_param: Option<BinanceSwapLogin>, ws_type: BinanceSwapWsType) -> BinanceSwapWs {
        return BinanceSwapWs::new_label("default-BinanceSwapWs".to_string(), is_colo, login_param, ws_type);
    }
    pub fn new_label(label: String, is_colo: bool, login_param: Option<BinanceSwapLogin>, ws_type: BinanceSwapWsType) -> BinanceSwapWs {
        /*******公共频道-私有频道数据组装*/
        let address_url = match ws_type {
            BinanceSwapWsType::PublicAndPrivate => {
                // "wss://fstream.binance.com/stream?streams=btcusdt@depth20@100ms".to_string(),
                "wss://fstream.binance.com/stream".to_string()
            }
        };

        if is_colo {
            info!("开启高速(未配置，走普通：{})通道",address_url);
        } else {
            info!("走普通通道：{}",address_url);
        }

        BinanceSwapWs {
            label,
            address_url,
            login_param,
            symbol_s: vec![],
            subscribe_types: vec![],
            heartbeat_time: 1000 * 20,
        }
    }

    /*******************************************************************************************************/
    /*****************************************订阅函数********************************************************/
    /*******************************************************************************************************/
    //手动添加订阅信息
    pub fn set_subscribe(&mut self, subscribe_types: Vec<BinanceSwapSubscribeType>) {
        self.subscribe_types.extend(subscribe_types);
    }
    //手动添加币对
    pub fn set_symbols(&mut self, mut b_array: Vec<String>) {
        for symbol in b_array.iter_mut() {
            // 小写
            *symbol = symbol.to_lowercase();
            // 字符串替换
            *symbol = symbol.replace("_", "");
            *symbol = symbol.replace("-", "");
        }
        self.symbol_s = b_array;
    }
    //频道是否需要登录
    fn contains_pr(&self) -> bool {
        for t in self.subscribe_types.clone() {
            if match t {
                BinanceSwapSubscribeType::PuBookTicker => false,
                BinanceSwapSubscribeType::PuAggTrade => false,
                BinanceSwapSubscribeType::PuDepth20levels100ms => false,
            } {
                return true;
            }
        }
        false
    }
    /*******************************************************************************************************/
    /*****************************************工具函数********************************************************/
    /*******************************************************************************************************/
    //订阅枚举解析
    pub fn enum_to_string(symbol: String, subscribe_type: BinanceSwapSubscribeType) -> String {
        match subscribe_type {
            BinanceSwapSubscribeType::PuAggTrade => {
                format!("{}@aggTrade", symbol)
            }
            BinanceSwapSubscribeType::PuDepth20levels100ms => {
                format!("{}@depth20@100ms", symbol)
            }
            BinanceSwapSubscribeType::PuBookTicker => {
                format!("{}@bookTicker", symbol)
            }
        }
    }
    //订阅信息生成
    pub fn get_subscription(&self) -> String {
        let mut params = vec![];
        for symbol in &self.symbol_s {
            for subscribe_type in &self.subscribe_types {
                let ty_str = Self::enum_to_string(symbol.clone(), subscribe_type.clone());
                params.push(ty_str);
            }
        }

        let str = json!({
            "method": "SUBSCRIBE",
            "params":  params,
            "id": 1
            });
        str.to_string()
    }
    /*******************************************************************************************************/
    /*****************************************socket基本*****************************************************/
    /*******************************************************************************************************/
    //链接
    pub async fn ws_connect_async(&mut self,
                                  bool_v1: Arc<AtomicBool>,
                                  write_tx_am: &Arc<Mutex<UnboundedSender<Message>>>,
                                  write_rx: UnboundedReceiver<Message>,
                                  read_tx: UnboundedSender<ResponseData>) -> Result<(), Error>
    {
        let login_is = self.contains_pr();
        let subscription = self.get_subscription();
        let address_url = self.address_url.clone();
        let label = self.label.clone();
        let heartbeat_time = self.heartbeat_time.clone();


        //心跳-- 方法内部线程启动
        let write_tx_clone1 = Arc::clone(write_tx_am);
        tokio::spawn(async move {
            trace!("线程-异步心跳-开始");
            AbstractWsMode::ping_or_pong(write_tx_clone1, HeartbeatType::Pong, heartbeat_time).await;
            trace!("线程-异步心跳-结束");
        });

        //设置订阅
        let mut subscribe_array = vec![];
        if login_is {
            //登录相关
        }
        subscribe_array.push(subscription.to_string());

        //链接
        let t2 = tokio::spawn(async move {
            trace!("线程-异步链接-开始");
            AbstractWsMode::ws_connect_async(bool_v1, address_url.clone(),
                                             label.clone(), subscribe_array,
                                             write_rx, read_tx,
                                             Self::message_text,
                                             Self::message_ping,
                                             Self::message_pong,
            ).await.expect("币安-期货");
            trace!("线程-异步链接-结束");
        });
        tokio::try_join!(t2).unwrap();
        trace!("线程-心跳与链接-结束");

        Ok(())
    }
    /*******************************************************************************************************/
    /*****************************************数据解析*****************************************************/
    /*******************************************************************************************************/
    //数据解析-Text
    pub fn message_text(text: String) -> Option<ResponseData> {
        let mut response_data = Self::ok_text(text);
        response_data.time = get_time_microsecond();
        match response_data.code.as_str() {
            "200" => Option::from(response_data),
            "-201" => {
                trace!("订阅成功：{:?}", response_data);
                None
            }
            _ => None
        }
    }
    //数据解析-ping
    pub fn message_ping(_pi: Vec<u8>) -> Option<ResponseData> {
        trace!("服务器响应-ping");
        return None;
    }
    //数据解析-pong
    pub fn message_pong(_po: Vec<u8>) -> Option<ResponseData> {
        trace!("服务器响应-pong");
        return None;
    }
    //数据解析
    pub fn ok_text(text: String) -> ResponseData {
        // trace!("原始数据");
        // trace!(?text);
        let mut res_data = ResponseData::new("".to_string(), "200".to_string(), "success".to_string(), "".to_string());
        let json_value: serde_json::Value = serde_json::from_str(&text).unwrap();

        if json_value.get("result").is_some() && json_value.get("id").is_some() &&
            json_value.get("id").unwrap() == 1
        {
            res_data.code = "-201".to_string();
            res_data.message = "订阅成功".to_string();
        } else if json_value.get("error").is_some() {//订阅返回
            res_data.code = json_value["error"]["code"].to_string();
            res_data.message = json_value["error"]["msg"].to_string();
        } else if json_value.get("stream").is_some() {//订阅返回
            res_data.data = format!("{}", json_value.get("data").as_ref().unwrap());
            res_data.code = "200".to_string();

            let channel = format!("{}", json_value.get("stream").as_ref().unwrap());
            if channel.contains("@aggTrade") {
                res_data.channel = "aggTrade".to_string();
            } else if channel.contains("@depth20@100ms") {
                res_data.channel = "depth".to_string();
            } else if channel.contains("@bookTicker") {
                res_data.channel = "bookTicker".to_string();
            } else {
                res_data.channel = "未知的频道".to_string();
            }
        } else {
            res_data.data = text
        }
        res_data
    }
}