nautilus_derive/websocket/
handler.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Inner I/O feed handler for the Derive WebSocket transport.
17//!
18//! The handler owns the [`WebSocketClient`] exclusively and runs in a dedicated
19//! Tokio task. The outer [`super::client::DeriveWebSocketClient`] talks to it
20//! via a command channel and consumes a stream of [`DeriveWsMessage`] events.
21//!
22//! Each outbound JSON-RPC request is registered in a `pending` map keyed by the
23//! correlator `id`. When the venue echoes the id on a response frame, the
24//! matching oneshot is fulfilled with `result` or the JSON-RPC error.
25
26use std::sync::{
27    Arc,
28    atomic::{AtomicBool, AtomicU64, Ordering},
29};
30
31use ahash::AHashMap;
32use nautilus_network::{
33    RECONNECTED,
34    websocket::{AuthTracker, WebSocketClient},
35};
36use serde_json::Value;
37use tokio_tungstenite::tungstenite::Message;
38
39use super::{
40    error::DeriveWsError,
41    messages::{DeriveWsChannel, DeriveWsFrame, WsSubscribeParams, WsSubscriptionPayload},
42};
43use crate::http::models::JsonRpcRequest;
44
45/// Outbound commands the outer client sends to the inner handler.
46#[derive(Debug)]
47pub(super) enum HandlerCommand {
48    /// Hand the active [`WebSocketClient`] to the handler.
49    SetClient(WebSocketClient),
50    /// Send a JSON-RPC request and resolve the oneshot when the venue replies.
51    /// `params` is a pre-serialized `Value` so the handler stays agnostic to the
52    /// per-method param types (login, subscribe, signed `private/*` bodies).
53    Request {
54        method: &'static str,
55        params: Value,
56        response_tx: tokio::sync::oneshot::Sender<Result<Value, DeriveWsError>>,
57    },
58    /// Gracefully tear down the WebSocket connection.
59    Disconnect,
60}
61
62/// Events emitted by the handler for the outer client and downstream consumers.
63#[derive(Debug, Clone)]
64pub enum DeriveWsMessage {
65    /// `public/login` succeeded. Consumed by the client's spawn loop to drive
66    /// resubscription; not forwarded to data/execution layers.
67    Authenticated,
68    /// Underlying transport reconnected; outer client triggers re-login and
69    /// resubscribes the tracked channels.
70    Reconnected,
71    /// Channel update pushed by the venue.
72    Subscription(WsSubscriptionPayload),
73}
74
75/// Inner I/O loop. Lives in a Tokio task spawned by
76/// [`super::client::DeriveWebSocketClient::connect`].
77pub(super) struct FeedHandler {
78    signal: Arc<AtomicBool>,
79    client: Option<WebSocketClient>,
80    cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
81    raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
82    next_id: Arc<AtomicU64>,
83    pending: AHashMap<u64, tokio::sync::oneshot::Sender<Result<Value, DeriveWsError>>>,
84    auth_tracker: AuthTracker,
85}
86
87impl FeedHandler {
88    pub(super) fn new(
89        signal: Arc<AtomicBool>,
90        cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
91        raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
92        next_id: Arc<AtomicU64>,
93        auth_tracker: AuthTracker,
94    ) -> Self {
95        Self {
96            signal,
97            client: None,
98            cmd_rx,
99            raw_rx,
100            next_id,
101            pending: AHashMap::new(),
102            auth_tracker,
103        }
104    }
105
106    /// Drains the next event from the underlying channels, processes it, and
107    /// returns the resulting outbound message (if any). Returns `None` when
108    /// the handler is shutting down or both channels closed.
109    pub(super) async fn next(&mut self) -> Option<DeriveWsMessage> {
110        loop {
111            tokio::select! {
112                Some(cmd) = self.cmd_rx.recv() => {
113                    match cmd {
114                        HandlerCommand::SetClient(client) => {
115                            log::debug!("Setting WebSocket client in Derive handler");
116                            self.client = Some(client);
117                        }
118                        HandlerCommand::Request { method, params, response_tx } => {
119                            self.dispatch_request(method, params, response_tx).await;
120                        }
121                        HandlerCommand::Disconnect => {
122                            log::debug!("Derive handler received disconnect command");
123                            if let Some(ref client) = self.client {
124                                client.disconnect().await;
125                            }
126                            self.signal.store(true, Ordering::SeqCst);
127                            return None;
128                        }
129                    }
130                }
131
132                Some(raw) = self.raw_rx.recv() => {
133                    match raw {
134                        Message::Text(text) => {
135                            if text.as_str() == RECONNECTED {
136                                log::info!("Derive WebSocket reconnected sentinel received");
137                                self.auth_tracker.invalidate();
138                                self.fail_pending("WebSocket reconnected before response was received");
139                                return Some(DeriveWsMessage::Reconnected);
140                            }
141
142                            match DeriveWsFrame::parse(&text) {
143                                Ok(DeriveWsFrame::Response { id, result, error }) => {
144                                    if let Some(sender) = self.pending.remove(&id) {
145                                        let outcome = match (result, error) {
146                                            (_, Some(err)) => Err(DeriveWsError::JsonRpc {
147                                                code: err.code,
148                                                message: err.message,
149                                                data: err.data,
150                                            }),
151                                            (Some(value), None) => Ok(value),
152                                            (None, None) => Ok(Value::Null),
153                                        };
154                                        let _ = sender.send(outcome);
155                                    } else {
156                                        log::debug!(
157                                            "Derive WebSocket response with unknown id={id} dropped",
158                                        );
159                                    }
160                                }
161                                Ok(DeriveWsFrame::Subscription(payload)) => {
162                                    return Some(DeriveWsMessage::Subscription(payload));
163                                }
164                                Ok(DeriveWsFrame::Unknown(value)) => {
165                                    log::debug!("Derive WebSocket unknown frame: {value}");
166                                }
167                                Err(e) => {
168                                    log::error!(
169                                        "Derive WebSocket frame parse error: {e}, text: {text}",
170                                    );
171                                }
172                            }
173                        }
174                        Message::Ping(data) => {
175                            if let Some(ref client) = self.client
176                                && let Err(e) = client.send_pong(data.to_vec()).await {
177                                log::error!("Derive WebSocket send_pong failed: {e}");
178                            }
179                        }
180                        Message::Close(_) => {
181                            log::debug!("Derive WebSocket close frame received");
182                            return None;
183                        }
184                        _ => {}
185                    }
186                }
187
188                else => {
189                    log::debug!("Derive handler shutting down: channels closed");
190                    return None;
191                }
192            }
193        }
194    }
195
196    async fn dispatch_request(
197        &mut self,
198        method: &'static str,
199        params: Value,
200        response_tx: tokio::sync::oneshot::Sender<Result<Value, DeriveWsError>>,
201    ) {
202        let Some(ref client) = self.client else {
203            let _ = response_tx.send(Err(DeriveWsError::NotConnected));
204            return;
205        };
206        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
207        let request = JsonRpcRequest::new(id, method, params);
208        let payload = match serde_json::to_string(&request) {
209            Ok(p) => p,
210            Err(e) => {
211                let _ = response_tx.send(Err(DeriveWsError::Serde(e)));
212                return;
213            }
214        };
215        self.pending.insert(id, response_tx);
216        log::debug!("Derive WebSocket sending `{method}` id={id}");
217        if let Err(e) = client.send_text(payload, None).await
218            && let Some(sender) = self.pending.remove(&id)
219        {
220            let _ = sender.send(Err(DeriveWsError::transport(e.to_string())));
221        }
222    }
223
224    fn fail_pending(&mut self, reason: &str) {
225        if self.pending.is_empty() {
226            return;
227        }
228        log::debug!(
229            "Failing {} pending Derive WebSocket request(s): {reason}",
230            self.pending.len(),
231        );
232
233        for (_, sender) in self.pending.drain() {
234            let _ = sender.send(Err(DeriveWsError::transport(reason.to_string())));
235        }
236    }
237}
238
239/// Builds `subscribe` params from a single channel topic.
240#[must_use]
241pub(super) fn subscribe_params(channel: DeriveWsChannel) -> WsSubscribeParams {
242    WsSubscribeParams {
243        channels: vec![channel],
244    }
245}
246
247/// Convenience wrapper that produces the `subscribe` params for the
248/// `ticker_slim.{instrument_name}.{interval}` channel.
249#[must_use]
250pub(super) fn ticker_subscribe_params(instrument_name: &str, interval: &str) -> WsSubscribeParams {
251    subscribe_params(DeriveWsChannel::ticker_slim(instrument_name, interval))
252}
253
254/// Convenience wrapper that produces the `subscribe` params for the
255/// `orderbook.{instrument_name}.{group}.{depth}` channel.
256#[must_use]
257pub(super) fn orderbook_subscribe_params(
258    instrument_name: &str,
259    group: &str,
260    depth: &str,
261) -> WsSubscribeParams {
262    subscribe_params(DeriveWsChannel::orderbook(instrument_name, group, depth))
263}
264
265/// Convenience wrapper that produces the `subscribe` params for the
266/// `trades.{instrument_type}.{currency}` channel.
267#[must_use]
268pub(super) fn trades_subscribe_params(instrument_type: &str, currency: &str) -> WsSubscribeParams {
269    subscribe_params(DeriveWsChannel::trades(instrument_type, currency))
270}
271
272#[cfg(test)]
273mod tests {
274    use rstest::rstest;
275
276    use super::*;
277
278    #[rstest]
279    fn test_subscribe_params_carries_single_channel() {
280        let params = subscribe_params(DeriveWsChannel::ticker_slim("ETH-PERP", "1000"));
281        assert_eq!(
282            params.channels,
283            vec![DeriveWsChannel::ticker_slim("ETH-PERP", "1000")],
284        );
285    }
286
287    #[rstest]
288    fn test_ticker_subscribe_params_formats_topic() {
289        let params = ticker_subscribe_params("ETH-PERP", "1000");
290        assert_eq!(
291            params.channels,
292            vec![DeriveWsChannel::ticker_slim("ETH-PERP", "1000")],
293        );
294    }
295
296    #[rstest]
297    fn test_orderbook_subscribe_params_formats_topic() {
298        let params = orderbook_subscribe_params("ETH-PERP", "1", "10");
299        assert_eq!(
300            params.channels,
301            vec![DeriveWsChannel::orderbook("ETH-PERP", "1", "10")],
302        );
303    }
304
305    #[rstest]
306    fn test_trades_subscribe_params_formats_topic() {
307        let params = trades_subscribe_params("perp", "ETH");
308        assert_eq!(
309            params.channels,
310            vec![DeriveWsChannel::trades("perp", "ETH")],
311        );
312    }
313
314    #[rstest]
315    #[tokio::test]
316    async fn test_dispatch_request_without_client_returns_not_connected() {
317        // Requests issued before SetClient must fail fast rather than hang.
318        let signal = Arc::new(AtomicBool::new(false));
319        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
320        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
321        let next_id = Arc::new(AtomicU64::new(1));
322        let auth_tracker = AuthTracker::new();
323        let mut handler = FeedHandler::new(signal, cmd_rx, raw_rx, next_id, auth_tracker);
324
325        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
326        let params = serde_json::to_value(WsSubscribeParams { channels: vec![] }).unwrap();
327        handler
328            .dispatch_request("public/login", params, response_tx)
329            .await;
330
331        let outcome = response_rx.await.expect("oneshot resolved");
332        match outcome {
333            Err(DeriveWsError::NotConnected) => {}
334            other => panic!("expected NotConnected, was {other:?}"),
335        }
336    }
337}
nautilus_derive/websocket/handler.rs

nautilus_derive/websocket/
handler.rs
