nautilus_bitmex/websocket/
client.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//! Provides the WebSocket client integration for the [BitMEX](https://bitmex.com) WebSocket API.
17//!
18//! This module defines and implements a [`BitmexWebSocketClient`] for
19//! connecting to BitMEX WebSocket streams. It handles authentication (when credentials
20//! are provided), manages subscriptions to market data and account update channels,
21//! and emits venue-specific message types for consumers to parse.
22
23use std::{
24    sync::{
25        Arc,
26        atomic::{AtomicBool, AtomicU8, Ordering},
27    },
28    time::Duration,
29};
30
31use arc_swap::ArcSwap;
32use dashmap::DashMap;
33use futures_util::Stream;
34use nautilus_common::live::get_runtime;
35use nautilus_core::{
36    consts::NAUTILUS_USER_AGENT,
37    env::{get_env_var, get_or_env_var_opt},
38};
39use nautilus_model::{
40    data::bar::BarType,
41    identifiers::{AccountId, InstrumentId},
42    instruments::{Instrument, InstrumentAny},
43};
44use nautilus_network::{
45    http::USER_AGENT,
46    mode::ConnectionMode,
47    websocket::{
48        AUTHENTICATION_TIMEOUT_SECS, AuthTracker, PingHandler, SubscriptionState, TransportBackend,
49        WebSocketClient, WebSocketConfig, channel_message_handler,
50    },
51};
52use tokio_tungstenite::tungstenite::Message;
53use ustr::Ustr;
54
55use super::{
56    enums::{BitmexWsAuthAction, BitmexWsAuthChannel, BitmexWsOperation, BitmexWsTopic},
57    error::BitmexWsError,
58    handler::{BitmexWsFeedHandler, HandlerCommand},
59    messages::{BitmexAuthentication, BitmexSubscription, BitmexWsMessage},
60    parse::{is_index_symbol, topic_from_bar_spec},
61};
62use crate::common::{
63    consts::{BITMEX_WS_TOPIC_DELIMITER, BITMEX_WS_URL},
64    credential::{Credential, credential_env_vars},
65    enums::BitmexEnvironment,
66};
67
68/// Provides a WebSocket client for connecting to the [BitMEX](https://bitmex.com) real-time API.
69///
70/// Key runtime patterns:
71/// - Authentication handshakes are managed by the internal auth tracker, ensuring resubscriptions
72///   occur only after BitMEX acknowledges `authKey` messages.
73/// - The subscription state maintains pending and confirmed topics so reconnection replay is
74///   deterministic and per-topic errors are surfaced.
75#[derive(Clone, Debug)]
76pub struct BitmexWebSocketClient {
77    url: String,
78    credential: Option<Credential>,
79    heartbeat: Option<u64>,
80    account_id: AccountId,
81    auth_tracker: AuthTracker,
82    signal: Arc<AtomicBool>,
83    connection_mode: Arc<ArcSwap<AtomicU8>>,
84    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
85    out_rx: Option<Arc<tokio::sync::mpsc::UnboundedReceiver<BitmexWsMessage>>>,
86    task_handle: Option<Arc<tokio::task::JoinHandle<()>>>,
87    subscriptions: SubscriptionState,
88    tracked_subscriptions: Arc<DashMap<String, ()>>,
89    instruments: Arc<DashMap<Ustr, InstrumentAny>>,
90    transport_backend: TransportBackend,
91    proxy_url: Option<String>,
92}
93
94impl BitmexWebSocketClient {
95    /// Creates a new [`BitmexWebSocketClient`] instance.
96    ///
97    /// # Errors
98    ///
99    /// Returns an error if only one of `api_key` or `api_secret` is provided (both or neither required).
100    pub fn new(
101        url: Option<String>,
102        api_key: Option<String>,
103        api_secret: Option<String>,
104        account_id: Option<AccountId>,
105        heartbeat: u64,
106        transport_backend: TransportBackend,
107        proxy_url: Option<String>,
108    ) -> anyhow::Result<Self> {
109        let credential = match (api_key, api_secret) {
110            (Some(key), Some(secret)) => Some(Credential::new(key, secret)),
111            (None, None) => None,
112            _ => anyhow::bail!("Both `api_key` and `api_secret` must be provided together"),
113        };
114
115        let account_id = account_id.unwrap_or(AccountId::from("BITMEX-master"));
116
117        let initial_mode = AtomicU8::new(ConnectionMode::Closed.as_u8());
118        let connection_mode = Arc::new(ArcSwap::from_pointee(initial_mode));
119
120        // Placeholder channel until connect() creates the real one
121        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
122
123        Ok(Self {
124            url: url.unwrap_or(BITMEX_WS_URL.to_string()),
125            credential,
126            heartbeat: Some(heartbeat),
127            account_id,
128            auth_tracker: AuthTracker::new(),
129            signal: Arc::new(AtomicBool::new(false)),
130            connection_mode,
131            cmd_tx: Arc::new(tokio::sync::RwLock::new(cmd_tx)),
132            out_rx: None,
133            task_handle: None,
134            subscriptions: SubscriptionState::new(BITMEX_WS_TOPIC_DELIMITER),
135            tracked_subscriptions: Arc::new(DashMap::new()),
136            instruments: Arc::new(DashMap::new()),
137            transport_backend,
138            proxy_url,
139        })
140    }
141
142    /// Creates a new [`BitmexWebSocketClient`] with environment variable credential resolution.
143    ///
144    /// If `api_key` or `api_secret` are not provided, they will be loaded from
145    /// environment variables based on the `environment`:
146    /// - Testnet: `BITMEX_TESTNET_API_KEY`, `BITMEX_TESTNET_API_SECRET`
147    /// - Mainnet: `BITMEX_API_KEY`, `BITMEX_API_SECRET`
148    ///
149    /// # Errors
150    ///
151    /// Returns an error if only one of `api_key` or `api_secret` is provided.
152    #[expect(clippy::too_many_arguments)]
153    pub fn new_with_env(
154        url: Option<String>,
155        api_key: Option<String>,
156        api_secret: Option<String>,
157        account_id: Option<AccountId>,
158        heartbeat: u64,
159        environment: BitmexEnvironment,
160        transport_backend: TransportBackend,
161        proxy_url: Option<String>,
162    ) -> anyhow::Result<Self> {
163        let (api_key_env, api_secret_env) = credential_env_vars(environment);
164
165        let key = get_or_env_var_opt(api_key, api_key_env);
166        let secret = get_or_env_var_opt(api_secret, api_secret_env);
167
168        Self::new(
169            url,
170            key,
171            secret,
172            account_id,
173            heartbeat,
174            transport_backend,
175            proxy_url,
176        )
177    }
178
179    /// Creates a new authenticated [`BitmexWebSocketClient`] using environment variables.
180    ///
181    /// # Errors
182    ///
183    /// Returns an error if environment variables are not set or credentials are invalid.
184    pub fn from_env() -> anyhow::Result<Self> {
185        let url = get_env_var("BITMEX_WS_URL")?;
186        let (key_var, secret_var) = credential_env_vars(BitmexEnvironment::Mainnet);
187        let api_key = get_env_var(key_var)?;
188        let api_secret = get_env_var(secret_var)?;
189
190        Self::new(
191            Some(url),
192            Some(api_key),
193            Some(api_secret),
194            None,
195            5,
196            TransportBackend::default(),
197            None,
198        )
199    }
200
201    /// Returns the websocket url being used by the client.
202    #[must_use]
203    pub const fn url(&self) -> &str {
204        self.url.as_str()
205    }
206
207    /// Returns the public API key being used by the client.
208    #[must_use]
209    pub fn api_key(&self) -> Option<&str> {
210        self.credential.as_ref().map(|c| c.api_key())
211    }
212
213    /// Returns a masked version of the API key for logging purposes.
214    #[must_use]
215    pub fn api_key_masked(&self) -> Option<String> {
216        self.credential.as_ref().map(|c| c.api_key_masked())
217    }
218
219    /// Returns a value indicating whether the client is active.
220    #[must_use]
221    pub fn is_active(&self) -> bool {
222        let connection_mode_arc = self.connection_mode.load();
223        ConnectionMode::from_atomic(&connection_mode_arc).is_active()
224            && !self.signal.load(Ordering::Relaxed)
225    }
226
227    /// Returns a value indicating whether the client is closed.
228    #[must_use]
229    pub fn is_closed(&self) -> bool {
230        let connection_mode_arc = self.connection_mode.load();
231        ConnectionMode::from_atomic(&connection_mode_arc).is_closed()
232            || self.signal.load(Ordering::Relaxed)
233    }
234
235    /// Returns the account ID.
236    #[must_use]
237    pub fn account_id(&self) -> AccountId {
238        self.account_id
239    }
240
241    /// Sets the account ID.
242    pub fn set_account_id(&mut self, account_id: AccountId) {
243        self.account_id = account_id;
244    }
245
246    /// Bulk-replaces the instrument cache with the given instruments.
247    pub fn cache_instruments(&self, instruments: &[InstrumentAny]) {
248        self.instruments.clear();
249        for inst in instruments {
250            self.instruments
251                .insert(inst.raw_symbol().inner(), inst.clone());
252        }
253    }
254
255    /// Upserts a single instrument into the cache.
256    pub fn cache_instrument(&self, instrument: InstrumentAny) {
257        self.instruments
258            .insert(instrument.raw_symbol().inner(), instrument);
259    }
260
261    /// Retrieves an instrument from the cache by symbol.
262    #[must_use]
263    pub fn get_instrument(&self, symbol: &Ustr) -> Option<InstrumentAny> {
264        self.instruments
265            .get(symbol)
266            .map(|entry| entry.value().clone())
267    }
268
269    /// Connect to the BitMEX WebSocket server.
270    ///
271    /// # Errors
272    ///
273    /// Returns an error if the WebSocket connection fails or authentication fails (if credentials provided).
274    ///
275    pub async fn connect(&mut self) -> Result<(), BitmexWsError> {
276        let (client, raw_rx) = self.connect_inner().await?;
277
278        // Reset shutdown signal so is_active() works after close+reconnect
279        self.signal.store(false, Ordering::Relaxed);
280
281        // Replace connection state so all clones see the underlying WebSocketClient's state
282        self.connection_mode.store(client.connection_mode_atomic());
283
284        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<BitmexWsMessage>();
285        self.out_rx = Some(Arc::new(out_rx));
286
287        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
288        *self.cmd_tx.write().await = cmd_tx.clone();
289
290        // Send WebSocketClient to handler
291        if let Err(e) = cmd_tx.send(HandlerCommand::SetClient(client)) {
292            return Err(BitmexWsError::ClientError(format!(
293                "Failed to send WebSocketClient to handler: {e}"
294            )));
295        }
296
297        let signal = self.signal.clone();
298        let credential = self.credential.clone();
299        let auth_tracker = self.auth_tracker.clone();
300        let subscriptions = self.subscriptions.clone();
301        let cmd_tx_for_reconnect = cmd_tx.clone();
302
303        let stream_handle = get_runtime().spawn(async move {
304            let mut handler = BitmexWsFeedHandler::new(
305                signal.clone(),
306                cmd_rx,
307                raw_rx,
308                out_tx,
309                auth_tracker.clone(),
310                subscriptions.clone(),
311            );
312
313            // Resubscribe all tracked subscriptions after reconnection.
314            let resubscribe_all = || {
315                // Use SubscriptionState as source of truth for what to restore
316                let topics = subscriptions.all_topics();
317
318                if topics.is_empty() {
319                    return;
320                }
321
322                log::debug!(
323                    "Resubscribing to confirmed subscriptions: count={}",
324                    topics.len()
325                );
326
327                for topic in &topics {
328                    subscriptions.mark_subscribe(topic.as_str());
329                }
330
331                // Serialize subscription messages
332                let mut payloads = Vec::with_capacity(topics.len());
333                for topic in &topics {
334                    let message = BitmexSubscription {
335                        op: BitmexWsOperation::Subscribe,
336                        args: vec![Ustr::from(topic.as_ref())],
337                    };
338
339                    if let Ok(payload) = serde_json::to_string(&message) {
340                        payloads.push(payload);
341                    }
342                }
343
344                if let Err(e) =
345                    cmd_tx_for_reconnect.send(HandlerCommand::Subscribe { topics: payloads })
346                {
347                    log::error!("Failed to send resubscribe command: {e}");
348                }
349            };
350
351            let mut waiting_for_reconnect_auth = false;
352
353            // Run message processing with reconnection handling
354            loop {
355                match handler.next().await {
356                    Some(BitmexWsMessage::Reconnected) => {
357                        if signal.load(Ordering::Relaxed) {
358                            continue;
359                        }
360
361                        log::info!("WebSocket reconnected");
362
363                        // Mark all confirmed subscriptions as failed so they transition to pending state
364                        let confirmed_topics: Vec<String> = {
365                            let confirmed = subscriptions.confirmed();
366                            let mut topics = Vec::new();
367
368                            for entry in confirmed.iter() {
369                                let (channel, symbols) = entry.pair();
370
371                                if *channel == BitmexWsTopic::Instrument.as_ref() {
372                                    continue;
373                                }
374
375                                for symbol in symbols {
376                                    if symbol.is_empty() {
377                                        topics.push(channel.to_string());
378                                    } else {
379                                        topics.push(format!("{channel}:{symbol}"));
380                                    }
381                                }
382                            }
383
384                            topics
385                        };
386
387                        if !confirmed_topics.is_empty() {
388                            log::debug!(
389                                "Marking confirmed subscriptions as pending for replay: count={}",
390                                confirmed_topics.len()
391                            );
392
393                            for topic in confirmed_topics {
394                                subscriptions.mark_failure(&topic);
395                            }
396                        }
397
398                        if let Some(cred) = &credential {
399                            log::debug!("Re-authenticating after reconnection");
400                            waiting_for_reconnect_auth = true;
401
402                            let expires =
403                                (chrono::Utc::now() + chrono::Duration::seconds(30)).timestamp();
404                            let signature = cred.sign("GET", "/realtime", expires, "");
405
406                            let auth_message = BitmexAuthentication {
407                                op: BitmexWsAuthAction::AuthKeyExpires,
408                                args: (cred.api_key().to_string(), expires, signature),
409                            };
410
411                            if let Ok(payload) = serde_json::to_string(&auth_message) {
412                                if let Err(e) = cmd_tx_for_reconnect
413                                    .send(HandlerCommand::Authenticate { payload })
414                                {
415                                    log::error!("Failed to send reconnection auth command: {e}");
416                                }
417                            } else {
418                                log::error!("Failed to serialize reconnection auth message");
419                            }
420                        }
421
422                        // Unauthenticated sessions resubscribe immediately after reconnection,
423                        // authenticated sessions wait for Authenticated message
424                        if credential.is_none() {
425                            log::debug!("No authentication required, resubscribing immediately");
426                            resubscribe_all();
427                        }
428
429                        if handler.send(BitmexWsMessage::Reconnected).is_err() {
430                            if handler.is_stopped() {
431                                log::debug!("Failed to forward reconnect event (receiver dropped)");
432                            } else {
433                                log::error!("Failed to forward reconnect event (receiver dropped)");
434                            }
435                            break;
436                        }
437                    }
438                    Some(BitmexWsMessage::Authenticated) => {
439                        if waiting_for_reconnect_auth {
440                            log::debug!("Authenticated after reconnection, resubscribing");
441                            resubscribe_all();
442                            waiting_for_reconnect_auth = false;
443                        }
444                    }
445                    Some(msg) => {
446                        if handler.send(msg).is_err() {
447                            if handler.is_stopped() {
448                                log::debug!("Failed to send message (receiver dropped)");
449                            } else {
450                                log::error!("Failed to send message (receiver dropped)");
451                            }
452                            break;
453                        }
454                    }
455                    None => {
456                        // Stream ended - check if it's a stop signal
457                        if handler.is_stopped() {
458                            log::debug!("Stop signal received, ending message processing");
459                            break;
460                        }
461                        // Otherwise it's an unexpected stream end
462                        log::warn!("WebSocket stream ended unexpectedly");
463                        break;
464                    }
465                }
466            }
467
468            log::debug!("Handler task exiting");
469        });
470
471        self.task_handle = Some(Arc::new(stream_handle));
472
473        if self.credential.is_some()
474            && let Err(e) = self.authenticate().await
475        {
476            if let Some(handle) = self.task_handle.take() {
477                handle.abort();
478            }
479            self.signal.store(true, Ordering::Relaxed);
480            return Err(e);
481        }
482
483        // Subscribe to instrument topic
484        let instrument_topic = BitmexWsTopic::Instrument.as_ref().to_string();
485        self.subscriptions.mark_subscribe(&instrument_topic);
486        self.tracked_subscriptions.insert(instrument_topic, ());
487
488        let subscribe_msg = BitmexSubscription {
489            op: BitmexWsOperation::Subscribe,
490            args: vec![Ustr::from(BitmexWsTopic::Instrument.as_ref())],
491        };
492
493        match serde_json::to_string(&subscribe_msg) {
494            Ok(subscribe_json) => {
495                if let Err(e) = self.cmd_tx.read().await.send(HandlerCommand::Subscribe {
496                    topics: vec![subscribe_json],
497                }) {
498                    log::error!("Failed to send subscribe command for instruments: {e}");
499                } else {
500                    log::debug!("Subscribed to all instruments");
501                }
502            }
503            Err(e) => {
504                log::error!("Failed to serialize subscribe message: {e}");
505            }
506        }
507
508        Ok(())
509    }
510
511    /// Connect to the WebSocket and return a message receiver.
512    ///
513    /// # Errors
514    ///
515    /// Returns an error if the WebSocket connection fails or if authentication fails (when credentials are provided).
516    async fn connect_inner(
517        &self,
518    ) -> Result<
519        (
520            WebSocketClient,
521            tokio::sync::mpsc::UnboundedReceiver<Message>,
522        ),
523        BitmexWsError,
524    > {
525        let (message_handler, rx) = channel_message_handler();
526
527        // No-op ping handler: handler owns the WebSocketClient and responds to pings directly
528        // in the message loop for minimal latency (see handler.rs pong response)
529        let ping_handler: PingHandler = Arc::new(move |_payload: Vec<u8>| {
530            // Handler responds to pings internally via select! loop
531        });
532
533        let config = WebSocketConfig {
534            url: self.url.clone(),
535            headers: vec![(USER_AGENT.to_string(), NAUTILUS_USER_AGENT.to_string())],
536            heartbeat: self.heartbeat,
537            heartbeat_msg: None,
538            reconnect_timeout_ms: Some(5_000),
539            reconnect_delay_initial_ms: None, // Use default
540            reconnect_delay_max_ms: None,     // Use default
541            reconnect_backoff_factor: None,   // Use default
542            reconnect_jitter_ms: None,        // Use default
543            reconnect_max_attempts: None,
544            idle_timeout_ms: None,
545            backend: self.transport_backend,
546            proxy_url: self.proxy_url.clone(),
547        };
548
549        let keyed_quotas = vec![];
550        let client = WebSocketClient::connect(
551            config,
552            Some(message_handler),
553            Some(ping_handler),
554            None, // post_reconnection
555            keyed_quotas,
556            None, // default_quota
557        )
558        .await
559        .map_err(|e| BitmexWsError::ClientError(e.to_string()))?;
560
561        Ok((client, rx))
562    }
563
564    /// Authenticate the WebSocket connection using the provided credentials.
565    ///
566    /// # Errors
567    ///
568    /// Returns an error if the WebSocket is not connected, if authentication fails,
569    /// or if credentials are not available.
570    async fn authenticate(&self) -> Result<(), BitmexWsError> {
571        let credential = match &self.credential {
572            Some(credential) => credential,
573            None => {
574                return Err(BitmexWsError::AuthenticationError(
575                    "API credentials not available to authenticate".to_string(),
576                ));
577            }
578        };
579
580        let receiver = self.auth_tracker.begin();
581
582        let expires = (chrono::Utc::now() + chrono::Duration::seconds(30)).timestamp();
583        let signature = credential.sign("GET", "/realtime", expires, "");
584
585        let auth_message = BitmexAuthentication {
586            op: BitmexWsAuthAction::AuthKeyExpires,
587            args: (credential.api_key().to_string(), expires, signature),
588        };
589
590        let auth_json = serde_json::to_string(&auth_message).map_err(|e| {
591            let msg = format!("Failed to serialize auth message: {e}");
592            self.auth_tracker.fail(msg.clone());
593            BitmexWsError::AuthenticationError(msg)
594        })?;
595
596        // Send Authenticate command to handler
597        self.cmd_tx
598            .read()
599            .await
600            .send(HandlerCommand::Authenticate { payload: auth_json })
601            .map_err(|e| {
602                let msg = format!("Failed to send authenticate command: {e}");
603                self.auth_tracker.fail(msg.clone());
604                BitmexWsError::AuthenticationError(msg)
605            })?;
606
607        self.auth_tracker
608            .wait_for_result::<BitmexWsError>(
609                Duration::from_secs(AUTHENTICATION_TIMEOUT_SECS),
610                receiver,
611            )
612            .await
613    }
614
615    /// Wait until the WebSocket connection is active.
616    ///
617    /// # Errors
618    ///
619    /// Returns an error if the connection times out.
620    pub async fn wait_until_active(&self, timeout_secs: f64) -> Result<(), BitmexWsError> {
621        let timeout = Duration::from_secs_f64(timeout_secs);
622
623        tokio::time::timeout(timeout, async {
624            while !self.is_active() {
625                tokio::time::sleep(Duration::from_millis(10)).await;
626            }
627        })
628        .await
629        .map_err(|_| {
630            BitmexWsError::ClientError(format!(
631                "WebSocket connection timeout after {timeout_secs} seconds"
632            ))
633        })?;
634
635        Ok(())
636    }
637
638    /// Provides the internal stream as a channel-based stream.
639    ///
640    /// # Panics
641    ///
642    /// This function panics:
643    /// - If the websocket is not connected.
644    /// - If `stream` has already been called somewhere else (stream receiver is then taken).
645    pub fn stream(&mut self) -> impl Stream<Item = BitmexWsMessage> + use<> {
646        let rx = self
647            .out_rx
648            .take()
649            .expect("Stream receiver already taken or not connected");
650        let mut rx = Arc::try_unwrap(rx).expect("Cannot take ownership - other references exist");
651        async_stream::stream! {
652            while let Some(msg) = rx.recv().await {
653                yield msg;
654            }
655        }
656    }
657
658    /// Closes the client.
659    ///
660    /// # Errors
661    ///
662    /// Returns an error if the WebSocket is not connected or if closing fails.
663    pub async fn close(&mut self) -> Result<(), BitmexWsError> {
664        log::debug!("Starting close process");
665
666        self.signal.store(true, Ordering::Relaxed);
667
668        // Send Disconnect command to handler
669        if let Err(e) = self.cmd_tx.read().await.send(HandlerCommand::Disconnect) {
670            log::debug!(
671                "Failed to send disconnect command (handler may already be shut down): {e}"
672            );
673        }
674
675        // Clean up task handle with timeout
676        if let Some(task_handle) = self.task_handle.take() {
677            match Arc::try_unwrap(task_handle) {
678                Ok(handle) => {
679                    log::debug!("Waiting for task handle to complete");
680                    match tokio::time::timeout(Duration::from_secs(2), handle).await {
681                        Ok(Ok(())) => log::debug!("Task handle completed successfully"),
682                        Ok(Err(e)) => log::error!("Task handle encountered an error: {e:?}"),
683                        Err(_) => {
684                            log::warn!(
685                                "Timeout waiting for task handle, task may still be running"
686                            );
687                            // The task will be dropped and should clean up automatically
688                        }
689                    }
690                }
691                Err(arc_handle) => {
692                    log::debug!(
693                        "Cannot take ownership of task handle - other references exist, aborting task"
694                    );
695                    arc_handle.abort();
696                }
697            }
698        } else {
699            log::debug!("No task handle to await");
700        }
701
702        log::debug!("Closed");
703
704        Ok(())
705    }
706
707    /// Subscribe to the specified topics.
708    ///
709    /// # Errors
710    ///
711    /// Returns an error if the WebSocket is not connected or if sending the subscription message fails.
712    pub async fn subscribe(&self, topics: Vec<String>) -> Result<(), BitmexWsError> {
713        log::debug!("Subscribing to topics: {topics:?}");
714
715        for topic in &topics {
716            self.subscriptions.mark_subscribe(topic.as_str());
717            self.tracked_subscriptions.insert(topic.clone(), ());
718        }
719
720        // Serialize subscription messages
721        let mut payloads = Vec::with_capacity(topics.len());
722        for topic in &topics {
723            let message = BitmexSubscription {
724                op: BitmexWsOperation::Subscribe,
725                args: vec![Ustr::from(topic.as_ref())],
726            };
727            let payload = serde_json::to_string(&message).map_err(|e| {
728                BitmexWsError::SubscriptionError(format!("Failed to serialize subscription: {e}"))
729            })?;
730            payloads.push(payload);
731        }
732
733        // Send Subscribe command to handler
734        let cmd = HandlerCommand::Subscribe { topics: payloads };
735
736        self.send_cmd(cmd).await.map_err(|e| {
737            BitmexWsError::SubscriptionError(format!("Failed to send subscribe command: {e}"))
738        })
739    }
740
741    /// Unsubscribe from the specified topics.
742    ///
743    /// # Errors
744    ///
745    /// Returns an error if the WebSocket is not connected or if sending the unsubscription message fails.
746    async fn unsubscribe(&self, topics: Vec<String>) -> Result<(), BitmexWsError> {
747        log::debug!("Attempting to unsubscribe from topics: {topics:?}");
748
749        if self.signal.load(Ordering::Relaxed) {
750            log::debug!("Shutdown signal detected, skipping unsubscribe");
751            return Ok(());
752        }
753
754        for topic in &topics {
755            self.subscriptions.mark_unsubscribe(topic.as_str());
756            self.tracked_subscriptions.remove(topic);
757        }
758
759        // Serialize unsubscription messages
760        let mut payloads = Vec::with_capacity(topics.len());
761        for topic in &topics {
762            let message = BitmexSubscription {
763                op: BitmexWsOperation::Unsubscribe,
764                args: vec![Ustr::from(topic.as_ref())],
765            };
766
767            if let Ok(payload) = serde_json::to_string(&message) {
768                payloads.push(payload);
769            }
770        }
771
772        // Send Unsubscribe command to handler
773        let cmd = HandlerCommand::Unsubscribe { topics: payloads };
774
775        if let Err(e) = self.send_cmd(cmd).await {
776            log::debug!("Failed to send unsubscribe command: {e}");
777        }
778
779        Ok(())
780    }
781
782    /// Get the current number of active subscriptions.
783    #[must_use]
784    pub fn subscription_count(&self) -> usize {
785        self.subscriptions.len()
786    }
787
788    pub fn get_subscriptions(&self, instrument_id: InstrumentId) -> Vec<String> {
789        let symbol = instrument_id.symbol.inner();
790        let confirmed = self.subscriptions.confirmed();
791        let mut channels = Vec::with_capacity(confirmed.len());
792
793        for entry in confirmed.iter() {
794            let (channel, symbols) = entry.pair();
795            if symbols.contains(&symbol) {
796                // Return the full topic string (e.g., "orderBookL2:XBTUSD")
797                channels.push(format!("{channel}:{symbol}"));
798            } else {
799                let has_channel_marker = symbols.iter().any(|s| s.is_empty());
800                if has_channel_marker
801                    && (*channel == BitmexWsAuthChannel::Execution.as_ref()
802                        || *channel == BitmexWsAuthChannel::Order.as_ref())
803                {
804                    // These are account-level subscriptions without symbols
805                    channels.push(channel.to_string());
806                }
807            }
808        }
809
810        channels
811    }
812
813    /// Subscribe to instrument updates for all instruments on the venue.
814    ///
815    /// # Errors
816    ///
817    /// Returns an error if the WebSocket is not connected or if the subscription fails.
818    pub async fn subscribe_instruments(&self) -> Result<(), BitmexWsError> {
819        // Already subscribed automatically on connection
820        log::debug!("Already subscribed to all instruments on connection, skipping");
821        Ok(())
822    }
823
824    /// Subscribe to instrument updates (mark/index prices) for the specified instrument.
825    ///
826    /// # Errors
827    ///
828    /// Returns an error if the WebSocket is not connected or if the subscription fails.
829    pub async fn subscribe_instrument(
830        &self,
831        instrument_id: InstrumentId,
832    ) -> Result<(), BitmexWsError> {
833        // Already subscribed to all instruments on connection
834        log::debug!(
835            "Already subscribed to all instruments on connection (includes {instrument_id}), skipping"
836        );
837        Ok(())
838    }
839
840    /// Subscribe to order book updates for the specified instrument.
841    ///
842    /// # Errors
843    ///
844    /// Returns an error if the WebSocket is not connected or if the subscription fails.
845    pub async fn subscribe_book(&self, instrument_id: InstrumentId) -> Result<(), BitmexWsError> {
846        let topic = BitmexWsTopic::OrderBookL2;
847        let symbol = instrument_id.symbol.inner();
848        self.subscribe(vec![format!("{topic}:{symbol}")]).await
849    }
850
851    /// Subscribe to order book L2 (25 levels) updates for the specified instrument.
852    ///
853    /// # Errors
854    ///
855    /// Returns an error if the WebSocket is not connected or if the subscription fails.
856    pub async fn subscribe_book_25(
857        &self,
858        instrument_id: InstrumentId,
859    ) -> Result<(), BitmexWsError> {
860        let topic = BitmexWsTopic::OrderBookL2_25;
861        let symbol = instrument_id.symbol.inner();
862        self.subscribe(vec![format!("{topic}:{symbol}")]).await
863    }
864
865    /// Subscribe to order book depth 10 updates for the specified instrument.
866    ///
867    /// # Errors
868    ///
869    /// Returns an error if the WebSocket is not connected or if the subscription fails.
870    pub async fn subscribe_book_depth10(
871        &self,
872        instrument_id: InstrumentId,
873    ) -> Result<(), BitmexWsError> {
874        let topic = BitmexWsTopic::OrderBook10;
875        let symbol = instrument_id.symbol.inner();
876        self.subscribe(vec![format!("{topic}:{symbol}")]).await
877    }
878
879    /// Subscribe to quote updates for the specified instrument.
880    ///
881    /// Note: Index symbols (starting with '.') do not have quotes and will be silently ignored.
882    ///
883    /// # Errors
884    ///
885    /// Returns an error if the WebSocket is not connected or if the subscription fails.
886    pub async fn subscribe_quotes(&self, instrument_id: InstrumentId) -> Result<(), BitmexWsError> {
887        let symbol = instrument_id.symbol.inner();
888
889        // Index symbols don't have quotes (bid/ask), only a single price
890        if is_index_symbol(&instrument_id.symbol.inner()) {
891            log::warn!("Ignoring quote subscription for index symbol: {symbol}");
892            return Ok(());
893        }
894
895        let topic = BitmexWsTopic::Quote;
896        self.subscribe(vec![format!("{topic}:{symbol}")]).await
897    }
898
899    /// Subscribe to trade updates for the specified instrument.
900    ///
901    /// Note: Index symbols (starting with '.') do not have trades and will be silently ignored.
902    ///
903    /// # Errors
904    ///
905    /// Returns an error if the WebSocket is not connected or if the subscription fails.
906    pub async fn subscribe_trades(&self, instrument_id: InstrumentId) -> Result<(), BitmexWsError> {
907        let symbol = instrument_id.symbol.inner();
908
909        // Index symbols don't have trades
910        if is_index_symbol(&symbol) {
911            log::warn!("Ignoring trade subscription for index symbol: {symbol}");
912            return Ok(());
913        }
914
915        let topic = BitmexWsTopic::Trade;
916        self.subscribe(vec![format!("{topic}:{symbol}")]).await
917    }
918
919    /// Subscribe to mark price updates for the specified instrument.
920    ///
921    /// # Errors
922    ///
923    /// Returns an error if the WebSocket is not connected or if the subscription fails.
924    pub async fn subscribe_mark_prices(
925        &self,
926        instrument_id: InstrumentId,
927    ) -> Result<(), BitmexWsError> {
928        self.subscribe_instrument(instrument_id).await
929    }
930
931    /// Subscribe to index price updates for the specified instrument.
932    ///
933    /// # Errors
934    ///
935    /// Returns an error if the WebSocket is not connected or if the subscription fails.
936    pub async fn subscribe_index_prices(
937        &self,
938        instrument_id: InstrumentId,
939    ) -> Result<(), BitmexWsError> {
940        self.subscribe_instrument(instrument_id).await
941    }
942
943    /// Subscribe to funding rate updates for the specified instrument.
944    ///
945    /// # Errors
946    ///
947    /// Returns an error if the WebSocket is not connected or if the subscription fails.
948    pub async fn subscribe_funding_rates(
949        &self,
950        instrument_id: InstrumentId,
951    ) -> Result<(), BitmexWsError> {
952        let topic = BitmexWsTopic::Funding;
953        let symbol = instrument_id.symbol.inner();
954        self.subscribe(vec![format!("{topic}:{symbol}")]).await
955    }
956
957    /// Subscribe to bar updates for the specified bar type.
958    ///
959    /// # Errors
960    ///
961    /// Returns an error if the WebSocket is not connected or if the subscription fails.
962    pub async fn subscribe_bars(&self, bar_type: BarType) -> Result<(), BitmexWsError> {
963        let topic = topic_from_bar_spec(bar_type.spec());
964        let symbol = bar_type.instrument_id().symbol.inner();
965        self.subscribe(vec![format!("{topic}:{symbol}")]).await
966    }
967
968    /// Unsubscribe from instrument updates for all instruments on the venue.
969    ///
970    /// # Errors
971    ///
972    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
973    pub async fn unsubscribe_instruments(&self) -> Result<(), BitmexWsError> {
974        // No-op: instruments are required for proper operation
975        log::debug!(
976            "Instruments subscription maintained for proper operation, skipping unsubscribe"
977        );
978        Ok(())
979    }
980
981    /// Unsubscribe from instrument updates (mark/index prices) for the specified instrument.
982    ///
983    /// # Errors
984    ///
985    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
986    pub async fn unsubscribe_instrument(
987        &self,
988        instrument_id: InstrumentId,
989    ) -> Result<(), BitmexWsError> {
990        // No-op: instruments are required for proper operation
991        log::debug!(
992            "Instruments subscription maintained for proper operation (includes {instrument_id}), skipping unsubscribe"
993        );
994        Ok(())
995    }
996
997    /// Unsubscribe from order book updates for the specified instrument.
998    ///
999    /// # Errors
1000    ///
1001    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1002    pub async fn unsubscribe_book(&self, instrument_id: InstrumentId) -> Result<(), BitmexWsError> {
1003        let topic = BitmexWsTopic::OrderBookL2;
1004        let symbol = instrument_id.symbol.inner();
1005        self.unsubscribe(vec![format!("{topic}:{symbol}")]).await
1006    }
1007
1008    /// Unsubscribe from order book L2 (25 levels) updates for the specified instrument.
1009    ///
1010    /// # Errors
1011    ///
1012    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1013    pub async fn unsubscribe_book_25(
1014        &self,
1015        instrument_id: InstrumentId,
1016    ) -> Result<(), BitmexWsError> {
1017        let topic = BitmexWsTopic::OrderBookL2_25;
1018        let symbol = instrument_id.symbol.inner();
1019        self.unsubscribe(vec![format!("{topic}:{symbol}")]).await
1020    }
1021
1022    /// Unsubscribe from order book depth 10 updates for the specified instrument.
1023    ///
1024    /// # Errors
1025    ///
1026    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1027    pub async fn unsubscribe_book_depth10(
1028        &self,
1029        instrument_id: InstrumentId,
1030    ) -> Result<(), BitmexWsError> {
1031        let topic = BitmexWsTopic::OrderBook10;
1032        let symbol = instrument_id.symbol.inner();
1033        self.unsubscribe(vec![format!("{topic}:{symbol}")]).await
1034    }
1035
1036    /// Unsubscribe from quote updates for the specified instrument.
1037    ///
1038    /// # Errors
1039    ///
1040    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1041    pub async fn unsubscribe_quotes(
1042        &self,
1043        instrument_id: InstrumentId,
1044    ) -> Result<(), BitmexWsError> {
1045        let symbol = instrument_id.symbol.inner();
1046
1047        // Index symbols don't have quotes
1048        if is_index_symbol(&symbol) {
1049            return Ok(());
1050        }
1051
1052        let topic = BitmexWsTopic::Quote;
1053        self.unsubscribe(vec![format!("{topic}:{symbol}")]).await
1054    }
1055
1056    /// Unsubscribe from trade updates for the specified instrument.
1057    ///
1058    /// # Errors
1059    ///
1060    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1061    pub async fn unsubscribe_trades(
1062        &self,
1063        instrument_id: InstrumentId,
1064    ) -> Result<(), BitmexWsError> {
1065        let symbol = instrument_id.symbol.inner();
1066
1067        // Index symbols don't have trades
1068        if is_index_symbol(&symbol) {
1069            return Ok(());
1070        }
1071
1072        let topic = BitmexWsTopic::Trade;
1073        self.unsubscribe(vec![format!("{topic}:{symbol}")]).await
1074    }
1075
1076    /// Unsubscribe from mark price updates for the specified instrument.
1077    ///
1078    /// # Errors
1079    ///
1080    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1081    pub async fn unsubscribe_mark_prices(
1082        &self,
1083        instrument_id: InstrumentId,
1084    ) -> Result<(), BitmexWsError> {
1085        // No-op: instrument channel shared with index prices
1086        log::debug!(
1087            "Mark prices for {instrument_id} uses shared instrument channel, skipping unsubscribe"
1088        );
1089        Ok(())
1090    }
1091
1092    /// Unsubscribe from index price updates for the specified instrument.
1093    ///
1094    /// # Errors
1095    ///
1096    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1097    pub async fn unsubscribe_index_prices(
1098        &self,
1099        instrument_id: InstrumentId,
1100    ) -> Result<(), BitmexWsError> {
1101        // No-op: instrument channel shared with mark prices
1102        log::debug!(
1103            "Index prices for {instrument_id} uses shared instrument channel, skipping unsubscribe"
1104        );
1105        Ok(())
1106    }
1107
1108    /// Unsubscribe from funding rate updates for the specified instrument.
1109    ///
1110    /// # Errors
1111    ///
1112    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1113    pub async fn unsubscribe_funding_rates(
1114        &self,
1115        instrument_id: InstrumentId,
1116    ) -> Result<(), BitmexWsError> {
1117        // No-op: unsubscribing during shutdown causes race conditions
1118        log::debug!(
1119            "Funding rates for {instrument_id}, skipping unsubscribe to avoid shutdown race"
1120        );
1121        Ok(())
1122    }
1123
1124    /// Unsubscribe from bar updates for the specified bar type.
1125    ///
1126    /// # Errors
1127    ///
1128    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1129    pub async fn unsubscribe_bars(&self, bar_type: BarType) -> Result<(), BitmexWsError> {
1130        let topic = topic_from_bar_spec(bar_type.spec());
1131        let symbol = bar_type.instrument_id().symbol.inner();
1132        self.unsubscribe(vec![format!("{topic}:{symbol}")]).await
1133    }
1134
1135    /// Subscribe to order updates for the authenticated account.
1136    ///
1137    /// # Errors
1138    ///
1139    /// Returns an error if the WebSocket is not connected, not authenticated, or if the subscription fails.
1140    pub async fn subscribe_orders(&self) -> Result<(), BitmexWsError> {
1141        if self.credential.is_none() {
1142            return Err(BitmexWsError::MissingCredentials);
1143        }
1144        self.subscribe(vec![BitmexWsAuthChannel::Order.to_string()])
1145            .await
1146    }
1147
1148    /// Subscribe to execution updates for the authenticated account.
1149    ///
1150    /// # Errors
1151    ///
1152    /// Returns an error if the WebSocket is not connected, not authenticated, or if the subscription fails.
1153    pub async fn subscribe_executions(&self) -> Result<(), BitmexWsError> {
1154        if self.credential.is_none() {
1155            return Err(BitmexWsError::MissingCredentials);
1156        }
1157        self.subscribe(vec![BitmexWsAuthChannel::Execution.to_string()])
1158            .await
1159    }
1160
1161    /// Subscribe to position updates for the authenticated account.
1162    ///
1163    /// # Errors
1164    ///
1165    /// Returns an error if the WebSocket is not connected, not authenticated, or if the subscription fails.
1166    pub async fn subscribe_positions(&self) -> Result<(), BitmexWsError> {
1167        if self.credential.is_none() {
1168            return Err(BitmexWsError::MissingCredentials);
1169        }
1170        self.subscribe(vec![BitmexWsAuthChannel::Position.to_string()])
1171            .await
1172    }
1173
1174    /// Subscribe to margin updates for the authenticated account.
1175    ///
1176    /// # Errors
1177    ///
1178    /// Returns an error if the WebSocket is not connected, not authenticated, or if the subscription fails.
1179    pub async fn subscribe_margin(&self) -> Result<(), BitmexWsError> {
1180        if self.credential.is_none() {
1181            return Err(BitmexWsError::MissingCredentials);
1182        }
1183        self.subscribe(vec![BitmexWsAuthChannel::Margin.to_string()])
1184            .await
1185    }
1186
1187    /// Subscribe to wallet updates for the authenticated account.
1188    ///
1189    /// # Errors
1190    ///
1191    /// Returns an error if the WebSocket is not connected, not authenticated, or if the subscription fails.
1192    pub async fn subscribe_wallet(&self) -> Result<(), BitmexWsError> {
1193        if self.credential.is_none() {
1194            return Err(BitmexWsError::MissingCredentials);
1195        }
1196        self.subscribe(vec![BitmexWsAuthChannel::Wallet.to_string()])
1197            .await
1198    }
1199
1200    /// Unsubscribe from order updates for the authenticated account.
1201    ///
1202    /// # Errors
1203    ///
1204    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1205    pub async fn unsubscribe_orders(&self) -> Result<(), BitmexWsError> {
1206        self.unsubscribe(vec![BitmexWsAuthChannel::Order.to_string()])
1207            .await
1208    }
1209
1210    /// Unsubscribe from execution updates for the authenticated account.
1211    ///
1212    /// # Errors
1213    ///
1214    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1215    pub async fn unsubscribe_executions(&self) -> Result<(), BitmexWsError> {
1216        self.unsubscribe(vec![BitmexWsAuthChannel::Execution.to_string()])
1217            .await
1218    }
1219
1220    /// Unsubscribe from position updates for the authenticated account.
1221    ///
1222    /// # Errors
1223    ///
1224    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1225    pub async fn unsubscribe_positions(&self) -> Result<(), BitmexWsError> {
1226        self.unsubscribe(vec![BitmexWsAuthChannel::Position.to_string()])
1227            .await
1228    }
1229
1230    /// Unsubscribe from margin updates for the authenticated account.
1231    ///
1232    /// # Errors
1233    ///
1234    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1235    pub async fn unsubscribe_margin(&self) -> Result<(), BitmexWsError> {
1236        self.unsubscribe(vec![BitmexWsAuthChannel::Margin.to_string()])
1237            .await
1238    }
1239
1240    /// Unsubscribe from wallet updates for the authenticated account.
1241    ///
1242    /// # Errors
1243    ///
1244    /// Returns an error if the WebSocket is not connected or if the unsubscription fails.
1245    pub async fn unsubscribe_wallet(&self) -> Result<(), BitmexWsError> {
1246        self.unsubscribe(vec![BitmexWsAuthChannel::Wallet.to_string()])
1247            .await
1248    }
1249
1250    /// Sends a command to the handler.
1251    async fn send_cmd(&self, cmd: HandlerCommand) -> Result<(), BitmexWsError> {
1252        self.cmd_tx
1253            .read()
1254            .await
1255            .send(cmd)
1256            .map_err(|e| BitmexWsError::ClientError(format!("Handler not available: {e}")))
1257    }
1258}
1259
1260#[cfg(test)]
1261mod tests {
1262    use ahash::AHashSet;
1263    use rstest::rstest;
1264    use ustr::Ustr;
1265
1266    use super::*;
1267
1268    #[rstest]
1269    fn test_reconnect_topics_restoration_logic() {
1270        // Create real client with credentials
1271        let client = BitmexWebSocketClient::new(
1272            Some("ws://test.com".to_string()),
1273            Some("test_key".to_string()),
1274            Some("test_secret".to_string()),
1275            Some(AccountId::new("BITMEX-TEST")),
1276            5,
1277            TransportBackend::default(),
1278            None,
1279        )
1280        .unwrap();
1281
1282        // Populate subscriptions like they would be during normal operation
1283        let subs = client.subscriptions.confirmed();
1284        subs.insert(Ustr::from(BitmexWsTopic::Trade.as_ref()), {
1285            let mut set = AHashSet::new();
1286            set.insert(Ustr::from("XBTUSD"));
1287            set.insert(Ustr::from("ETHUSD"));
1288            set
1289        });
1290
1291        subs.insert(Ustr::from(BitmexWsTopic::OrderBookL2.as_ref()), {
1292            let mut set = AHashSet::new();
1293            set.insert(Ustr::from("XBTUSD"));
1294            set
1295        });
1296
1297        // Private channels (no symbols)
1298        subs.insert(Ustr::from(BitmexWsAuthChannel::Order.as_ref()), {
1299            let mut set = AHashSet::new();
1300            set.insert(Ustr::from(""));
1301            set
1302        });
1303        subs.insert(Ustr::from(BitmexWsAuthChannel::Position.as_ref()), {
1304            let mut set = AHashSet::new();
1305            set.insert(Ustr::from(""));
1306            set
1307        });
1308
1309        // Test the actual reconnection topic building logic
1310        let mut topics_to_restore = Vec::new();
1311
1312        for entry in subs.iter() {
1313            let (channel, symbols) = entry.pair();
1314            for symbol in symbols {
1315                if symbol.is_empty() {
1316                    topics_to_restore.push(channel.to_string());
1317                } else {
1318                    topics_to_restore.push(format!("{channel}:{symbol}"));
1319                }
1320            }
1321        }
1322
1323        // Verify it builds the correct restoration topics
1324        assert!(topics_to_restore.contains(&format!("{}:XBTUSD", BitmexWsTopic::Trade.as_ref())));
1325        assert!(topics_to_restore.contains(&format!("{}:ETHUSD", BitmexWsTopic::Trade.as_ref())));
1326        assert!(
1327            topics_to_restore.contains(&format!("{}:XBTUSD", BitmexWsTopic::OrderBookL2.as_ref()))
1328        );
1329        assert!(topics_to_restore.contains(&BitmexWsAuthChannel::Order.as_ref().to_string()));
1330        assert!(topics_to_restore.contains(&BitmexWsAuthChannel::Position.as_ref().to_string()));
1331        assert_eq!(topics_to_restore.len(), 5);
1332    }
1333
1334    #[rstest]
1335    fn test_reconnect_auth_message_building() {
1336        // Test with credentials
1337        let client_with_creds = BitmexWebSocketClient::new(
1338            Some("ws://test.com".to_string()),
1339            Some("test_key".to_string()),
1340            Some("test_secret".to_string()),
1341            Some(AccountId::new("BITMEX-TEST")),
1342            5,
1343            TransportBackend::default(),
1344            None,
1345        )
1346        .unwrap();
1347
1348        // Test the actual auth message building logic from lines 220-228
1349        if let Some(cred) = &client_with_creds.credential {
1350            let expires = (chrono::Utc::now() + chrono::Duration::seconds(30)).timestamp();
1351            let signature = cred.sign("GET", "/realtime", expires, "");
1352
1353            let auth_message = BitmexAuthentication {
1354                op: BitmexWsAuthAction::AuthKeyExpires,
1355                args: (cred.api_key().to_string(), expires, signature),
1356            };
1357
1358            // Verify auth message structure
1359            assert_eq!(auth_message.op, BitmexWsAuthAction::AuthKeyExpires);
1360            assert_eq!(auth_message.args.0, "test_key");
1361            assert!(auth_message.args.1 > 0); // expires should be positive
1362            assert!(!auth_message.args.2.is_empty()); // signature should exist
1363        } else {
1364            panic!("Client should have credentials");
1365        }
1366
1367        // Test without credentials
1368        let client_no_creds = BitmexWebSocketClient::new(
1369            Some("ws://test.com".to_string()),
1370            None,
1371            None,
1372            Some(AccountId::new("BITMEX-TEST")),
1373            5,
1374            TransportBackend::default(),
1375            None,
1376        )
1377        .unwrap();
1378
1379        assert!(client_no_creds.credential.is_none());
1380    }
1381
1382    #[rstest]
1383    fn test_subscription_state_after_unsubscribe() {
1384        let client = BitmexWebSocketClient::new(
1385            Some("ws://test.com".to_string()),
1386            Some("test_key".to_string()),
1387            Some("test_secret".to_string()),
1388            Some(AccountId::new("BITMEX-TEST")),
1389            5,
1390            TransportBackend::default(),
1391            None,
1392        )
1393        .unwrap();
1394
1395        // Set up initial subscriptions
1396        let subs = client.subscriptions.confirmed();
1397        subs.insert(Ustr::from(BitmexWsTopic::Trade.as_ref()), {
1398            let mut set = AHashSet::new();
1399            set.insert(Ustr::from("XBTUSD"));
1400            set.insert(Ustr::from("ETHUSD"));
1401            set
1402        });
1403
1404        subs.insert(Ustr::from(BitmexWsTopic::OrderBookL2.as_ref()), {
1405            let mut set = AHashSet::new();
1406            set.insert(Ustr::from("XBTUSD"));
1407            set
1408        });
1409
1410        // Simulate unsubscribe logic (like from unsubscribe() method lines 586-599)
1411        let topic = format!("{}:ETHUSD", BitmexWsTopic::Trade.as_ref());
1412        if let Some((channel, symbol)) = topic.split_once(':')
1413            && let Some(mut entry) = subs.get_mut(&Ustr::from(channel))
1414        {
1415            entry.remove(&Ustr::from(symbol));
1416            if entry.is_empty() {
1417                drop(entry);
1418                subs.remove(&Ustr::from(channel));
1419            }
1420        }
1421
1422        // Build restoration topics after unsubscribe
1423        let mut topics_to_restore = Vec::new();
1424
1425        for entry in subs.iter() {
1426            let (channel, symbols) = entry.pair();
1427            for symbol in symbols {
1428                if symbol.is_empty() {
1429                    topics_to_restore.push(channel.to_string());
1430                } else {
1431                    topics_to_restore.push(format!("{channel}:{symbol}"));
1432                }
1433            }
1434        }
1435
1436        // Should have XBTUSD trade but not ETHUSD trade
1437        let trade_xbt = format!("{}:XBTUSD", BitmexWsTopic::Trade.as_ref());
1438        let trade_eth = format!("{}:ETHUSD", BitmexWsTopic::Trade.as_ref());
1439        let book_xbt = format!("{}:XBTUSD", BitmexWsTopic::OrderBookL2.as_ref());
1440
1441        assert!(topics_to_restore.contains(&trade_xbt));
1442        assert!(!topics_to_restore.contains(&trade_eth));
1443        assert!(topics_to_restore.contains(&book_xbt));
1444        assert_eq!(topics_to_restore.len(), 2);
1445    }
1446
1447    #[rstest]
1448    fn test_race_unsubscribe_failure_recovery() {
1449        // Simulates the race condition where venue rejects an unsubscribe request.
1450        // The adapter must perform the 3-step recovery:
1451        // 1. confirm_unsubscribe() - clear pending_unsubscribe
1452        // 2. mark_subscribe() - mark as subscribing again
1453        // 3. confirm_subscribe() - restore to confirmed state
1454        let client = BitmexWebSocketClient::new(
1455            Some("ws://test.com".to_string()),
1456            None,
1457            None,
1458            Some(AccountId::new("BITMEX-TEST")),
1459            5,
1460            TransportBackend::default(),
1461            None,
1462        )
1463        .unwrap();
1464
1465        let topic = format!("{}:XBTUSD", BitmexWsTopic::Trade.as_ref());
1466
1467        // Initial subscribe flow
1468        client.subscriptions.mark_subscribe(&topic);
1469        client.subscriptions.confirm_subscribe(&topic);
1470        assert_eq!(client.subscriptions.len(), 1);
1471
1472        // User unsubscribes
1473        client.subscriptions.mark_unsubscribe(&topic);
1474        assert_eq!(client.subscriptions.len(), 0);
1475        assert_eq!(
1476            client.subscriptions.pending_unsubscribe_topics(),
1477            vec![topic.clone()]
1478        );
1479
1480        // Venue REJECTS the unsubscribe (error message)
1481        // Adapter must perform 3-step recovery (from lines 1884-1891)
1482        client.subscriptions.confirm_unsubscribe(&topic); // Step 1: clear pending_unsubscribe
1483        client.subscriptions.mark_subscribe(&topic); // Step 2: mark as subscribing
1484        client.subscriptions.confirm_subscribe(&topic); // Step 3: confirm subscription
1485
1486        // Verify recovery: topic should be back in confirmed state
1487        assert_eq!(client.subscriptions.len(), 1);
1488        assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
1489        assert!(client.subscriptions.pending_subscribe_topics().is_empty());
1490
1491        // Verify topic is in all_topics() for reconnect
1492        let all = client.subscriptions.all_topics();
1493        assert_eq!(all.len(), 1);
1494        assert!(all.contains(&topic));
1495    }
1496
1497    #[rstest]
1498    fn test_race_resubscribe_before_unsubscribe_ack() {
1499        // Simulates: User unsubscribes, then immediately resubscribes before
1500        // the unsubscribe ACK arrives from the venue.
1501        // This is the race condition fixed in the subscription tracker.
1502        let client = BitmexWebSocketClient::new(
1503            Some("ws://test.com".to_string()),
1504            None,
1505            None,
1506            Some(AccountId::new("BITMEX-TEST")),
1507            5,
1508            TransportBackend::default(),
1509            None,
1510        )
1511        .unwrap();
1512
1513        let topic = format!("{}:XBTUSD", BitmexWsTopic::OrderBookL2.as_ref());
1514
1515        // Initial subscribe
1516        client.subscriptions.mark_subscribe(&topic);
1517        client.subscriptions.confirm_subscribe(&topic);
1518        assert_eq!(client.subscriptions.len(), 1);
1519
1520        // User unsubscribes
1521        client.subscriptions.mark_unsubscribe(&topic);
1522        assert_eq!(client.subscriptions.len(), 0);
1523        assert_eq!(
1524            client.subscriptions.pending_unsubscribe_topics(),
1525            vec![topic.clone()]
1526        );
1527
1528        // User immediately changes mind and resubscribes (before unsubscribe ACK)
1529        client.subscriptions.mark_subscribe(&topic);
1530        assert_eq!(
1531            client.subscriptions.pending_subscribe_topics(),
1532            vec![topic.clone()]
1533        );
1534
1535        // NOW the unsubscribe ACK arrives - should NOT clear pending_subscribe
1536        client.subscriptions.confirm_unsubscribe(&topic);
1537        assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
1538        assert_eq!(
1539            client.subscriptions.pending_subscribe_topics(),
1540            vec![topic.clone()]
1541        );
1542
1543        // Subscribe ACK arrives
1544        client.subscriptions.confirm_subscribe(&topic);
1545        assert_eq!(client.subscriptions.len(), 1);
1546        assert!(client.subscriptions.pending_subscribe_topics().is_empty());
1547
1548        // Verify final state is correct
1549        let all = client.subscriptions.all_topics();
1550        assert_eq!(all.len(), 1);
1551        assert!(all.contains(&topic));
1552    }
1553
1554    #[rstest]
1555    fn test_race_channel_level_reconnection_with_pending_states() {
1556        // Simulates reconnection with mixed pending states including channel-level subscriptions.
1557        let client = BitmexWebSocketClient::new(
1558            Some("ws://test.com".to_string()),
1559            Some("test_key".to_string()),
1560            Some("test_secret".to_string()),
1561            Some(AccountId::new("BITMEX-TEST")),
1562            5,
1563            TransportBackend::default(),
1564            None,
1565        )
1566        .unwrap();
1567
1568        // Set up mixed state before reconnection
1569        // Confirmed: trade:XBTUSD
1570        let trade_xbt = format!("{}:XBTUSD", BitmexWsTopic::Trade.as_ref());
1571        client.subscriptions.mark_subscribe(&trade_xbt);
1572        client.subscriptions.confirm_subscribe(&trade_xbt);
1573
1574        // Confirmed: order (channel-level, no symbol)
1575        let order_channel = BitmexWsAuthChannel::Order.as_ref();
1576        client.subscriptions.mark_subscribe(order_channel);
1577        client.subscriptions.confirm_subscribe(order_channel);
1578
1579        // Pending subscribe: trade:ETHUSD
1580        let trade_eth = format!("{}:ETHUSD", BitmexWsTopic::Trade.as_ref());
1581        client.subscriptions.mark_subscribe(&trade_eth);
1582
1583        // Pending unsubscribe: orderBookL2:XBTUSD (user cancelled)
1584        let book_xbt = format!("{}:XBTUSD", BitmexWsTopic::OrderBookL2.as_ref());
1585        client.subscriptions.mark_subscribe(&book_xbt);
1586        client.subscriptions.confirm_subscribe(&book_xbt);
1587        client.subscriptions.mark_unsubscribe(&book_xbt);
1588
1589        // Get topics for reconnection
1590        let topics_to_restore = client.subscriptions.all_topics();
1591
1592        // Should include: confirmed + pending_subscribe (NOT pending_unsubscribe)
1593        assert_eq!(topics_to_restore.len(), 3);
1594        assert!(topics_to_restore.contains(&trade_xbt));
1595        assert!(topics_to_restore.contains(&order_channel.to_string()));
1596        assert!(topics_to_restore.contains(&trade_eth));
1597        assert!(!topics_to_restore.contains(&book_xbt)); // Excluded
1598
1599        // Verify channel-level marker is handled correctly
1600        // order channel should not have ':' delimiter
1601        for topic in &topics_to_restore {
1602            if topic == order_channel {
1603                assert!(
1604                    !topic.contains(':'),
1605                    "Channel-level topic should not have delimiter"
1606                );
1607            }
1608        }
1609    }
1610}
nautilus_bitmex/websocket/client.rs

nautilus_bitmex/websocket/
client.rs
