df_hyperliquid_core/client/native/
socket.rs

1//! The websocket for native programs: a [`Feed`] carries one [`crate::live::Session`]
2//! over tokio-tungstenite. Subscriptions made while down are sent on the next open; a
3//! post made while down is refused, and one cut off by a drop is reported, never resent.
4
5use std::time::{Duration, SystemTime, UNIX_EPOCH};
6
7use crate::live::{BadFrame, Frame, PING, Post, Refused, Retry, Session, Subscription, Tick};
8use futures_util::{SinkExt, StreamExt};
9use tokio::{
10    net::TcpStream,
11    sync::{mpsc, oneshot},
12    task::JoinHandle,
13    time::{Instant, sleep_until},
14};
15use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async, tungstenite::Message};
16
17use crate::client::native::clock;
18
19/// What a feed hands back, in the order it happened.
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum Event {
22    /// The socket opened and everything wanted was subscribed again.
23    Opened,
24    /// Boxed: a frame is large and most events are not frames.
25    Frame(Box<Frame>),
26    /// A frame that did not read. The socket is still alive.
27    Unreadable(BadFrame),
28    /// The socket closed, or could not open. The feed dials again after
29    /// `retry.after_ms`; the posts in `retry.unanswered` may or may not have
30    /// landed, so read the account to find out.
31    Closed { reason: String, retry: Retry },
32}
33
34#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
35pub enum FeedError {
36    #[error(transparent)]
37    Refused(#[from] Refused),
38    #[error("the feed has stopped")]
39    Stopped,
40}
41
42enum Command {
43    Subscribe(Subscription, oneshot::Sender<Result<(), Refused>>),
44    Unsubscribe(Subscription),
45    Post(Post, oneshot::Sender<Result<u64, Refused>>),
46}
47
48/// One websocket connection, kept open. Dropping the feed closes it.
49pub struct Feed {
50    commands: mpsc::UnboundedSender<Command>,
51    events: mpsc::UnboundedReceiver<Event>,
52    task: JoinHandle<()>,
53}
54
55impl Feed {
56    /// Connects to the session's network on the current tokio runtime.
57    #[must_use]
58    pub fn start(session: Session) -> Feed {
59        let url = session.url();
60        Feed::start_at(url, session)
61    }
62
63    /// Connects to `url` instead of the session's network: a proxy, or a test.
64    #[must_use]
65    pub fn start_at(url: impl Into<String>, session: Session) -> Feed {
66        let (commands, inbox) = mpsc::unbounded_channel();
67        let (outbox, events) = mpsc::unbounded_channel();
68        let task = tokio::spawn(run(url.into(), session, inbox, outbox));
69        Feed { commands, events, task }
70    }
71
72    /// Wants `subscription` from now on, across reconnects. Refused past the
73    /// session's limits.
74    pub async fn subscribe(&self, subscription: Subscription) -> Result<(), FeedError> {
75        let (reply, answer) = oneshot::channel();
76        self.commands.send(Command::Subscribe(subscription, reply)).map_err(|_| FeedError::Stopped)?;
77        Ok(answer.await.map_err(|_| FeedError::Stopped)??)
78    }
79
80    pub fn unsubscribe(&self, subscription: Subscription) -> Result<(), FeedError> {
81        self.commands.send(Command::Unsubscribe(subscription)).map_err(|_| FeedError::Stopped)
82    }
83
84    /// Sends an info request or a signed action over the socket, and gives
85    /// its id; the answer arrives as [`Frame::Post`] with the same id.
86    pub async fn post(&self, request: Post) -> Result<u64, FeedError> {
87        let (reply, answer) = oneshot::channel();
88        self.commands.send(Command::Post(request, reply)).map_err(|_| FeedError::Stopped)?;
89        Ok(answer.await.map_err(|_| FeedError::Stopped)??)
90    }
91
92    /// The next event; `None` once the feed has stopped.
93    pub async fn next(&mut self) -> Option<Event> {
94        self.events.recv().await
95    }
96
97    /// Closes the socket and stops.
98    pub fn close(self) {
99        drop(self);
100    }
101}
102
103impl Drop for Feed {
104    fn drop(&mut self) {
105        self.task.abort();
106    }
107}
108
109type Socket = WebSocketStream<MaybeTlsStream<TcpStream>>;
110
111/// Whether the connection was lost, and why, or the feed is done.
112enum Outcome {
113    Lost(String),
114    Stop,
115}
116
117async fn run(
118    url: String,
119    mut session: Session,
120    mut commands: mpsc::UnboundedReceiver<Command>,
121    events: mpsc::UnboundedSender<Event>,
122) {
123    loop {
124        // While it dials, the session still answers: a subscription is kept for
125        // the open, a post is refused. Nothing waits on a handshake.
126        let dial = connect_async(url.as_str());
127        tokio::pin!(dial);
128        let dialled = loop {
129            tokio::select! {
130                dialled = &mut dial => break dialled,
131                command = commands.recv() => match command {
132                    Some(command) => {
133                        handle(&mut session, command);
134                    }
135                    None => return,
136                },
137            }
138        };
139        let reason = match dialled {
140            Ok((mut socket, _)) => match open(&mut socket, &mut session, &mut commands, &events).await {
141                Outcome::Lost(reason) => reason,
142                Outcome::Stop => return,
143            },
144            Err(error) => error.to_string(),
145        };
146        let retry = session.closed(clock(), random());
147        let until = Instant::now() + Duration::from_millis(retry.after_ms);
148        if events.send(Event::Closed { reason, retry }).is_err() {
149            return;
150        }
151        // Until the next dial, subscriptions are kept for it and posts refused.
152        loop {
153            tokio::select! {
154                () = sleep_until(until) => break,
155                command = commands.recv() => match command {
156                    Some(command) => {
157                        handle(&mut session, command);
158                    }
159                    None => return,
160                },
161            }
162        }
163    }
164}
165
166async fn open(
167    socket: &mut Socket,
168    session: &mut Session,
169    commands: &mut mpsc::UnboundedReceiver<Command>,
170    events: &mpsc::UnboundedSender<Event>,
171) -> Outcome {
172    for frame in session.opened(clock()) {
173        if let Err(error) = socket.send(Message::text(frame)).await {
174            return Outcome::Lost(error.to_string());
175        }
176    }
177    if events.send(Event::Opened).is_err() {
178        return Outcome::Stop;
179    }
180    loop {
181        let wake = session.due().map(|at| Instant::now() + Duration::from_millis(at.saturating_sub(clock())));
182        tokio::select! {
183            message = socket.next() => match message {
184                Some(Ok(Message::Text(text))) => {
185                    let event = match session.received(clock(), &text) {
186                        Ok(frame) => Event::Frame(Box::new(frame)),
187                        Err(bad) => Event::Unreadable(bad),
188                    };
189                    if events.send(event).is_err() {
190                        return Outcome::Stop;
191                    }
192                }
193                Some(Ok(Message::Close(frame))) => {
194                    return Outcome::Lost(frame.map_or_else(|| "the venue closed the socket".to_owned(), |frame| frame.reason.to_string()));
195                }
196                // Pings are answered by tungstenite; the venue's own heartbeat is text.
197                Some(Ok(_)) => {}
198                Some(Err(error)) => return Outcome::Lost(error.to_string()),
199                None => return Outcome::Lost("the socket ended".to_owned()),
200            },
201            command = commands.recv() => match command {
202                Some(command) => {
203                    if let Some(frame) = handle(session, command)
204                        && let Err(error) = socket.send(Message::text(frame)).await
205                    {
206                        return Outcome::Lost(error.to_string());
207                    }
208                }
209                None => {
210                    let _ = socket.close(None).await;
211                    return Outcome::Stop;
212                }
213            },
214            () = until(wake) => match session.tick(clock()) {
215                Tick::Wait => {}
216                Tick::Ping => {
217                    if let Err(error) = socket.send(Message::text(PING)).await {
218                        return Outcome::Lost(error.to_string());
219                    }
220                }
221                Tick::Reconnect => {
222                    let _ = socket.close(None).await;
223                    return Outcome::Lost("nothing arrived for too long".to_owned());
224                }
225            },
226        }
227    }
228}
229
230/// Puts a command to the session and answers its caller; gives the frame to
231/// send, if the socket is open and there is one.
232fn handle(session: &mut Session, command: Command) -> Option<String> {
233    match command {
234        Command::Subscribe(subscription, reply) => match session.subscribe(subscription) {
235            Ok(frame) => {
236                let _ = reply.send(Ok(()));
237                frame
238            }
239            Err(refused) => {
240                let _ = reply.send(Err(refused));
241                None
242            }
243        },
244        Command::Unsubscribe(subscription) => session.unsubscribe(&subscription),
245        Command::Post(request, reply) => match session.post(&request) {
246            Ok((id, frame)) => {
247                let _ = reply.send(Ok(id));
248                Some(frame)
249            }
250            Err(refused) => {
251                let _ = reply.send(Err(refused));
252                None
253            }
254        },
255    }
256}
257
258async fn until(wake: Option<Instant>) {
259    match wake {
260        Some(at) => sleep_until(at).await,
261        None => std::future::pending().await,
262    }
263}
264
265/// Enough to spread reconnects: many clients dropped at once must not return
266/// at once. Not for anything secret.
267fn random() -> u64 {
268    let nanos = SystemTime::now().duration_since(UNIX_EPOCH).map_or(0, |since| u64::from(since.subsec_nanos()));
269    nanos.wrapping_mul(0x9e37_79b9_7f4a_7c15) ^ clock().rotate_left(29)
270}
271
272#[cfg(test)]
273#[path = "tests/socket.rs"]
274mod tests;
df_hyperliquid_core/client/native/socket.rs

df_hyperliquid_core/client/native/
socket.rs
