1use std::{
19 collections::HashMap,
20 sync::{
21 Arc,
22 atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering},
23 },
24};
25
26use arc_swap::ArcSwap;
27#[cfg(test)]
28use nautilus_core::string::secret::REDACTED;
29use nautilus_core::{AtomicMap, string::secret::SecretString};
30use nautilus_live::{
31 SocketControl,
32 task::{TaskGroup, TaskShutdownError},
33};
34use nautilus_model::{
35 data::BarType,
36 enums::BarAggregation,
37 identifiers::{AccountId, ClientOrderId, InstrumentId, StrategyId, TraderId, VenueOrderId},
38 instruments::{Instrument, InstrumentAny},
39};
40use nautilus_network::{
41 ConnectionUrl,
42 http::create_standard_nautilus_headers,
43 mode::ConnectionMode,
44 websocket::{
45 AuthTracker, SubscriptionState, TransportBackend, WebSocketClient, WebSocketConfig,
46 channel_message_handler,
47 },
48};
49use parking_lot::RwLock;
50use tokio_util::sync::CancellationToken;
51use ustr::Ustr;
52
53pub const KRAKEN_SPOT_WS_TOPIC_DELIMITER: char = ':';
57
58use super::{
59 enums::{KrakenWsChannel, KrakenWsMethod},
60 handler::{SpotFeedHandler, SpotHandlerCommand},
61 level_2::{L2BookRequest, L2BookRequests, L2Depths},
62 messages::{KrakenSpotWsMessage, KrakenWsChannelParams, KrakenWsParams, KrakenWsRequest},
63};
64use crate::{
65 common::{
66 consts::{
67 KRAKEN_RATE_LIMIT_KEY_ORDER, KRAKEN_RATE_LIMIT_KEY_SUBSCRIPTION,
68 KRAKEN_SPOT_WS_ORDER_QUOTA, KRAKEN_SPOT_WS_SUBSCRIPTION_QUOTA,
69 },
70 parse::normalize_spot_symbol,
71 },
72 config::KrakenDataClientConfig,
73 http::{KrakenSpotHttpClient, spot::client::KRAKEN_SPOT_DEFAULT_RATE_LIMIT_PER_SECOND},
74 websocket::error::KrakenWsError,
75};
76
77const WS_PING_MSG: &str = r#"{"method":"ping"}"#;
78
79#[derive(Debug)]
81pub struct KrakenSpotWebSocketClient {
82 url: ConnectionUrl,
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 handler_tasks: Arc<TaskGroup>,
89 connect_lock: Arc<tokio::sync::Mutex<()>>,
90 subscriptions: SubscriptionState,
91 subscription_payloads: Arc<tokio::sync::RwLock<HashMap<String, SecretString>>>,
92 auth_tracker: AuthTracker,
93 cancellation_token: CancellationToken,
94 req_id_counter: Arc<AtomicU64>,
95 auth_token: Arc<tokio::sync::RwLock<Option<SecretString>>>,
96 account_id: Arc<RwLock<Option<AccountId>>>,
97 truncated_id_map: Arc<AtomicMap<String, ClientOrderId>>,
98 instruments: Arc<AtomicMap<InstrumentId, InstrumentAny>>,
99 l2_depths: L2Depths,
100 l2_book_commands: Arc<tokio::sync::Mutex<()>>,
103 book_requests: L2BookRequests,
106 l3_depths: Arc<parking_lot::Mutex<ahash::AHashMap<String, u32>>>,
107 transport_backend: TransportBackend,
108 proxy_url: Option<SecretString>,
109 socket_control: Option<SocketControl>,
110}
111
112impl Clone for KrakenSpotWebSocketClient {
113 fn clone(&self) -> Self {
114 Self {
115 url: self.url.clone(),
116 config: self.config.clone(),
117 signal: Arc::clone(&self.signal),
118 connection_mode: Arc::clone(&self.connection_mode),
119 cmd_tx: Arc::clone(&self.cmd_tx),
120 out_rx: self.out_rx.clone(),
121 handler_tasks: Arc::clone(&self.handler_tasks),
122 connect_lock: Arc::clone(&self.connect_lock),
123 subscriptions: self.subscriptions.clone(),
124 subscription_payloads: Arc::clone(&self.subscription_payloads),
125 auth_tracker: self.auth_tracker.clone(),
126 cancellation_token: self.cancellation_token.clone(),
127 req_id_counter: self.req_id_counter.clone(),
128 auth_token: self.auth_token.clone(),
129 account_id: Arc::clone(&self.account_id),
130 truncated_id_map: Arc::clone(&self.truncated_id_map),
131 instruments: Arc::clone(&self.instruments),
132 l2_depths: self.l2_depths.clone(),
133 l2_book_commands: Arc::clone(&self.l2_book_commands),
134 book_requests: Arc::clone(&self.book_requests),
135 l3_depths: Arc::clone(&self.l3_depths),
136 transport_backend: self.transport_backend,
137 proxy_url: self.proxy_url.clone(),
138 socket_control: self.socket_control.clone(),
139 }
140 }
141}
142
143impl KrakenSpotWebSocketClient {
144 pub fn new(
150 config: KrakenDataClientConfig,
151 cancellation_token: CancellationToken,
152 proxy_url: Option<String>,
153 ) -> anyhow::Result<Self> {
154 config.validate()?;
155 let url = if config.ws_private_url.is_some() {
156 config.ws_private_url()
157 } else {
158 config.ws_public_url()
159 };
160 Self::new_with_url(url, config, cancellation_token, proxy_url)
161 }
162
163 pub fn l3(
172 config: KrakenDataClientConfig,
173 cancellation_token: CancellationToken,
174 proxy_url: Option<String>,
175 ) -> anyhow::Result<Self> {
176 config.validate()?;
177
178 let url = config.ws_l3_url();
179 Self::new_with_url(url, config, cancellation_token, proxy_url)
180 }
181
182 fn new_with_url(
183 url: String,
184 mut config: KrakenDataClientConfig,
185 cancellation_token: CancellationToken,
186 proxy_url: Option<String>,
187 ) -> anyhow::Result<Self> {
188 config.validate()?;
189
190 let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel::<SpotHandlerCommand>();
191 let initial_mode = AtomicU8::new(ConnectionMode::Closed.as_u8());
192 let connection_mode = Arc::new(ArcSwap::from_pointee(initial_mode));
193
194 let transport_backend = config.transport_backend;
195 config.proxy_url = proxy_url.clone().map(SecretString::from);
196
197 Ok(Self {
198 url: url.into(),
199 config,
200 signal: Arc::new(AtomicBool::new(false)),
201 connection_mode,
202 cmd_tx: Arc::new(tokio::sync::RwLock::new(cmd_tx)),
203 out_rx: None,
204 handler_tasks: Arc::new(TaskGroup::new()),
205 connect_lock: Arc::new(tokio::sync::Mutex::new(())),
206 subscriptions: SubscriptionState::new(KRAKEN_SPOT_WS_TOPIC_DELIMITER),
207 subscription_payloads: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
208 auth_tracker: AuthTracker::new(),
209 cancellation_token,
210 req_id_counter: Arc::new(AtomicU64::new(0)),
211 auth_token: Arc::new(tokio::sync::RwLock::new(None)),
212 account_id: Arc::new(RwLock::new(None)),
213 truncated_id_map: Arc::new(AtomicMap::new()),
214 instruments: Arc::new(AtomicMap::new()),
215 l2_depths: L2Depths::default(),
216 l2_book_commands: Arc::new(tokio::sync::Mutex::new(())),
217 book_requests: Arc::new(parking_lot::Mutex::new(ahash::AHashMap::new())),
218 l3_depths: Arc::new(parking_lot::Mutex::new(ahash::AHashMap::new())),
219 transport_backend,
220 proxy_url: proxy_url.map(SecretString::from),
221 socket_control: None,
222 })
223 }
224
225 pub(crate) fn begin_shutdown(&self) {
226 self.handler_tasks.begin_shutdown();
227 self.cancellation_token.cancel();
228 self.signal.store(true, Ordering::Relaxed);
229 }
230
231 #[must_use]
233 pub fn with_socket_control(mut self, control: SocketControl) -> Self {
234 self.socket_control = Some(control);
235 self
236 }
237
238 fn get_next_req_id(&self) -> u64 {
239 self.req_id_counter.fetch_add(1, Ordering::Relaxed) + 1
240 }
241
242 pub fn req_id_counter(&self) -> Arc<AtomicU64> {
244 self.req_id_counter.clone()
245 }
246
247 pub async fn handler_command_sender(
249 &self,
250 ) -> tokio::sync::mpsc::UnboundedSender<SpotHandlerCommand> {
251 self.cmd_tx.read().await.clone()
252 }
253
254 pub fn handler_command_handle(
259 &self,
260 ) -> Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<SpotHandlerCommand>>> {
261 self.cmd_tx.clone()
262 }
263
264 pub async fn auth_token(&self) -> Option<SecretString> {
266 self.auth_token.read().await.clone()
267 }
268
269 pub fn auth_token_blocking(&self) -> Option<SecretString> {
275 self.auth_token
276 .try_read()
277 .ok()
278 .and_then(|guard| guard.clone())
279 }
280
281 pub(crate) fn auth_token_handle(&self) -> Arc<tokio::sync::RwLock<Option<SecretString>>> {
284 self.auth_token.clone()
285 }
286
287 pub async fn connect(&mut self) -> Result<(), KrakenWsError> {
289 let connect_lock = Arc::clone(&self.connect_lock);
290 let _connect_guard = connect_lock.lock().await;
291
292 log::debug!("Connecting to {}", self.url.origin());
293
294 if !self.handler_tasks.is_open() || !self.handler_tasks.is_empty() {
295 self.disconnect_locked().await?;
296 self.handler_tasks.start_generation().map_err(|e| {
297 KrakenWsError::ConnectionError(format!(
298 "Failed to start WebSocket handler task generation: {e}"
299 ))
300 })?;
301 }
302 let handler_spawner = self.handler_tasks.spawner().map_err(|e| {
303 KrakenWsError::ConnectionError(format!(
304 "Failed to acquire WebSocket handler task spawner: {e}"
305 ))
306 })?;
307
308 if self.cancellation_token.is_cancelled() {
309 self.cancellation_token = CancellationToken::new();
310 }
311
312 self.signal.store(false, Ordering::Relaxed);
313
314 let (raw_handler, raw_rx) = channel_message_handler();
315 let headers = create_standard_nautilus_headers();
316
317 let ws_config = WebSocketConfig {
318 url: self.url.expose().to_owned(),
319 headers,
320 heartbeat_interval_secs: Some(self.config.heartbeat_interval_secs),
321 heartbeat_payload: Some(WS_PING_MSG.to_string()),
322 connect_timeout_ms: Some(5_000),
323 reconnect_delay_initial_ms: Some(500),
324 reconnect_delay_max_ms: Some(5_000),
325 reconnect_backoff_factor: Some(1.5),
326 reconnect_jitter_ms: Some(250),
327 reconnect_max_attempts: None,
328 heartbeat_timeout_secs: None,
331 idle_timeout_ms: (self.config.ws_idle_timeout_ms != 0)
332 .then_some(self.config.ws_idle_timeout_ms),
333 writer_capacity: None,
334 backend: self.transport_backend,
335 proxy_url: self
336 .proxy_url
337 .as_ref()
338 .map(|value| value.expose_secret().to_owned()),
339 max_message_size_bytes: None,
340 max_frame_size_bytes: None,
341 };
342
343 let keyed_quotas = vec![
344 (
345 KRAKEN_RATE_LIMIT_KEY_SUBSCRIPTION[0].to_string(),
346 *KRAKEN_SPOT_WS_SUBSCRIPTION_QUOTA,
347 ),
348 (
349 KRAKEN_RATE_LIMIT_KEY_ORDER[0].to_string(),
350 *KRAKEN_SPOT_WS_ORDER_QUOTA,
351 ),
352 ];
353
354 let ws_client = WebSocketClient::builder()
355 .config(ws_config)
356 .message_handler(raw_handler)
357 .keyed_quotas(keyed_quotas)
358 .maybe_state_sink(self.socket_control.as_ref().map(SocketControl::sink))
359 .connect()
360 .await
361 .map_err(|e| KrakenWsError::ConnectionError(e.to_string()))?;
362
363 let connection_mode = ws_client.connection_mode_atomic();
364 let reconnect_handle = ws_client.reconnect_handle();
365
366 let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<KrakenSpotWsMessage>();
367 self.out_rx = Some(Arc::new(out_rx));
368
369 let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<SpotHandlerCommand>();
370
371 if let Err(e) = cmd_tx.send(SpotHandlerCommand::SetClient(ws_client)) {
372 return Err(KrakenWsError::ConnectionError(format!(
373 "Failed to send WebSocketClient to handler: {e}"
374 )));
375 }
376
377 {
378 let mut command_sender = self.cmd_tx.write().await;
379 *command_sender = cmd_tx.clone();
380 self.connection_mode.store(connection_mode);
381 }
382
383 if let Some(control) = &self.socket_control {
384 control.register(move || reconnect_handle.request_reconnect());
385 }
386
387 let signal = self.signal.clone();
388 let subscriptions = self.subscriptions.clone();
389 let subscription_payloads = self.subscription_payloads.clone();
390 let config_for_reconnect = self.config.clone();
391 let auth_token_for_reconnect = self.auth_token.clone();
392 let auth_tracker_for_reconnect = self.auth_tracker.clone();
393 let cmd_tx_for_reconnect = cmd_tx.clone();
394 let book_requests_for_reconnect = Arc::clone(&self.book_requests);
395 let l2_depths_for_reconnect = self.l2_depths.clone();
396
397 let handler_task = async move {
398 let mut handler =
399 SpotFeedHandler::new(signal.clone(), cmd_rx, raw_rx, subscriptions.clone());
400
401 loop {
402 match handler.next().await {
403 Some(KrakenSpotWsMessage::Reconnected) => {
404 if signal.load(Ordering::Relaxed) {
405 continue;
406 }
407 log::info!("WebSocket reconnected, resubscribing");
408
409 begin_resubscribe(
410 &subscriptions,
411 &book_requests_for_reconnect,
412 &l2_depths_for_reconnect,
413 );
414
415 let payloads = subscription_payloads.read().await;
416 if payloads.is_empty() {
417 log::debug!("No subscriptions to restore after reconnection");
418 } else {
419 let had_auth = auth_token_for_reconnect.read().await.is_some();
420
421 if had_auth && config_for_reconnect.has_api_credentials() {
422 log::debug!("Re-authenticating after reconnect");
423
424 auth_tracker_for_reconnect.invalidate();
425 let _rx = auth_tracker_for_reconnect.begin();
426
427 match refresh_auth_token(&config_for_reconnect).await {
428 Ok(new_token) => {
429 *auth_token_for_reconnect.write().await = Some(new_token);
430 auth_tracker_for_reconnect.succeed();
431 log::debug!("Re-authentication successful");
432 }
433 Err(e) => {
434 log::error!(
435 "Failed to re-authenticate after reconnect: {e}"
436 );
437 *auth_token_for_reconnect.write().await = None;
438 auth_tracker_for_reconnect.fail(e.to_string());
439 }
440 }
441 }
442
443 log::debug!(
444 "Resubscribing after reconnection: count={}",
445 payloads.len()
446 );
447
448 for (topic, payload) in payloads.iter() {
449 let needs_token =
450 topic == "executions" || topic.starts_with("level3:");
451 let payload = if needs_token {
452 let auth_token = auth_token_for_reconnect.read().await.clone();
453 match auth_token {
454 Some(token) => {
455 match update_auth_token_in_payload(
456 payload.expose_secret(),
457 token.expose_secret(),
458 ) {
459 Ok(p) => p,
460 Err(e) => {
461 log::error!("Failed to update auth token: {e}");
462 continue;
463 }
464 }
465 }
466 None => {
467 log::warn!(
468 "Cannot resubscribe to {topic}: no auth token"
469 );
470 continue;
471 }
472 }
473 } else {
474 payload.clone()
475 };
476
477 if let Err(e) = cmd_tx_for_reconnect
478 .send(SpotHandlerCommand::Subscribe { payload })
479 {
480 log::error!(
481 "Failed to send resubscribe command: error={e}, \
482 topic={topic}"
483 );
484 }
485
486 subscriptions.mark_subscribe(topic);
487 }
488 }
489
490 if out_tx.send(KrakenSpotWsMessage::Reconnected).is_err() {
491 if handler.is_stopped() {
492 log::debug!("Failed to send message (receiver dropped)");
493 } else {
494 log::error!("Failed to send message (receiver dropped)");
495 }
496 break;
497 }
498 }
499 Some(msg) => {
500 if out_tx.send(msg).is_err() {
501 if handler.is_stopped() {
502 log::debug!("Failed to send message (receiver dropped)");
503 } else {
504 log::error!("Failed to send message (receiver dropped)");
505 }
506 break;
507 }
508 }
509 None => {
510 if handler.is_stopped() {
511 log::debug!("Stop signal received, ending message processing");
512 break;
513 }
514 log::warn!("WebSocket stream ended unexpectedly");
515 break;
516 }
517 }
518 }
519
520 log::debug!("Handler task exiting");
521 };
522
523 if let Err(e) = handler_spawner.spawn(handler_task) {
524 if let Some(control) = &self.socket_control {
525 control.deregister();
526 }
527 self.out_rx = None;
528 return Err(KrakenWsError::ConnectionError(format!(
529 "Failed to register WebSocket handler task: {e}"
530 )));
531 }
532
533 log::debug!("WebSocket connected successfully");
534 Ok(())
535 }
536
537 pub async fn disconnect(&mut self) -> Result<(), KrakenWsError> {
539 let connect_lock = Arc::clone(&self.connect_lock);
540 let _connect_guard = connect_lock.lock().await;
541 self.disconnect_locked().await
542 }
543
544 async fn disconnect_locked(&self) -> Result<(), KrakenWsError> {
545 log::debug!("Disconnecting WebSocket");
546
547 self.handler_tasks.begin_shutdown();
548 self.signal.store(true, Ordering::Relaxed);
549
550 if let Err(e) = self
551 .cmd_tx
552 .read()
553 .await
554 .send(SpotHandlerCommand::Disconnect)
555 {
556 log::debug!(
557 "Failed to send disconnect command (handler may already be shut down): {e}"
558 );
559 }
560
561 let task_result = self
562 .handler_tasks
563 .finish_shutdown(
564 tokio::time::Duration::from_secs(2),
565 tokio::time::Duration::from_secs(2),
566 )
567 .await;
568
569 self.subscriptions.clear();
570 self.subscription_payloads.write().await.clear();
571 self.auth_tracker.fail("Disconnected");
572
573 self.l3_depths.lock().clear();
574 self.l2_depths.clear();
575 self.book_requests.lock().clear();
576
577 if let Some(control) = &self.socket_control {
578 control.deregister();
579 }
580
581 match task_result {
582 Ok(()) => Ok(()),
583 Err(error @ TaskShutdownError::Timeout { .. }) => Err(KrakenWsError::Timeout(format!(
584 "Spot WebSocket handler shutdown timed out: {error}"
585 ))),
586 Err(e) => Err(KrakenWsError::Disconnected(format!(
587 "Spot WebSocket handler shutdown failed: {e}"
588 ))),
589 }
590 }
591
592 pub async fn close(&mut self) -> Result<(), KrakenWsError> {
594 self.disconnect().await
595 }
596
597 pub async fn wait_until_active(&self, timeout_secs: f64) -> Result<(), KrakenWsError> {
599 let timeout = tokio::time::Duration::from_secs_f64(timeout_secs);
600
601 tokio::time::timeout(timeout, async {
602 while !self.is_active() {
603 tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
604 }
605 })
606 .await
607 .map_err(|_| {
608 KrakenWsError::ConnectionError(format!(
609 "WebSocket connection timeout after {timeout_secs} seconds"
610 ))
611 })?;
612
613 Ok(())
614 }
615
616 #[must_use]
618 pub fn is_authenticated(&self) -> bool {
619 self.auth_tracker.is_authenticated()
620 }
621
622 pub async fn wait_until_authenticated(&self, timeout_secs: f64) -> Result<(), KrakenWsError> {
626 let timeout = tokio::time::Duration::from_secs_f64(timeout_secs);
627
628 if self.auth_tracker.wait_for_authenticated(timeout).await {
629 Ok(())
630 } else {
631 Err(KrakenWsError::AuthenticationError(format!(
632 "Authentication not completed within {timeout_secs} seconds"
633 )))
634 }
635 }
636
637 pub async fn authenticate(&self) -> Result<(), KrakenWsError> {
639 if !self.config.has_api_credentials() {
640 return Err(KrakenWsError::AuthenticationError(
641 "API credentials required for authentication".to_string(),
642 ));
643 }
644
645 let _receiver = self.auth_tracker.begin();
646
647 match refresh_auth_token(&self.config).await {
648 Ok(token) => {
649 *self.auth_token.write().await = Some(token);
650 self.auth_tracker.succeed();
651 Ok(())
652 }
653 Err(e) => {
654 *self.auth_token.write().await = None;
655 self.auth_tracker.fail(e.to_string());
656 Err(e)
657 }
658 }
659 }
660
661 pub fn cancel_all_requests(&self) {
663 self.cancellation_token.cancel();
664 }
665
666 pub fn cancellation_token(&self) -> &CancellationToken {
668 &self.cancellation_token
669 }
670
671 pub async fn subscribe(
673 &self,
674 channel: KrakenWsChannel,
675 symbols: Vec<Ustr>,
676 depth: Option<u32>,
677 ) -> Result<(), KrakenWsError> {
678 if matches!(channel, KrakenWsChannel::Level3) {
679 return Err(KrakenWsError::InvalidMessage(
680 "Use subscribe_book_l3 / unsubscribe_book_l3 for the Level3 channel".to_string(),
681 ));
682 }
683 let mut symbols_to_subscribe = Vec::new();
684 let channel_str = channel.as_ref();
685 for symbol in &symbols {
686 let key = format!("{channel_str}:{symbol}");
687 if self.subscriptions.add_reference(&key) {
688 self.subscriptions.mark_subscribe(&key);
689 symbols_to_subscribe.push(*symbol);
690 }
691 }
692
693 if symbols_to_subscribe.is_empty() {
694 return Ok(());
695 }
696
697 let is_private = matches!(
698 channel,
699 KrakenWsChannel::Executions | KrakenWsChannel::Balances
700 );
701 let token = if is_private {
702 Some(self.auth_token().await.ok_or_else(|| {
703 KrakenWsError::AuthenticationError(
704 "Authentication token required for private channels. Call authenticate() first"
705 .to_string(),
706 )
707 })?)
708 } else {
709 None
710 };
711
712 let req_id = self.get_next_req_id();
713 let request = KrakenWsRequest {
714 method: KrakenWsMethod::Subscribe,
715 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
716 channel,
717 symbol: Some(symbols_to_subscribe.clone()),
718 snapshot: None,
719 depth,
720 interval: None,
721 event_trigger: None,
722 token,
723 snap_orders: None,
724 snap_trades: None,
725 })),
726 req_id: Some(req_id),
727 };
728
729 let payload = self.send_command(&request).await?;
730
731 for symbol in &symbols_to_subscribe {
732 let key = format!("{channel_str}:{symbol}");
733 self.subscriptions.confirm_subscribe(&key);
734 self.subscription_payloads
735 .write()
736 .await
737 .insert(key, payload.clone());
738 }
739
740 Ok(())
741 }
742
743 async fn subscribe_with_interval(
745 &self,
746 channel: KrakenWsChannel,
747 symbols: Vec<Ustr>,
748 interval: u32,
749 ) -> Result<(), KrakenWsError> {
750 let mut symbols_to_subscribe = Vec::new();
751 let channel_str = channel.as_ref();
752 for symbol in &symbols {
753 let key = format!("{channel_str}:{symbol}:{interval}");
754 if self.subscriptions.add_reference(&key) {
755 self.subscriptions.mark_subscribe(&key);
756 symbols_to_subscribe.push(*symbol);
757 }
758 }
759
760 if symbols_to_subscribe.is_empty() {
761 return Ok(());
762 }
763
764 let req_id = self.get_next_req_id();
765 let request = KrakenWsRequest {
766 method: KrakenWsMethod::Subscribe,
767 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
768 channel,
769 symbol: Some(symbols_to_subscribe.clone()),
770 snapshot: Some(false),
771 depth: None,
772 interval: Some(interval),
773 event_trigger: None,
774 token: None,
775 snap_orders: None,
776 snap_trades: None,
777 })),
778 req_id: Some(req_id),
779 };
780
781 let payload = self.send_command(&request).await?;
782
783 for symbol in &symbols_to_subscribe {
784 let key = format!("{channel_str}:{symbol}:{interval}");
785 self.subscriptions.confirm_subscribe(&key);
786 self.subscription_payloads
787 .write()
788 .await
789 .insert(key, payload.clone());
790 }
791
792 Ok(())
793 }
794
795 async fn unsubscribe_with_interval(
797 &self,
798 channel: KrakenWsChannel,
799 symbols: Vec<Ustr>,
800 interval: u32,
801 ) -> Result<(), KrakenWsError> {
802 let mut symbols_to_unsubscribe = Vec::new();
803 let channel_str = channel.as_ref();
804 for symbol in &symbols {
805 let key = format!("{channel_str}:{symbol}:{interval}");
806 if self.subscriptions.remove_reference(&key) {
807 self.subscriptions.mark_unsubscribe(&key);
808 symbols_to_unsubscribe.push(*symbol);
809 }
810 }
811
812 if symbols_to_unsubscribe.is_empty() {
813 return Ok(());
814 }
815
816 let req_id = self.get_next_req_id();
817 let request = KrakenWsRequest {
818 method: KrakenWsMethod::Unsubscribe,
819 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
820 channel,
821 symbol: Some(symbols_to_unsubscribe.clone()),
822 snapshot: None,
823 depth: None,
824 interval: Some(interval),
825 event_trigger: None,
826 token: None,
827 snap_orders: None,
828 snap_trades: None,
829 })),
830 req_id: Some(req_id),
831 };
832
833 self.send_command(&request).await?;
834
835 for symbol in &symbols_to_unsubscribe {
836 let key = format!("{channel_str}:{symbol}:{interval}");
837 self.subscriptions.confirm_unsubscribe(&key);
838 self.subscription_payloads.write().await.remove(&key);
839 }
840
841 Ok(())
842 }
843
844 pub async fn unsubscribe(
846 &self,
847 channel: KrakenWsChannel,
848 symbols: Vec<Ustr>,
849 ) -> Result<(), KrakenWsError> {
850 if matches!(channel, KrakenWsChannel::Level3) {
851 return Err(KrakenWsError::InvalidMessage(
852 "Use subscribe_book_l3 / unsubscribe_book_l3 for the Level3 channel".to_string(),
853 ));
854 }
855 let mut symbols_to_unsubscribe = Vec::new();
856 let channel_str = channel.as_ref();
857 for symbol in &symbols {
858 let key = format!("{channel_str}:{symbol}");
859 if self.subscriptions.remove_reference(&key) {
860 self.subscriptions.mark_unsubscribe(&key);
861 symbols_to_unsubscribe.push(*symbol);
862 } else {
863 log::debug!(
864 "Channel {channel_str} symbol {symbol} still has active subscriptions, not unsubscribing"
865 );
866 }
867 }
868
869 if symbols_to_unsubscribe.is_empty() {
870 return Ok(());
871 }
872
873 let is_private = matches!(
874 channel,
875 KrakenWsChannel::Executions | KrakenWsChannel::Balances
876 );
877 let token = if is_private {
878 Some(self.auth_token().await.ok_or_else(|| {
879 KrakenWsError::AuthenticationError(
880 "Authentication token required for private channels. Call authenticate() first"
881 .to_string(),
882 )
883 })?)
884 } else {
885 None
886 };
887
888 let req_id = self.get_next_req_id();
889 let request = KrakenWsRequest {
890 method: KrakenWsMethod::Unsubscribe,
891 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
892 channel,
893 symbol: Some(symbols_to_unsubscribe.clone()),
894 snapshot: None,
895 depth: None,
896 interval: None,
897 event_trigger: None,
898 token,
899 snap_orders: None,
900 snap_trades: None,
901 })),
902 req_id: Some(req_id),
903 };
904
905 self.send_command(&request).await?;
906
907 for symbol in &symbols_to_unsubscribe {
908 let key = format!("{channel_str}:{symbol}");
909 self.subscriptions.confirm_unsubscribe(&key);
910 self.subscription_payloads.write().await.remove(&key);
911 }
912
913 Ok(())
914 }
915
916 pub async fn send_ping(&self) -> Result<(), KrakenWsError> {
918 let req_id = self.get_next_req_id();
919
920 let request = KrakenWsRequest {
921 method: KrakenWsMethod::Ping,
922 params: None,
923 req_id: Some(req_id),
924 };
925
926 self.send_command(&request).await?;
927 Ok(())
928 }
929
930 async fn send_command(&self, request: &KrakenWsRequest) -> Result<SecretString, KrakenWsError> {
931 let payload = SecretString::from(
932 serde_json::to_string(request).map_err(|e| KrakenWsError::JsonError(e.to_string()))?,
933 );
934
935 log::trace!(
936 "Sending WebSocket request: method={:?} ({} bytes)",
937 request.method,
938 payload.expose_secret().len(),
939 );
940
941 let cmd = match request.method {
942 KrakenWsMethod::Subscribe => SpotHandlerCommand::Subscribe {
943 payload: payload.clone(),
944 },
945 KrakenWsMethod::Unsubscribe => SpotHandlerCommand::Unsubscribe {
946 payload: payload.clone(),
947 },
948 KrakenWsMethod::Ping | KrakenWsMethod::Pong => SpotHandlerCommand::Ping {
949 payload: payload.clone(),
950 },
951 KrakenWsMethod::AddOrder
952 | KrakenWsMethod::AmendOrder
953 | KrakenWsMethod::CancelOrder
954 | KrakenWsMethod::BatchAdd => {
955 return Err(KrakenWsError::InvalidMessage(
956 "Order methods must not be sent via send_command; use the dedicated order submission path".to_string()
957 ));
958 }
959 };
960
961 self.cmd_tx
962 .read()
963 .await
964 .send(cmd)
965 .map_err(|e| KrakenWsError::ConnectionError(format!("Failed to send request: {e}")))?;
966
967 Ok(payload)
968 }
969
970 pub fn is_connected(&self) -> bool {
972 let connection_mode_arc = self.connection_mode.load();
973 !ConnectionMode::from_atomic(&connection_mode_arc).is_closed()
974 }
975
976 pub fn is_active(&self) -> bool {
978 let connection_mode_arc = self.connection_mode.load();
979 ConnectionMode::from_atomic(&connection_mode_arc).is_active()
980 && !self.signal.load(Ordering::Relaxed)
981 }
982
983 pub fn is_closed(&self) -> bool {
985 let connection_mode_arc = self.connection_mode.load();
986 ConnectionMode::from_atomic(&connection_mode_arc).is_closed()
987 || self.signal.load(Ordering::Relaxed)
988 }
989
990 pub fn url(&self) -> &str {
992 self.url.expose()
993 }
994
995 pub fn get_subscriptions(&self) -> Vec<String> {
997 self.subscriptions.all_topics()
998 }
999
1000 pub fn subscriptions_contains(&self, topic: &str) -> bool {
1002 self.subscriptions.all_topics().iter().any(|t| t == topic)
1003 }
1004
1005 pub fn set_account_id(&self, account_id: AccountId) {
1007 *self.account_id.write() = Some(account_id);
1008 }
1009
1010 #[must_use]
1012 pub fn account_id(&self) -> Option<AccountId> {
1013 *self.account_id.read()
1014 }
1015
1016 pub fn cache_instrument(&self, instrument: InstrumentAny) {
1018 self.instruments.insert(instrument.id(), instrument);
1019 }
1020
1021 pub fn account_id_shared(&self) -> &Arc<RwLock<Option<AccountId>>> {
1023 &self.account_id
1024 }
1025
1026 pub fn truncated_id_map(&self) -> &Arc<AtomicMap<String, ClientOrderId>> {
1028 &self.truncated_id_map
1029 }
1030
1031 pub fn cache_client_order(
1033 &self,
1034 client_order_id: ClientOrderId,
1035 _venue_order_id: Option<VenueOrderId>,
1036 _instrument_id: InstrumentId,
1037 _trader_id: TraderId,
1038 _strategy_id: StrategyId,
1039 ) {
1040 let truncated = crate::common::parse::truncate_cl_ord_id(&client_order_id);
1041
1042 if truncated != client_order_id.as_str() {
1043 self.truncated_id_map.insert(truncated, client_order_id);
1044 }
1045 }
1046
1047 pub fn stream(
1055 &mut self,
1056 ) -> Result<impl futures_util::Stream<Item = KrakenSpotWsMessage> + use<>, KrakenWsError> {
1057 let rx = self.out_rx.take().ok_or_else(|| {
1058 KrakenWsError::ChannelError(
1059 "Stream receiver already taken or client not connected".to_string(),
1060 )
1061 })?;
1062 let mut rx = Arc::try_unwrap(rx).map_err(|_| {
1063 KrakenWsError::ChannelError(
1064 "Cannot take ownership of stream - other client clones still hold references"
1065 .to_string(),
1066 )
1067 })?;
1068 Ok(async_stream::stream! {
1069 while let Some(msg) = rx.recv().await {
1070 yield msg;
1071 }
1072 })
1073 }
1074
1075 pub async fn subscribe_book(
1077 &self,
1078 instrument_id: InstrumentId,
1079 depth: Option<u32>,
1080 ) -> Result<(), KrakenWsError> {
1081 let _book_commands = self.l2_book_commands.lock().await;
1082 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1083 let depth = depth.unwrap_or(10);
1084
1085 if !matches!(depth, 10 | 25 | 100 | 500 | 1000) {
1086 return Err(KrakenWsError::InvalidMessage(format!(
1087 "Invalid L2 depth {depth}, valid values: 10, 25, 100, 500, 1000",
1088 )));
1089 }
1090
1091 let channel_str = KrakenWsChannel::Book.as_ref();
1092 let key = format!("{channel_str}:{symbol}");
1093 let is_first_reference = self.subscriptions.add_reference(&key);
1094
1095 if !is_first_reference {
1096 let existing_depth = self.l2_depths.get(symbol.as_str());
1097
1098 if existing_depth != Some(depth) {
1099 self.subscriptions.remove_reference(&key);
1100 return Err(KrakenWsError::InvalidMessage(format!(
1101 "L2 subscription for {symbol} already exists with depth \
1102 {existing_depth:?}, cannot resubscribe with depth {depth}",
1103 )));
1104 }
1105 return Ok(());
1106 }
1107
1108 self.subscriptions.mark_subscribe(&key);
1109 let req_id = self.get_next_req_id();
1110 self.l2_depths.insert(symbol.as_str(), depth, req_id);
1111
1112 let request = KrakenWsRequest {
1113 method: KrakenWsMethod::Subscribe,
1114 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1115 channel: KrakenWsChannel::Book,
1116 symbol: Some(vec![symbol]),
1117 snapshot: None,
1118 depth: Some(depth),
1119 interval: None,
1120 event_trigger: None,
1121 token: None,
1122 snap_orders: None,
1123 snap_trades: None,
1124 })),
1125 req_id: Some(req_id),
1126 };
1127 self.record_book_request(req_id, symbol);
1128
1129 let payload = match self.send_command(&request).await {
1130 Ok(payload) => payload,
1131 Err(e) => {
1132 self.book_requests.lock().remove(&req_id);
1133 self.l2_depths.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 subscribe_book_l3(&self, symbol: Ustr, depth: u32) -> Result<(), KrakenWsError> {
1168 if !matches!(depth, 10 | 100 | 1000) {
1169 return Err(KrakenWsError::InvalidMessage(format!(
1170 "Invalid L3 depth {depth}, valid values: 10, 100, 1000",
1171 )));
1172 }
1173
1174 let token = self.auth_token().await.ok_or_else(|| {
1175 KrakenWsError::AuthenticationError(
1176 "Authentication token required for level3. Call authenticate() first".to_string(),
1177 )
1178 })?;
1179
1180 let channel_str = KrakenWsChannel::Level3.as_ref();
1181 let key = format!("{channel_str}:{symbol}");
1182
1183 let is_first_reference = self.subscriptions.add_reference(&key);
1184
1185 if !is_first_reference {
1186 let existing_depth = self.l3_depths.lock().get(symbol.as_str()).copied();
1187
1188 if existing_depth != Some(depth) {
1189 self.subscriptions.remove_reference(&key);
1190 return Err(KrakenWsError::InvalidMessage(format!(
1191 "L3 subscription for {symbol} already exists with depth \
1192 {existing_depth:?}, cannot resubscribe with depth {depth}",
1193 )));
1194 }
1195 return Ok(());
1196 }
1197
1198 self.subscriptions.mark_subscribe(&key);
1199
1200 self.l3_depths.lock().insert(symbol.to_string(), depth);
1201
1202 let req_id = self.get_next_req_id();
1203 let request = KrakenWsRequest {
1204 method: KrakenWsMethod::Subscribe,
1205 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1206 channel: KrakenWsChannel::Level3,
1207 symbol: Some(vec![symbol]),
1208 snapshot: Some(true),
1209 depth: Some(depth),
1210 interval: None,
1211 event_trigger: None,
1212 token: Some(token),
1213 snap_orders: None,
1214 snap_trades: None,
1215 })),
1216 req_id: Some(req_id),
1217 };
1218
1219 let payload = match self.send_command(&request).await {
1220 Ok(p) => p,
1221 Err(e) => {
1222 self.l3_depths.lock().remove(symbol.as_str());
1223 self.subscriptions.remove_reference(&key);
1224 self.subscriptions.mark_unsubscribe(&key);
1225 self.subscriptions.confirm_unsubscribe(&key);
1226 return Err(e);
1227 }
1228 };
1229
1230 self.subscriptions.confirm_subscribe(&key);
1231 self.subscription_payloads
1232 .write()
1233 .await
1234 .insert(key, payload);
1235 Ok(())
1236 }
1237
1238 pub async fn unsubscribe_book_l3(&self, symbol: Ustr) -> Result<(), KrakenWsError> {
1253 let channel_str = KrakenWsChannel::Level3.as_ref();
1254 let key = format!("{channel_str}:{symbol}");
1255 if !self.subscriptions.remove_reference(&key) {
1256 return Ok(());
1257 }
1258 self.subscriptions.mark_unsubscribe(&key);
1259
1260 let token = self.auth_token().await;
1261 let req_id = self.get_next_req_id();
1262 let request = KrakenWsRequest {
1263 method: KrakenWsMethod::Unsubscribe,
1264 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1265 channel: KrakenWsChannel::Level3,
1266 symbol: Some(vec![symbol]),
1267 snapshot: None,
1268 depth: None,
1269 interval: None,
1270 event_trigger: None,
1271 token,
1272 snap_orders: None,
1273 snap_trades: None,
1274 })),
1275 req_id: Some(req_id),
1276 };
1277
1278 self.send_command(&request).await?;
1279 self.subscriptions.confirm_unsubscribe(&key);
1280 self.subscription_payloads.write().await.remove(&key);
1281 self.l3_depths.lock().remove(symbol.as_str());
1282 Ok(())
1283 }
1284
1285 pub async fn resync_book_l3(&self, symbol: Ustr, depth: u32) -> Result<(), KrakenWsError> {
1299 let channel_str = KrakenWsChannel::Level3.as_ref();
1300 let key = format!("{channel_str}:{symbol}");
1301
1302 if !self.subscriptions_contains(&key) {
1308 log::debug!("Skipping L3 resync: subscription cancelled mid-retry, symbol={symbol}",);
1309 return Ok(());
1310 }
1311
1312 let new_token = refresh_auth_token(&self.config).await?;
1313 *self.auth_token.write().await = Some(new_token.clone());
1314
1315 if !self.subscriptions_contains(&key) {
1318 log::debug!(
1319 "Skipping L3 resync: subscription cancelled after token refresh, symbol={symbol}",
1320 );
1321 return Ok(());
1322 }
1323
1324 let unsub_req_id = self.get_next_req_id();
1325 let unsub = KrakenWsRequest {
1326 method: KrakenWsMethod::Unsubscribe,
1327 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1328 channel: KrakenWsChannel::Level3,
1329 symbol: Some(vec![symbol]),
1330 snapshot: None,
1331 depth: None,
1332 interval: None,
1333 event_trigger: None,
1334 token: Some(new_token.clone()),
1335 snap_orders: None,
1336 snap_trades: None,
1337 })),
1338 req_id: Some(unsub_req_id),
1339 };
1340 self.send_command(&unsub).await?;
1341
1342 if !self.subscriptions_contains(&key) {
1344 log::debug!("Skipping L3 resync resubscribe: cancelled before send, symbol={symbol}",);
1345 return Ok(());
1346 }
1347
1348 let sub_req_id = self.get_next_req_id();
1349 let sub = KrakenWsRequest {
1350 method: KrakenWsMethod::Subscribe,
1351 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1352 channel: KrakenWsChannel::Level3,
1353 symbol: Some(vec![symbol]),
1354 snapshot: Some(true),
1355 depth: Some(depth),
1356 interval: None,
1357 event_trigger: None,
1358 token: Some(new_token.clone()),
1359 snap_orders: None,
1360 snap_trades: None,
1361 })),
1362 req_id: Some(sub_req_id),
1363 };
1364 let payload = self.send_command(&sub).await?;
1365
1366 if self.subscriptions_contains(&key) {
1371 self.subscription_payloads
1372 .write()
1373 .await
1374 .insert(key, payload);
1375 self.l3_depths.lock().insert(symbol.to_string(), depth);
1376 }
1377
1378 Ok(())
1379 }
1380
1381 pub fn validate_l2_checksum(&self) -> bool {
1383 self.config.validate_l2_checksum
1384 }
1385
1386 pub async fn resync_book(
1407 &self,
1408 instrument_id: InstrumentId,
1409 generation: u64,
1410 epoch: u64,
1411 ) -> Result<(), KrakenWsError> {
1412 let _book_commands = self.l2_book_commands.lock().await;
1413 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1414 let channel_str = KrakenWsChannel::Book.as_ref();
1415 let key = format!("{channel_str}:{symbol}");
1416 let req_id = self.get_next_req_id();
1417 let live = self
1418 .subscriptions_contains(&key)
1419 .then(|| {
1420 self.l2_depths
1421 .begin_resync(symbol.as_str(), generation, epoch, req_id)
1422 })
1423 .flatten();
1424
1425 let Some(live) = live else {
1426 log::debug!(
1430 "Skipping L2 resync: subscription canceled or replaced, or its snapshot served \
1431 since, symbol={symbol}, generation={generation}, epoch={epoch}"
1432 );
1433 return Ok(());
1434 };
1435 let depth = live.depth;
1436
1437 let unsub = KrakenWsRequest {
1438 method: KrakenWsMethod::Unsubscribe,
1439 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1440 channel: KrakenWsChannel::Book,
1441 symbol: Some(vec![symbol]),
1442 snapshot: None,
1443 depth: Some(depth),
1446 interval: None,
1447 event_trigger: None,
1448 token: None,
1449 snap_orders: None,
1450 snap_trades: None,
1451 })),
1452 req_id: Some(self.get_next_req_id()),
1453 };
1454 self.send_command(&unsub).await?;
1455
1456 let sub = KrakenWsRequest {
1457 method: KrakenWsMethod::Subscribe,
1458 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1459 channel: KrakenWsChannel::Book,
1460 symbol: Some(vec![symbol]),
1461 snapshot: Some(true),
1463 depth: Some(depth),
1464 interval: None,
1465 event_trigger: None,
1466 token: None,
1467 snap_orders: None,
1468 snap_trades: None,
1469 })),
1470 req_id: Some(req_id),
1471 };
1472 self.record_book_request(req_id, symbol);
1473
1474 let payload = match self.send_command(&sub).await {
1475 Ok(payload) => payload,
1476 Err(e) => {
1477 self.book_requests.lock().remove(&req_id);
1478 return Err(e);
1479 }
1480 };
1481 self.subscription_payloads
1482 .write()
1483 .await
1484 .insert(key, payload);
1485
1486 Ok(())
1487 }
1488
1489 fn record_book_request(&self, req_id: u64, symbol: Ustr) {
1490 self.book_requests
1491 .lock()
1492 .insert(req_id, L2BookRequest { symbol });
1493 }
1494
1495 pub fn validate_l3_checksum(&self) -> bool {
1497 self.config.validate_l3_checksum
1498 }
1499
1500 pub fn has_credentials(&self) -> bool {
1503 self.config.has_api_credentials()
1504 }
1505
1506 pub fn instruments_handle(&self) -> Arc<AtomicMap<InstrumentId, InstrumentAny>> {
1512 Arc::clone(&self.instruments)
1513 }
1514
1515 pub fn l3_depths_handle(&self) -> Arc<parking_lot::Mutex<ahash::AHashMap<String, u32>>> {
1520 Arc::clone(&self.l3_depths)
1521 }
1522
1523 pub(crate) fn l2_depths_handle(&self) -> L2Depths {
1525 self.l2_depths.clone()
1526 }
1527
1528 pub(crate) fn book_requests_handle(&self) -> L2BookRequests {
1530 Arc::clone(&self.book_requests)
1531 }
1532
1533 pub async fn subscribe_quotes(&self, instrument_id: InstrumentId) -> Result<(), KrakenWsError> {
1538 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1539 let key = format!("quotes:{symbol}");
1540
1541 if !self.subscriptions.add_reference(&key) {
1542 return Ok(());
1543 }
1544
1545 self.subscriptions.mark_subscribe(&key);
1546
1547 let req_id = self.get_next_req_id();
1548 let request = KrakenWsRequest {
1549 method: KrakenWsMethod::Subscribe,
1550 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1551 channel: KrakenWsChannel::Ticker,
1552 symbol: Some(vec![symbol]),
1553 snapshot: None,
1554 depth: None,
1555 interval: None,
1556 event_trigger: Some("bbo".to_string()),
1557 token: None,
1558 snap_orders: None,
1559 snap_trades: None,
1560 })),
1561 req_id: Some(req_id),
1562 };
1563
1564 let payload = self.send_command(&request).await?;
1565 self.subscriptions.confirm_subscribe(&key);
1566 self.subscription_payloads
1567 .write()
1568 .await
1569 .insert(key, payload);
1570 Ok(())
1571 }
1572
1573 pub async fn subscribe_trades(&self, instrument_id: InstrumentId) -> Result<(), KrakenWsError> {
1575 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1576 self.subscribe(KrakenWsChannel::Trade, vec![symbol], None)
1577 .await
1578 }
1579
1580 pub async fn subscribe_bars(&self, bar_type: BarType) -> Result<(), KrakenWsError> {
1586 let symbol = to_ws_v2_symbol(bar_type.instrument_id().symbol.inner());
1587 let interval = bar_type_to_ws_interval(bar_type)?;
1588 self.subscribe_with_interval(KrakenWsChannel::Ohlc, vec![symbol], interval)
1589 .await
1590 }
1591
1592 pub async fn subscribe_executions(
1596 &self,
1597 snap_orders: bool,
1598 snap_trades: bool,
1599 ) -> Result<(), KrakenWsError> {
1600 let req_id = self.get_next_req_id();
1601
1602 let token = self.auth_token().await.ok_or_else(|| {
1603 KrakenWsError::AuthenticationError(
1604 "Authentication token required for executions channel. Call authenticate() first"
1605 .to_string(),
1606 )
1607 })?;
1608
1609 let request = KrakenWsRequest {
1610 method: KrakenWsMethod::Subscribe,
1611 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1612 channel: KrakenWsChannel::Executions,
1613 symbol: None,
1614 snapshot: None,
1615 depth: None,
1616 interval: None,
1617 event_trigger: None,
1618 token: Some(token),
1619 snap_orders: Some(snap_orders),
1620 snap_trades: Some(snap_trades),
1621 })),
1622 req_id: Some(req_id),
1623 };
1624
1625 let payload = self.send_command(&request).await?;
1626
1627 let key = "executions";
1628 if self.subscriptions.add_reference(key) {
1629 self.subscriptions.mark_subscribe(key);
1630 self.subscriptions.confirm_subscribe(key);
1631 self.subscription_payloads
1632 .write()
1633 .await
1634 .insert(key.to_string(), payload);
1635 }
1636
1637 Ok(())
1638 }
1639
1640 pub async fn unsubscribe_book(&self, instrument_id: InstrumentId) -> Result<(), KrakenWsError> {
1645 let _book_commands = self.l2_book_commands.lock().await;
1646 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1647 let channel_str = KrakenWsChannel::Book.as_ref();
1648 let key = format!("{channel_str}:{symbol}");
1649
1650 if !self.subscriptions.remove_reference(&key) {
1651 return Ok(());
1652 }
1653
1654 self.subscriptions.mark_unsubscribe(&key);
1655
1656 let req_id = self.get_next_req_id();
1657 let request = KrakenWsRequest {
1658 method: KrakenWsMethod::Unsubscribe,
1659 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1660 channel: KrakenWsChannel::Book,
1661 symbol: Some(vec![symbol]),
1662 snapshot: None,
1663 depth: self.l2_depths.get(symbol.as_str()),
1664 interval: None,
1665 event_trigger: None,
1666 token: None,
1667 snap_orders: None,
1668 snap_trades: None,
1669 })),
1670 req_id: Some(req_id),
1671 };
1672
1673 self.send_command(&request).await?;
1674 self.subscriptions.confirm_unsubscribe(&key);
1675 self.subscription_payloads.write().await.remove(&key);
1676 self.l2_depths.remove(symbol.as_str());
1677 Ok(())
1678 }
1679
1680 pub async fn unsubscribe_quotes(
1682 &self,
1683 instrument_id: InstrumentId,
1684 ) -> Result<(), KrakenWsError> {
1685 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1686 let key = format!("quotes:{symbol}");
1687
1688 if !self.subscriptions.remove_reference(&key) {
1689 return Ok(());
1690 }
1691
1692 self.subscriptions.mark_unsubscribe(&key);
1693
1694 let req_id = self.get_next_req_id();
1695 let request = KrakenWsRequest {
1696 method: KrakenWsMethod::Unsubscribe,
1697 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1698 channel: KrakenWsChannel::Ticker,
1699 symbol: Some(vec![symbol]),
1700 snapshot: None,
1701 depth: None,
1702 interval: None,
1703 event_trigger: Some("bbo".to_string()),
1704 token: None,
1705 snap_orders: None,
1706 snap_trades: None,
1707 })),
1708 req_id: Some(req_id),
1709 };
1710
1711 self.send_command(&request).await?;
1712 self.subscriptions.confirm_unsubscribe(&key);
1713 self.subscription_payloads.write().await.remove(&key);
1714 Ok(())
1715 }
1716
1717 pub async fn unsubscribe_trades(
1719 &self,
1720 instrument_id: InstrumentId,
1721 ) -> Result<(), KrakenWsError> {
1722 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1723 self.unsubscribe(KrakenWsChannel::Trade, vec![symbol]).await
1724 }
1725
1726 pub async fn unsubscribe_bars(&self, bar_type: BarType) -> Result<(), KrakenWsError> {
1732 let symbol = to_ws_v2_symbol(bar_type.instrument_id().symbol.inner());
1733 let interval = bar_type_to_ws_interval(bar_type)?;
1734 self.unsubscribe_with_interval(KrakenWsChannel::Ohlc, vec![symbol], interval)
1735 .await
1736 }
1737}
1738
1739fn begin_resubscribe(
1748 subscriptions: &SubscriptionState,
1749 book_requests: &L2BookRequests,
1750 l2_depths: &L2Depths,
1751) {
1752 subscriptions.reset_after_reconnect();
1753 book_requests.lock().clear();
1754 l2_depths.advance_epochs();
1755}
1756
1757async fn refresh_auth_token(
1759 config: &KrakenDataClientConfig,
1760) -> Result<SecretString, KrakenWsError> {
1761 let api_key = config
1762 .api_key
1763 .clone()
1764 .ok_or_else(|| KrakenWsError::AuthenticationError("Missing API key".to_string()))?;
1765 let api_secret = config
1766 .api_secret
1767 .clone()
1768 .ok_or_else(|| KrakenWsError::AuthenticationError("Missing API secret".to_string()))?;
1769
1770 let http_client = KrakenSpotHttpClient::with_credentials(
1771 api_key.into_inner(),
1772 api_secret.into_inner(),
1773 config.environment,
1774 Some(config.http_base_url()),
1775 config.timeout_secs,
1776 None,
1777 None,
1778 None,
1779 config
1780 .proxy_url
1781 .as_ref()
1782 .map(|value| value.expose_secret().to_owned()),
1783 config
1784 .max_requests_per_second
1785 .unwrap_or(KRAKEN_SPOT_DEFAULT_RATE_LIMIT_PER_SECOND),
1786 )
1787 .map_err(|e| {
1788 KrakenWsError::AuthenticationError(format!("Failed to create HTTP client: {e}"))
1789 })?;
1790
1791 let ws_token = http_client.get_websockets_token().await.map_err(|e| {
1792 KrakenWsError::AuthenticationError(format!("Failed to get WebSocket token: {e}"))
1793 })?;
1794
1795 log::debug!(
1796 "WebSocket authentication token refreshed: token_length={}, expires={}",
1797 ws_token.token.expose_secret().len(),
1798 ws_token.expires
1799 );
1800
1801 Ok(ws_token.into_token())
1802}
1803
1804fn update_auth_token_in_payload(
1805 payload: &str,
1806 new_token: &str,
1807) -> Result<SecretString, KrakenWsError> {
1808 let mut value: serde_json::Value =
1809 serde_json::from_str(payload).map_err(|e| KrakenWsError::JsonError(e.to_string()))?;
1810
1811 if let Some(params) = value.get_mut("params") {
1812 params["token"] = serde_json::Value::String(new_token.to_string());
1813 }
1814
1815 serde_json::to_string(&value)
1816 .map(SecretString::from)
1817 .map_err(|e| KrakenWsError::JsonError(e.to_string()))
1818}
1819
1820#[inline]
1821fn to_ws_v2_symbol(symbol: Ustr) -> Ustr {
1822 Ustr::from(&normalize_spot_symbol(symbol.as_str()))
1823}
1824
1825fn bar_type_to_ws_interval(bar_type: BarType) -> Result<u32, KrakenWsError> {
1826 const VALID_INTERVALS: [u32; 9] = [1, 5, 15, 30, 60, 240, 1440, 10080, 21600];
1827
1828 let spec = bar_type.spec();
1829 let step = spec.step.get() as u32;
1830
1831 let base_minutes = match spec.aggregation {
1832 BarAggregation::Minute => 1,
1833 BarAggregation::Hour => 60,
1834 BarAggregation::Day => 1440,
1835 BarAggregation::Week => 10080,
1836 other => {
1837 return Err(KrakenWsError::SubscriptionError(format!(
1838 "Unsupported bar aggregation for Kraken OHLC streaming: {other:?}"
1839 )));
1840 }
1841 };
1842
1843 let interval = base_minutes * step;
1844
1845 if !VALID_INTERVALS.contains(&interval) {
1846 return Err(KrakenWsError::SubscriptionError(format!(
1847 "Invalid bar interval {interval} minutes for Kraken OHLC streaming. \
1848 Supported intervals: 1, 5, 15, 30, 60, 240, 1440, 10080, 21600"
1849 )));
1850 }
1851
1852 Ok(interval)
1853}
1854
1855#[cfg(test)]
1856mod tests {
1857 use std::{
1858 sync::{Arc, atomic::Ordering},
1859 time::Duration,
1860 };
1861
1862 use futures_util::StreamExt;
1863 use log::{Level, LevelFilter, Log, Metadata, Record};
1864 use nautilus_testkit::connection::race_command_with_publication;
1865 use parking_lot::Mutex;
1866 use rstest::rstest;
1867 use tokio_util::sync::CancellationToken;
1868
1869 use super::*;
1870 use crate::config::KrakenDataClientConfig;
1871
1872 #[tokio::test]
1873 async fn command_during_sender_publication_reaches_initialized_handler() {
1874 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1875 let addr = listener.local_addr().unwrap();
1876 let (received_tx, mut received_rx) = tokio::sync::mpsc::unbounded_channel();
1877
1878 let server = tokio::spawn(async move {
1879 let (socket, _) = listener.accept().await.unwrap();
1880 let mut socket = tokio_tungstenite::accept_async(socket).await.unwrap();
1881 while let Some(Ok(message)) = socket.next().await {
1882 if let tokio_tungstenite::tungstenite::Message::Text(text) = message {
1883 let payload: serde_json::Value = serde_json::from_str(text.as_str()).unwrap();
1884 let _ = received_tx.send(payload);
1885 }
1886 }
1887 });
1888
1889 let url = format!("ws://{addr}");
1890
1891 let mut ws_client = KrakenSpotWebSocketClient::new(
1892 KrakenDataClientConfig {
1893 ws_public_url: Some(url),
1894 ..Default::default()
1895 },
1896 CancellationToken::new(),
1897 None,
1898 )
1899 .unwrap();
1900
1901 let observer = ws_client.clone();
1902 let expected = serde_json::json!({
1903 "method": "subscribe",
1904 "params": {
1905 "channel": "trade",
1906 "symbol": [
1907 "BTC/USD"
1908 ]
1909 },
1910 "req_id": 79
1911 });
1912
1913 let command = SpotHandlerCommand::Subscribe {
1914 payload: SecretString::from(expected.to_string()),
1915 };
1916
1917 race_command_with_publication(
1918 Arc::clone(&ws_client.cmd_tx),
1919 ws_client.connect(),
1920 move || command,
1921 || observer.is_active(),
1922 )
1923 .await
1924 .expect("connect failed");
1925 let received = tokio::time::timeout(Duration::from_secs(2), received_rx.recv()).await;
1926 ws_client.close().await.expect("disconnect failed");
1927 server.abort();
1928
1929 let payload = received
1930 .expect("startup command must reach server")
1931 .unwrap();
1932 assert_eq!(payload, expected);
1933 }
1934
1935 #[rstest]
1936 #[case("wss://stream.example.com/feed", "wss://stream.example.com")]
1937 #[case(
1938 "wss://SYN-user:SYN-password@stream.example.com:9443/SYN-path?token=SYN-query#SYN-fragment",
1939 "wss://stream.example.com:9443"
1940 )]
1941 fn test_debug_preserves_url_origin(#[case] url: &str, #[case] origin: &str) {
1942 let client = KrakenSpotWebSocketClient::new(
1943 KrakenDataClientConfig {
1944 ws_public_url: Some(url.to_string()),
1945 ..Default::default()
1946 },
1947 CancellationToken::new(),
1948 None,
1949 )
1950 .unwrap();
1951
1952 let debug = format!("{client:?}");
1953 assert!(debug.contains(&format!("url: {origin:?}")));
1954 assert!(!debug.contains("SYN-"));
1955 assert_eq!(client.url.expose(), url);
1956 }
1957
1958 #[rstest]
1959 #[tokio::test]
1960 async fn test_debug_redacts_auth_state_and_proxy_url() {
1961 let client = KrakenSpotWebSocketClient::new(
1962 KrakenDataClientConfig::default(),
1963 CancellationToken::new(),
1964 Some("http://user:proxy-secret@localhost".to_string()),
1965 )
1966 .unwrap();
1967 *client.auth_token.write().await = Some(SecretString::from("auth-token".to_string()));
1968
1969 let debug = format!("{client:?}");
1970
1971 assert!(debug.contains(REDACTED));
1972 assert!(!debug.contains("proxy-secret"));
1973 assert!(!debug.contains("auth-token"));
1974 }
1975
1976 const SECRET_MARKER: &str = "OUTBOUND_SECRET_MARKER";
1977
1978 struct OutboundLogCapture {
1979 messages: Mutex<Vec<String>>,
1980 }
1981
1982 static OUTBOUND_LOG_CAPTURE: OutboundLogCapture = OutboundLogCapture {
1983 messages: Mutex::new(Vec::new()),
1984 };
1985
1986 impl OutboundLogCapture {
1987 fn clear(&self) {
1988 self.messages.lock().clear();
1989 }
1990
1991 fn messages(&self) -> Vec<String> {
1992 self.messages.lock().clone()
1993 }
1994 }
1995
1996 impl Log for OutboundLogCapture {
1997 fn enabled(&self, metadata: &Metadata<'_>) -> bool {
1998 metadata.level() == Level::Trace
1999 && metadata.target() == "nautilus_kraken::websocket::spot_v2::client"
2000 }
2001
2002 fn log(&self, record: &Record<'_>) {
2003 if self.enabled(record.metadata()) {
2004 let message = record.args().to_string();
2005 if message.starts_with("Sending WebSocket request") {
2006 self.messages.lock().push(message);
2007 }
2008 }
2009 }
2010
2011 fn flush(&self) {}
2012 }
2013
2014 #[rstest]
2018 fn test_a_reconnect_replay_retires_the_unanswered_book_requests() {
2019 let client = test_client_without_credentials();
2020 client.record_book_request(7, Ustr::from("BTC/USD"));
2021 assert_eq!(client.book_requests.lock().len(), 1);
2022
2023 begin_resubscribe(
2024 &client.subscriptions,
2025 &client.book_requests,
2026 &client.l2_depths,
2027 );
2028
2029 assert!(client.book_requests.lock().is_empty());
2030 }
2031
2032 #[rstest]
2033 fn test_req_id_counter_is_shared_arc_and_monotonic() {
2034 let cfg = KrakenDataClientConfig::default();
2035 let client = KrakenSpotWebSocketClient::new(cfg, CancellationToken::new(), None).unwrap();
2036 let counter = client.req_id_counter();
2037 let a = counter.fetch_add(1, Ordering::Relaxed);
2038 let b = counter.fetch_add(1, Ordering::Relaxed);
2039 assert!(b > a);
2040 #[allow(clippy::redundant_clone)]
2041 let cloned = client.clone();
2042 let cloned_counter = cloned.req_id_counter();
2043 assert!(Arc::ptr_eq(&counter, &cloned_counter));
2044 }
2045
2046 #[rstest]
2047 #[case("XBT/EUR", "BTC/EUR")]
2048 #[case("XBT/USD", "BTC/USD")]
2049 #[case("XBT/USDT", "BTC/USDT")]
2050 #[case("ETH/USD", "ETH/USD")]
2051 #[case("ETH/XBT", "ETH/BTC")]
2052 #[case("SOL/XBT", "SOL/BTC")]
2053 #[case("SOL/USD", "SOL/USD")]
2054 #[case("BTC/USD", "BTC/USD")]
2055 #[case("ETH/BTC", "ETH/BTC")]
2056 #[case("XDG/USD", "DOGE/USD")]
2057 #[case("XDG/EUR", "DOGE/EUR")]
2058 fn test_to_kraken_ws_v2_symbol(#[case] input: &str, #[case] expected: &str) {
2059 let symbol = Ustr::from(input);
2060 let result = to_ws_v2_symbol(symbol);
2061 assert_eq!(result, expected);
2062 }
2063
2064 fn test_client_without_credentials() -> KrakenSpotWebSocketClient {
2065 KrakenSpotWebSocketClient::new(
2066 KrakenDataClientConfig::default(),
2067 CancellationToken::new(),
2068 None,
2069 )
2070 .unwrap()
2071 }
2072
2073 #[rstest]
2074 #[tokio::test]
2075 async fn test_outbound_logs_omit_payload_bodies() {
2076 log::set_logger(&OUTBOUND_LOG_CAPTURE).expect("test logger already installed");
2077 log::set_max_level(LevelFilter::Trace);
2078
2079 let client = test_client_without_credentials();
2080 let request = KrakenWsRequest {
2081 method: KrakenWsMethod::Subscribe,
2082 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
2083 channel: KrakenWsChannel::Executions,
2084 symbol: None,
2085 snapshot: None,
2086 depth: None,
2087 interval: None,
2088 event_trigger: None,
2089 token: Some(SecretString::from(SECRET_MARKER)),
2090 snap_orders: Some(true),
2091 snap_trades: Some(false),
2092 })),
2093 req_id: Some(426),
2094 };
2095 let payload_len = serde_json::to_string(&request).unwrap().len();
2096 OUTBOUND_LOG_CAPTURE.clear();
2097
2098 let error = client.send_command(&request).await.unwrap_err();
2099 let messages = OUTBOUND_LOG_CAPTURE.messages();
2100
2101 assert!(matches!(error, KrakenWsError::ConnectionError(_)));
2102 assert!(
2103 messages
2104 .iter()
2105 .all(|message| !message.contains(SECRET_MARKER)),
2106 "outbound logs exposed the secret marker: {messages:?}"
2107 );
2108 assert!(
2109 messages.iter().any(|message| {
2110 message
2111 == &format!("Sending WebSocket request: method=Subscribe ({payload_len} bytes)")
2112 }),
2113 "subscribe metadata missing or inaccurate: {messages:?}"
2114 );
2115 }
2116
2117 #[rstest]
2118 #[tokio::test]
2119 async fn test_authenticate_without_credentials_errors() {
2120 let client = test_client_without_credentials();
2121
2122 let err = client.authenticate().await.expect_err("should fail");
2123 assert!(
2124 matches!(err, KrakenWsError::AuthenticationError(ref msg) if msg.contains("API credentials required")),
2125 "unexpected error: {err:?}"
2126 );
2127 assert!(!client.is_authenticated());
2128 }
2129
2130 #[rstest]
2131 #[tokio::test]
2132 async fn test_wait_until_authenticated_times_out() {
2133 let client = test_client_without_credentials();
2134
2135 let err = client
2136 .wait_until_authenticated(0.05)
2137 .await
2138 .expect_err("should time out");
2139 assert!(matches!(err, KrakenWsError::AuthenticationError(_)));
2140 }
2141
2142 #[rstest]
2143 #[tokio::test]
2144 async fn test_wait_until_authenticated_resolves_after_succeed() {
2145 let client = test_client_without_credentials();
2146
2147 let tracker = client.auth_tracker.clone();
2148 let _rx = tracker.begin();
2149
2150 tokio::spawn(async move {
2151 tokio::time::sleep(tokio::time::Duration::from_millis(20)).await;
2152 tracker.succeed();
2153 });
2154
2155 client
2156 .wait_until_authenticated(1.0)
2157 .await
2158 .expect("should resolve once tracker succeeds");
2159 assert!(client.is_authenticated());
2160 }
2161
2162 #[rstest]
2163 #[tokio::test]
2164 async fn test_is_authenticated_flips_on_fail() {
2165 let client = test_client_without_credentials();
2166
2167 let _rx = client.auth_tracker.begin();
2168 client.auth_tracker.succeed();
2169 assert!(client.is_authenticated());
2170
2171 client.auth_tracker.fail("test failure");
2172 assert!(!client.is_authenticated());
2173 }
2174
2175 #[rstest]
2176 fn test_l3_factory_uses_ws_l3_url() {
2177 let cfg = KrakenDataClientConfig::default();
2178 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None).unwrap();
2179 assert_eq!(client.url(), "wss://ws-l3.kraken.com/v2");
2180 }
2181
2182 #[rstest]
2183 fn test_l3_factory_respects_override() {
2184 let cfg = KrakenDataClientConfig {
2185 ws_l3_url: Some("wss://override.example/v2".to_string()),
2186 ..Default::default()
2187 };
2188
2189 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None).unwrap();
2190 assert_eq!(client.url(), "wss://override.example/v2");
2191 }
2192
2193 #[rstest]
2194 #[tokio::test]
2195 async fn test_subscribe_book_l3_without_auth_errors_and_leaves_clean_state() {
2196 let cfg = KrakenDataClientConfig::default();
2197 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None).unwrap();
2198
2199 let err = client
2200 .subscribe_book_l3(Ustr::from("BTC/USD"), 1000)
2201 .await
2202 .expect_err("should fail without auth token");
2203
2204 assert!(matches!(err, KrakenWsError::AuthenticationError(_)));
2205 assert!(
2206 client.subscriptions.is_empty(),
2207 "no state must leak on auth failure"
2208 );
2209 }
2210
2211 #[rstest]
2212 #[tokio::test]
2213 async fn test_subscribe_book_l3_invalid_depth_errors() {
2214 let cfg = KrakenDataClientConfig::default();
2215 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None).unwrap();
2216
2217 let err = client
2218 .subscribe_book_l3(Ustr::from("BTC/USD"), 50)
2219 .await
2220 .expect_err("should fail on invalid depth");
2221
2222 assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
2223 assert!(!client.subscriptions_contains("level3:BTC/USD"));
2224 }
2225
2226 #[rstest]
2227 #[tokio::test]
2228 async fn test_subscribe_book_invalid_depth_errors() {
2229 let client = test_client_without_credentials();
2230 let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
2231
2232 let err = client
2233 .subscribe_book(instrument_id, Some(50))
2234 .await
2235 .expect_err("should fail on invalid depth");
2236
2237 assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
2238 assert!(!client.subscriptions_contains("book:BTC/USD"));
2239 assert_eq!(client.l2_depths.get("BTC/USD"), None);
2240 }
2241
2242 #[rstest]
2243 #[tokio::test]
2244 async fn test_subscribe_book_defaults_to_depth_10_and_stores_state() {
2245 let client = test_client_without_credentials();
2246 let key = "book:BTC/USD";
2247 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2248 *client.cmd_tx.write().await = cmd_tx;
2249
2250 client
2251 .subscribe_book(InstrumentId::from("BTC/USD.KRAKEN"), None)
2252 .await
2253 .expect("subscribe should succeed");
2254
2255 let cmd = cmd_rx.try_recv().expect("expected subscribe command");
2256 let SpotHandlerCommand::Subscribe { payload } = cmd else {
2257 panic!("expected subscribe command");
2258 };
2259 assert_book_subscribe_payload(payload.expose_secret(), "BTC/USD", 10);
2260
2261 assert_eq!(client.subscriptions.get_reference_count(key), 1);
2262 assert!(client.subscriptions_contains(key));
2263 assert_eq!(client.l2_depths.get("BTC/USD"), Some(10));
2264 assert_eq!(
2265 client.subscription_payloads.read().await.get(key),
2266 Some(&payload)
2267 );
2268 }
2269
2270 #[rstest]
2271 #[tokio::test]
2272 async fn test_subscribe_book_send_failure_rolls_back_l2_state() {
2273 let client = test_client_without_credentials();
2274 let key = "book:BTC/USD";
2275
2276 let err = client
2277 .subscribe_book(InstrumentId::from("BTC/USD.KRAKEN"), Some(10))
2278 .await
2279 .expect_err("should fail when command receiver is closed");
2280
2281 assert!(matches!(err, KrakenWsError::ConnectionError(_)));
2282 assert_eq!(client.subscriptions.get_reference_count(key), 0);
2283 assert!(!client.subscriptions_contains(key));
2284 assert!(client.subscriptions.pending_subscribe_topics().is_empty());
2285 assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
2286 assert_eq!(client.l2_depths.get("BTC/USD"), None);
2287 assert!(!client.subscription_payloads.read().await.contains_key(key));
2288 }
2289
2290 #[rstest]
2291 fn test_subscribe_book_l3_refcount_idempotent() {
2292 let cfg = KrakenDataClientConfig::default();
2293 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None).unwrap();
2294 let key = "level3:BTC/USD";
2295
2296 assert!(client.subscriptions.add_reference(key));
2297 assert!(!client.subscriptions.add_reference(key));
2298 assert!(!client.subscriptions.remove_reference(key));
2299 assert!(client.subscriptions.remove_reference(key));
2300 }
2301
2302 #[rstest]
2303 #[tokio::test]
2304 async fn test_subscribe_book_l3_rejects_depth_mismatch() {
2305 let cfg = KrakenDataClientConfig::default();
2306 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None).unwrap();
2307
2308 let key = "level3:BTC/USD";
2309 client.subscriptions.add_reference(key);
2310 client.subscriptions.mark_subscribe(key);
2311 client.subscriptions.confirm_subscribe(key);
2312 client.l3_depths.lock().insert("BTC/USD".to_string(), 1000);
2313
2314 *client.auth_token.write().await = Some(SecretString::from("test-token".to_string()));
2315
2316 let err = client
2317 .subscribe_book_l3(Ustr::from("BTC/USD"), 10)
2318 .await
2319 .expect_err("should reject depth mismatch");
2320 assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
2321
2322 assert!(client.subscriptions.remove_reference(key));
2323 }
2324
2325 #[rstest]
2326 #[tokio::test]
2327 async fn test_subscribe_book_rejects_depth_mismatch() {
2328 let client = test_client_without_credentials();
2329 let key = "book:BTC/USD";
2330 client.subscriptions.add_reference(key);
2331 client.subscriptions.mark_subscribe(key);
2332 client.subscriptions.confirm_subscribe(key);
2333 client.l2_depths.insert("BTC/USD", 10, 1);
2334
2335 let err = client
2336 .subscribe_book(InstrumentId::from("BTC/USD.KRAKEN"), Some(25))
2337 .await
2338 .expect_err("should reject depth mismatch");
2339 assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
2340
2341 assert!(client.subscriptions.remove_reference(key));
2342 }
2343
2344 #[rstest]
2345 #[tokio::test]
2346 async fn test_unsubscribe_book_removes_l2_depth_on_last_reference() {
2347 let client = test_client_without_credentials();
2348 let key = "book:BTC/USD";
2349 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2350 *client.cmd_tx.write().await = cmd_tx;
2351
2352 client.subscriptions.add_reference(key);
2353 client.subscriptions.mark_subscribe(key);
2354 client.subscriptions.confirm_subscribe(key);
2355 client.l2_depths.insert("BTC/USD", 10, 1);
2356 client
2357 .subscription_payloads
2358 .write()
2359 .await
2360 .insert(key.to_string(), SecretString::from("payload".to_string()));
2361
2362 client
2363 .unsubscribe_book(InstrumentId::from("BTC/USD.KRAKEN"))
2364 .await
2365 .expect("unsubscribe should succeed");
2366
2367 assert_eq!(client.subscriptions.get_reference_count(key), 0);
2368 assert!(!client.subscriptions_contains(key));
2369 assert_eq!(client.l2_depths.get("BTC/USD"), None);
2370 assert!(!client.subscription_payloads.read().await.contains_key(key));
2371
2372 let cmd = cmd_rx.try_recv().expect("expected unsubscribe command");
2373 let SpotHandlerCommand::Unsubscribe { payload } = cmd else {
2374 panic!("expected unsubscribe command");
2375 };
2376 assert_book_unsubscribe_payload(payload.expose_secret(), "BTC/USD", 10);
2377 }
2378
2379 #[rstest]
2380 #[tokio::test]
2381 async fn test_generic_subscribe_rejects_level3() {
2382 let cfg = KrakenDataClientConfig::default();
2383 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None).unwrap();
2384
2385 let err = client
2386 .subscribe(
2387 KrakenWsChannel::Level3,
2388 vec![Ustr::from("BTC/USD")],
2389 Some(1000),
2390 )
2391 .await
2392 .expect_err("generic subscribe must reject Level3");
2393 assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
2394
2395 let err = client
2396 .unsubscribe(KrakenWsChannel::Level3, vec![Ustr::from("BTC/USD")])
2397 .await
2398 .expect_err("generic unsubscribe must reject Level3");
2399 assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
2400 }
2401
2402 #[rstest]
2403 fn test_update_auth_token_in_payload_for_level3() {
2404 let original = r#"{"method":"subscribe","params":{"channel":"level3","symbol":["BTC/USD"],"depth":1000,"snapshot":true,"token":"OLD"},"req_id":1}"#;
2405 let rewritten = update_auth_token_in_payload(original, "NEW").unwrap();
2406 assert!(rewritten.expose_secret().contains(r#""token":"NEW""#));
2407 assert!(!rewritten.expose_secret().contains(r#""token":"OLD""#));
2408 }
2409
2410 #[rstest]
2411 fn test_l3_depths_shared_between_subscribe_and_handle() {
2412 let cfg = KrakenDataClientConfig::default();
2413 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None).unwrap();
2414
2415 let handle = client.l3_depths_handle();
2416 client.l3_depths.lock().insert("BTC/USD".to_string(), 100);
2417
2418 assert_eq!(handle.lock().get("BTC/USD").copied(), Some(100));
2419 }
2420
2421 #[rstest]
2422 fn test_l2_depths_shared_between_subscribe_and_handle() {
2423 let client = test_client_without_credentials();
2424
2425 let handle = client.l2_depths_handle();
2426 client.l2_depths.insert("BTC/USD", 10, 1);
2427
2428 assert_eq!(handle.get("BTC/USD"), Some(10));
2429 }
2430
2431 #[rstest]
2434 #[tokio::test]
2435 async fn test_subscribe_and_resync_book_record_their_requests() {
2436 let client = test_client_without_credentials();
2437 let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
2438 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2439 *client.cmd_tx.write().await = cmd_tx;
2440 let handle = client.book_requests_handle();
2441
2442 client
2443 .subscribe_book(instrument_id, Some(10))
2444 .await
2445 .unwrap();
2446 let subscription = client.l2_depths.subscription("BTC/USD").unwrap();
2447 let SpotHandlerCommand::Subscribe { payload } = cmd_rx.try_recv().expect("subscribe")
2448 else {
2449 panic!("expected a subscribe command");
2450 };
2451 let subscribe_req_id = payload_req_id(payload.expose_secret());
2452
2453 assert_eq!(
2454 handle.lock().get(&subscribe_req_id),
2455 Some(&L2BookRequest {
2456 symbol: Ustr::from("BTC/USD"),
2457 })
2458 );
2459 assert_eq!(subscription.latest_request, subscribe_req_id);
2460
2461 client
2462 .resync_book(
2463 instrument_id,
2464 subscription.generation,
2465 subscription.snapshot_epoch,
2466 )
2467 .await
2468 .unwrap();
2469 let SpotHandlerCommand::Unsubscribe { payload } = cmd_rx.try_recv().expect("unsubscribe")
2470 else {
2471 panic!("expected an unsubscribe command");
2472 };
2473 assert_book_unsubscribe_payload(payload.expose_secret(), "BTC/USD", 10);
2474 let unsubscribe_req_id = payload_req_id(payload.expose_secret());
2475 let SpotHandlerCommand::Subscribe { payload } = cmd_rx.try_recv().expect("subscribe")
2476 else {
2477 panic!("expected a subscribe command");
2478 };
2479 let resync_req_id = payload_req_id(payload.expose_secret());
2480
2481 let requests = handle.lock();
2482 assert_eq!(
2483 requests.get(&resync_req_id),
2484 Some(&L2BookRequest {
2485 symbol: Ustr::from("BTC/USD"),
2486 })
2487 );
2488 assert_eq!(
2489 client
2490 .l2_depths
2491 .subscription("BTC/USD")
2492 .unwrap()
2493 .latest_request,
2494 resync_req_id
2495 );
2496 assert!(
2497 !requests.contains_key(&unsubscribe_req_id),
2498 "only subscribes are on record"
2499 );
2500 assert_eq!(requests.len(), 2);
2501 }
2502
2503 #[rstest]
2506 #[tokio::test]
2507 async fn test_disconnect_clears_the_book_requests() {
2508 let client = test_client_without_credentials();
2509 let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2510 *client.cmd_tx.write().await = cmd_tx;
2511 client
2512 .subscribe_book(InstrumentId::from("BTC/USD.KRAKEN"), Some(10))
2513 .await
2514 .unwrap();
2515 assert_eq!(client.book_requests.lock().len(), 1);
2516
2517 let _ = client.disconnect_locked().await;
2518
2519 assert!(client.book_requests.lock().is_empty());
2520 }
2521
2522 fn payload_req_id(payload: &str) -> u64 {
2523 let value: serde_json::Value = serde_json::from_str(payload).unwrap();
2524 value["req_id"].as_u64().expect("a req_id")
2525 }
2526
2527 #[rstest]
2531 #[tokio::test]
2532 async fn test_resync_book_skips_once_a_snapshot_is_accepted() {
2533 let client = test_client_without_credentials();
2534 let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
2535 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2536 *client.cmd_tx.write().await = cmd_tx;
2537 client
2538 .subscribe_book(instrument_id, Some(10))
2539 .await
2540 .unwrap();
2541 while cmd_rx.try_recv().is_ok() {}
2542 let issued = client.l2_depths.subscription("BTC/USD").unwrap();
2543
2544 let accepted = client
2545 .l2_depths
2546 .accept_snapshot("BTC/USD", issued.latest_request)
2547 .expect("the live stream's snapshot is accepted");
2548 client
2549 .resync_book(instrument_id, issued.generation, issued.snapshot_epoch)
2550 .await
2551 .unwrap();
2552
2553 assert!(
2554 cmd_rx.try_recv().is_err(),
2555 "a served recovery must send nothing"
2556 );
2557 assert_eq!(
2558 client
2559 .l2_depths
2560 .subscription("BTC/USD")
2561 .unwrap()
2562 .latest_request,
2563 issued.latest_request
2564 );
2565
2566 client
2567 .resync_book(instrument_id, accepted.generation, accepted.snapshot_epoch)
2568 .await
2569 .unwrap();
2570 assert_book_recovery_commands(&mut cmd_rx, 10);
2571 }
2572
2573 #[rstest]
2576 #[tokio::test]
2577 async fn test_a_recovery_issued_before_a_reconnect_sends_nothing() {
2578 let client = test_client_without_credentials();
2579 let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
2580 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2581 *client.cmd_tx.write().await = cmd_tx;
2582 client
2583 .subscribe_book(instrument_id, Some(10))
2584 .await
2585 .unwrap();
2586 while cmd_rx.try_recv().is_ok() {}
2587 let issued = client.l2_depths.subscription("BTC/USD").unwrap();
2588
2589 begin_resubscribe(
2590 &client.subscriptions,
2591 &client.book_requests,
2592 &client.l2_depths,
2593 );
2594 client
2595 .resync_book(instrument_id, issued.generation, issued.snapshot_epoch)
2596 .await
2597 .unwrap();
2598
2599 assert!(
2600 cmd_rx.try_recv().is_err(),
2601 "a recovery from before the reconnect must send nothing"
2602 );
2603
2604 let replayed = client.l2_depths.subscription("BTC/USD").unwrap();
2605 client
2606 .resync_book(instrument_id, replayed.generation, replayed.snapshot_epoch)
2607 .await
2608 .unwrap();
2609 assert_book_recovery_commands(&mut cmd_rx, 10);
2610 }
2611
2612 #[rstest]
2616 #[tokio::test]
2617 async fn test_resync_book_keeps_its_latest_request_when_a_send_fails() {
2618 let client = test_client_without_credentials();
2619 let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
2620 let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2621 *client.cmd_tx.write().await = cmd_tx;
2622 client
2623 .subscribe_book(instrument_id, Some(10))
2624 .await
2625 .unwrap();
2626 let issued = client.l2_depths.subscription("BTC/USD").unwrap();
2627 drop(cmd_rx);
2628
2629 client
2630 .resync_book(instrument_id, issued.generation, issued.snapshot_epoch)
2631 .await
2632 .expect_err("the command channel is closed");
2633
2634 let latest = client
2635 .l2_depths
2636 .subscription("BTC/USD")
2637 .unwrap()
2638 .latest_request;
2639 assert_ne!(latest, issued.latest_request);
2640 assert!(
2641 !client.book_requests.lock().contains_key(&latest),
2642 "an unsent subscribe has no answer on record"
2643 );
2644 }
2645
2646 #[rstest]
2647 fn test_resync_book_l3_does_not_touch_refcount() {
2648 let cfg = KrakenDataClientConfig::default();
2649 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None).unwrap();
2650 let key = "level3:BTC/USD";
2651
2652 assert!(client.subscriptions.add_reference(key));
2653 assert!(!client.subscriptions.add_reference(key));
2654 client.subscriptions.mark_subscribe(key);
2655 client.subscriptions.confirm_subscribe(key);
2656
2657 assert!(client.subscriptions_contains(key));
2658 }
2659
2660 #[rstest]
2663 #[tokio::test]
2664 async fn test_resync_book_leaves_a_replacement_subscription_alone() {
2665 let client = test_client_without_credentials();
2666 let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
2667 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2668 *client.cmd_tx.write().await = cmd_tx;
2669
2670 client
2671 .subscribe_book(instrument_id, Some(10))
2672 .await
2673 .unwrap();
2674 let retired = client.l2_depths.subscription("BTC/USD").unwrap().generation;
2675 client.unsubscribe_book(instrument_id).await.unwrap();
2676 client
2677 .subscribe_book(instrument_id, Some(100))
2678 .await
2679 .unwrap();
2680
2681 while cmd_rx.try_recv().is_ok() {}
2682
2683 client.resync_book(instrument_id, retired, 0).await.unwrap();
2684
2685 assert!(
2686 cmd_rx.try_recv().is_err(),
2687 "a retired subscription's recovery must send nothing"
2688 );
2689 assert_eq!(client.l2_depths.get("BTC/USD"), Some(100));
2690
2691 let live = client.l2_depths.subscription("BTC/USD").unwrap();
2693 client
2694 .resync_book(instrument_id, live.generation, live.snapshot_epoch)
2695 .await
2696 .unwrap();
2697
2698 let SpotHandlerCommand::Unsubscribe { payload } = cmd_rx.try_recv().expect("unsubscribe")
2699 else {
2700 panic!("expected an unsubscribe command");
2701 };
2702 assert_book_unsubscribe_payload(payload.expose_secret(), "BTC/USD", 100);
2703 let SpotHandlerCommand::Subscribe { payload } = cmd_rx.try_recv().expect("subscribe")
2704 else {
2705 panic!("expected a subscribe command");
2706 };
2707 assert_book_subscribe_payload(payload.expose_secret(), "BTC/USD", 100);
2708 assert_eq!(client.l2_depths.get("BTC/USD"), Some(100));
2709 }
2710
2711 async fn spawn_parked<F>(command: F) -> tokio::task::JoinHandle<Result<(), KrakenWsError>>
2714 where
2715 F: Future<Output = Result<(), KrakenWsError>> + Send + 'static,
2716 {
2717 let handle = tokio::spawn(command);
2718
2719 for _ in 0..10 {
2720 tokio::task::yield_now().await;
2721 }
2722
2723 handle
2724 }
2725
2726 fn assert_book_recovery_commands(
2727 cmd_rx: &mut tokio::sync::mpsc::UnboundedReceiver<SpotHandlerCommand>,
2728 depth: u32,
2729 ) {
2730 let SpotHandlerCommand::Unsubscribe { payload } = cmd_rx.try_recv().expect("unsubscribe")
2731 else {
2732 panic!("expected an unsubscribe command");
2733 };
2734 assert_book_unsubscribe_payload(payload.expose_secret(), "BTC/USD", depth);
2735 let SpotHandlerCommand::Subscribe { payload } = cmd_rx.try_recv().expect("subscribe")
2736 else {
2737 panic!("expected a subscribe command");
2738 };
2739 assert_book_subscribe_payload(payload.expose_secret(), "BTC/USD", depth);
2740 let value: serde_json::Value = serde_json::from_str(payload.expose_secret()).unwrap();
2741 assert_eq!(value["params"]["snapshot"], serde_json::json!(true));
2742 assert!(cmd_rx.try_recv().is_err(), "nothing else is sent");
2743 }
2744
2745 #[rstest]
2748 #[tokio::test]
2749 async fn test_resync_book_waits_for_the_book_command_lock() {
2750 let client = test_client_without_credentials();
2751 let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
2752 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2753 *client.cmd_tx.write().await = cmd_tx;
2754 client
2755 .subscribe_book(instrument_id, Some(25))
2756 .await
2757 .unwrap();
2758 while cmd_rx.try_recv().is_ok() {}
2759 let live = client.l2_depths.subscription("BTC/USD").unwrap();
2760
2761 let guard = client.l2_book_commands.lock().await;
2762 let task_client = client.clone();
2763 let recovery = spawn_parked(async move {
2764 task_client
2765 .resync_book(instrument_id, live.generation, live.snapshot_epoch)
2766 .await
2767 })
2768 .await;
2769 assert!(
2770 cmd_rx.try_recv().is_err(),
2771 "no command while a subscription change holds the lock"
2772 );
2773
2774 drop(guard);
2775 recovery.await.unwrap().unwrap();
2776 assert_book_recovery_commands(&mut cmd_rx, 25);
2777 }
2778
2779 #[rstest]
2782 #[tokio::test]
2783 async fn test_subscribe_and_unsubscribe_book_wait_for_the_book_command_lock() {
2784 let client = test_client_without_credentials();
2785 let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
2786 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2787 *client.cmd_tx.write().await = cmd_tx;
2788
2789 let guard = client.l2_book_commands.lock().await;
2790 let task_client = client.clone();
2791 let subscribe =
2792 spawn_parked(async move { task_client.subscribe_book(instrument_id, Some(10)).await })
2793 .await;
2794 assert!(
2795 cmd_rx.try_recv().is_err(),
2796 "no subscribe while the lock is held"
2797 );
2798
2799 drop(guard);
2800 subscribe.await.unwrap().unwrap();
2801 let SpotHandlerCommand::Subscribe { payload } = cmd_rx.try_recv().expect("subscribe")
2802 else {
2803 panic!("expected a subscribe command");
2804 };
2805 assert_book_subscribe_payload(payload.expose_secret(), "BTC/USD", 10);
2806 assert!(cmd_rx.try_recv().is_err());
2807
2808 let guard = client.l2_book_commands.lock().await;
2809 let task_client = client.clone();
2810 let unsubscribe =
2811 spawn_parked(async move { task_client.unsubscribe_book(instrument_id).await }).await;
2812 assert!(
2813 cmd_rx.try_recv().is_err(),
2814 "no unsubscribe while the lock is held"
2815 );
2816
2817 drop(guard);
2818 unsubscribe.await.unwrap().unwrap();
2819 let SpotHandlerCommand::Unsubscribe { payload } = cmd_rx.try_recv().expect("unsubscribe")
2820 else {
2821 panic!("expected an unsubscribe command");
2822 };
2823 assert_book_unsubscribe_payload(payload.expose_secret(), "BTC/USD", 10);
2824 assert!(cmd_rx.try_recv().is_err());
2825 assert_eq!(client.l2_depths.get("BTC/USD"), None);
2826 }
2827
2828 #[rstest]
2831 #[tokio::test]
2832 async fn test_resync_book_sends_both_commands_once_admitted() {
2833 let client = test_client_without_credentials();
2834 let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
2835 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2836 *client.cmd_tx.write().await = cmd_tx;
2837 client
2838 .subscribe_book(instrument_id, Some(10))
2839 .await
2840 .unwrap();
2841 while cmd_rx.try_recv().is_ok() {}
2842 let live = client.l2_depths.subscription("BTC/USD").unwrap();
2843
2844 let admission = client.cmd_tx.write().await;
2846 let task_client = client.clone();
2847 let recovery = spawn_parked(async move {
2848 task_client
2849 .resync_book(instrument_id, live.generation, live.snapshot_epoch)
2850 .await
2851 })
2852 .await;
2853 assert!(cmd_rx.try_recv().is_err());
2854
2855 client.l2_depths.insert("BTC/USD", 25, 99);
2857 drop(admission);
2858 recovery.await.unwrap().unwrap();
2859
2860 assert_book_recovery_commands(&mut cmd_rx, 10);
2861 }
2862
2863 fn assert_book_subscribe_payload(payload: &str, symbol: &str, depth: u32) {
2864 let value: serde_json::Value =
2865 serde_json::from_str(payload).expect("payload should parse as JSON");
2866
2867 assert_eq!(value["method"], serde_json::json!("subscribe"));
2868 assert_eq!(value["params"]["channel"], serde_json::json!("book"));
2869 assert_eq!(value["params"]["symbol"], serde_json::json!([symbol]));
2870 assert_eq!(value["params"]["depth"], serde_json::json!(depth));
2871 }
2872
2873 fn assert_book_unsubscribe_payload(payload: &str, symbol: &str, depth: u32) {
2875 let value: serde_json::Value =
2876 serde_json::from_str(payload).expect("payload should parse as JSON");
2877
2878 assert_eq!(value["method"], serde_json::json!("unsubscribe"));
2879 assert_eq!(value["params"]["channel"], serde_json::json!("book"));
2880 assert_eq!(value["params"]["symbol"], serde_json::json!([symbol]));
2881 assert_eq!(value["params"]["depth"], serde_json::json!(depth));
2882 }
2883}