1use std::{
22 num::NonZeroU32,
23 str::FromStr,
24 sync::{
25 Arc, LazyLock,
26 atomic::{AtomicBool, AtomicU8, Ordering},
27 },
28 time::Duration,
29};
30
31use arc_swap::ArcSwap;
32use nautilus_core::AtomicMap;
33use nautilus_live::{
34 SocketControl,
35 task::{TaskJoinOutcome, TaskSlot, finish_task},
36};
37use nautilus_model::{
38 data::BarType,
39 identifiers::{AccountId, InstrumentId},
40 instruments::{Instrument, InstrumentAny},
41};
42use nautilus_network::{
43 mode::ConnectionMode,
44 ratelimiter::quota::Quota,
45 websocket::{
46 SubscriptionState, TransportBackend, WebSocketClient, WebSocketConfig,
47 channel_message_handler,
48 },
49};
50use ustr::Ustr;
51
52use crate::{
53 common::{
54 consts::{
55 RECONNECT_BACKOFF_FACTOR, RECONNECT_BASE_BACKOFF, RECONNECT_JITTER_MS,
56 RECONNECT_MAX_BACKOFF, RECONNECT_TIMEOUT, WS_DISCONNECT_TIMEOUT, WS_HEARTBEAT_SECS,
57 },
58 credential::CoinbaseCredential,
59 enums::CoinbaseWsChannel,
60 },
61 websocket::{
62 handler::{FeedHandler, HandlerCommand, NautilusWsMessage},
63 messages::{CoinbaseWsAction, CoinbaseWsSubscription},
64 },
65};
66
67pub static COINBASE_WS_CONNECTION_QUOTA: LazyLock<Quota> = LazyLock::new(|| {
69 Quota::per_second(NonZeroU32::new(8).expect("non-zero")).expect("valid constant")
70});
71
72pub static COINBASE_WS_SUBSCRIPTION_QUOTA: LazyLock<Quota> = LazyLock::new(|| {
74 Quota::per_second(NonZeroU32::new(8).expect("non-zero")).expect("valid constant")
75});
76
77pub const COINBASE_RATE_LIMIT_KEY_SUBSCRIPTION: &str = "subscription";
79
80pub static COINBASE_WS_SUBSCRIPTION_KEYS: LazyLock<[Ustr; 1]> =
82 LazyLock::new(|| [Ustr::from(COINBASE_RATE_LIMIT_KEY_SUBSCRIPTION)]);
83
84#[derive(Debug)]
89pub struct CoinbaseWebSocketClient {
90 url: String,
91 connection_mode: Arc<ArcSwap<AtomicU8>>,
92 signal: Arc<AtomicBool>,
93 cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
94 out_rx: Option<tokio::sync::mpsc::UnboundedReceiver<NautilusWsMessage>>,
95 instruments: Arc<AtomicMap<InstrumentId, InstrumentAny>>,
96 subscription_aliases: Arc<AtomicMap<Ustr, Ustr>>,
101 bar_types: ahash::AHashMap<String, BarType>,
102 subscriptions: SubscriptionState,
103 credential: Option<CoinbaseCredential>,
104 account_id: Option<AccountId>,
105 task_handle: TaskSlot<()>,
106 shutdown_errors: Vec<String>,
107 transport_backend: TransportBackend,
108 proxy_url: Option<String>,
109 socket_control: Option<SocketControl>,
110}
111
112impl Clone for CoinbaseWebSocketClient {
113 fn clone(&self) -> Self {
114 Self {
115 url: self.url.clone(),
116 connection_mode: Arc::clone(&self.connection_mode),
117 signal: Arc::clone(&self.signal),
118 cmd_tx: Arc::clone(&self.cmd_tx),
119 out_rx: None,
120 instruments: Arc::clone(&self.instruments),
121 subscription_aliases: Arc::clone(&self.subscription_aliases),
122 bar_types: self.bar_types.clone(),
123 subscriptions: self.subscriptions.clone(),
124 credential: self.credential.clone(),
125 account_id: self.account_id,
126 task_handle: TaskSlot::new(),
127 shutdown_errors: Vec::new(),
128 transport_backend: self.transport_backend,
129 proxy_url: self.proxy_url.clone(),
130 socket_control: self.socket_control.clone(),
131 }
132 }
133}
134
135impl CoinbaseWebSocketClient {
136 pub fn new(url: &str, transport_backend: TransportBackend, proxy_url: Option<String>) -> Self {
138 let (placeholder_tx, _) = tokio::sync::mpsc::unbounded_channel();
139
140 Self {
141 url: url.to_string(),
142 connection_mode: Arc::new(ArcSwap::from_pointee(AtomicU8::new(
143 ConnectionMode::Closed.as_u8(),
144 ))),
145 signal: Arc::new(AtomicBool::new(false)),
146 cmd_tx: Arc::new(tokio::sync::RwLock::new(placeholder_tx)),
147 out_rx: None,
148 instruments: Arc::new(AtomicMap::new()),
149 subscription_aliases: Arc::new(AtomicMap::new()),
150 bar_types: ahash::AHashMap::new(),
151 subscriptions: SubscriptionState::new('|'),
152 credential: None,
153 account_id: None,
154 task_handle: TaskSlot::new(),
155 shutdown_errors: Vec::new(),
156 transport_backend,
157 proxy_url,
158 socket_control: None,
159 }
160 }
161
162 #[must_use]
164 pub fn with_socket_control(mut self, control: SocketControl) -> Self {
165 self.socket_control = Some(control);
166 self
167 }
168
169 pub fn with_credential(
171 url: &str,
172 credential: CoinbaseCredential,
173 transport_backend: TransportBackend,
174 proxy_url: Option<String>,
175 ) -> Self {
176 let mut client = Self::new(url, transport_backend, proxy_url);
177 client.credential = Some(credential);
178 client
179 }
180
181 pub async fn set_account_id(&mut self, account_id: AccountId) {
186 self.account_id = Some(account_id);
187
188 let cmd_tx = self.cmd_tx.read().await;
189 if let Err(e) = cmd_tx.send(HandlerCommand::SetAccountId(account_id)) {
190 log::debug!("Failed to send SetAccountId: {e}");
191 }
192 }
193
194 pub async fn initialize_instruments(&self, instruments: Vec<InstrumentAny>) {
201 for instrument in &instruments {
202 self.instruments.insert(instrument.id(), instrument.clone());
203 }
204
205 let cmd_tx = self.cmd_tx.read().await;
206 if let Err(e) = cmd_tx.send(HandlerCommand::InitializeInstruments(instruments)) {
207 log::debug!("Failed to send InitializeInstruments: {e}");
208 }
209 }
210
211 fn prime_default_subscriptions(&self) {
215 self.subscriptions
216 .mark_subscribe(CoinbaseWsChannel::Heartbeats.as_ref());
217 }
218
219 pub async fn connect(&mut self) -> anyhow::Result<()> {
221 if self.is_active() || self.is_reconnecting() {
222 log::warn!("WebSocket already connected or reconnecting");
223 return Ok(());
224 }
225
226 if let Some(outcome) = finish_task(
227 &mut self.task_handle,
228 WS_DISCONNECT_TIMEOUT,
229 WS_DISCONNECT_TIMEOUT,
230 )
231 .await
232 {
233 match outcome {
234 TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => {}
235 TaskJoinOutcome::Failed(error) => {
236 anyhow::bail!("Coinbase WebSocket handler failed: {error}");
237 }
238 TaskJoinOutcome::Incomplete => {
239 anyhow::bail!("Coinbase WebSocket handler did not stop after abort");
240 }
241 }
242 }
243
244 self.signal.store(false, Ordering::Relaxed);
246
247 let (message_handler, raw_rx) = channel_message_handler();
248 let cfg = WebSocketConfig {
249 url: self.url.clone(),
250 headers: vec![],
251 heartbeat_interval_secs: Some(WS_HEARTBEAT_SECS),
254 heartbeat_payload: None,
255 connect_timeout_ms: Some(RECONNECT_TIMEOUT.as_millis() as u64),
256 reconnect_delay_initial_ms: Some(RECONNECT_BASE_BACKOFF.as_millis() as u64),
257 reconnect_delay_max_ms: Some(RECONNECT_MAX_BACKOFF.as_millis() as u64),
258 reconnect_backoff_factor: Some(RECONNECT_BACKOFF_FACTOR),
259 reconnect_jitter_ms: Some(RECONNECT_JITTER_MS),
260 reconnect_max_attempts: None,
261 heartbeat_timeout_secs: None,
262 idle_timeout_ms: None,
263 backend: self.transport_backend,
264 proxy_url: self.proxy_url.clone(),
265 };
266
267 let keyed_quotas = vec![(
268 COINBASE_RATE_LIMIT_KEY_SUBSCRIPTION.to_string(),
269 *COINBASE_WS_SUBSCRIPTION_QUOTA,
270 )];
271
272 let client = WebSocketClient::builder()
273 .config(cfg)
274 .message_handler(message_handler)
275 .keyed_quotas(keyed_quotas)
276 .default_quota(*COINBASE_WS_CONNECTION_QUOTA)
277 .maybe_state_sink(self.socket_control.as_ref().map(SocketControl::sink))
278 .connect()
279 .await?;
280
281 let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
282 let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<NautilusWsMessage>();
283
284 *self.cmd_tx.write().await = cmd_tx.clone();
285 self.out_rx = Some(out_rx);
286 self.connection_mode.store(client.connection_mode_atomic());
287 let reconnect_handle = client.reconnect_handle();
288 log::debug!("Coinbase WebSocket connected: {}", self.url);
289
290 if let Err(e) = cmd_tx.send(HandlerCommand::SetClient(client)) {
291 anyhow::bail!("Failed to send SetClient command: {e}");
292 }
293
294 if let Some(control) = &self.socket_control {
295 control.register(move || reconnect_handle.request_reconnect());
296 }
297
298 let instruments_vec: Vec<InstrumentAny> =
299 self.instruments.load().values().cloned().collect();
300
301 if !instruments_vec.is_empty()
302 && let Err(e) = cmd_tx.send(HandlerCommand::InitializeInstruments(instruments_vec))
303 {
304 log::error!("Failed to send InitializeInstruments: {e}");
305 }
306
307 for (key, bar_type) in &self.bar_types {
309 if let Err(e) = cmd_tx.send(HandlerCommand::AddBarType {
310 key: key.clone(),
311 bar_type: *bar_type,
312 }) {
313 log::error!("Failed to restore bar type {key}: {e}");
314 }
315 }
316
317 if let Some(account_id) = self.account_id
318 && let Err(e) = cmd_tx.send(HandlerCommand::SetAccountId(account_id))
319 {
320 log::error!("Failed to restore account_id: {e}");
321 }
322
323 self.prime_default_subscriptions();
324
325 resubscribe_all(
327 &self.subscriptions,
328 &self.credential,
329 &cmd_tx,
330 Some(&out_tx),
331 );
332
333 let signal = Arc::clone(&self.signal);
334 let subscriptions = self.subscriptions.clone();
335 let credential = self.credential.clone();
336 let cmd_tx_reconnect = cmd_tx.clone();
337 let aliases_for_handler = Arc::clone(&self.subscription_aliases);
338
339 if let Err(e) = self.task_handle.spawn(async move {
340 let mut handler = FeedHandler::new(signal, cmd_rx, raw_rx, aliases_for_handler);
341
342 loop {
343 match handler.next().await {
344 Some(NautilusWsMessage::Reconnected) => {
345 subscriptions.reset_after_reconnect();
346 resubscribe_all(
347 &subscriptions,
348 &credential,
349 &cmd_tx_reconnect,
350 Some(&out_tx),
351 );
352
353 if let Err(e) = out_tx.send(NautilusWsMessage::Reconnected) {
354 log::debug!("Output channel closed: {e}");
355 break;
356 }
357 }
358 Some(msg) => {
359 if let Err(e) = out_tx.send(msg) {
360 log::debug!("Output channel closed: {e}");
361 break;
362 }
363 }
364 None => {
365 log::debug!("Feed handler stopped");
366 break;
367 }
368 }
369 }
370 }) {
371 self.out_rx = None;
372 anyhow::bail!("Failed to start Coinbase WebSocket handler task: {e}");
373 }
374
375 Ok(())
376 }
377
378 pub async fn subscribe(
380 &self,
381 channel: CoinbaseWsChannel,
382 product_ids: &[Ustr],
383 ) -> anyhow::Result<()> {
384 let jwt = if channel.requires_auth() {
385 let credential = self
386 .credential
387 .as_ref()
388 .ok_or_else(|| anyhow::anyhow!("Credentials required for {channel}"))?;
389 Some(credential.build_ws_jwt()?)
390 } else {
391 self.credential.as_ref().and_then(|c| c.build_ws_jwt().ok())
392 };
393
394 let sub = CoinbaseWsSubscription {
395 msg_type: CoinbaseWsAction::Subscribe,
396 product_ids: product_ids.to_vec(),
397 channel,
398 jwt,
399 };
400
401 let channel_str = channel.as_ref();
402
403 if product_ids.is_empty() {
404 self.subscriptions.mark_subscribe(channel_str);
405 } else {
406 for product_id in product_ids {
407 let topic = format!("{channel_str}|{product_id}");
408 self.subscriptions.mark_subscribe(&topic);
409 }
410 }
411
412 let cmd_tx = self.cmd_tx.read().await;
413 cmd_tx
414 .send(HandlerCommand::Subscribe(sub))
415 .map_err(|e| anyhow::anyhow!("Failed to send Subscribe command: {e}"))
416 }
417
418 pub async fn unsubscribe(
420 &self,
421 channel: CoinbaseWsChannel,
422 product_ids: &[Ustr],
423 ) -> anyhow::Result<()> {
424 let jwt = self.credential.as_ref().and_then(|c| c.build_ws_jwt().ok());
425
426 let unsub = CoinbaseWsSubscription {
427 msg_type: CoinbaseWsAction::Unsubscribe,
428 product_ids: product_ids.to_vec(),
429 channel,
430 jwt,
431 };
432
433 let channel_str = channel.as_ref();
434
435 if product_ids.is_empty() {
436 self.subscriptions.mark_unsubscribe(channel_str);
437 } else {
438 for product_id in product_ids {
439 let topic = format!("{channel_str}|{product_id}");
440 self.subscriptions.mark_unsubscribe(&topic);
441 }
442 }
443
444 let cmd_tx = self.cmd_tx.read().await;
445 cmd_tx
446 .send(HandlerCommand::Unsubscribe(unsub))
447 .map_err(|e| anyhow::anyhow!("Failed to send Unsubscribe command: {e}"))
448 }
449
450 pub async fn next_message(&mut self) -> Option<NautilusWsMessage> {
452 self.out_rx.as_mut()?.recv().await
453 }
454
455 pub(crate) fn begin_shutdown(&self) {
457 self.signal.store(true, Ordering::Release);
458 }
459
460 pub async fn disconnect(&mut self) -> anyhow::Result<()> {
462 let cmd_tx = self.cmd_tx.read().await;
465
466 if let Err(e) = cmd_tx.send(HandlerCommand::Disconnect) {
467 log::debug!("Failed to send Disconnect command: {e}");
468 }
469 drop(cmd_tx);
470
471 self.begin_shutdown();
474
475 if let Some(outcome) = finish_task(
476 &mut self.task_handle,
477 WS_DISCONNECT_TIMEOUT,
478 WS_DISCONNECT_TIMEOUT,
479 )
480 .await
481 {
482 match outcome {
483 TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => {}
484 TaskJoinOutcome::Failed(error) => {
485 self.shutdown_errors
486 .push(format!("Coinbase WebSocket handler failed: {error}"));
487 }
488 TaskJoinOutcome::Incomplete => {
489 self.shutdown_errors
490 .push("Coinbase WebSocket handler did not stop after abort".to_string());
491 }
492 }
493 }
494
495 let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
500
501 loop {
502 let mode_ptr = self.connection_mode.load();
503
504 if ConnectionMode::from_u8(mode_ptr.load(Ordering::Relaxed)).is_closed() {
505 break;
506 }
507
508 if tokio::time::Instant::now() >= deadline {
509 self.shutdown_errors
510 .push("Timed out waiting for WebSocket to reach Closed state".to_string());
511 break;
512 }
513
514 tokio::time::sleep(Duration::from_millis(20)).await;
515 }
516
517 if let Some(control) = &self.socket_control {
518 control.deregister();
519 }
520
521 if self.shutdown_errors.is_empty() {
522 Ok(())
523 } else {
524 let errors = std::mem::take(&mut self.shutdown_errors);
525 anyhow::bail!(errors.join("; "))
526 }
527 }
528
529 #[must_use]
531 pub fn is_active(&self) -> bool {
532 let mode_ptr = self.connection_mode.load();
533 let mode_val = mode_ptr.load(Ordering::Relaxed);
534 ConnectionMode::from_u8(mode_val).is_active()
535 }
536
537 #[must_use]
539 pub fn is_reconnecting(&self) -> bool {
540 let mode_ptr = self.connection_mode.load();
541 let mode_val = mode_ptr.load(Ordering::Relaxed);
542 ConnectionMode::from_u8(mode_val).is_reconnect()
543 }
544
545 #[must_use]
547 pub fn instruments(&self) -> &Arc<AtomicMap<InstrumentId, InstrumentAny>> {
548 &self.instruments
549 }
550
551 #[must_use]
553 pub fn subscription_aliases(&self) -> &Arc<AtomicMap<Ustr, Ustr>> {
554 &self.subscription_aliases
555 }
556
557 pub fn register_subscription_alias(&self, canonical: Ustr, subscribed: Ustr) {
561 self.subscription_aliases.insert(canonical, subscribed);
562 }
563
564 pub fn unregister_subscription_alias(&self, canonical: &Ustr) {
566 self.subscription_aliases.remove(canonical);
567 }
568
569 #[must_use]
571 pub fn subscriptions(&self) -> &SubscriptionState {
572 &self.subscriptions
573 }
574
575 pub async fn update_instrument(&self, instrument: InstrumentAny) {
577 let id = instrument.id();
578 self.instruments.insert(id, instrument.clone());
579
580 let cmd_tx = self.cmd_tx.read().await;
581
582 if let Err(e) = cmd_tx.send(HandlerCommand::UpdateInstrument(Box::new(instrument))) {
583 log::debug!("Failed to send UpdateInstrument: {e}");
584 }
585 }
586
587 pub fn take_out_rx(
591 &mut self,
592 ) -> Option<tokio::sync::mpsc::UnboundedReceiver<NautilusWsMessage>> {
593 self.out_rx.take()
594 }
595
596 pub fn register_bar_type(&mut self, key: String, bar_type: BarType) {
601 self.bar_types.insert(key, bar_type);
602 }
603
604 pub async fn add_bar_type(&mut self, key: String, bar_type: BarType) {
606 self.bar_types.insert(key.clone(), bar_type);
607
608 let cmd_tx = self.cmd_tx.read().await;
609
610 if let Err(e) = cmd_tx.send(HandlerCommand::AddBarType { key, bar_type }) {
611 log::debug!("Failed to send AddBarType: {e}");
612 }
613 }
614}
615
616impl Drop for CoinbaseWebSocketClient {
617 fn drop(&mut self) {
618 if let Some(handle) = self.task_handle.as_ref() {
619 self.signal.store(true, Ordering::Release);
620 handle.abort();
621 }
622 }
623}
624
625fn resubscribe_all(
626 subscriptions: &SubscriptionState,
627 credential: &Option<CoinbaseCredential>,
628 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
629 out_tx: Option<&tokio::sync::mpsc::UnboundedSender<NautilusWsMessage>>,
630) {
631 let topics = subscriptions.all_topics();
632
633 if topics.is_empty() {
634 log::debug!("No active subscriptions to restore");
635 return;
636 }
637
638 log::info!(
639 "Resubscribing to {} topics after reconnection",
640 topics.len()
641 );
642
643 for topic in topics {
644 let (channel, product_id) = match topic.split_once('|') {
645 Some((ch, pid)) => (ch, Some(pid)),
646 None => (topic.as_str(), None),
647 };
648
649 let channel_enum = match CoinbaseWsChannel::from_str(channel) {
650 Ok(ch) => ch,
651 Err(_) => {
652 log::warn!("Unknown channel in topic: {topic}");
653 continue;
654 }
655 };
656
657 let jwt = match credential.as_ref() {
658 Some(c) => match c.build_ws_jwt() {
659 Ok(token) => Some(token),
660 Err(e) => {
661 if channel_enum.requires_auth() {
662 let msg = format!(
663 "JWT required for {channel} but build failed: {e}; topic {topic} not restored"
664 );
665 log::error!("{msg}");
666 if let Some(tx) = out_tx {
667 let _ = tx.send(NautilusWsMessage::Error(msg));
668 }
669 continue;
670 }
671 None
672 }
673 },
674 None => {
675 if channel_enum.requires_auth() {
676 let msg = format!(
677 "JWT required for {channel} but no credentials configured; topic {topic} not restored"
678 );
679 log::error!("{msg}");
680 if let Some(tx) = out_tx {
681 let _ = tx.send(NautilusWsMessage::Error(msg));
682 }
683 continue;
684 }
685 None
686 }
687 };
688
689 let product_ids = match product_id {
690 Some(pid) => vec![Ustr::from(pid)],
691 None => vec![],
692 };
693
694 let sub = CoinbaseWsSubscription {
695 msg_type: CoinbaseWsAction::Subscribe,
696 product_ids,
697 channel: channel_enum,
698 jwt,
699 };
700
701 if let Err(e) = cmd_tx.send(HandlerCommand::Subscribe(sub)) {
702 log::error!("Failed to resubscribe {topic}: {e}");
703 }
704 }
705}
706
707#[cfg(test)]
708mod tests {
709 use nautilus_network::websocket::SubscriptionState;
710 use rstest::rstest;
711
712 use super::*;
713
714 #[rstest]
715 fn test_drop_clone_does_not_signal_handler() {
716 let client = CoinbaseWebSocketClient::new("wss://test", TransportBackend::default(), None);
717 let clone = client.clone();
718
719 drop(clone);
720
721 assert!(!client.signal.load(Ordering::Acquire));
722 }
723
724 #[rstest]
725 fn test_resubscribe_all_product_level_topic() {
726 let subs = SubscriptionState::new('|');
727 subs.mark_subscribe("level2|BTC-USD");
728
729 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
730 resubscribe_all(&subs, &None, &tx, None);
731
732 let cmd = rx.try_recv().unwrap();
733
734 match cmd {
735 HandlerCommand::Subscribe(sub) => {
736 assert_eq!(sub.channel, CoinbaseWsChannel::Level2);
737 assert_eq!(sub.product_ids.len(), 1);
738 assert_eq!(sub.product_ids[0], "BTC-USD");
739 assert!(sub.jwt.is_none());
740 }
741 other => panic!("Expected Subscribe, was {other:?}"),
742 }
743 }
744
745 #[rstest]
746 fn test_resubscribe_all_channel_level_topic() {
747 let subs = SubscriptionState::new('|');
748 subs.mark_subscribe("heartbeats");
749
750 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
751 resubscribe_all(&subs, &None, &tx, None);
752
753 let cmd = rx.try_recv().unwrap();
754
755 match cmd {
756 HandlerCommand::Subscribe(sub) => {
757 assert_eq!(sub.channel, CoinbaseWsChannel::Heartbeats);
758 assert!(sub.product_ids.is_empty());
759 }
760 other => panic!("Expected Subscribe, was {other:?}"),
761 }
762 }
763
764 #[rstest]
765 fn test_resubscribe_all_multiple_topics() {
766 let subs = SubscriptionState::new('|');
767 subs.mark_subscribe("market_trades|BTC-USD");
768 subs.mark_subscribe("ticker|ETH-USD");
769
770 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
771 resubscribe_all(&subs, &None, &tx, None);
772
773 let cmd1 = rx.try_recv().unwrap();
774 let cmd2 = rx.try_recv().unwrap();
775
776 assert!(matches!(cmd1, HandlerCommand::Subscribe(_)));
777 assert!(matches!(cmd2, HandlerCommand::Subscribe(_)));
778 assert!(rx.try_recv().is_err());
779 }
780
781 #[rstest]
782 fn test_resubscribe_all_empty_subscriptions() {
783 let subs = SubscriptionState::new('|');
784
785 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
786 resubscribe_all(&subs, &None, &tx, None);
787
788 assert!(rx.try_recv().is_err());
789 }
790
791 #[rstest]
792 fn test_resubscribe_all_unknown_channel_skipped() {
793 let subs = SubscriptionState::new('|');
794 subs.mark_subscribe("nonexistent_channel|BTC-USD");
795
796 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
797 resubscribe_all(&subs, &None, &tx, None);
798
799 assert!(rx.try_recv().is_err());
800 }
801
802 #[rstest]
803 #[case("level2|BTC-USD", CoinbaseWsChannel::Level2)]
804 #[case("market_trades|ETH-USD", CoinbaseWsChannel::MarketTrades)]
805 #[case("ticker|BTC-USD", CoinbaseWsChannel::Ticker)]
806 #[case("ticker_batch|BTC-USD", CoinbaseWsChannel::TickerBatch)]
807 #[case("candles|BTC-USD", CoinbaseWsChannel::Candles)]
808 #[case("heartbeats", CoinbaseWsChannel::Heartbeats)]
809 #[case("status", CoinbaseWsChannel::Status)]
810 fn test_resubscribe_all_channel_mapping(
811 #[case] topic: &str,
812 #[case] expected_channel: CoinbaseWsChannel,
813 ) {
814 let subs = SubscriptionState::new('|');
815 subs.mark_subscribe(topic);
816
817 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
818 resubscribe_all(&subs, &None, &tx, None);
819
820 let cmd = rx.try_recv().unwrap();
821
822 match cmd {
823 HandlerCommand::Subscribe(sub) => {
824 assert_eq!(sub.channel, expected_channel);
825 }
826 other => panic!("Expected Subscribe, was {other:?}"),
827 }
828 }
829
830 #[rstest]
831 #[case("user|BTC-USD")]
832 #[case("futures_balance_summary")]
833 fn test_resubscribe_all_auth_channel_skipped_without_credentials(#[case] topic: &str) {
834 let subs = SubscriptionState::new('|');
835 subs.mark_subscribe(topic);
836
837 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
838 resubscribe_all(&subs, &None, &tx, None);
839
840 assert!(rx.try_recv().is_err());
842 }
843
844 #[rstest]
845 #[case("user|BTC-USD", "user")]
846 #[case("futures_balance_summary", "futures_balance_summary")]
847 fn test_resubscribe_all_emits_error_for_auth_channel_without_credentials(
848 #[case] topic: &str,
849 #[case] channel: &str,
850 ) {
851 let subs = SubscriptionState::new('|');
852 subs.mark_subscribe(topic);
853
854 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
855 let (out_tx, mut out_rx) = tokio::sync::mpsc::unbounded_channel();
856 resubscribe_all(&subs, &None, &cmd_tx, Some(&out_tx));
857
858 assert!(cmd_rx.try_recv().is_err());
860
861 let msg = out_rx
862 .try_recv()
863 .expect("Error event must be emitted when auth channel cannot resubscribe");
864 match msg {
865 NautilusWsMessage::Error(text) => {
866 assert!(
867 text.contains(channel),
868 "error must mention the channel, was: {text}"
869 );
870 assert!(
871 text.contains(topic),
872 "error must mention the topic, was: {text}"
873 );
874 }
875 other => panic!("expected Error variant, was {other:?}"),
876 }
877 }
878
879 #[rstest]
880 fn test_resubscribe_all_emits_error_when_jwt_build_fails() {
881 let subs = SubscriptionState::new('|');
882 let topic = "user|BTC-USD";
883 subs.mark_subscribe(topic);
884
885 let bad_credential = Some(CoinbaseCredential::new(
888 "organizations/test/apiKeys/test".to_string(),
889 "not-a-pem-key".to_string(),
890 ));
891
892 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
893 let (out_tx, mut out_rx) = tokio::sync::mpsc::unbounded_channel();
894 resubscribe_all(&subs, &bad_credential, &cmd_tx, Some(&out_tx));
895
896 assert!(cmd_rx.try_recv().is_err(), "no subscribe should be sent");
897 let msg = out_rx
898 .try_recv()
899 .expect("Error event must be emitted when JWT build fails for an auth channel");
900 match msg {
901 NautilusWsMessage::Error(text) => {
902 assert!(text.contains("user"), "error must mention channel: {text}");
903 assert!(text.contains(topic), "error must mention topic: {text}");
904 }
905 other => panic!("expected Error variant, was {other:?}"),
906 }
907 }
908
909 #[rstest]
910 fn test_prime_default_subscriptions_marks_heartbeats() {
911 let client = CoinbaseWebSocketClient::new("wss://test", TransportBackend::default(), None);
912 assert!(client.subscriptions.all_topics().is_empty());
913
914 client.prime_default_subscriptions();
915
916 let topics = client.subscriptions.all_topics();
917 assert!(topics.iter().any(|t| t == "heartbeats"), "{topics:?}");
918 }
919
920 #[rstest]
921 fn test_ws_quotas_match_documented_limits() {
922 assert_eq!(COINBASE_WS_CONNECTION_QUOTA.burst_size().get(), 8);
923 assert_eq!(COINBASE_WS_SUBSCRIPTION_QUOTA.burst_size().get(), 8);
924 }
925
926 #[rstest]
927 fn test_ws_subscription_rate_limit_key_is_stable() {
928 assert_eq!(COINBASE_RATE_LIMIT_KEY_SUBSCRIPTION, "subscription");
929 assert_eq!(
930 COINBASE_WS_SUBSCRIPTION_KEYS[0].as_str(),
931 COINBASE_RATE_LIMIT_KEY_SUBSCRIPTION,
932 );
933 }
934}