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

df_hyperliquid_core/client/native/
socket.rs
