df-hyperliquid-core 0.0.8-beta

//! The websocket in the browser and Workers: a [`Feed`] carries one [`crate::live::Session`]
//! over the global `WebSocket`. Subscriptions made while down are sent on the next open; a
//! post made while down is refused, and one cut off by a drop is reported, never resent.
//!
//! Everything runs on the one JS thread: the socket's events and the timer call into the
//! feed as tasks of their own, and [`Feed::next`] waits for the events they leave.

use std::{
    cell::RefCell,
    collections::VecDeque,
    future::poll_fn,
    rc::{Rc, Weak},
    task::{Poll, Waker},
};

use crate::client::web::clock;
use crate::live::{BadFrame, Frame, PING, Post, Refused, Retry, Session, Subscription, Tick};
use js_sys::{Array, Function, Math, Reflect};
use wasm_bindgen::{JsCast, JsValue, closure::Closure};
use web_sys::{CloseEvent, MessageEvent, WebSocket};

/// What a feed hands back, in the order it happened.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Event {
    /// The socket opened and everything wanted was subscribed again.
    Opened,
    /// Boxed: a frame is large and most events are not frames.
    Frame(Box<Frame>),
    /// A frame that did not read. The socket is still alive.
    Unreadable(BadFrame),
    /// The socket closed, or could not open. The feed dials again after
    /// `retry.after_ms`; the posts in `retry.unanswered` may or may not have
    /// landed, so read the account to find out.
    Closed { reason: String, retry: Retry },
}

#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum FeedError {
    #[error(transparent)]
    Refused(#[from] Refused),
    /// The runtime gave no timer (no `setTimeout`, or one that threw), so the
    /// feed could not keep its time and closed for good.
    #[error("the feed has stopped")]
    Stopped,
}

/// One websocket connection, kept open. Dropping the feed closes it.
pub struct Feed {
    driver: Rc<RefCell<Driver>>,
}

impl Feed {
    /// Connects to the session's network.
    #[must_use]
    pub fn start(session: Session) -> Feed {
        let url = session.url();
        Feed::start_at(url, session)
    }

    /// Connects to `url` instead of the session's network: a proxy, or a test.
    #[must_use]
    pub fn start_at(url: impl Into<String>, session: Session) -> Feed {
        let global: JsValue = js_sys::global().into();
        let timers = function(&global, "setTimeout").zip(function(&global, "clearTimeout"));
        let driver = Rc::new_cyclic(|driver| {
            RefCell::new(Driver {
                url: url.into(),
                session,
                global,
                timers,
                socket: None,
                timer: None,
                calls: Calls::new(driver),
                events: VecDeque::new(),
                waker: None,
                stopped: false,
            })
        });
        if let Ok(mut driver) = driver.try_borrow_mut() {
            if driver.timers.is_some() {
                driver.dial();
            } else {
                driver.stop();
            }
        }
        Feed { driver }
    }

    /// Wants `subscription` from now on, across reconnects. Refused past the
    /// session's limits.
    pub async fn subscribe(&self, subscription: Subscription) -> Result<(), FeedError> {
        self.with(|driver| {
            if let Some(frame) = driver.session.subscribe(subscription)? {
                driver.carry(&frame);
            }
            Ok(())
        })
    }

    pub fn unsubscribe(&self, subscription: Subscription) -> Result<(), FeedError> {
        self.with(|driver| {
            if let Some(frame) = driver.session.unsubscribe(&subscription) {
                driver.carry(&frame);
            }
            Ok(())
        })
    }

    /// Sends an info request or a signed action over the socket, and gives
    /// its id; the answer arrives as [`Frame::Post`] with the same id.
    pub async fn post(&self, request: Post) -> Result<u64, FeedError> {
        self.with(|driver| {
            let (id, frame) = driver.session.post(&request)?;
            driver.carry(&frame);
            Ok(id)
        })
    }

    /// The next event; `None` once the feed has stopped.
    pub async fn next(&mut self) -> Option<Event> {
        poll_fn(|cx| {
            let Ok(mut driver) = self.driver.try_borrow_mut() else {
                // Never: the driver is only held inside a call, never across one.
                cx.waker().wake_by_ref();
                return Poll::Pending;
            };
            if let Some(event) = driver.events.pop_front() {
                return Poll::Ready(Some(event));
            }
            if driver.stopped {
                return Poll::Ready(None);
            }
            driver.waker = Some(cx.waker().clone());
            Poll::Pending
        })
        .await
    }

    /// Closes the socket and stops.
    pub fn close(self) {
        drop(self);
    }

