xapi_okx/common/ws/
stream.rs

1use crate::common::{
2    payload::{OkxWsStreamOperation, OkxWsStreamPayload},
3    response::{OkxWsStreamData, OkxWsStreamResponse},
4};
5use ezsockets::{Bytes, Client, ClientConfig, ClientExt, Error, Utf8Bytes};
6use std::collections::HashMap;
7use tokio::sync::{mpsc, oneshot};
8use ulid::Ulid;
9use xapi_shared::ws::{api::SharedWsApiTrait, error::SharedWsError, stream::SharedWsStreamTrait};
10
11pub struct OkxWsStream {
12    client: Client<Self>,
13    on_connect_tx: Option<oneshot::Sender<()>>,
14    oneshot_tx_map: HashMap<String, oneshot::Sender<Result<OkxWsStreamResponse, SharedWsError>>>,
15    stream_tx_map: HashMap<serde_json::Value, mpsc::Sender<Result<OkxWsStreamData, SharedWsError>>>,
16}
17
18pub enum OkxWsStreamCall {
19    SubscribeStream {
20        args: Vec<(
21            serde_json::Value,
22            mpsc::Sender<Result<OkxWsStreamData, SharedWsError>>,
23        )>,
24        tx: oneshot::Sender<Result<OkxWsStreamResponse, SharedWsError>>,
25    },
26}
27
28#[async_trait::async_trait]
29impl ClientExt for OkxWsStream {
30    type Call = OkxWsStreamCall;
31
32    async fn on_text(&mut self, text: Utf8Bytes) -> Result<(), Error> {
33        let msg = text.to_string();
34
35        if let Some(result) = self.recv_stream_resp(&msg).await {
36            return result.map_err(|err| err.into());
37        }
38
39        if let Some(result) = self.recv_oneshot_resp(&msg) {
40            return result.map_err(|err| err.into());
41        }
42
43        tracing::error!(?msg, "unhandled bn ws message");
44        Err(SharedWsError::AppError("unhandled bn ws message".to_string()).into())
45    }
46
47    async fn on_binary(&mut self, _bytes: Bytes) -> Result<(), Error> {
48        unimplemented!()
49    }
50
51    async fn on_call(&mut self, call: Self::Call) -> Result<(), Error> {
52        match call {
53            OkxWsStreamCall::SubscribeStream { args, tx } => self.subscribe_streams(args, tx)?,
54        }
55
56        Ok(())
57    }
58
59    async fn on_connect(&mut self) -> Result<(), Error> {
60        if let Some(tx) = self.on_connect_tx.take() {
61            tx.send(())
62                .inspect_err(|err| {
63                    tracing::error!(?err, "failed to send on_connect signal");
64                })
65                .map_err(|_| {
66                    SharedWsError::ChannelClosedError("first on connect channel closed".to_string())
67                })?;
68        }
69        Ok(())
70    }
71}
72
73impl SharedWsApiTrait<String, OkxWsStreamPayload, OkxWsStreamResponse> for OkxWsStream {
74    fn get_client(&self) -> &Client<Self> {
75        &self.client
76    }
77
78    fn get_oneshot_tx_map(
79        &mut self,
80    ) -> &mut HashMap<String, oneshot::Sender<Result<OkxWsStreamResponse, SharedWsError>>> {
81        &mut self.oneshot_tx_map
82    }
83}
84
85#[async_trait::async_trait]
86impl SharedWsStreamTrait<serde_json::Value, OkxWsStreamData> for OkxWsStream {
87    fn get_stream_tx_map(
88        &mut self,
89    ) -> &mut HashMap<serde_json::Value, mpsc::Sender<Result<OkxWsStreamData, SharedWsError>>> {
90        &mut self.stream_tx_map
91    }
92}
93
94impl OkxWsStream {
95    pub async fn connect(config: ClientConfig) -> Client<Self> {
96        let (on_connect_tx, on_connect_rx) = oneshot::channel();
97
98        let (client, future) = ezsockets::connect(
99            |client| Self {
100                client,
101                on_connect_tx: Some(on_connect_tx),
102                oneshot_tx_map: Default::default(),
103                stream_tx_map: Default::default(),
104            },
105            config,
106        )
107        .await;
108
109        tokio::spawn(async move {
110            future.await.inspect_err(|err| {
111                tracing::error!(?err, "okx ws client connection error");
112            })
113        });
114
115        _ = on_connect_rx.await;
116
117        client
118    }
119
120    fn subscribe_streams(
121        &mut self,
122        args: Vec<(
123            serde_json::Value,
124            mpsc::Sender<Result<OkxWsStreamData, SharedWsError>>,
125        )>,
126        tx: oneshot::Sender<Result<OkxWsStreamResponse, SharedWsError>>,
127    ) -> Result<(), SharedWsError> {
128        if args.is_empty() {
129            tracing::warn!("no args to subscribe");
130            return Ok(());
131        }
132
133        for (arg, _) in &args {
134            if self.stream_tx_map.contains_key(arg) {
135                tracing::error!(?arg, "stream already subscribed");
136                return Err(SharedWsError::AppError(
137                    "stream already subscribed".to_string(),
138                ));
139            }
140        }
141
142        let id = Ulid::new().to_string();
143
144        let payload = OkxWsStreamPayload {
145            id: id.clone(),
146            operation: OkxWsStreamOperation::Subscribe,
147            args: serde_json::Value::Array(
148                args.iter().map(|(arg, _)| arg.clone()).collect::<Vec<_>>(),
149            ),
150        };
151
152        for (arg, tx) in args {
153            self.stream_tx_map.insert(arg, tx);
154        }
155
156        self.send_oneshot(payload, tx)
157    }
158}
xapi_okx/common/ws/stream.rs

xapi_okx/common/ws/
stream.rs
