1use std::{
19 collections::HashMap,
20 sync::{
21 Arc, RwLock,
22 atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering},
23 },
24};
25
26use arc_swap::ArcSwap;
27use nautilus_common::live::get_runtime;
28use nautilus_core::AtomicMap;
29use nautilus_model::{
30 data::BarType,
31 enums::BarAggregation,
32 identifiers::{AccountId, ClientOrderId, InstrumentId, StrategyId, TraderId, VenueOrderId},
33 instruments::{Instrument, InstrumentAny},
34};
35use nautilus_network::{
36 mode::ConnectionMode,
37 websocket::{
38 AuthTracker, SubscriptionState, TransportBackend, WebSocketClient, WebSocketConfig,
39 channel_message_handler,
40 },
41};
42use tokio_util::sync::CancellationToken;
43use ustr::Ustr;
44
45pub const KRAKEN_SPOT_WS_TOPIC_DELIMITER: char = ':';
49
50use super::{
51 enums::{KrakenWsChannel, KrakenWsMethod},
52 handler::{SpotFeedHandler, SpotHandlerCommand},
53 level_2::L2Depths,
54 messages::{KrakenSpotWsMessage, KrakenWsChannelParams, KrakenWsParams, KrakenWsRequest},
55};
56use crate::{
57 common::{
58 consts::{
59 KRAKEN_RATE_LIMIT_KEY_ORDER, KRAKEN_RATE_LIMIT_KEY_SUBSCRIPTION,
60 KRAKEN_SPOT_WS_ORDER_QUOTA, KRAKEN_SPOT_WS_SUBSCRIPTION_QUOTA,
61 },
62 parse::normalize_spot_symbol,
63 },
64 config::KrakenDataClientConfig,
65 http::{KrakenSpotHttpClient, spot::client::KRAKEN_SPOT_DEFAULT_RATE_LIMIT_PER_SECOND},
66 websocket::error::KrakenWsError,
67};
68
69const WS_PING_MSG: &str = r#"{"method":"ping"}"#;
70
71#[derive(Debug)]
73#[cfg_attr(
74 feature = "python",
75 pyo3::pyclass(module = "nautilus_trader.core.nautilus_pyo3.kraken", from_py_object)
76)]
77#[cfg_attr(
78 feature = "python",
79 pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.adapters.kraken")
80)]
81pub struct KrakenSpotWebSocketClient {
82 url: String,
83 config: KrakenDataClientConfig,
84 signal: Arc<AtomicBool>,
85 connection_mode: Arc<ArcSwap<AtomicU8>>,
86 cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<SpotHandlerCommand>>>,
87 out_rx: Option<Arc<tokio::sync::mpsc::UnboundedReceiver<KrakenSpotWsMessage>>>,
88 task_handle: Option<Arc<tokio::task::JoinHandle<()>>>,
89 subscriptions: SubscriptionState,
90 subscription_payloads: Arc<tokio::sync::RwLock<HashMap<String, String>>>,
91 auth_tracker: AuthTracker,
92 cancellation_token: CancellationToken,
93 req_id_counter: Arc<AtomicU64>,
94 auth_token: Arc<tokio::sync::RwLock<Option<String>>>,
95 account_id: Arc<RwLock<Option<AccountId>>>,
96 truncated_id_map: Arc<AtomicMap<String, ClientOrderId>>,
97 instruments: Arc<AtomicMap<InstrumentId, InstrumentAny>>,
98 l2_depths: L2Depths,
99 l3_depths: Arc<std::sync::Mutex<ahash::AHashMap<String, u32>>>,
100 transport_backend: TransportBackend,
101 proxy_url: Option<String>,
102}
103
104impl Clone for KrakenSpotWebSocketClient {
105 fn clone(&self) -> Self {
106 Self {
107 url: self.url.clone(),
108 config: self.config.clone(),
109 signal: Arc::clone(&self.signal),
110 connection_mode: Arc::clone(&self.connection_mode),
111 cmd_tx: Arc::clone(&self.cmd_tx),
112 out_rx: self.out_rx.clone(),
113 task_handle: self.task_handle.clone(),
114 subscriptions: self.subscriptions.clone(),
115 subscription_payloads: Arc::clone(&self.subscription_payloads),
116 auth_tracker: self.auth_tracker.clone(),
117 cancellation_token: self.cancellation_token.clone(),
118 req_id_counter: self.req_id_counter.clone(),
119 auth_token: self.auth_token.clone(),
120 account_id: Arc::clone(&self.account_id),
121 truncated_id_map: Arc::clone(&self.truncated_id_map),
122 instruments: Arc::clone(&self.instruments),
123 l2_depths: self.l2_depths.clone(),
124 l3_depths: Arc::clone(&self.l3_depths),
125 transport_backend: self.transport_backend,
126 proxy_url: self.proxy_url.clone(),
127 }
128 }
129}
130
131impl KrakenSpotWebSocketClient {
132 pub fn new(
134 config: KrakenDataClientConfig,
135 cancellation_token: CancellationToken,
136 proxy_url: Option<String>,
137 ) -> Self {
138 let url = if config.ws_private_url.is_some() {
139 config.ws_private_url()
140 } else {
141 config.ws_public_url()
142 };
143 Self::new_with_url(url, config, cancellation_token, proxy_url)
144 }
145
146 pub fn l3(
151 config: KrakenDataClientConfig,
152 cancellation_token: CancellationToken,
153 proxy_url: Option<String>,
154 ) -> Self {
155 let url = config.ws_l3_url();
156 Self::new_with_url(url, config, cancellation_token, proxy_url)
157 }
158
159 fn new_with_url(
160 url: String,
161 mut config: KrakenDataClientConfig,
162 cancellation_token: CancellationToken,
163 proxy_url: Option<String>,
164 ) -> Self {
165 let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel::<SpotHandlerCommand>();
166 let initial_mode = AtomicU8::new(ConnectionMode::Closed.as_u8());
167 let connection_mode = Arc::new(ArcSwap::from_pointee(initial_mode));
168
169 let transport_backend = config.transport_backend;
170 config.proxy_url = proxy_url.clone();
171
172 Self {
173 url,
174 config,
175 signal: Arc::new(AtomicBool::new(false)),
176 connection_mode,
177 cmd_tx: Arc::new(tokio::sync::RwLock::new(cmd_tx)),
178 out_rx: None,
179 task_handle: None,
180 subscriptions: SubscriptionState::new(KRAKEN_SPOT_WS_TOPIC_DELIMITER),
181 subscription_payloads: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
182 auth_tracker: AuthTracker::new(),
183 cancellation_token,
184 req_id_counter: Arc::new(AtomicU64::new(0)),
185 auth_token: Arc::new(tokio::sync::RwLock::new(None)),
186 account_id: Arc::new(RwLock::new(None)),
187 truncated_id_map: Arc::new(AtomicMap::new()),
188 instruments: Arc::new(AtomicMap::new()),
189 l2_depths: L2Depths::default(),
190 l3_depths: Arc::new(std::sync::Mutex::new(ahash::AHashMap::new())),
191 transport_backend,
192 proxy_url,
193 }
194 }
195
196 fn get_next_req_id(&self) -> u64 {
197 self.req_id_counter.fetch_add(1, Ordering::Relaxed) + 1
198 }
199
200 pub fn req_id_counter(&self) -> Arc<AtomicU64> {
202 self.req_id_counter.clone()
203 }
204
205 pub async fn handler_command_sender(
207 &self,
208 ) -> tokio::sync::mpsc::UnboundedSender<SpotHandlerCommand> {
209 self.cmd_tx.read().await.clone()
210 }
211
212 pub fn handler_command_handle(
217 &self,
218 ) -> Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<SpotHandlerCommand>>> {
219 self.cmd_tx.clone()
220 }
221
222 pub async fn auth_token(&self) -> Option<String> {
224 self.auth_token.read().await.clone()
225 }
226
227 pub fn auth_token_blocking(&self) -> Option<String> {
233 self.auth_token.try_read().ok().and_then(|g| g.clone())
234 }
235
236 pub fn auth_token_handle(&self) -> Arc<tokio::sync::RwLock<Option<String>>> {
240 self.auth_token.clone()
241 }
242
243 pub async fn connect(&mut self) -> Result<(), KrakenWsError> {
245 log::debug!("Connecting to {}", self.url);
246
247 self.signal.store(false, Ordering::Relaxed);
248
249 let (raw_handler, raw_rx) = channel_message_handler();
250
251 let ws_config = WebSocketConfig {
252 url: self.url.clone(),
253 headers: vec![],
254 heartbeat: Some(self.config.heartbeat_interval_secs),
255 heartbeat_msg: Some(WS_PING_MSG.to_string()),
256 reconnect_timeout_ms: Some(5_000),
257 reconnect_delay_initial_ms: Some(500),
258 reconnect_delay_max_ms: Some(5_000),
259 reconnect_backoff_factor: Some(1.5),
260 reconnect_jitter_ms: Some(250),
261 reconnect_max_attempts: None,
262 idle_timeout_ms: (self.config.ws_idle_timeout_ms != 0)
265 .then_some(self.config.ws_idle_timeout_ms),
266 backend: self.transport_backend,
267 proxy_url: self.proxy_url.clone(),
268 };
269
270 let keyed_quotas = vec![
271 (
272 KRAKEN_RATE_LIMIT_KEY_SUBSCRIPTION[0].to_string(),
273 *KRAKEN_SPOT_WS_SUBSCRIPTION_QUOTA,
274 ),
275 (
276 KRAKEN_RATE_LIMIT_KEY_ORDER[0].to_string(),
277 *KRAKEN_SPOT_WS_ORDER_QUOTA,
278 ),
279 ];
280
281 let ws_client = WebSocketClient::connect(
282 ws_config,
283 Some(raw_handler),
284 None, None, keyed_quotas,
287 None,
288 )
289 .await
290 .map_err(|e| KrakenWsError::ConnectionError(e.to_string()))?;
291
292 self.connection_mode
294 .store(ws_client.connection_mode_atomic());
295
296 let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<KrakenSpotWsMessage>();
297 self.out_rx = Some(Arc::new(out_rx));
298
299 let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<SpotHandlerCommand>();
300 *self.cmd_tx.write().await = cmd_tx.clone();
301
302 if let Err(e) = cmd_tx.send(SpotHandlerCommand::SetClient(ws_client)) {
303 return Err(KrakenWsError::ConnectionError(format!(
304 "Failed to send WebSocketClient to handler: {e}"
305 )));
306 }
307
308 let signal = self.signal.clone();
309 let subscriptions = self.subscriptions.clone();
310 let subscription_payloads = self.subscription_payloads.clone();
311 let config_for_reconnect = self.config.clone();
312 let auth_token_for_reconnect = self.auth_token.clone();
313 let auth_tracker_for_reconnect = self.auth_tracker.clone();
314 let cmd_tx_for_reconnect = cmd_tx.clone();
315
316 let stream_handle = get_runtime().spawn(async move {
317 let mut handler =
318 SpotFeedHandler::new(signal.clone(), cmd_rx, raw_rx, subscriptions.clone());
319
320 loop {
321 match handler.next().await {
322 Some(KrakenSpotWsMessage::Reconnected) => {
323 if signal.load(Ordering::Relaxed) {
324 continue;
325 }
326 log::info!("WebSocket reconnected, resubscribing");
327
328 let confirmed_topics = subscriptions.all_topics();
329 for topic in &confirmed_topics {
330 subscriptions.mark_failure(topic);
331 }
332
333 let payloads = subscription_payloads.read().await;
334 if payloads.is_empty() {
335 log::debug!("No subscriptions to restore after reconnection");
336 } else {
337 let had_auth = auth_token_for_reconnect.read().await.is_some();
338
339 if had_auth && config_for_reconnect.has_api_credentials() {
340 log::debug!("Re-authenticating after reconnect");
341
342 auth_tracker_for_reconnect.invalidate();
343 let _rx = auth_tracker_for_reconnect.begin();
344
345 match refresh_auth_token(&config_for_reconnect).await {
346 Ok(new_token) => {
347 *auth_token_for_reconnect.write().await = Some(new_token);
348 auth_tracker_for_reconnect.succeed();
349 log::debug!("Re-authentication successful");
350 }
351 Err(e) => {
352 log::error!(
353 "Failed to re-authenticate after reconnect: {e}"
354 );
355 *auth_token_for_reconnect.write().await = None;
356 auth_tracker_for_reconnect.fail(e.to_string());
357 }
358 }
359 }
360
361 log::debug!(
362 "Resubscribing after reconnection: count={}",
363 payloads.len()
364 );
365
366 for (topic, payload) in payloads.iter() {
367 let needs_token =
368 topic == "executions" || topic.starts_with("level3:");
369 let payload = if needs_token {
370 let auth_token = auth_token_for_reconnect.read().await.clone();
371 match auth_token {
372 Some(token) => {
373 match update_auth_token_in_payload(payload, &token) {
374 Ok(p) => p,
375 Err(e) => {
376 log::error!("Failed to update auth token: {e}");
377 continue;
378 }
379 }
380 }
381 None => {
382 log::warn!(
383 "Cannot resubscribe to {topic}: no auth token"
384 );
385 continue;
386 }
387 }
388 } else {
389 payload.clone()
390 };
391
392 if let Err(e) = cmd_tx_for_reconnect
393 .send(SpotHandlerCommand::Subscribe { payload })
394 {
395 log::error!(
396 "Failed to send resubscribe command: error={e}, \
397 topic={topic}"
398 );
399 }
400
401 subscriptions.mark_subscribe(topic);
402 }
403 }
404
405 if out_tx.send(KrakenSpotWsMessage::Reconnected).is_err() {
406 if handler.is_stopped() {
407 log::debug!("Failed to send message (receiver dropped)");
408 } else {
409 log::error!("Failed to send message (receiver dropped)");
410 }
411 break;
412 }
413 }
414 Some(msg) => {
415 if out_tx.send(msg).is_err() {
416 if handler.is_stopped() {
417 log::debug!("Failed to send message (receiver dropped)");
418 } else {
419 log::error!("Failed to send message (receiver dropped)");
420 }
421 break;
422 }
423 }
424 None => {
425 if handler.is_stopped() {
426 log::debug!("Stop signal received, ending message processing");
427 break;
428 }
429 log::warn!("WebSocket stream ended unexpectedly");
430 break;
431 }
432 }
433 }
434
435 log::debug!("Handler task exiting");
436 });
437
438 self.task_handle = Some(Arc::new(stream_handle));
439
440 log::debug!("WebSocket connected successfully");
441 Ok(())
442 }
443
444 pub async fn disconnect(&mut self) -> Result<(), KrakenWsError> {
446 log::debug!("Disconnecting WebSocket");
447
448 self.signal.store(true, Ordering::Relaxed);
449
450 if let Err(e) = self
451 .cmd_tx
452 .read()
453 .await
454 .send(SpotHandlerCommand::Disconnect)
455 {
456 log::debug!(
457 "Failed to send disconnect command (handler may already be shut down): {e}"
458 );
459 }
460
461 if let Some(task_handle) = self.task_handle.take() {
462 match Arc::try_unwrap(task_handle) {
463 Ok(handle) => {
464 log::debug!("Waiting for task handle to complete");
465 match tokio::time::timeout(tokio::time::Duration::from_secs(2), handle).await {
466 Ok(Ok(())) => log::debug!("Task handle completed successfully"),
467 Ok(Err(e)) => log::error!("Task handle encountered an error: {e:?}"),
468 Err(_) => {
469 log::warn!(
470 "Timeout waiting for task handle, task may still be running"
471 );
472 }
473 }
474 }
475 Err(arc_handle) => {
476 log::debug!(
477 "Cannot take ownership of task handle - other references exist, aborting task"
478 );
479 arc_handle.abort();
480 }
481 }
482 } else {
483 log::debug!("No task handle to await");
484 }
485
486 self.subscriptions.clear();
487 self.subscription_payloads.write().await.clear();
488 self.auth_tracker.fail("Disconnected");
489
490 if let Ok(mut depths) = self.l3_depths.lock() {
491 depths.clear();
492 }
493 self.l2_depths.clear();
494
495 Ok(())
496 }
497
498 pub async fn close(&mut self) -> Result<(), KrakenWsError> {
500 self.disconnect().await
501 }
502
503 pub async fn wait_until_active(&self, timeout_secs: f64) -> Result<(), KrakenWsError> {
505 let timeout = tokio::time::Duration::from_secs_f64(timeout_secs);
506
507 tokio::time::timeout(timeout, async {
508 while !self.is_active() {
509 tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
510 }
511 })
512 .await
513 .map_err(|_| {
514 KrakenWsError::ConnectionError(format!(
515 "WebSocket connection timeout after {timeout_secs} seconds"
516 ))
517 })?;
518
519 Ok(())
520 }
521
522 #[must_use]
524 pub fn is_authenticated(&self) -> bool {
525 self.auth_tracker.is_authenticated()
526 }
527
528 pub async fn wait_until_authenticated(&self, timeout_secs: f64) -> Result<(), KrakenWsError> {
532 let timeout = tokio::time::Duration::from_secs_f64(timeout_secs);
533
534 if self.auth_tracker.wait_for_authenticated(timeout).await {
535 Ok(())
536 } else {
537 Err(KrakenWsError::AuthenticationError(format!(
538 "Authentication not completed within {timeout_secs} seconds"
539 )))
540 }
541 }
542
543 pub async fn authenticate(&self) -> Result<(), KrakenWsError> {
545 if !self.config.has_api_credentials() {
546 return Err(KrakenWsError::AuthenticationError(
547 "API credentials required for authentication".to_string(),
548 ));
549 }
550
551 let _receiver = self.auth_tracker.begin();
552
553 match refresh_auth_token(&self.config).await {
554 Ok(token) => {
555 *self.auth_token.write().await = Some(token);
556 self.auth_tracker.succeed();
557 Ok(())
558 }
559 Err(e) => {
560 *self.auth_token.write().await = None;
561 self.auth_tracker.fail(e.to_string());
562 Err(e)
563 }
564 }
565 }
566
567 pub fn cancel_all_requests(&self) {
569 self.cancellation_token.cancel();
570 }
571
572 pub fn cancellation_token(&self) -> &CancellationToken {
574 &self.cancellation_token
575 }
576
577 pub async fn subscribe(
579 &self,
580 channel: KrakenWsChannel,
581 symbols: Vec<Ustr>,
582 depth: Option<u32>,
583 ) -> Result<(), KrakenWsError> {
584 if matches!(channel, KrakenWsChannel::Level3) {
585 return Err(KrakenWsError::InvalidMessage(
586 "Use subscribe_book_l3 / unsubscribe_book_l3 for the Level3 channel".to_string(),
587 ));
588 }
589 let mut symbols_to_subscribe = Vec::new();
590 let channel_str = channel.as_ref();
591 for symbol in &symbols {
592 let key = format!("{channel_str}:{symbol}");
593 if self.subscriptions.add_reference(&key) {
594 self.subscriptions.mark_subscribe(&key);
595 symbols_to_subscribe.push(*symbol);
596 }
597 }
598
599 if symbols_to_subscribe.is_empty() {
600 return Ok(());
601 }
602
603 let is_private = matches!(
604 channel,
605 KrakenWsChannel::Executions | KrakenWsChannel::Balances
606 );
607 let token = if is_private {
608 Some(self.auth_token.read().await.clone().ok_or_else(|| {
609 KrakenWsError::AuthenticationError(
610 "Authentication token required for private channels. Call authenticate() first"
611 .to_string(),
612 )
613 })?)
614 } else {
615 None
616 };
617
618 let req_id = self.get_next_req_id();
619 let request = KrakenWsRequest {
620 method: KrakenWsMethod::Subscribe,
621 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
622 channel,
623 symbol: Some(symbols_to_subscribe.clone()),
624 snapshot: None,
625 depth,
626 interval: None,
627 event_trigger: None,
628 token,
629 snap_orders: None,
630 snap_trades: None,
631 })),
632 req_id: Some(req_id),
633 };
634
635 let payload = self.send_command(&request).await?;
636
637 for symbol in &symbols_to_subscribe {
638 let key = format!("{channel_str}:{symbol}");
639 self.subscriptions.confirm_subscribe(&key);
640 self.subscription_payloads
641 .write()
642 .await
643 .insert(key, payload.clone());
644 }
645
646 Ok(())
647 }
648
649 async fn subscribe_with_interval(
651 &self,
652 channel: KrakenWsChannel,
653 symbols: Vec<Ustr>,
654 interval: u32,
655 ) -> Result<(), KrakenWsError> {
656 let mut symbols_to_subscribe = Vec::new();
657 let channel_str = channel.as_ref();
658 for symbol in &symbols {
659 let key = format!("{channel_str}:{symbol}:{interval}");
660 if self.subscriptions.add_reference(&key) {
661 self.subscriptions.mark_subscribe(&key);
662 symbols_to_subscribe.push(*symbol);
663 }
664 }
665
666 if symbols_to_subscribe.is_empty() {
667 return Ok(());
668 }
669
670 let req_id = self.get_next_req_id();
671 let request = KrakenWsRequest {
672 method: KrakenWsMethod::Subscribe,
673 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
674 channel,
675 symbol: Some(symbols_to_subscribe.clone()),
676 snapshot: Some(false),
677 depth: None,
678 interval: Some(interval),
679 event_trigger: None,
680 token: None,
681 snap_orders: None,
682 snap_trades: None,
683 })),
684 req_id: Some(req_id),
685 };
686
687 let payload = self.send_command(&request).await?;
688
689 for symbol in &symbols_to_subscribe {
690 let key = format!("{channel_str}:{symbol}:{interval}");
691 self.subscriptions.confirm_subscribe(&key);
692 self.subscription_payloads
693 .write()
694 .await
695 .insert(key, payload.clone());
696 }
697
698 Ok(())
699 }
700
701 async fn unsubscribe_with_interval(
703 &self,
704 channel: KrakenWsChannel,
705 symbols: Vec<Ustr>,
706 interval: u32,
707 ) -> Result<(), KrakenWsError> {
708 let mut symbols_to_unsubscribe = Vec::new();
709 let channel_str = channel.as_ref();
710 for symbol in &symbols {
711 let key = format!("{channel_str}:{symbol}:{interval}");
712 if self.subscriptions.remove_reference(&key) {
713 self.subscriptions.mark_unsubscribe(&key);
714 symbols_to_unsubscribe.push(*symbol);
715 }
716 }
717
718 if symbols_to_unsubscribe.is_empty() {
719 return Ok(());
720 }
721
722 let req_id = self.get_next_req_id();
723 let request = KrakenWsRequest {
724 method: KrakenWsMethod::Unsubscribe,
725 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
726 channel,
727 symbol: Some(symbols_to_unsubscribe.clone()),
728 snapshot: None,
729 depth: None,
730 interval: Some(interval),
731 event_trigger: None,
732 token: None,
733 snap_orders: None,
734 snap_trades: None,
735 })),
736 req_id: Some(req_id),
737 };
738
739 self.send_command(&request).await?;
740
741 for symbol in &symbols_to_unsubscribe {
742 let key = format!("{channel_str}:{symbol}:{interval}");
743 self.subscriptions.confirm_unsubscribe(&key);
744 self.subscription_payloads.write().await.remove(&key);
745 }
746
747 Ok(())
748 }
749
750 pub async fn unsubscribe(
752 &self,
753 channel: KrakenWsChannel,
754 symbols: Vec<Ustr>,
755 ) -> Result<(), KrakenWsError> {
756 if matches!(channel, KrakenWsChannel::Level3) {
757 return Err(KrakenWsError::InvalidMessage(
758 "Use subscribe_book_l3 / unsubscribe_book_l3 for the Level3 channel".to_string(),
759 ));
760 }
761 let mut symbols_to_unsubscribe = Vec::new();
762 let channel_str = channel.as_ref();
763 for symbol in &symbols {
764 let key = format!("{channel_str}:{symbol}");
765 if self.subscriptions.remove_reference(&key) {
766 self.subscriptions.mark_unsubscribe(&key);
767 symbols_to_unsubscribe.push(*symbol);
768 } else {
769 log::debug!(
770 "Channel {channel_str} symbol {symbol} still has active subscriptions, not unsubscribing"
771 );
772 }
773 }
774
775 if symbols_to_unsubscribe.is_empty() {
776 return Ok(());
777 }
778
779 let is_private = matches!(
780 channel,
781 KrakenWsChannel::Executions | KrakenWsChannel::Balances
782 );
783 let token = if is_private {
784 Some(self.auth_token.read().await.clone().ok_or_else(|| {
785 KrakenWsError::AuthenticationError(
786 "Authentication token required for private channels. Call authenticate() first"
787 .to_string(),
788 )
789 })?)
790 } else {
791 None
792 };
793
794 let req_id = self.get_next_req_id();
795 let request = KrakenWsRequest {
796 method: KrakenWsMethod::Unsubscribe,
797 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
798 channel,
799 symbol: Some(symbols_to_unsubscribe.clone()),
800 snapshot: None,
801 depth: None,
802 interval: None,
803 event_trigger: None,
804 token,
805 snap_orders: None,
806 snap_trades: None,
807 })),
808 req_id: Some(req_id),
809 };
810
811 self.send_command(&request).await?;
812
813 for symbol in &symbols_to_unsubscribe {
814 let key = format!("{channel_str}:{symbol}");
815 self.subscriptions.confirm_unsubscribe(&key);
816 self.subscription_payloads.write().await.remove(&key);
817 }
818
819 Ok(())
820 }
821
822 pub async fn send_ping(&self) -> Result<(), KrakenWsError> {
824 let req_id = self.get_next_req_id();
825
826 let request = KrakenWsRequest {
827 method: KrakenWsMethod::Ping,
828 params: None,
829 req_id: Some(req_id),
830 };
831
832 self.send_command(&request).await?;
833 Ok(())
834 }
835
836 async fn send_command(&self, request: &KrakenWsRequest) -> Result<String, KrakenWsError> {
837 let payload =
838 serde_json::to_string(request).map_err(|e| KrakenWsError::JsonError(e.to_string()))?;
839
840 log::trace!("Sending message: {payload}");
841
842 let cmd = match request.method {
843 KrakenWsMethod::Subscribe => SpotHandlerCommand::Subscribe {
844 payload: payload.clone(),
845 },
846 KrakenWsMethod::Unsubscribe => SpotHandlerCommand::Unsubscribe {
847 payload: payload.clone(),
848 },
849 KrakenWsMethod::Ping | KrakenWsMethod::Pong => SpotHandlerCommand::Ping {
850 payload: payload.clone(),
851 },
852 KrakenWsMethod::AddOrder
853 | KrakenWsMethod::AmendOrder
854 | KrakenWsMethod::CancelOrder
855 | KrakenWsMethod::BatchAdd => {
856 return Err(KrakenWsError::InvalidMessage(
857 "Order methods must not be sent via send_command; use the dedicated order submission path".to_string()
858 ));
859 }
860 };
861
862 self.cmd_tx
863 .read()
864 .await
865 .send(cmd)
866 .map_err(|e| KrakenWsError::ConnectionError(format!("Failed to send request: {e}")))?;
867
868 Ok(payload)
869 }
870
871 pub fn is_connected(&self) -> bool {
873 let connection_mode_arc = self.connection_mode.load();
874 !ConnectionMode::from_atomic(&connection_mode_arc).is_closed()
875 }
876
877 pub fn is_active(&self) -> bool {
879 let connection_mode_arc = self.connection_mode.load();
880 ConnectionMode::from_atomic(&connection_mode_arc).is_active()
881 && !self.signal.load(Ordering::Relaxed)
882 }
883
884 pub fn is_closed(&self) -> bool {
886 let connection_mode_arc = self.connection_mode.load();
887 ConnectionMode::from_atomic(&connection_mode_arc).is_closed()
888 || self.signal.load(Ordering::Relaxed)
889 }
890
891 pub fn url(&self) -> &str {
893 &self.url
894 }
895
896 pub fn get_subscriptions(&self) -> Vec<String> {
898 self.subscriptions.all_topics()
899 }
900
901 pub fn subscriptions_contains(&self, topic: &str) -> bool {
903 self.subscriptions.all_topics().iter().any(|t| t == topic)
904 }
905
906 pub fn set_account_id(&self, account_id: AccountId) {
908 if let Ok(mut guard) = self.account_id.write() {
909 *guard = Some(account_id);
910 }
911 }
912
913 #[must_use]
915 pub fn account_id(&self) -> Option<AccountId> {
916 self.account_id.read().ok().and_then(|g| *g)
917 }
918
919 pub fn cache_instrument(&self, instrument: InstrumentAny) {
921 self.instruments.insert(instrument.id(), instrument);
922 }
923
924 pub fn account_id_shared(&self) -> &Arc<RwLock<Option<AccountId>>> {
926 &self.account_id
927 }
928
929 pub fn truncated_id_map(&self) -> &Arc<AtomicMap<String, ClientOrderId>> {
931 &self.truncated_id_map
932 }
933
934 pub fn cache_client_order(
936 &self,
937 client_order_id: ClientOrderId,
938 _venue_order_id: Option<VenueOrderId>,
939 _instrument_id: InstrumentId,
940 _trader_id: TraderId,
941 _strategy_id: StrategyId,
942 ) {
943 let truncated = crate::common::parse::truncate_cl_ord_id(&client_order_id);
944
945 if truncated != client_order_id.as_str() {
946 self.truncated_id_map.insert(truncated, client_order_id);
947 }
948 }
949
950 pub fn stream(
958 &mut self,
959 ) -> Result<impl futures_util::Stream<Item = KrakenSpotWsMessage> + use<>, KrakenWsError> {
960 let rx = self.out_rx.take().ok_or_else(|| {
961 KrakenWsError::ChannelError(
962 "Stream receiver already taken or client not connected".to_string(),
963 )
964 })?;
965 let mut rx = Arc::try_unwrap(rx).map_err(|_| {
966 KrakenWsError::ChannelError(
967 "Cannot take ownership of stream - other client clones still hold references"
968 .to_string(),
969 )
970 })?;
971 Ok(async_stream::stream! {
972 while let Some(msg) = rx.recv().await {
973 yield msg;
974 }
975 })
976 }
977
978 pub async fn subscribe_book(
980 &self,
981 instrument_id: InstrumentId,
982 depth: Option<u32>,
983 ) -> Result<(), KrakenWsError> {
984 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
985 let depth = depth.unwrap_or(10);
986
987 if !matches!(depth, 10 | 25 | 100 | 500 | 1000) {
988 return Err(KrakenWsError::InvalidMessage(format!(
989 "Invalid L2 depth {depth}, valid values: 10, 25, 100, 500, 1000",
990 )));
991 }
992
993 let channel_str = KrakenWsChannel::Book.as_ref();
994 let key = format!("{channel_str}:{symbol}");
995 let is_first_reference = self.subscriptions.add_reference(&key);
996
997 if !is_first_reference {
998 let existing_depth = self.l2_depths.get(symbol.as_str());
999
1000 if existing_depth != Some(depth) {
1001 self.subscriptions.remove_reference(&key);
1002 return Err(KrakenWsError::InvalidMessage(format!(
1003 "L2 subscription for {symbol} already exists with depth \
1004 {existing_depth:?}, cannot resubscribe with depth {depth}",
1005 )));
1006 }
1007 return Ok(());
1008 }
1009
1010 self.subscriptions.mark_subscribe(&key);
1011 self.l2_depths.insert(symbol.as_str(), depth);
1012
1013 let req_id = self.get_next_req_id();
1014 let request = KrakenWsRequest {
1015 method: KrakenWsMethod::Subscribe,
1016 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1017 channel: KrakenWsChannel::Book,
1018 symbol: Some(vec![symbol]),
1019 snapshot: None,
1020 depth: Some(depth),
1021 interval: None,
1022 event_trigger: None,
1023 token: None,
1024 snap_orders: None,
1025 snap_trades: None,
1026 })),
1027 req_id: Some(req_id),
1028 };
1029
1030 let payload = match self.send_command(&request).await {
1031 Ok(payload) => payload,
1032 Err(e) => {
1033 self.l2_depths.remove(symbol.as_str());
1034 self.subscriptions.remove_reference(&key);
1035 self.subscriptions.mark_unsubscribe(&key);
1036 self.subscriptions.confirm_unsubscribe(&key);
1037 return Err(e);
1038 }
1039 };
1040
1041 self.subscriptions.confirm_subscribe(&key);
1042 self.subscription_payloads
1043 .write()
1044 .await
1045 .insert(key, payload);
1046 Ok(())
1047 }
1048
1049 pub async fn subscribe_book_l3(&self, symbol: Ustr, depth: u32) -> Result<(), KrakenWsError> {
1068 if !matches!(depth, 10 | 100 | 1000) {
1069 return Err(KrakenWsError::InvalidMessage(format!(
1070 "Invalid L3 depth {depth}, valid values: 10, 100, 1000",
1071 )));
1072 }
1073
1074 let token = self.auth_token.read().await.clone().ok_or_else(|| {
1075 KrakenWsError::AuthenticationError(
1076 "Authentication token required for level3. Call authenticate() first".to_string(),
1077 )
1078 })?;
1079
1080 let channel_str = KrakenWsChannel::Level3.as_ref();
1081 let key = format!("{channel_str}:{symbol}");
1082
1083 let is_first_reference = self.subscriptions.add_reference(&key);
1084
1085 if !is_first_reference {
1086 let existing_depth = self
1087 .l3_depths
1088 .lock()
1089 .expect("L3 depth map mutex poisoned")
1090 .get(symbol.as_str())
1091 .copied();
1092
1093 if existing_depth != Some(depth) {
1094 self.subscriptions.remove_reference(&key);
1095 return Err(KrakenWsError::InvalidMessage(format!(
1096 "L3 subscription for {symbol} already exists with depth \
1097 {existing_depth:?}, cannot resubscribe with depth {depth}",
1098 )));
1099 }
1100 return Ok(());
1101 }
1102
1103 self.subscriptions.mark_subscribe(&key);
1104
1105 self.l3_depths
1106 .lock()
1107 .expect("L3 depth map mutex poisoned")
1108 .insert(symbol.to_string(), depth);
1109
1110 let req_id = self.get_next_req_id();
1111 let request = KrakenWsRequest {
1112 method: KrakenWsMethod::Subscribe,
1113 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1114 channel: KrakenWsChannel::Level3,
1115 symbol: Some(vec![symbol]),
1116 snapshot: Some(true),
1117 depth: Some(depth),
1118 interval: None,
1119 event_trigger: None,
1120 token: Some(token),
1121 snap_orders: None,
1122 snap_trades: None,
1123 })),
1124 req_id: Some(req_id),
1125 };
1126
1127 let payload = match self.send_command(&request).await {
1128 Ok(p) => p,
1129 Err(e) => {
1130 self.l3_depths
1131 .lock()
1132 .expect("L3 depth map mutex poisoned")
1133 .remove(symbol.as_str());
1134 self.subscriptions.remove_reference(&key);
1135 self.subscriptions.mark_unsubscribe(&key);
1136 self.subscriptions.confirm_unsubscribe(&key);
1137 return Err(e);
1138 }
1139 };
1140
1141 self.subscriptions.confirm_subscribe(&key);
1142 self.subscription_payloads
1143 .write()
1144 .await
1145 .insert(key, payload);
1146 Ok(())
1147 }
1148
1149 pub async fn unsubscribe_book_l3(&self, symbol: Ustr) -> Result<(), KrakenWsError> {
1164 let channel_str = KrakenWsChannel::Level3.as_ref();
1165 let key = format!("{channel_str}:{symbol}");
1166 if !self.subscriptions.remove_reference(&key) {
1167 return Ok(());
1168 }
1169 self.subscriptions.mark_unsubscribe(&key);
1170
1171 let token = self.auth_token.read().await.clone();
1172 let req_id = self.get_next_req_id();
1173 let request = KrakenWsRequest {
1174 method: KrakenWsMethod::Unsubscribe,
1175 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1176 channel: KrakenWsChannel::Level3,
1177 symbol: Some(vec![symbol]),
1178 snapshot: None,
1179 depth: None,
1180 interval: None,
1181 event_trigger: None,
1182 token,
1183 snap_orders: None,
1184 snap_trades: None,
1185 })),
1186 req_id: Some(req_id),
1187 };
1188
1189 self.send_command(&request).await?;
1190 self.subscriptions.confirm_unsubscribe(&key);
1191 self.subscription_payloads.write().await.remove(&key);
1192 self.l3_depths
1193 .lock()
1194 .expect("L3 depth map mutex poisoned")
1195 .remove(symbol.as_str());
1196 Ok(())
1197 }
1198
1199 pub async fn resync_book_l3(&self, symbol: Ustr, depth: u32) -> Result<(), KrakenWsError> {
1213 let channel_str = KrakenWsChannel::Level3.as_ref();
1214 let key = format!("{channel_str}:{symbol}");
1215
1216 if !self.subscriptions_contains(&key) {
1222 log::debug!("Skipping L3 resync: subscription cancelled mid-retry, symbol={symbol}",);
1223 return Ok(());
1224 }
1225
1226 let new_token = refresh_auth_token(&self.config).await?;
1227 *self.auth_token.write().await = Some(new_token.clone());
1228
1229 if !self.subscriptions_contains(&key) {
1232 log::debug!(
1233 "Skipping L3 resync: subscription cancelled after token refresh, symbol={symbol}",
1234 );
1235 return Ok(());
1236 }
1237
1238 let unsub_req_id = self.get_next_req_id();
1239 let unsub = KrakenWsRequest {
1240 method: KrakenWsMethod::Unsubscribe,
1241 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1242 channel: KrakenWsChannel::Level3,
1243 symbol: Some(vec![symbol]),
1244 snapshot: None,
1245 depth: None,
1246 interval: None,
1247 event_trigger: None,
1248 token: Some(new_token.clone()),
1249 snap_orders: None,
1250 snap_trades: None,
1251 })),
1252 req_id: Some(unsub_req_id),
1253 };
1254 self.send_command(&unsub).await?;
1255
1256 if !self.subscriptions_contains(&key) {
1258 log::debug!("Skipping L3 resync resubscribe: cancelled before send, symbol={symbol}",);
1259 return Ok(());
1260 }
1261
1262 let sub_req_id = self.get_next_req_id();
1263 let sub = KrakenWsRequest {
1264 method: KrakenWsMethod::Subscribe,
1265 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1266 channel: KrakenWsChannel::Level3,
1267 symbol: Some(vec![symbol]),
1268 snapshot: Some(true),
1269 depth: Some(depth),
1270 interval: None,
1271 event_trigger: None,
1272 token: Some(new_token),
1273 snap_orders: None,
1274 snap_trades: None,
1275 })),
1276 req_id: Some(sub_req_id),
1277 };
1278 let payload = self.send_command(&sub).await?;
1279
1280 if self.subscriptions_contains(&key) {
1285 self.subscription_payloads
1286 .write()
1287 .await
1288 .insert(key, payload);
1289 self.l3_depths
1290 .lock()
1291 .expect("L3 depth map mutex poisoned")
1292 .insert(symbol.to_string(), depth);
1293 }
1294
1295 Ok(())
1296 }
1297
1298 pub fn validate_l3_checksum(&self) -> bool {
1300 self.config.validate_l3_checksum
1301 }
1302
1303 pub fn has_credentials(&self) -> bool {
1306 self.config.has_api_credentials()
1307 }
1308
1309 pub fn instruments_handle(&self) -> Arc<AtomicMap<InstrumentId, InstrumentAny>> {
1315 Arc::clone(&self.instruments)
1316 }
1317
1318 pub fn l3_depths_handle(&self) -> Arc<std::sync::Mutex<ahash::AHashMap<String, u32>>> {
1323 Arc::clone(&self.l3_depths)
1324 }
1325
1326 pub(crate) fn l2_depths_handle(&self) -> L2Depths {
1328 self.l2_depths.clone()
1329 }
1330
1331 pub async fn subscribe_quotes(&self, instrument_id: InstrumentId) -> Result<(), KrakenWsError> {
1336 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1337 let key = format!("quotes:{symbol}");
1338
1339 if !self.subscriptions.add_reference(&key) {
1340 return Ok(());
1341 }
1342
1343 self.subscriptions.mark_subscribe(&key);
1344
1345 let req_id = self.get_next_req_id();
1346 let request = KrakenWsRequest {
1347 method: KrakenWsMethod::Subscribe,
1348 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1349 channel: KrakenWsChannel::Ticker,
1350 symbol: Some(vec![symbol]),
1351 snapshot: None,
1352 depth: None,
1353 interval: None,
1354 event_trigger: Some("bbo".to_string()),
1355 token: None,
1356 snap_orders: None,
1357 snap_trades: None,
1358 })),
1359 req_id: Some(req_id),
1360 };
1361
1362 let payload = self.send_command(&request).await?;
1363 self.subscriptions.confirm_subscribe(&key);
1364 self.subscription_payloads
1365 .write()
1366 .await
1367 .insert(key, payload);
1368 Ok(())
1369 }
1370
1371 pub async fn subscribe_trades(&self, instrument_id: InstrumentId) -> Result<(), KrakenWsError> {
1373 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1374 self.subscribe(KrakenWsChannel::Trade, vec![symbol], None)
1375 .await
1376 }
1377
1378 pub async fn subscribe_bars(&self, bar_type: BarType) -> Result<(), KrakenWsError> {
1384 let symbol = to_ws_v2_symbol(bar_type.instrument_id().symbol.inner());
1385 let interval = bar_type_to_ws_interval(bar_type)?;
1386 self.subscribe_with_interval(KrakenWsChannel::Ohlc, vec![symbol], interval)
1387 .await
1388 }
1389
1390 pub async fn subscribe_executions(
1394 &self,
1395 snap_orders: bool,
1396 snap_trades: bool,
1397 ) -> Result<(), KrakenWsError> {
1398 let req_id = self.get_next_req_id();
1399
1400 let token = self.auth_token.read().await.clone().ok_or_else(|| {
1401 KrakenWsError::AuthenticationError(
1402 "Authentication token required for executions channel. Call authenticate() first"
1403 .to_string(),
1404 )
1405 })?;
1406
1407 let request = KrakenWsRequest {
1408 method: KrakenWsMethod::Subscribe,
1409 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1410 channel: KrakenWsChannel::Executions,
1411 symbol: None,
1412 snapshot: None,
1413 depth: None,
1414 interval: None,
1415 event_trigger: None,
1416 token: Some(token),
1417 snap_orders: Some(snap_orders),
1418 snap_trades: Some(snap_trades),
1419 })),
1420 req_id: Some(req_id),
1421 };
1422
1423 let payload = self.send_command(&request).await?;
1424
1425 let key = "executions";
1426 if self.subscriptions.add_reference(key) {
1427 self.subscriptions.mark_subscribe(key);
1428 self.subscriptions.confirm_subscribe(key);
1429 self.subscription_payloads
1430 .write()
1431 .await
1432 .insert(key.to_string(), payload);
1433 }
1434
1435 Ok(())
1436 }
1437
1438 pub async fn unsubscribe_book(&self, instrument_id: InstrumentId) -> Result<(), KrakenWsError> {
1440 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1441 let channel_str = KrakenWsChannel::Book.as_ref();
1442 let key = format!("{channel_str}:{symbol}");
1443
1444 if !self.subscriptions.remove_reference(&key) {
1445 return Ok(());
1446 }
1447
1448 self.subscriptions.mark_unsubscribe(&key);
1449
1450 let req_id = self.get_next_req_id();
1451 let request = KrakenWsRequest {
1452 method: KrakenWsMethod::Unsubscribe,
1453 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1454 channel: KrakenWsChannel::Book,
1455 symbol: Some(vec![symbol]),
1456 snapshot: None,
1457 depth: None,
1458 interval: None,
1459 event_trigger: None,
1460 token: None,
1461 snap_orders: None,
1462 snap_trades: None,
1463 })),
1464 req_id: Some(req_id),
1465 };
1466
1467 self.send_command(&request).await?;
1468 self.subscriptions.confirm_unsubscribe(&key);
1469 self.subscription_payloads.write().await.remove(&key);
1470 self.l2_depths.remove(symbol.as_str());
1471 Ok(())
1472 }
1473
1474 pub async fn unsubscribe_quotes(
1476 &self,
1477 instrument_id: InstrumentId,
1478 ) -> Result<(), KrakenWsError> {
1479 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1480 let key = format!("quotes:{symbol}");
1481
1482 if !self.subscriptions.remove_reference(&key) {
1483 return Ok(());
1484 }
1485
1486 self.subscriptions.mark_unsubscribe(&key);
1487
1488 let req_id = self.get_next_req_id();
1489 let request = KrakenWsRequest {
1490 method: KrakenWsMethod::Unsubscribe,
1491 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1492 channel: KrakenWsChannel::Ticker,
1493 symbol: Some(vec![symbol]),
1494 snapshot: None,
1495 depth: None,
1496 interval: None,
1497 event_trigger: Some("bbo".to_string()),
1498 token: None,
1499 snap_orders: None,
1500 snap_trades: None,
1501 })),
1502 req_id: Some(req_id),
1503 };
1504
1505 self.send_command(&request).await?;
1506 self.subscriptions.confirm_unsubscribe(&key);
1507 self.subscription_payloads.write().await.remove(&key);
1508 Ok(())
1509 }
1510
1511 pub async fn unsubscribe_trades(
1513 &self,
1514 instrument_id: InstrumentId,
1515 ) -> Result<(), KrakenWsError> {
1516 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1517 self.unsubscribe(KrakenWsChannel::Trade, vec![symbol]).await
1518 }
1519
1520 pub async fn unsubscribe_bars(&self, bar_type: BarType) -> Result<(), KrakenWsError> {
1526 let symbol = to_ws_v2_symbol(bar_type.instrument_id().symbol.inner());
1527 let interval = bar_type_to_ws_interval(bar_type)?;
1528 self.unsubscribe_with_interval(KrakenWsChannel::Ohlc, vec![symbol], interval)
1529 .await
1530 }
1531}
1532
1533async fn refresh_auth_token(config: &KrakenDataClientConfig) -> Result<String, KrakenWsError> {
1535 let api_key = config
1536 .api_key
1537 .clone()
1538 .ok_or_else(|| KrakenWsError::AuthenticationError("Missing API key".to_string()))?;
1539 let api_secret = config
1540 .api_secret
1541 .clone()
1542 .ok_or_else(|| KrakenWsError::AuthenticationError("Missing API secret".to_string()))?;
1543
1544 let http_client = KrakenSpotHttpClient::with_credentials(
1545 api_key,
1546 api_secret,
1547 config.environment,
1548 Some(config.http_base_url()),
1549 config.timeout_secs,
1550 None,
1551 None,
1552 None,
1553 config.proxy_url.clone(),
1554 config
1555 .max_requests_per_second
1556 .unwrap_or(KRAKEN_SPOT_DEFAULT_RATE_LIMIT_PER_SECOND),
1557 )
1558 .map_err(|e| {
1559 KrakenWsError::AuthenticationError(format!("Failed to create HTTP client: {e}"))
1560 })?;
1561
1562 let ws_token = http_client.get_websockets_token().await.map_err(|e| {
1563 KrakenWsError::AuthenticationError(format!("Failed to get WebSocket token: {e}"))
1564 })?;
1565
1566 log::debug!(
1567 "WebSocket authentication token refreshed: token_length={}, expires={}",
1568 ws_token.token.len(),
1569 ws_token.expires
1570 );
1571
1572 Ok(ws_token.token)
1573}
1574
1575fn update_auth_token_in_payload(payload: &str, new_token: &str) -> Result<String, KrakenWsError> {
1576 let mut value: serde_json::Value =
1577 serde_json::from_str(payload).map_err(|e| KrakenWsError::JsonError(e.to_string()))?;
1578
1579 if let Some(params) = value.get_mut("params") {
1580 params["token"] = serde_json::Value::String(new_token.to_string());
1581 }
1582
1583 serde_json::to_string(&value).map_err(|e| KrakenWsError::JsonError(e.to_string()))
1584}
1585
1586#[inline]
1587fn to_ws_v2_symbol(symbol: Ustr) -> Ustr {
1588 Ustr::from(&normalize_spot_symbol(symbol.as_str()))
1589}
1590
1591fn bar_type_to_ws_interval(bar_type: BarType) -> Result<u32, KrakenWsError> {
1592 const VALID_INTERVALS: [u32; 9] = [1, 5, 15, 30, 60, 240, 1440, 10080, 21600];
1593
1594 let spec = bar_type.spec();
1595 let step = spec.step.get() as u32;
1596
1597 let base_minutes = match spec.aggregation {
1598 BarAggregation::Minute => 1,
1599 BarAggregation::Hour => 60,
1600 BarAggregation::Day => 1440,
1601 BarAggregation::Week => 10080,
1602 other => {
1603 return Err(KrakenWsError::SubscriptionError(format!(
1604 "Unsupported bar aggregation for Kraken OHLC streaming: {other:?}"
1605 )));
1606 }
1607 };
1608
1609 let interval = base_minutes * step;
1610
1611 if !VALID_INTERVALS.contains(&interval) {
1612 return Err(KrakenWsError::SubscriptionError(format!(
1613 "Invalid bar interval {interval} minutes for Kraken OHLC streaming. \
1614 Supported intervals: 1, 5, 15, 30, 60, 240, 1440, 10080, 21600"
1615 )));
1616 }
1617
1618 Ok(interval)
1619}
1620
1621#[cfg(test)]
1622mod tests {
1623 use std::sync::{Arc, atomic::Ordering};
1624
1625 use rstest::rstest;
1626 use tokio_util::sync::CancellationToken;
1627
1628 use super::*;
1629 use crate::config::KrakenDataClientConfig;
1630
1631 #[rstest]
1632 fn test_req_id_counter_is_shared_arc_and_monotonic() {
1633 let cfg = KrakenDataClientConfig::default();
1634 let client = KrakenSpotWebSocketClient::new(cfg, CancellationToken::new(), None);
1635 let counter = client.req_id_counter();
1636 let a = counter.fetch_add(1, Ordering::Relaxed);
1637 let b = counter.fetch_add(1, Ordering::Relaxed);
1638 assert!(b > a);
1639 #[allow(clippy::redundant_clone)]
1640 let cloned = client.clone();
1641 let cloned_counter = cloned.req_id_counter();
1642 assert!(Arc::ptr_eq(&counter, &cloned_counter));
1643 }
1644
1645 #[rstest]
1646 #[case("XBT/EUR", "BTC/EUR")]
1647 #[case("XBT/USD", "BTC/USD")]
1648 #[case("XBT/USDT", "BTC/USDT")]
1649 #[case("ETH/USD", "ETH/USD")]
1650 #[case("ETH/XBT", "ETH/BTC")]
1651 #[case("SOL/XBT", "SOL/BTC")]
1652 #[case("SOL/USD", "SOL/USD")]
1653 #[case("BTC/USD", "BTC/USD")]
1654 #[case("ETH/BTC", "ETH/BTC")]
1655 #[case("XDG/USD", "DOGE/USD")]
1656 #[case("XDG/EUR", "DOGE/EUR")]
1657 fn test_to_kraken_ws_v2_symbol(#[case] input: &str, #[case] expected: &str) {
1658 let symbol = Ustr::from(input);
1659 let result = to_ws_v2_symbol(symbol);
1660 assert_eq!(result.as_str(), expected);
1661 }
1662
1663 fn test_client_without_credentials() -> KrakenSpotWebSocketClient {
1664 KrakenSpotWebSocketClient::new(
1665 KrakenDataClientConfig::default(),
1666 CancellationToken::new(),
1667 None,
1668 )
1669 }
1670
1671 #[rstest]
1672 #[tokio::test]
1673 async fn test_authenticate_without_credentials_errors() {
1674 let client = test_client_without_credentials();
1675
1676 let err = client.authenticate().await.expect_err("should fail");
1677 assert!(
1678 matches!(err, KrakenWsError::AuthenticationError(ref msg) if msg.contains("API credentials required")),
1679 "unexpected error: {err:?}"
1680 );
1681 assert!(!client.is_authenticated());
1682 }
1683
1684 #[rstest]
1685 #[tokio::test]
1686 async fn test_wait_until_authenticated_times_out() {
1687 let client = test_client_without_credentials();
1688
1689 let err = client
1690 .wait_until_authenticated(0.05)
1691 .await
1692 .expect_err("should time out");
1693 assert!(matches!(err, KrakenWsError::AuthenticationError(_)));
1694 }
1695
1696 #[rstest]
1697 #[tokio::test]
1698 async fn test_wait_until_authenticated_resolves_after_succeed() {
1699 let client = test_client_without_credentials();
1700
1701 let tracker = client.auth_tracker.clone();
1702 let _rx = tracker.begin();
1703
1704 tokio::spawn(async move {
1705 tokio::time::sleep(tokio::time::Duration::from_millis(20)).await;
1706 tracker.succeed();
1707 });
1708
1709 client
1710 .wait_until_authenticated(1.0)
1711 .await
1712 .expect("should resolve once tracker succeeds");
1713 assert!(client.is_authenticated());
1714 }
1715
1716 #[rstest]
1717 #[tokio::test]
1718 async fn test_is_authenticated_flips_on_fail() {
1719 let client = test_client_without_credentials();
1720
1721 let _rx = client.auth_tracker.begin();
1722 client.auth_tracker.succeed();
1723 assert!(client.is_authenticated());
1724
1725 client.auth_tracker.fail("test failure");
1726 assert!(!client.is_authenticated());
1727 }
1728
1729 #[rstest]
1730 fn test_l3_factory_uses_ws_l3_url() {
1731 let cfg = KrakenDataClientConfig::default();
1732 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
1733 assert_eq!(client.url(), "wss://ws-l3.kraken.com/v2");
1734 }
1735
1736 #[rstest]
1737 fn test_l3_factory_respects_override() {
1738 let cfg = KrakenDataClientConfig {
1739 ws_l3_url: Some("wss://override.example/v2".to_string()),
1740 ..Default::default()
1741 };
1742 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
1743 assert_eq!(client.url(), "wss://override.example/v2");
1744 }
1745
1746 #[rstest]
1747 #[tokio::test]
1748 async fn test_subscribe_book_l3_without_auth_errors_and_leaves_clean_state() {
1749 let cfg = KrakenDataClientConfig::default();
1750 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
1751
1752 let err = client
1753 .subscribe_book_l3(Ustr::from("BTC/USD"), 1000)
1754 .await
1755 .expect_err("should fail without auth token");
1756
1757 assert!(matches!(err, KrakenWsError::AuthenticationError(_)));
1758 assert!(
1759 client.subscriptions.is_empty(),
1760 "no state must leak on auth failure"
1761 );
1762 }
1763
1764 #[rstest]
1765 #[tokio::test]
1766 async fn test_subscribe_book_l3_invalid_depth_errors() {
1767 let cfg = KrakenDataClientConfig::default();
1768 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
1769
1770 let err = client
1771 .subscribe_book_l3(Ustr::from("BTC/USD"), 50)
1772 .await
1773 .expect_err("should fail on invalid depth");
1774
1775 assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
1776 assert!(!client.subscriptions_contains("level3:BTC/USD"));
1777 }
1778
1779 #[rstest]
1780 #[tokio::test]
1781 async fn test_subscribe_book_invalid_depth_errors() {
1782 let client = test_client_without_credentials();
1783 let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
1784
1785 let err = client
1786 .subscribe_book(instrument_id, Some(50))
1787 .await
1788 .expect_err("should fail on invalid depth");
1789
1790 assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
1791 assert!(!client.subscriptions_contains("book:BTC/USD"));
1792 assert_eq!(client.l2_depths.get("BTC/USD"), None);
1793 }
1794
1795 #[rstest]
1796 #[tokio::test]
1797 async fn test_subscribe_book_defaults_to_depth_10_and_stores_state() {
1798 let client = test_client_without_credentials();
1799 let key = "book:BTC/USD";
1800 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1801 *client.cmd_tx.write().await = cmd_tx;
1802
1803 client
1804 .subscribe_book(InstrumentId::from("BTC/USD.KRAKEN"), None)
1805 .await
1806 .expect("subscribe should succeed");
1807
1808 let cmd = cmd_rx.try_recv().expect("expected subscribe command");
1809 let SpotHandlerCommand::Subscribe { payload } = cmd else {
1810 panic!("expected subscribe command");
1811 };
1812 assert_book_subscribe_payload(&payload, "BTC/USD", 10);
1813
1814 assert_eq!(client.subscriptions.get_reference_count(key), 1);
1815 assert!(client.subscriptions_contains(key));
1816 assert_eq!(client.l2_depths.get("BTC/USD"), Some(10));
1817 assert_eq!(
1818 client.subscription_payloads.read().await.get(key),
1819 Some(&payload)
1820 );
1821 }
1822
1823 #[rstest]
1824 #[tokio::test]
1825 async fn test_subscribe_book_send_failure_rolls_back_l2_state() {
1826 let client = test_client_without_credentials();
1827 let key = "book:BTC/USD";
1828
1829 let err = client
1830 .subscribe_book(InstrumentId::from("BTC/USD.KRAKEN"), Some(10))
1831 .await
1832 .expect_err("should fail when command receiver is closed");
1833
1834 assert!(matches!(err, KrakenWsError::ConnectionError(_)));
1835 assert_eq!(client.subscriptions.get_reference_count(key), 0);
1836 assert!(!client.subscriptions_contains(key));
1837 assert!(client.subscriptions.pending_subscribe_topics().is_empty());
1838 assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
1839 assert_eq!(client.l2_depths.get("BTC/USD"), None);
1840 assert!(!client.subscription_payloads.read().await.contains_key(key));
1841 }
1842
1843 #[rstest]
1844 fn test_subscribe_book_l3_refcount_idempotent() {
1845 let cfg = KrakenDataClientConfig::default();
1846 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
1847 let key = "level3:BTC/USD";
1848
1849 assert!(client.subscriptions.add_reference(key));
1850 assert!(!client.subscriptions.add_reference(key));
1851 assert!(!client.subscriptions.remove_reference(key));
1852 assert!(client.subscriptions.remove_reference(key));
1853 }
1854
1855 #[rstest]
1856 #[tokio::test]
1857 async fn test_subscribe_book_l3_rejects_depth_mismatch() {
1858 let cfg = KrakenDataClientConfig::default();
1859 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
1860
1861 let key = "level3:BTC/USD";
1862 client.subscriptions.add_reference(key);
1863 client.subscriptions.mark_subscribe(key);
1864 client.subscriptions.confirm_subscribe(key);
1865 client
1866 .l3_depths
1867 .lock()
1868 .unwrap()
1869 .insert("BTC/USD".to_string(), 1000);
1870
1871 *client.auth_token.write().await = Some("test-token".to_string());
1872
1873 let err = client
1874 .subscribe_book_l3(Ustr::from("BTC/USD"), 10)
1875 .await
1876 .expect_err("should reject depth mismatch");
1877 assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
1878
1879 assert!(client.subscriptions.remove_reference(key));
1880 }
1881
1882 #[rstest]
1883 #[tokio::test]
1884 async fn test_subscribe_book_rejects_depth_mismatch() {
1885 let client = test_client_without_credentials();
1886 let key = "book:BTC/USD";
1887 client.subscriptions.add_reference(key);
1888 client.subscriptions.mark_subscribe(key);
1889 client.subscriptions.confirm_subscribe(key);
1890 client.l2_depths.insert("BTC/USD", 10);
1891
1892 let err = client
1893 .subscribe_book(InstrumentId::from("BTC/USD.KRAKEN"), Some(25))
1894 .await
1895 .expect_err("should reject depth mismatch");
1896 assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
1897
1898 assert!(client.subscriptions.remove_reference(key));
1899 }
1900
1901 #[rstest]
1902 #[tokio::test]
1903 async fn test_unsubscribe_book_removes_l2_depth_on_last_reference() {
1904 let client = test_client_without_credentials();
1905 let key = "book:BTC/USD";
1906 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1907 *client.cmd_tx.write().await = cmd_tx;
1908
1909 client.subscriptions.add_reference(key);
1910 client.subscriptions.mark_subscribe(key);
1911 client.subscriptions.confirm_subscribe(key);
1912 client.l2_depths.insert("BTC/USD", 10);
1913 client
1914 .subscription_payloads
1915 .write()
1916 .await
1917 .insert(key.to_string(), "payload".to_string());
1918
1919 client
1920 .unsubscribe_book(InstrumentId::from("BTC/USD.KRAKEN"))
1921 .await
1922 .expect("unsubscribe should succeed");
1923
1924 assert_eq!(client.subscriptions.get_reference_count(key), 0);
1925 assert!(!client.subscriptions_contains(key));
1926 assert_eq!(client.l2_depths.get("BTC/USD"), None);
1927 assert!(!client.subscription_payloads.read().await.contains_key(key));
1928
1929 let cmd = cmd_rx.try_recv().expect("expected unsubscribe command");
1930 let SpotHandlerCommand::Unsubscribe { payload } = cmd else {
1931 panic!("expected unsubscribe command");
1932 };
1933 assert!(payload.contains(r#""method":"unsubscribe""#));
1934 assert!(payload.contains(r#""channel":"book""#));
1935 assert!(payload.contains(r#""BTC/USD""#));
1936 }
1937
1938 #[rstest]
1939 #[tokio::test]
1940 async fn test_generic_subscribe_rejects_level3() {
1941 let cfg = KrakenDataClientConfig::default();
1942 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
1943
1944 let err = client
1945 .subscribe(
1946 KrakenWsChannel::Level3,
1947 vec![Ustr::from("BTC/USD")],
1948 Some(1000),
1949 )
1950 .await
1951 .expect_err("generic subscribe must reject Level3");
1952 assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
1953
1954 let err = client
1955 .unsubscribe(KrakenWsChannel::Level3, vec![Ustr::from("BTC/USD")])
1956 .await
1957 .expect_err("generic unsubscribe must reject Level3");
1958 assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
1959 }
1960
1961 #[rstest]
1962 fn test_update_auth_token_in_payload_for_level3() {
1963 let original = r#"{"method":"subscribe","params":{"channel":"level3","symbol":["BTC/USD"],"depth":1000,"snapshot":true,"token":"OLD"},"req_id":1}"#;
1964 let rewritten = update_auth_token_in_payload(original, "NEW").unwrap();
1965 assert!(rewritten.contains(r#""token":"NEW""#));
1966 assert!(!rewritten.contains(r#""token":"OLD""#));
1967 }
1968
1969 #[rstest]
1970 fn test_l3_depths_shared_between_subscribe_and_handle() {
1971 let cfg = KrakenDataClientConfig::default();
1972 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
1973
1974 let handle = client.l3_depths_handle();
1975 client
1976 .l3_depths
1977 .lock()
1978 .unwrap()
1979 .insert("BTC/USD".to_string(), 100);
1980
1981 assert_eq!(handle.lock().unwrap().get("BTC/USD").copied(), Some(100));
1982 }
1983
1984 #[rstest]
1985 fn test_l2_depths_shared_between_subscribe_and_handle() {
1986 let client = test_client_without_credentials();
1987
1988 let handle = client.l2_depths_handle();
1989 client.l2_depths.insert("BTC/USD", 10);
1990
1991 assert_eq!(handle.get("BTC/USD"), Some(10));
1992 }
1993
1994 #[rstest]
1995 fn test_resync_book_l3_does_not_touch_refcount() {
1996 let cfg = KrakenDataClientConfig::default();
1997 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
1998 let key = "level3:BTC/USD";
1999
2000 assert!(client.subscriptions.add_reference(key));
2001 assert!(!client.subscriptions.add_reference(key));
2002 client.subscriptions.confirm_subscribe(key);
2003
2004 assert!(client.subscriptions_contains(key));
2005 }
2006
2007 fn assert_book_subscribe_payload(payload: &str, symbol: &str, depth: u32) {
2008 let value: serde_json::Value =
2009 serde_json::from_str(payload).expect("payload should parse as JSON");
2010
2011 assert_eq!(value["method"], serde_json::json!("subscribe"));
2012 assert_eq!(value["params"]["channel"], serde_json::json!("book"));
2013 assert_eq!(value["params"]["symbol"], serde_json::json!([symbol]));
2014 assert_eq!(value["params"]["depth"], serde_json::json!(depth));
2015 }
2016}