    /// Runs `act` on a feed that has not stopped. An event it leaves needs no
    /// wake: [`Feed::next`] takes the feed mutably, so nothing waits on it.
    fn with<T>(&self, act: impl FnOnce(&mut Driver) -> Result<T, FeedError>) -> Result<T, FeedError> {
        let mut driver = self.driver.try_borrow_mut().map_err(|_| FeedError::Stopped)?;
        if driver.stopped {
            return Err(FeedError::Stopped);
        }
        act(&mut driver)
    }
}

impl Drop for Feed {
    fn drop(&mut self) {
        if let Ok(mut driver) = self.driver.try_borrow_mut() {
            driver.stop();
        }
    }
}

/// The feed's state. The socket's handlers and the timer hold it weakly, so
/// the feed alone keeps it, and with it every closure JS can call.
struct Driver {
    url: String,
    session: Session,
    global: JsValue,
    /// The global `setTimeout` and `clearTimeout`.
    timers: Option<(Function, Function)>,
    /// The socket dialling or open; none while waiting to dial again.
    socket: Option<WebSocket>,
    /// The pending timeout: the next tick while open, the next dial while
    /// closed, none while dialling. A browser gives a number, Node an object.
    timer: Option<JsValue>,
    calls: Calls,
    events: VecDeque<Event>,
    waker: Option<Waker>,
    stopped: bool,
}

impl Driver {
    /// A new socket. While it dials, the session still answers: a subscription
    /// is kept for the open, a post is refused. Nothing waits on a handshake.
    fn dial(&mut self) {
        match connect(&self.global, &self.url) {
            Ok(socket) => {
                socket.set_onopen(Some(self.calls.open.as_ref().unchecked_ref()));
                socket.set_onmessage(Some(self.calls.message.as_ref().unchecked_ref()));
                socket.set_onclose(Some(self.calls.close.as_ref().unchecked_ref()));
                socket.set_onerror(Some(self.calls.error.as_ref().unchecked_ref()));
                self.socket = Some(socket);
            }
            Err(reason) => self.lost(reason),
        }
    }

    fn opened(&mut self) {
        for frame in self.session.opened(clock()) {
            if let Err(reason) = self.send(&frame) {
                return self.lost(reason);
            }
        }
        self.events.push_back(Event::Opened);
        self.arm_tick();
    }

    fn received(&mut self, message: &MessageEvent) {
        // The venue sends text only; anything else is not its.
        let Some(text) = message.data().as_string() else { return };
        let event = match self.session.received(clock(), &text) {
            Ok(frame) => Event::Frame(Box::new(frame)),
            Err(bad) => Event::Unreadable(bad),
        };
        self.events.push_back(event);
    }

    /// The timer fired: a tick while open, the next dial while closed.
    fn woke(&mut self) {
        self.timer = None;
        if self.session.is_open() {
            match self.session.tick(clock()) {
                Tick::Wait => {}
                Tick::Ping => {
                    if let Err(reason) = self.send(PING) {
                        return self.lost(reason);
                    }
                }
                Tick::Reconnect => return self.lost("nothing arrived for too long".to_owned()),
            }
            self.arm_tick();
        } else if self.socket.is_none() {
            self.dial();
        }
    }

    /// Sends `frame`, or drops a socket that will not take it.
    fn carry(&mut self, frame: &str) {
        if let Err(reason) = self.send(frame) {
            self.lost(reason);
        }
    }

    fn send(&self, frame: &str) -> Result<(), String> {
        let Some(socket) = &self.socket else { return Err("no socket to send on".to_owned()) };
        socket.send_with_str(frame).map_err(|error| said(&error).unwrap_or_else(|| "the send failed".to_owned()))
    }

    /// The socket is gone, or never came: report it with the session's delay,
    /// and dial again after it.
    fn lost(&mut self, reason: String) {
        self.hang_up();
        let retry = self.session.closed(clock(), random());
        let after = retry.after_ms;
        self.events.push_back(Event::Closed { reason, retry });
        self.arm(after);
    }

    /// Closes the socket after taking its handlers off, so nothing it still
    /// has queued reaches the feed. A half-open socket's close may never
    /// finish; nothing waits for it.
    fn hang_up(&mut self) {
        if let Some(socket) = self.socket.take() {
            socket.set_onopen(None);
            socket.set_onmessage(None);
            socket.set_onclose(None);
            socket.set_onerror(None);
            let _ = socket.close();
        }
    }

    fn stop(&mut self) {
        self.hang_up();
        self.disarm();
        self.stopped = true;
    }

    /// The timer for the session's next tick, while open.
    fn arm_tick(&mut self) {
        if let Some(at) = self.session.due() {
            self.arm(at.saturating_sub(clock()));
        }
    }

