xtb_client/
connection.rs

1use std::collections::HashMap;
2use std::future::Future;
3use std::pin::{Pin, pin};
4use std::sync::Arc;
5use std::task::{Context, Poll, Waker};
6
7use async_trait::async_trait;
8use futures_util::{SinkExt, StreamExt};
9use futures_util::stream::{SplitSink};
10use serde_json::Value;
11use thiserror::Error;
12use tokio::sync::Mutex;
13use tokio::task::JoinHandle;
14use tokio_tungstenite::{connect_async};
15use tokio_tungstenite::tungstenite::Message;
16use tracing::{error, warn};
17use url::Url;
18
19use crate::schema::Request;
20use crate::listener::{listen_for_responses, ResponseHandler, Stream};
21use crate::message_processing::ProcessedMessage;
22
23/// Interface for XTB servers connectors.
24#[async_trait]
25pub trait XtbConnection {
26
27    type Error;
28
29    type Response: Future<Output=Result<ProcessedMessage, BasicXtbConnectionError>>;
30
31    /// Send standard command to the server.
32    async fn send_command(&mut self, command: &str, payload: Option<Value>) -> Result<Self::Response, Self::Error>;
33}
34
35
36#[derive(Debug, Error)]
37pub enum BasicXtbConnectionError {
38    #[error("Cannot connect to server ({0}")]
39    CannotConnect(String),
40    #[error("Cannot serialize command payload")]
41    SerializationError(serde_json::Error),
42    #[error("Cannot send request to the XTB server.")]
43    CannotSendRequest(tokio_tungstenite::tungstenite::Error),
44}
45
46
47/// Common implementation of the `XtbConnection` trait.
48pub struct BasicXtbConnection {
49    sink: SplitSink<Stream, Message>,
50    tag_maker: TagMaker,
51    promise_state_by_tag: Arc<Mutex<HashMap<String, Arc<Mutex<ResponsePromiseState>>>>>,
52    listener_join: JoinHandle<()>,
53}
54
55
56impl BasicXtbConnection {
57    /// Create new instance from server url
58    pub async fn new(url: Url) -> Result<Self, BasicXtbConnectionError> {
59        let host_clone = url.as_str().to_owned();
60        let (conn, _) = connect_async(url).await.map_err(|err| {
61            error!("Cannot connect to server {}: {:?}", host_clone, err);
62            BasicXtbConnectionError::CannotConnect(host_clone)
63        })?;
64
65        let (sink, stream) = conn.split();
66        let lookup = Arc::new(Mutex::new(HashMap::new()));
67        let listener_join = listen_for_responses(stream, BasicConnectionResponseHandler(lookup.clone()));
68        let instance = Self {
69            sink,
70            tag_maker: TagMaker::default(),
71            promise_state_by_tag: lookup,
72            listener_join
73        };
74        Ok(instance)
75    }
76
77    /// Build a request from command and payload.
78    /// Return request and its tag.
79    fn build_request(&mut self, command: &str, mut payload: Option<Value>) -> (Request, String) {
80        let tag = self.tag_maker.next();
81
82        if let Some(p) = &payload {
83            if p.is_null() {
84                payload = None;
85            }
86        }
87
88        let r = Request::default()
89            .with_command(command)
90            .with_maybe_arguments(payload)
91            .with_custom_tag(&tag);
92        (r, tag)
93    }
94}
95
96
97
98#[async_trait]
99impl XtbConnection for BasicXtbConnection {
100
101    type Error = BasicXtbConnectionError;
102
103    type Response = ResponsePromise;
104
105    async fn send_command(&mut self, command: &str, payload: Option<Value>) -> Result<Self::Response, Self::Error> {
106        let (request, tag) = self.build_request(command, payload);
107        let request_json = serde_json::to_string(&request).map_err(BasicXtbConnectionError::SerializationError)?;
108        let message = Message::Text(request_json);
109
110        let (promise, state) = ResponsePromise::new();
111        self.promise_state_by_tag.lock().await.insert(tag, state);
112        self.sink.send(message).await.map_err(BasicXtbConnectionError::CannotSendRequest)?;
113
114        Ok(promise)
115    }
116}
117
118
119impl Drop for BasicXtbConnection {
120    fn drop(&mut self) {
121        // Stop the listening task
122        self.listener_join.abort();
123    }
124}
125
126
127/// Internal state shared between the ResponsePromise and BasicXtbConnection instance.
128/// This state is used to deliver response to the consumer.
129#[derive(Default, Debug)]
130pub struct ResponsePromiseState {
131    /// The response.
132    ///
133    /// * `None` - the response is not ready yet.
134    /// * `Some(response)` - the response is ready to be delivered.
135    result: Option<Result<ProcessedMessage, BasicXtbConnectionError>>,
136    /// If the `ResponsePromise` was palled, the `Waker` is stored here.
137    /// When response is set and the waker is set, the waker is called.
138    waker: Option<Waker>,
139}
140
141
142impl ResponsePromiseState {
143    /// Set response. If a waker is set in the state, it is notified.
144    pub fn set_result(&mut self, result: Result<ProcessedMessage, BasicXtbConnectionError>) {
145        self.result = Some(result);
146        if let Some(waker) = self.waker.take() {
147            waker.wake();
148        }
149    }
150}
151
152
153/// Handle messages delivered by XTB server
154struct BasicConnectionResponseHandler(Arc<Mutex<HashMap<String, Arc<Mutex<ResponsePromiseState>>>>>);
155
156#[async_trait]
157impl ResponseHandler for BasicConnectionResponseHandler {
158    async fn handle_response(&self, response: ProcessedMessage) {
159        let maybe_tag = match &response {
160            ProcessedMessage::Response(resp) => resp.custom_tag.as_ref(),
161            ProcessedMessage::ErrorResponse(resp) => resp.custom_tag.as_ref(),
162        };
163
164        // if there is no tag, continue (the message cannot be routed to consumer)
165        let tag = match maybe_tag {
166            Some(t) => t,
167            _ => {
168                warn!("Response has no tag and cannot be routed: {:?}", response);
169                return;
170            }
171        };
172
173        // try to deliver message to its consumer
174        if let Some(state) = self.0.lock().await.remove(tag) {
175            state.lock().await.set_result(Ok(response));
176        }
177    }
178}
179
180
181/// Represent promise of a response delivery in a future.
182///
183/// Implements the `Future` trait and when the future is awaited, it is resolved by response
184/// returned from a server. The response is type of `Result<Response, ErrorResponse>`.
185#[derive(Debug)]
186pub struct ResponsePromise {
187    /// Shared internal state. The second "point" is in the source connection.
188    state: Arc<Mutex<ResponsePromiseState>>,
189}
190
191
192impl ResponsePromise {
193    /// Create new instance and return tuple:
194    ///
195    /// 1. instance of `Self`
196    /// 2. thread safe `ResponsePromiseState` for response delivery.
197    pub fn new() -> (Self, Arc<Mutex<ResponsePromiseState>>) {
198        let state = ResponsePromiseState::default();
199        let wrapped_state = Arc::new(Mutex::new(state));
200        (Self { state: wrapped_state.clone() }, wrapped_state)
201    }
202}
203
204
205impl Future for ResponsePromise {
206    type Output = Result<ProcessedMessage, BasicXtbConnectionError>;
207
208    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
209        // Try to get the lock
210        if let Poll::Ready(mut guard) = pin!(self.state.lock()).poll(cx) {
211            // If response is set, return it as `Poll::Ready`
212            if let Some(response) = guard.result.take() {
213                return Poll::Ready(response);
214            }
215            // If response is not ready yet, register the waker.
216            guard.waker = Some(cx.waker().clone());
217        }
218        // Wait until response is ready
219        Poll::Pending
220    }
221}
222
223
224/// Helper struct generating message tags.
225///
226/// It generates unique tags with prefix "message_" followed by incremented positive integer number.
227/// The increment step is 1 and the first number is 1.
228///
229/// Example of series is: "message_1", "message_2", "message_3", ...
230#[derive(Default, Debug)]
231struct TagMaker(u64);
232
233
234impl TagMaker {
235    fn next(&mut self) -> String {
236        self.0 += 1;
237        format!("message_{}", self.0)
238    }
239}
240
241
242#[cfg(test)]
243mod tests {
244    mod response_promise {
245        use std::sync::Arc;
246        use std::time::Duration;
247
248        use rstest::*;
249        use serde_json::to_value;
250        use tokio::spawn;
251        use tokio::sync::Mutex;
252        use tokio::time::sleep;
253
254        use crate::schema::Response;
255        use crate::connection::ResponsePromiseState;
256        use crate::message_processing::ProcessedMessage;
257        use crate::ResponsePromise;
258
259        #[rstest]
260        #[case(0)]
261        #[case(1)]
262        #[case(50)]
263        #[case(100)]
264        #[timeout(Duration::from_millis(500))]
265        #[tokio::test]
266        async fn deliver_data(#[case] delay_ms: u64) {
267            let (instance, target) = ResponsePromise::new();
268            spawn(write_data(target, delay_ms));
269            let result = instance.await;
270        }
271
272        async fn write_data(target: Arc<Mutex<ResponsePromiseState>>, delay: u64) {
273            if delay > 0 {
274                sleep(Duration::from_millis(delay)).await;
275            }
276            let mut lock = target.lock().await;
277            let mut response = Response::default();
278            response.return_data = Some(to_value(42).unwrap());
279            lock.set_result(Ok(ProcessedMessage::Response(response)));
280        }
281    }
282
283    mod tag_maker {
284        use crate::connection::TagMaker;
285
286        #[test]
287        fn make_series() {
288            let mut maker = TagMaker::default();
289
290            let tag = maker.next();
291            assert_eq!(tag, "message_1");
292            let tag = maker.next();
293            assert_eq!(tag, "message_2");
294            let tag = maker.next();
295            assert_eq!(tag, "message_3");
296        }
297    }
298}
xtb_client/connection.rs

xtb_client/
connection.rs
