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#[async_trait]
19pub trait XtbStreamConnection {
20 type MessageStream: MessageStream;
22
23 type Error;
24
25 async fn subscribe(&mut self, command: &str, arguments: Option<Value>) -> Result<(), Self::Error>;
29
30 async fn unsubscribe(&mut self, command: &str, arguments: Option<Value>) -> Result<(), Self::Error>;
34
35 async fn make_message_stream(&mut self, filter: DataMessageFilter) -> Self::MessageStream;
37}
38
39
40#[derive(Debug)]
41pub struct BasicXtbStreamConnection {
42 stream_session_id: String,
44 sender: Sender<StreamDataMessage>,
46 sink: SplitSink<Stream, Message>,
48 listener_join: JoinHandle<()>,
50}
51
52
53impl BasicXtbStreamConnection {
54 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 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 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 fn prepare_arguments(arguments: Option<Value>) -> Result<Option<Map<String, Value>>, BasicXtbStreamConnectionError> {
87 match arguments {
88 None => Ok(None),
90 Some(Value::Object(obj)) => Ok(Some(obj)),
92 Some(Value::Null) => Ok(None),
93 _ => 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
135struct MessageHandler {
137 sender: Sender<StreamDataMessage>,
139}
140
141
142impl MessageHandler {
143 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 #[default]
168 Always,
169 Never,
171 Command(String),
173 FieldValue { name: String, value: Value },
177 Custom(Box<dyn Fn(&StreamDataMessage) -> bool + Send + Sync>),
179 All(Vec<DataMessageFilter>),
181 Any(Vec<DataMessageFilter>),
183}
184
185
186impl DataMessageFilter {
187 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 fn resolve_always(_: &StreamDataMessage) -> bool {
202 true
203 }
204
205 fn resolve_never(_: &StreamDataMessage) -> bool {
207 false
208 }
209
210 fn resolve_command(msg: &StreamDataMessage, command: &str) -> bool {
212 return msg.command.as_str() == command;
213 }
214
215 fn resolve_all(msg: &StreamDataMessage, ops: &Vec<DataMessageFilter>) -> bool {
217 ops.iter().all(|f| f.test_message(msg))
218 }
219
220 fn resolve_any(msg: &StreamDataMessage, ops: &Vec<DataMessageFilter>) -> bool {
222 ops.iter().any(|f| f.test_message(msg))
223 }
224
225 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 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 async fn next(&mut self) -> Option<StreamDataMessage>;
268}
269
270
271pub struct BasicMessageStream {
272 filter: DataMessageFilter,
274 stream: Receiver<StreamDataMessage>,
276}
277
278
279impl BasicMessageStream {
280 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}