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