df_hyperliquid_core/client/native/
socket.rs

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

df_hyperliquid_core/client/native/
socket.rs