    /// One timeout for `ms`, in place of any pending. Without a timer the feed
    /// can neither tell a dead socket nor dial again, so it stops.
    fn arm(&mut self, ms: u64) {
        self.disarm();
        let Some((set_timeout, _)) = &self.timers else { return self.stop() };
        let delay = JsValue::from(u32::try_from(ms).map_or(LONGEST_DELAY, |ms| ms.min(LONGEST_DELAY)));
        match set_timeout.call2(&self.global, self.calls.wake.as_ref(), &delay) {
            Ok(timer) => self.timer = Some(timer),
            Err(_) => self.stop(),
        }
    }

    fn disarm(&mut self) {
        if let (Some(timer), Some((_, clear_timeout))) = (self.timer.take(), &self.timers) {
            let _ = clear_timeout.call1(&self.global, &timer);
        }
    }

    /// The waker to call once the driver is let go, when `next` has something.
    fn wakeable(&mut self) -> Option<Waker> {
        if self.events.is_empty() && !self.stopped { None } else { self.waker.take() }
    }
}

/// The longest delay `setTimeout` keeps; a longer one fires at once.
const LONGEST_DELAY: u32 = 0x7fff_ffff;

/// The closures JS calls, made once a feed and dropped with it. Each holds the
/// driver weakly and is handed to every socket the feed dials in turn.
struct Calls {
    open: Closure<dyn FnMut()>,
    message: Closure<dyn FnMut(MessageEvent)>,
    close: Closure<dyn FnMut(CloseEvent)>,
    error: Closure<dyn FnMut(JsValue)>,
    wake: Closure<dyn FnMut()>,
}

impl Calls {
    fn new(driver: &Weak<RefCell<Driver>>) -> Calls {
        let (open, message, close, error, wake) =
            (driver.clone(), driver.clone(), driver.clone(), driver.clone(), driver.clone());
        Calls {
            open: Closure::new(move || on(&open, Driver::opened)),
            message: Closure::new(move |event: MessageEvent| on(&message, |driver| driver.received(&event))),
            close: Closure::new(move |event: CloseEvent| {
                let reason = match event.reason() {
                    reason if reason.is_empty() => format!("the socket closed ({})", event.code()),
                    reason => reason,
                };
                on(&close, |driver| driver.lost(reason));
            }),
            // An error always ends the socket; the close that follows is not
            // waited for.
            error: Closure::new(move |event: JsValue| {
                let reason = said(&event).unwrap_or_else(|| "the socket failed".to_owned());
                on(&error, |driver| driver.lost(reason));
            }),
            wake: Closure::new(move || on(&wake, Driver::woke)),
        }
    }
}

/// Runs `act` on the driver for JS, then wakes `next` if it has something. The
/// waker is called once the driver is let go, so one that polls at once finds
/// it free.
fn on(driver: &Weak<RefCell<Driver>>, act: impl FnOnce(&mut Driver)) {
    let Some(driver) = driver.upgrade() else { return };
    let waker = {
        let Ok(mut driver) = driver.try_borrow_mut() else { return };
        act(&mut driver);
        driver.wakeable()
    };
    if let Some(waker) = waker {
        waker.wake();
    }
}

/// A new socket to `url`, from the global `WebSocket`.
fn connect(global: &JsValue, url: &str) -> Result<WebSocket, String> {
    let class = function(global, "WebSocket").ok_or_else(|| "this runtime has no WebSocket".to_owned())?;
    let socket = Reflect::construct(&class, &Array::of1(&JsValue::from_str(url)))
        .map_err(|error| said(&error).unwrap_or_else(|| "the socket could not be made".to_owned()))?;
    Ok(socket.unchecked_into())
}

/// The global function `name`, if the runtime has it.
fn function(global: &JsValue, name: &str) -> Option<Function> {
    Reflect::get(global, &JsValue::from_str(name)).ok()?.dyn_into().ok()
}

/// The `message` an error or an error event carries, if any. A browser's
/// socket error carries none.
fn said(error: &JsValue) -> Option<String> {
    let message = Reflect::get(error, &JsValue::from_str("message")).ok()?.as_string()?;
    (!message.is_empty()).then_some(message)
}

/// Enough to spread reconnects: many clients dropped at once must not return
/// at once. Not for anything secret. `Math.random()`'s bits, mixed so the low
/// ones the session's jitter reads are not a small float's trailing zeros.
fn random() -> u64 {
    let bits = Math::random().to_bits().wrapping_mul(0x9e37_79b9_7f4a_7c15);
    bits ^ (bits >> 29) ^ clock().rotate_left(29)
}

#[cfg(test)]
#[cfg(target_arch = "wasm32")]
#[path = "tests/socket.rs"]
mod tests;