Skip to main content

xtb_client/
stream_connection.rs

1use async_trait::async_trait;
2use futures_util::stream::SplitSink;
3use futures_util::{SinkExt, StreamExt};
4use serde::Serialize;
5use serde_json::{Map, to_string, to_value, Value};
6use thiserror::Error;
7use tokio::sync::broadcast::{channel, Receiver, Sender};
8use tokio::task::JoinHandle;
9use tokio_tungstenite::connect_async;
10use tokio_tungstenite::tungstenite::Message;
11use tracing::{debug, error, info};
12use url::Url;
13use crate::schema::{StreamDataMessage, SubscribeRequest, UnsubscribeRequest};
14
15use crate::listener::{listen_for_stream_data, Stream, StreamDataMessageHandler};
16
17/// Common interface for stream command api of the XTB.
18#[async_trait]
19pub trait XtbStreamConnection {
20    /// Type of message stream returned by the `make_message_stream` method.
21    type MessageStream: MessageStream;
22
23    type Error;
24
25    /// Subscribe for data stream from the XTB server
26    ///
27    /// The `arguments` must be `Value::Object`. Any other variants causes an error
28    async fn subscribe(&mut self, command: &str, arguments: Option<Value>) -> Result<(), Self::Error>;
29
30    /// Unsubscribe from data stream from the XTB server
31    ///
32    /// The `arguments` value must be `Value::Object` or `Value::Null`. Any other variants causes an error
33    async fn unsubscribe(&mut self, command: &str, arguments: Option<Value>) -> Result<(), Self::Error>;
34
35    /// Create message stream builder
36    async fn make_message_stream(&mut self, filter: DataMessageFilter) -> Self::MessageStream;
37}
38
39
40#[derive(Debug)]
41pub struct BasicXtbStreamConnection {
42    /// Stream session id used to identify for the stream server
43    stream_session_id: String,
44    /// Sender of messages used for delivering messages to `MessageStream` implementors
45    sender: Sender<StreamDataMessage>,
46    /// Sink used for sending messages to the XTB server
47    sink: SplitSink<Stream, Message>,
48    /// Handle used for join of listening task
49    listener_join: JoinHandle<()>,
50}
51
52
53impl BasicXtbStreamConnection {
54    /// Create new instance of the stream connection.
55    pub async fn new(url: Url, stream_session_id: String) -> Result<Self, BasicXtbStreamConnectionError> {
56        let (sender, _) = channel(64usize);
57        let host_clone = url.as_str().to_owned();
58        let (conn, _) = connect_async(url).await.map_err(|_| BasicXtbStreamConnectionError::CannotConnect(host_clone))?;
59        let (sink, stream) = conn.split();
60        let listener_join = listen_for_stream_data(stream, MessageHandler::new(sender.clone()));
61        Ok(Self {
62            stream_session_id,
63            sender,
64            sink,
65            listener_join,
66        })
67    }
68
69    /// Build message from request and arguments and send it to the server.
70    async fn assemble_and_send<T: Serialize>(&mut self, request: T, arguments: Option<Value>) -> Result<(), BasicXtbStreamConnectionError> {
71        let mut obj = to_value(request).map_err(|err| BasicXtbStreamConnectionError::SerializationFailed(err))?;
72        let prepared_arguments = Self::prepare_arguments(arguments)?;
73
74        if let Some(mut prepared_obj) = prepared_arguments {
75            // unwrap is safe here.
76            obj.as_object_mut().unwrap().append(&mut prepared_obj);
77        }
78        let serialized = to_string(&obj).map_err(|err| BasicXtbStreamConnectionError::SerializationFailed(err))?;
79        let message = Message::text(serialized);
80        self.sink.send(message).await.map_err(|err| BasicXtbStreamConnectionError::CannotSend(err))
81    }
82
83    /// Check and prepare arguments.
84    ///
85    /// The arguments must be None or Some(Value::Object) or Some(Value::Null). Otherwise, an error is returned.
86    fn prepare_arguments(arguments: Option<Value>) -> Result<Option<Map<String, Value>>, BasicXtbStreamConnectionError> {
87        match arguments {
88            // No arguments are provided
89            None => Ok(None),
90            // There is arguments object
91            Some(Value::Object(obj)) => Ok(Some(obj)),
92            Some(Value::Null) => Ok(None),
93            // Invalid input data
94            _ => Err(BasicXtbStreamConnectionError::InvalidArgumentsType)
95        }
96    }
97}
98
99
100impl Drop for BasicXtbStreamConnection {
101    fn drop(&mut self) {
102        self.listener_join.abort();
103    }
104}
105
106
107#[async_trait]
108impl XtbStreamConnection for BasicXtbStreamConnection {
109    type MessageStream = BasicMessageStream;
110
111    type Error = BasicXtbStreamConnectionError;
112
113    async fn subscribe(&mut self, command: &str, arguments: Option<Value>) -> Result<(), Self::Error> {
114        let request = SubscribeRequest::default()
115            .with_command(command)
116            .with_stream_session_id(&self.stream_session_id);
117        info!("Subscribing for {command}");
118        debug!("Subscription arguments are {arguments:?}");
119        self.assemble_and_send(request, arguments).await
120    }
121
122    async fn unsubscribe(&mut self, command: &str, arguments: Option<Value>) -> Result<(), Self::Error> {
123        let request = UnsubscribeRequest::default().with_command(command);
124        info!("Unsubscribing from {command}");
125        debug!("Unsubscription arguments are {arguments:?}");
126        self.assemble_and_send(request, arguments).await
127    }
128
129    async fn make_message_stream(&mut self, filter: DataMessageFilter) -> Self::MessageStream {
130        BasicMessageStream::new(filter, self.sender.subscribe())
131    }
132}
133
134
135/// Handle incoming data messages from stream
136struct MessageHandler {
137    /// Broadcast sender for messages
138    sender: Sender<StreamDataMessage>,
139}
140
141
142impl MessageHandler {
143    /// Create new instance of the MessageHandler
144    pub fn new(sender: Sender<StreamDataMessage>) -> Self {
145        Self { sender }
146    }
147}
148
149
150#[async_trait]
151impl StreamDataMessageHandler for MessageHandler {
152    async fn handle_message(&self, message: StreamDataMessage) {
153        let cmd = message.command.to_owned();
154        info!("Handling incoming message {cmd}");
155        debug!("Incoming message: {message:?}");
156        match self.sender.send(message) {
157            Err(err) => error!("Cannot broadcast message: {}", err),
158            _ => debug!("Message {cmd} was broadcast to the {} receivers", self.sender.len())
159        }
160    }
161}
162
163
164#[derive(Default)]
165pub enum DataMessageFilter {
166    /// Always true
167    #[default]
168    Always,
169    /// Always false
170    Never,
171    /// Command name must match
172    Command(String),
173    /// Value of field in `data` must match
174    /// Return true if and only if the `data` field is type of `Object::Value`, contains key
175    /// defined by `name` and the field is equal to `value`.
176    FieldValue { name: String, value: Value },
177    /// Apply custom filter fn
178    Custom(Box<dyn Fn(&StreamDataMessage) -> bool + Send + Sync>),
179    /// All inner filters must match. If list of predicates is empty, return true.
180    All(Vec<DataMessageFilter>),
181    /// Any inner filter must match. If list of predicates is empty, return false
182    Any(Vec<DataMessageFilter>),
183}
184
185
186impl DataMessageFilter {
187    /// Return true if the filter match, return false otherwise.
188    pub fn test_message(&self, msg: &StreamDataMessage) -> bool {
189        match self {
190            Self::Always => Self::resolve_always(msg),
191            Self::Never => Self::resolve_never(msg),
192            Self::Command(cmd) => Self::resolve_command(msg, cmd),
193            Self::All(ops) => Self::resolve_all(msg, ops),
194            Self::Any(ops) => Self::resolve_any(msg, ops),
195            Self::FieldValue { name, value } => Self::resolve_field_value(msg, name, value),
196            Self::Custom(cbk) => Self::resolve_custom(msg, cbk),
197        }
198    }
199
200    /// resolve StreamFilter::Always
201    fn resolve_always(_: &StreamDataMessage) -> bool {
202        true
203    }
204
205    /// resolve StreamFilter::Never
206    fn resolve_never(_: &StreamDataMessage) -> bool {
207        false
208    }
209
210    /// resolve StreamFilter::Command
211    fn resolve_command(msg: &StreamDataMessage, command: &str) -> bool {
212        return msg.command.as_str() == command;
213    }
214
215    /// resolve StreamFilter::All
216    fn resolve_all(msg: &StreamDataMessage, ops: &Vec<DataMessageFilter>) -> bool {
217        ops.iter().all(|f| f.test_message(msg))
218    }
219
220    /// resolve StreamFilter::Any
221    fn resolve_any(msg: &StreamDataMessage, ops: &Vec<DataMessageFilter>) -> bool {
222        ops.iter().any(|f| f.test_message(msg))
223    }
224
225    /// resolve StreamFilter::FieldValue
226    fn resolve_field_value(msg: &StreamDataMessage, field_name: &str, field_value: &Value) -> bool {
227        match &msg.data {
228            Value::Object(data_obj) => {
229                if let Some(field_content) = data_obj.get(field_name) {
230                    field_content == field_value
231                } else {
232                    false
233                }
234            }
235            _ => false
236        }
237    }
238
239    /// resolve StreamFilter::Custom
240    fn resolve_custom(msg: &StreamDataMessage, cbk: &Box<dyn Fn(&StreamDataMessage) -> bool + Send + Sync>) -> bool {
241        (*cbk)(msg)
242    }
243}
244
245
246#[derive(Debug, Error)]
247pub enum BasicXtbStreamConnectionError {
248    #[error("Cannot connect to server ({0}")]
249    CannotConnect(String),
250    #[error("Cannot send message")]
251    CannotSend(tokio_tungstenite::tungstenite::Error),
252    #[error("Cannot serialize data")]
253    SerializationFailed(serde_json::Error),
254    #[error("Only Value::Object can be used for the arguments")]
255    InvalidArgumentsType,
256}
257
258
259#[async_trait]
260pub trait MessageStream {
261    /// Get next message from stream
262    ///
263    /// # Returns
264    ///
265    /// `Some(x)` - next message in stream
266    /// `None` - there is no more message
267    async fn next(&mut self) -> Option<StreamDataMessage>;
268}
269
270
271pub struct BasicMessageStream {
272    /// The filter for messages
273    filter: DataMessageFilter,
274    /// Stream with incoming messages
275    stream: Receiver<StreamDataMessage>,
276}
277
278
279impl BasicMessageStream {
280    /// Create new instance
281    pub fn new(filter: DataMessageFilter, stream: Receiver<StreamDataMessage>) -> Self {
282        BasicMessageStream {
283            filter,
284            stream,
285        }
286    }
287}
288
289
290#[async_trait]
291impl MessageStream for BasicMessageStream {
292    async fn next(&mut self) -> Option<StreamDataMessage> {
293        while let Some(msg) = self.stream.recv().await.ok() {
294            if self.filter.test_message(&msg) {
295                return Some(msg);
296            }
297        }
298        None
299    }
300}
301
302
303#[cfg(test)]
304mod tests {
305    mod data_message_filter {
306        use rstest::rstest;
307        use serde_json::{from_str, Value};
308        use crate::schema::StreamDataMessage;
309        use crate::DataMessageFilter;
310
311        #[test]
312        fn always() {
313            let msg = StreamDataMessage::default();
314            assert!(DataMessageFilter::Always.test_message(&msg));
315        }
316
317        #[test]
318        fn never() {
319            let msg = StreamDataMessage::default();
320            assert!(!DataMessageFilter::Never.test_message(&msg));
321        }
322
323        #[rstest]
324        #[case("command", true)]
325        #[case("other_command", false)]
326        fn command(#[case] cmd: &str, #[case] expected_result: bool) {
327            let msg = StreamDataMessage { command: "command".to_string(), data: Value::Null };
328            assert_eq!(DataMessageFilter::Command(cmd.to_string()).test_message(&msg), expected_result);
329        }
330
331        #[rstest]
332        #[case(vec ! [], true)]
333        #[case(vec ! [DataMessageFilter::Always], true)]
334        #[case(vec ! [DataMessageFilter::Never], false)]
335        #[case(vec ! [DataMessageFilter::Command("command".to_string()), DataMessageFilter::Always], true)]
336        #[case(vec ! [DataMessageFilter::Command("command".to_string()), DataMessageFilter::Never], false)]
337        #[case(vec ! [DataMessageFilter::Command("other_command".to_string()), DataMessageFilter::Never], false)]
338        #[case(vec ! [DataMessageFilter::Command("other_command".to_string()), DataMessageFilter::Always], false)]
339        fn all(#[case] filters: Vec<DataMessageFilter>, #[case] expected_result: bool) {
340            let msg = StreamDataMessage { command: "command".to_owned(), data: Value::Null };
341            let f = DataMessageFilter::All(filters);
342            assert_eq!(f.test_message(&msg), expected_result);
343        }
344
345        #[rstest]
346        #[case(vec ! [], false)]
347        #[case(vec ! [DataMessageFilter::Always], true)]
348        #[case(vec ! [DataMessageFilter::Never], false)]
349        #[case(vec ! [DataMessageFilter::Command("command".to_string()), DataMessageFilter::Always], true)]
350        #[case(vec ! [DataMessageFilter::Command("command".to_string()), DataMessageFilter::Never], true)]
351        #[case(vec ! [DataMessageFilter::Command("other_command".to_string()), DataMessageFilter::Never], false)]
352        #[case(vec ! [DataMessageFilter::Command("other_command".to_string()), DataMessageFilter::Always], true)]
353        fn any(#[case] filters: Vec<DataMessageFilter>, #[case] expected_result: bool) {
354            let msg = StreamDataMessage { command: "command".to_owned(), data: Value::Null };
355            let f = DataMessageFilter::Any(filters);
356            assert_eq!(f.test_message(&msg), expected_result);
357        }
358
359        #[rstest]
360        #[case(r#"{"field": "value"}"#, true)]
361        #[case(r#"{"field": 10}"#, false)]
362        #[case(r#"{"other_field": 10}"#, false)]
363        #[case(r#"null"#, false)]
364        fn filed_value(#[case] source_data: &str, #[case] expected_value: bool) {
365            let data: Value = from_str(source_data).unwrap();
366            let msg = StreamDataMessage { data, command: "".to_owned() };
367            let f = DataMessageFilter::FieldValue { name: "field".to_owned(), value: Value::String("value".to_owned()) };
368            assert_eq!(f.test_message(&msg), expected_value)
369        }
370
371        #[test]
372        fn custom_true() {
373            let msg = StreamDataMessage::default();
374            let f = DataMessageFilter::Custom(Box::new(|msg| true));
375            assert_eq!(f.test_message(&msg), true)
376        }
377
378        #[test]
379        fn custom_false() {
380            let msg = StreamDataMessage::default();
381            let f = DataMessageFilter::Custom(Box::new(|msg| false));
382            assert_eq!(f.test_message(&msg), false)
383        }
384    }
385}