1use std::{
25 fmt::Debug,
26 sync::{
27 Arc,
28 atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering},
29 },
30 time::Duration,
31};
32
33use alloy::signers::local::PrivateKeySigner;
34use arc_swap::ArcSwap;
35use dashmap::DashMap;
36#[cfg(test)]
37use nautilus_common::live::get_runtime;
38use nautilus_core::{
39 UUID4,
40 string::secret::{REDACTED, SecretString},
41};
42use nautilus_live::{
43 SocketControl,
44 task::{SharedTaskSlot, TaskJoinOutcome, TaskSlot, finish_task},
45};
46use nautilus_network::{
47 http::create_standard_nautilus_headers,
48 mode::ConnectionMode,
49 ratelimiter::clock::MonotonicClock,
50 websocket::{
51 AuthTracker, TransportBackend, WebSocketClient, WebSocketConfig, channel_message_handler,
52 },
53};
54use serde::{Serialize, de::DeserializeOwned};
55use serde_json::Value;
56use ustr::Ustr;
57use zeroize::Zeroizing;
58
59use super::{
60 error::{DeriveWsError, Result},
61 handler::{
62 DeriveWsMessage, FeedHandler, HandlerCommand, orderbook_subscribe_params,
63 ticker_subscribe_params, trades_subscribe_params,
64 },
65 messages::{
66 DeriveWsChannel, WsLoginParams, WsLoginResult, WsSubscribeParams, WsSubscribeResult,
67 WsUnsubscribeParams, WsUnsubscribeResult, methods, orderbook_channel, ticker_channel,
68 trades_channel,
69 },
70};
71use crate::{
72 common::{
73 consts::{
74 RECONNECT_BACKOFF_FACTOR, RECONNECT_BASE_BACKOFF, RECONNECT_JITTER_MS,
75 RECONNECT_MAX_BACKOFF, RECONNECT_TIMEOUT, WS_HEARTBEAT_SECS, WS_HEARTBEAT_TIMEOUT,
76 WS_REQUEST_TIMEOUT,
77 },
78 enums::DeriveEnvironment,
79 rate_limit::{
80 DeriveRateLimiter, FixedWindowLimiter, FixedWindowLimits, RateClass,
81 rate_class_for_method,
82 },
83 urls,
84 },
85 http::{
86 models::{
87 DeriveCancelByInstrumentResult, DeriveCancelByLabelResult, DeriveEmptyResult,
88 DeriveOpenOrdersResult, DeriveOrder, DeriveOrderResult, DeriveReplaceOutcome,
89 DeriveReplaceResult,
90 },
91 query::{
92 DeriveCancelAllParams, DeriveCancelByInstrumentParams, DeriveCancelByLabelParams,
93 DeriveCancelParams, DeriveCancelTriggerOrderParams, DeriveGetTriggerOrdersParams,
94 DeriveOrderParams, DeriveReplaceParams, DeriveTriggerOrderParams,
95 },
96 },
97 signing::auth::build_ws_login,
98};
99
100#[derive(Clone)]
103pub struct DeriveWsCredentials {
104 pub wallet_address: String,
106 pub signer: PrivateKeySigner,
108}
109
110impl DeriveWsCredentials {
111 pub fn new(wallet_address: impl Into<String>, session_key_hex: &str) -> Result<Self> {
117 let signer: PrivateKeySigner = session_key_hex
118 .parse()
119 .map_err(|e| DeriveWsError::transport(format!("invalid session key: {e}")))?;
120 Ok(Self {
121 wallet_address: wallet_address.into(),
122 signer,
123 })
124 }
125}
126
127impl Debug for DeriveWsCredentials {
128 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
129 f.debug_struct(stringify!(DeriveWsCredentials))
130 .field("wallet_address", &self.wallet_address)
131 .field("signer", &REDACTED)
132 .finish()
133 }
134}
135
136type WsRateLimiter = DeriveRateLimiter;
139
140const MAX_SESSION_RECOVERY_ATTEMPTS: u32 = 3;
141const SUBSCRIPTION_ACCEPTED_STATUSES: &[&str] = &["ok"];
142const SUBSCRIPTION_REPLAY_ACCEPTED_STATUSES: &[&str] = &["ok", "already subscribed"];
143pub(super) const UNAUTHENTICATED_CONNECTION_EPOCH: u64 = u64::MAX;
144
145#[derive(Debug)]
151pub struct DeriveWebSocketClient {
152 url: String,
153 transport_backend: TransportBackend,
154 proxy_url: Option<SecretString>,
155 connection_mode: Arc<ArcSwap<AtomicU8>>,
156 connection_epoch: Arc<ArcSwap<AtomicU64>>,
157 signal: Arc<AtomicBool>,
158 auth_tracker: AuthTracker,
159 authenticated_epoch: Arc<AtomicU64>,
160 credentials: Option<DeriveWsCredentials>,
161 next_id: Arc<AtomicU64>,
162 cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
163 out_rx: Option<tokio::sync::mpsc::UnboundedReceiver<DeriveWsMessage>>,
164 subscriptions: Arc<DashMap<String, ()>>,
165 subscription_lock: Arc<tokio::sync::Mutex<()>>,
166 task_handle: TaskSlot<()>,
167 send_task: Arc<SharedTaskSlot<()>>,
168 shutdown_errors: Vec<String>,
169 request_timeout: Duration,
170 conn_id: Arc<ArcSwap<String>>,
171 rate_limiter: Arc<WsRateLimiter>,
172 socket_control: Option<SocketControl>,
173}
174
175#[derive(Clone, Debug)]
176pub(crate) struct DeriveWebSocketShutdownHandle {
177 signal: Arc<AtomicBool>,
178}
179
180impl DeriveWebSocketShutdownHandle {
181 pub(crate) fn begin_shutdown(&self) {
182 self.signal.store(true, Ordering::Release);
183 }
184}
185
186struct DeriveWebSocketSetupGuard {
187 shutdown: DeriveWebSocketShutdownHandle,
188 armed: bool,
189}
190
191impl DeriveWebSocketSetupGuard {
192 fn new(shutdown: DeriveWebSocketShutdownHandle) -> Self {
193 Self {
194 shutdown,
195 armed: true,
196 }
197 }
198
199 fn disarm(mut self) {
200 self.armed = false;
201 }
202}
203
204impl Drop for DeriveWebSocketSetupGuard {
205 fn drop(&mut self) {
206 if self.armed {
207 self.shutdown.begin_shutdown();
208 }
209 }
210}
211
212#[derive(Debug, Clone)]
214pub struct DeriveWebSocketSubscriptionHandle {
215 cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
216 subscriptions: Arc<DashMap<String, ()>>,
217 subscription_lock: Arc<tokio::sync::Mutex<()>>,
218 request_timeout: Duration,
219 rate_limiter: Arc<WsRateLimiter>,
220}
221
222#[derive(Debug, Clone)]
231pub struct DeriveWsExecutionHandle {
232 cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
233 auth_tracker: AuthTracker,
234 request_timeout: Duration,
235 conn_id: Arc<ArcSwap<String>>,
236 rate_limiter: Arc<WsRateLimiter>,
237}
238
239#[derive(Debug)]
240pub(crate) struct MatchingRateLimitReservation {
241 method: &'static str,
242 instrument_name: Ustr,
243 window: u32,
244}
245
246impl DeriveWebSocketClient {
247 #[must_use]
250 pub fn new(
251 url: Option<String>,
252 environment: DeriveEnvironment,
253 transport_backend: TransportBackend,
254 proxy_url: Option<String>,
255 ) -> Self {
256 let url = url.unwrap_or_else(|| urls::ws_url(environment).to_string());
257 Self::build(
258 url,
259 transport_backend,
260 proxy_url,
261 None,
262 FixedWindowLimits::websocket(None, None),
263 )
264 }
265
266 #[must_use]
274 pub fn with_credentials(
275 url: Option<String>,
276 environment: DeriveEnvironment,
277 transport_backend: TransportBackend,
278 proxy_url: Option<String>,
279 credentials: DeriveWsCredentials,
280 max_matching_requests_per_second: Option<u32>,
281 max_per_instrument_matching_requests_per_second: Option<u32>,
282 ) -> Self {
283 let url = url.unwrap_or_else(|| urls::ws_url(environment).to_string());
284 let limits = FixedWindowLimits::websocket(
285 max_matching_requests_per_second,
286 max_per_instrument_matching_requests_per_second,
287 );
288 Self::build(url, transport_backend, proxy_url, Some(credentials), limits)
289 }
290
291 fn build(
292 url: String,
293 transport_backend: TransportBackend,
294 proxy_url: Option<String>,
295 credentials: Option<DeriveWsCredentials>,
296 limits: FixedWindowLimits,
297 ) -> Self {
298 let connection_mode = Arc::new(ArcSwap::new(Arc::new(AtomicU8::new(
299 ConnectionMode::Closed as u8,
300 ))));
301 let connection_epoch = Arc::new(ArcSwap::new(Arc::new(AtomicU64::new(0))));
302
303 let (placeholder_tx, _) = tokio::sync::mpsc::unbounded_channel();
305
306 let rate_limiter = Arc::new(FixedWindowLimiter::new(limits, MonotonicClock {}));
312 Self {
313 url,
314 transport_backend,
315 proxy_url: proxy_url.map(SecretString::from),
316 connection_mode,
317 connection_epoch,
318 signal: Arc::new(AtomicBool::new(false)),
319 auth_tracker: AuthTracker::new(),
320 authenticated_epoch: Arc::new(AtomicU64::new(UNAUTHENTICATED_CONNECTION_EPOCH)),
321 credentials,
322 next_id: Arc::new(AtomicU64::new(1)),
323 cmd_tx: Arc::new(tokio::sync::RwLock::new(placeholder_tx)),
324 out_rx: None,
325 subscriptions: Arc::new(DashMap::new()),
326 subscription_lock: Arc::new(tokio::sync::Mutex::new(())),
327 task_handle: TaskSlot::new(),
328 send_task: Arc::new(SharedTaskSlot::new()),
329 shutdown_errors: Vec::new(),
330 request_timeout: WS_REQUEST_TIMEOUT,
331 conn_id: Arc::new(ArcSwap::from_pointee(UUID4::new().to_string())),
332 rate_limiter,
333 socket_control: None,
334 }
335 }
336
337 #[must_use]
339 pub fn with_socket_control(mut self, control: SocketControl) -> Self {
340 self.socket_control = Some(control);
341 self
342 }
343
344 #[must_use]
346 pub fn url(&self) -> &str {
347 &self.url
348 }
349
350 pub fn set_request_timeout(&mut self, timeout: Duration) {
354 self.request_timeout = timeout;
355 }
356
357 #[must_use]
360 pub fn is_authenticated(&self) -> bool {
361 self.auth_tracker.is_authenticated()
362 && self.is_active()
363 && self.authenticated_epoch.load(Ordering::Acquire)
364 == self.connection_epoch.load().load(Ordering::Acquire)
365 }
366
367 #[must_use]
369 pub fn is_active(&self) -> bool {
370 self.connection_mode.load().load(Ordering::Relaxed) == ConnectionMode::Active as u8
371 }
372
373 pub async fn connect(&mut self) -> Result<()> {
384 let auth_ok = self.credentials.is_none() || self.is_authenticated();
387 if self.is_active() && auth_ok && self.task_handle.is_some() {
388 log::warn!("Derive WebSocket already connected");
389 return Ok(());
390 }
391
392 if self.task_handle.is_some() || !self.send_task.is_empty() {
394 log::debug!("Tearing down stale Derive WebSocket state before connect");
395 self.teardown().await?;
396 }
397
398 self.signal.store(false, Ordering::Release);
399 let setup_guard = DeriveWebSocketSetupGuard::new(self.shutdown_handle());
400
401 self.authenticated_epoch
402 .store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
403
404 let (message_handler, raw_rx) = channel_message_handler();
405 let headers = create_standard_nautilus_headers();
406
407 let cfg = WebSocketConfig {
408 url: self.url.clone(),
409 headers,
410 heartbeat_interval_secs: Some(WS_HEARTBEAT_SECS),
411 heartbeat_payload: None,
412 connect_timeout_ms: Some(RECONNECT_TIMEOUT.as_millis() as u64),
413 reconnect_delay_initial_ms: Some(RECONNECT_BASE_BACKOFF.as_millis() as u64),
414 reconnect_delay_max_ms: Some(RECONNECT_MAX_BACKOFF.as_millis() as u64),
415 reconnect_backoff_factor: Some(RECONNECT_BACKOFF_FACTOR),
416 reconnect_jitter_ms: Some(RECONNECT_JITTER_MS),
417 reconnect_max_attempts: None,
418 heartbeat_timeout_secs: Some(WS_HEARTBEAT_TIMEOUT.as_secs()),
419 idle_timeout_ms: None,
420 backend: self.transport_backend,
421 proxy_url: self
422 .proxy_url
423 .as_ref()
424 .map(|value| value.expose_secret().to_owned()),
425 };
426 let client = WebSocketClient::builder()
430 .config(cfg)
431 .message_handler(message_handler)
432 .maybe_state_sink(self.socket_control.as_ref().map(SocketControl::sink))
433 .connect()
434 .await
435 .map_err(|e| DeriveWsError::transport(e.to_string()))?;
436
437 client.set_auth_tracker(self.auth_tracker.clone(), false);
441
442 let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
443 let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<DeriveWsMessage>();
444
445 *self.cmd_tx.write().await = cmd_tx.clone();
446 self.out_rx = Some(out_rx);
447 self.conn_id.store(Arc::new(UUID4::new().to_string()));
448
449 let connection_mode = client.connection_mode_atomic();
450 let connection_epoch = client.connection_epoch_atomic();
451 let reconnect_handle = client.reconnect_handle();
452 self.connection_mode.store(Arc::clone(&connection_mode));
453 self.connection_epoch.store(Arc::clone(&connection_epoch));
454 log::debug!("Derive WebSocket connected: {}", self.url);
455
456 if let Err(e) = cmd_tx.send(HandlerCommand::SetClient(client)) {
457 return Err(DeriveWsError::transport(format!(
458 "failed to send SetClient command: {e}",
459 )));
460 }
461
462 let signal = Arc::clone(&self.signal);
463 let auth_tracker = self.auth_tracker.clone();
464 let authenticated_epoch = Arc::clone(&self.authenticated_epoch);
465 let next_id = Arc::clone(&self.next_id);
466 let credentials = self.credentials.clone();
467 let subscriptions = Arc::clone(&self.subscriptions);
468 let subscription_lock = Arc::clone(&self.subscription_lock);
469 let conn_id = Arc::clone(&self.conn_id);
470 let cmd_tx_for_loop = cmd_tx.clone();
471 let rate_limiter = Arc::clone(&self.rate_limiter);
472 let request_timeout = self.request_timeout;
473 let recovery_connection_mode = Arc::clone(&connection_mode);
474 let recovery_connection_epoch = Arc::clone(&connection_epoch);
475 let send_task = Arc::clone(&self.send_task);
476
477 if let Err(e) = self.task_handle.spawn(async move {
478 let (recovery_tx, mut recovery_rx) = tokio::sync::mpsc::unbounded_channel::<u64>();
479 let recovery_out_tx = out_tx.clone();
480 let recovery_cmd_tx = cmd_tx_for_loop.clone();
481 let recovery_auth_tracker = auth_tracker.clone();
482 let recovery_authenticated_epoch = Arc::clone(&authenticated_epoch);
483 let recovery_subscriptions = Arc::clone(&subscriptions);
484 let recovery_subscription_lock = Arc::clone(&subscription_lock);
485 let recovery_rate_limiter = Arc::clone(&rate_limiter);
486 let recovery_credentials = credentials.clone();
487 let recovery_mode = Arc::clone(&recovery_connection_mode);
488 let recovery_epoch = Arc::clone(&recovery_connection_epoch);
489 let mut recovery_tasks = tokio::task::JoinSet::new();
490
491 recovery_tasks.spawn(async move {
492 while let Some(mut requested_epoch) = recovery_rx.recv().await {
493 while let Ok(epoch) = recovery_rx.try_recv() {
494 requested_epoch = requested_epoch.max(epoch);
495 }
496
497 loop {
498 match recover_session(
499 &recovery_rate_limiter,
500 &recovery_cmd_tx,
501 &recovery_auth_tracker,
502 &recovery_authenticated_epoch,
503 &recovery_mode,
504 &recovery_epoch,
505 recovery_credentials.as_ref(),
506 &recovery_subscriptions,
507 &recovery_subscription_lock,
508 request_timeout,
509 )
510 .await
511 {
512 Ok(recovered_epoch) => {
513 while let Ok(epoch) = recovery_rx.try_recv() {
514 requested_epoch = requested_epoch.max(epoch);
515 }
516
517 if requested_epoch > recovered_epoch
518 || !session_connection_is_active(
519 &recovery_mode,
520 &recovery_epoch,
521 recovered_epoch,
522 )
523 {
524 continue;
525 }
526
527 if recovery_out_tx.send(DeriveWsMessage::Reconnected).is_err() {
528 log::debug!("Derive outer receiver dropped during recovery");
529 }
530 break;
531 }
532 Err(e) => {
533 let mut retry = false;
534
535 while let Ok(epoch) = recovery_rx.try_recv() {
536 requested_epoch = requested_epoch.max(epoch);
537 retry = true;
538 }
539
540 if retry {
541 continue;
542 }
543 log::error!("Derive WebSocket session recovery failed: {e}");
544 let _ = recovery_out_tx
545 .send(DeriveWsMessage::SessionRecoveryFailed(e.to_string()));
546 let _ = recovery_cmd_tx.send(HandlerCommand::Disconnect);
547 return;
548 }
549 }
550 }
551 }
552 });
553
554 let mut handler = FeedHandler::new_with_send_task(
555 signal,
556 cmd_rx,
557 raw_rx,
558 next_id,
559 auth_tracker.clone(),
560 Arc::clone(&authenticated_epoch),
561 send_task,
562 );
563
564 loop {
565 match handler.next().await {
566 Some(DeriveWsMessage::Reconnected) => {
567 log::info!("Derive WebSocket re-establishing session after reconnect");
568 conn_id.store(Arc::new(UUID4::new().to_string()));
569 let epoch = recovery_connection_epoch.load(Ordering::Acquire);
570 if recovery_tx.send(epoch).is_err() {
571 log::error!("Derive WebSocket recovery task stopped unexpectedly");
572 let _ = cmd_tx_for_loop.send(HandlerCommand::Disconnect);
573 }
574 }
575 Some(msg) => {
576 if out_tx.send(msg).is_err() {
577 log::debug!("Derive outer receiver dropped, exiting stream loop");
578 break;
579 }
580 }
581 None => {
582 log::debug!("Derive handler task ended");
583 break;
584 }
585 }
586 }
587 }) {
588 let shutdown_result = self.teardown().await;
589 return Err(DeriveWsError::transport(match shutdown_result {
590 Ok(()) => format!("failed to start WebSocket handler task: {e}"),
591 Err(shutdown_error) => format!(
592 "failed to start WebSocket handler task: {e}; startup rollback failed: \
593 {shutdown_error}"
594 ),
595 }));
596 }
597
598 if let Some(control) = &self.socket_control {
599 control.register(move || reconnect_handle.request_reconnect());
600 }
601
602 if let Some(creds) = self.credentials.clone()
603 && let Err(e) = login_via_handler(
604 &self.rate_limiter,
605 &cmd_tx,
606 &self.auth_tracker,
607 &self.authenticated_epoch,
608 &connection_mode,
609 &connection_epoch,
610 &creds,
611 self.request_timeout,
612 )
613 .await
614 {
615 log::warn!("Derive WebSocket login failed; tearing down transport: {e}");
618 self.teardown().await?;
619 return Err(e);
620 }
621
622 setup_guard.disarm();
623 Ok(())
624 }
625
626 pub(crate) fn begin_shutdown(&self) {
627 self.signal.store(true, Ordering::Release);
628 }
629
630 pub(crate) fn shutdown_handle(&self) -> DeriveWebSocketShutdownHandle {
631 DeriveWebSocketShutdownHandle {
632 signal: Arc::clone(&self.signal),
633 }
634 }
635
636 async fn teardown(&mut self) -> Result<()> {
640 self.begin_shutdown();
641
642 if let Err(e) = self.cmd_tx.read().await.send(HandlerCommand::Disconnect) {
643 log::debug!(
644 "Failed to enqueue Disconnect command (handler may already be shut down): {e}",
645 );
646 }
647
648 match finish_task(
649 &mut self.task_handle,
650 Duration::from_secs(2),
651 Duration::from_secs(2),
652 )
653 .await
654 {
655 None | Some(TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted) => {}
656 Some(TaskJoinOutcome::Failed(error)) => self
657 .shutdown_errors
658 .push(format!("WebSocket handler task failed: {error}")),
659 Some(TaskJoinOutcome::Incomplete) => self
660 .shutdown_errors
661 .push("WebSocket handler task did not stop after abort".to_string()),
662 }
663
664 if self.task_handle.is_some() {
665 if let Some(control) = &self.socket_control {
666 control.deregister();
667 }
668 return self.take_shutdown_result();
669 }
670
671 match self
672 .send_task
673 .finish(Duration::from_secs(1), Duration::from_secs(2))
674 .await
675 {
676 None | Some(TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted) => {}
677 Some(TaskJoinOutcome::Failed(error)) => self
678 .shutdown_errors
679 .push(format!("WebSocket send worker failed: {error}")),
680 Some(TaskJoinOutcome::Incomplete) => self
681 .shutdown_errors
682 .push("WebSocket send worker did not stop after abort".to_string()),
683 }
684
685 if !self.send_task.is_empty() {
686 if let Some(control) = &self.socket_control {
687 control.deregister();
688 }
689 return self.take_shutdown_result();
690 }
691
692 let (placeholder_tx, _) = tokio::sync::mpsc::unbounded_channel();
695 *self.cmd_tx.write().await = placeholder_tx;
696 self.out_rx = None;
697 self.connection_mode
698 .store(Arc::new(AtomicU8::new(ConnectionMode::Closed as u8)));
699 self.connection_epoch.store(Arc::new(AtomicU64::new(0)));
700 self.auth_tracker.invalidate();
701 self.authenticated_epoch
702 .store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
703 self.subscriptions.clear();
704 self.signal.store(false, Ordering::Relaxed);
705
706 if let Some(control) = &self.socket_control {
707 control.deregister();
708 }
709
710 self.take_shutdown_result()
711 }
712
713 fn take_shutdown_result(&mut self) -> Result<()> {
714 if self.shutdown_errors.is_empty() {
715 Ok(())
716 } else {
717 Err(DeriveWsError::transport(
718 std::mem::take(&mut self.shutdown_errors).join("; "),
719 ))
720 }
721 }
722
723 pub async fn disconnect(&mut self) -> Result<()> {
730 log::debug!("Disconnecting Derive WebSocket");
731 self.begin_shutdown();
732 self.teardown().await
733 }
734
735 pub async fn subscribe_ticker(&self, instrument_name: &str, interval: &str) -> Result<()> {
743 self.subscription_handle()
744 .subscribe_ticker(instrument_name, interval)
745 .await
746 }
747
748 pub async fn unsubscribe_ticker(&self, instrument_name: &str, interval: &str) -> Result<()> {
755 self.subscription_handle()
756 .unsubscribe_ticker(instrument_name, interval)
757 .await
758 }
759
760 pub async fn subscribe_orderbook(
767 &self,
768 instrument_name: &str,
769 group: &str,
770 depth: &str,
771 ) -> Result<()> {
772 self.subscription_handle()
773 .subscribe_orderbook(instrument_name, group, depth)
774 .await
775 }
776
777 pub async fn unsubscribe_orderbook(
784 &self,
785 instrument_name: &str,
786 group: &str,
787 depth: &str,
788 ) -> Result<()> {
789 self.subscription_handle()
790 .unsubscribe_orderbook(instrument_name, group, depth)
791 .await
792 }
793
794 pub async fn subscribe_trades(&self, instrument_type: &str, currency: &str) -> Result<()> {
801 self.subscription_handle()
802 .subscribe_trades(instrument_type, currency)
803 .await
804 }
805
806 pub async fn unsubscribe_trades(&self, instrument_type: &str, currency: &str) -> Result<()> {
813 self.subscription_handle()
814 .unsubscribe_trades(instrument_type, currency)
815 .await
816 }
817
818 pub async fn subscribe_channels<C>(&self, channels: Vec<C>) -> Result<()>
829 where
830 C: Into<DeriveWsChannel>,
831 {
832 self.subscription_handle()
833 .subscribe_channels(channels)
834 .await
835 }
836
837 pub async fn unsubscribe_channels<C>(&self, channels: Vec<C>) -> Result<()>
845 where
846 C: Into<DeriveWsChannel>,
847 {
848 self.subscription_handle()
849 .unsubscribe_channels(channels)
850 .await
851 }
852
853 pub async fn next_event(&mut self) -> Option<DeriveWsMessage> {
855 if let Some(rx) = self.out_rx.as_mut() {
856 rx.recv().await
857 } else {
858 None
859 }
860 }
861
862 #[must_use]
865 pub fn subscription_count(&self) -> usize {
866 self.subscriptions.len()
867 }
868
869 #[must_use]
871 pub fn subscription_handle(&self) -> DeriveWebSocketSubscriptionHandle {
872 DeriveWebSocketSubscriptionHandle {
873 cmd_tx: Arc::clone(&self.cmd_tx),
874 subscriptions: Arc::clone(&self.subscriptions),
875 subscription_lock: Arc::clone(&self.subscription_lock),
876 request_timeout: self.request_timeout,
877 rate_limiter: Arc::clone(&self.rate_limiter),
878 }
879 }
880
881 #[must_use]
888 pub fn execution_handle(&self) -> DeriveWsExecutionHandle {
889 DeriveWsExecutionHandle {
890 cmd_tx: Arc::clone(&self.cmd_tx),
891 auth_tracker: self.auth_tracker.clone(),
892 request_timeout: self.request_timeout,
893 conn_id: Arc::clone(&self.conn_id),
894 rate_limiter: Arc::clone(&self.rate_limiter),
895 }
896 }
897
898 pub fn take_event_receiver(
903 &mut self,
904 ) -> Option<tokio::sync::mpsc::UnboundedReceiver<DeriveWsMessage>> {
905 self.out_rx.take()
906 }
907}
908
909impl Drop for DeriveWebSocketClient {
910 fn drop(&mut self) {
911 self.signal.store(true, Ordering::Relaxed);
912
913 if let Some(handle) = self.task_handle.as_ref() {
914 handle.abort();
915 }
916
917 if let Some(control) = &self.socket_control {
918 control.deregister();
919 }
920 }
921}
922
923impl DeriveWebSocketSubscriptionHandle {
924 pub(crate) fn has_subscription(&self, channel: &str) -> bool {
925 self.subscriptions.contains_key(channel)
926 }
927
928 pub(crate) fn forget_subscription(&self, channel: &str) {
929 self.subscriptions.remove(channel);
930 }
931
932 pub(crate) fn remember_subscription(&self, channel: &str) {
933 self.subscriptions.insert(channel.to_string(), ());
934 }
935
936 pub async fn subscribe_ticker(&self, instrument_name: &str, interval: &str) -> Result<()> {
943 let channel = ticker_channel(instrument_name, interval);
944 let params = ticker_subscribe_params(instrument_name, interval);
945 self.send_subscribe(channel, ¶ms).await
946 }
947
948 pub async fn unsubscribe_ticker(&self, instrument_name: &str, interval: &str) -> Result<()> {
955 let channel = ticker_channel(instrument_name, interval);
956 self.send_unsubscribe(channel).await
957 }
958
959 pub async fn subscribe_orderbook(
966 &self,
967 instrument_name: &str,
968 group: &str,
969 depth: &str,
970 ) -> Result<()> {
971 let channel = orderbook_channel(instrument_name, group, depth);
972 let params = orderbook_subscribe_params(instrument_name, group, depth);
973 self.send_subscribe(channel, ¶ms).await
974 }
975
976 pub async fn unsubscribe_orderbook(
983 &self,
984 instrument_name: &str,
985 group: &str,
986 depth: &str,
987 ) -> Result<()> {
988 let channel = orderbook_channel(instrument_name, group, depth);
989 self.send_unsubscribe(channel).await
990 }
991
992 pub async fn subscribe_trades(&self, instrument_type: &str, currency: &str) -> Result<()> {
999 let channel = trades_channel(instrument_type, currency);
1000 let params = trades_subscribe_params(instrument_type, currency);
1001 self.send_subscribe(channel, ¶ms).await
1002 }
1003
1004 pub async fn unsubscribe_trades(&self, instrument_type: &str, currency: &str) -> Result<()> {
1011 let channel = trades_channel(instrument_type, currency);
1012 self.send_unsubscribe(channel).await
1013 }
1014
1015 pub async fn subscribe_channels<C>(&self, channels: Vec<C>) -> Result<()>
1022 where
1023 C: Into<DeriveWsChannel>,
1024 {
1025 let channels = channels.into_iter().map(Into::into).collect::<Vec<_>>();
1026 if channels.is_empty() {
1027 return Ok(());
1028 }
1029 let _guard = self.subscription_lock.lock().await;
1030 let params = WsSubscribeParams { channels };
1031 let cmd_tx = self.cmd_tx.read().await.clone();
1032 let result: WsSubscribeResult = send_request(
1033 &self.rate_limiter,
1034 &cmd_tx,
1035 methods::PUBLIC_SUBSCRIBE,
1036 ¶ms,
1037 self.request_timeout,
1038 )
1039 .await?;
1040
1041 let (confirmed, failure) =
1042 subscription_outcome(¶ms.channels, &result, SUBSCRIPTION_ACCEPTED_STATUSES);
1043
1044 for channel in confirmed {
1045 self.subscriptions.insert(channel, ());
1046 }
1047 failure.map_or(Ok(()), Err)
1048 }
1049
1050 pub async fn unsubscribe_channels<C>(&self, channels: Vec<C>) -> Result<()>
1058 where
1059 C: Into<DeriveWsChannel>,
1060 {
1061 let channels = channels.into_iter().map(Into::into).collect::<Vec<_>>();
1062 if channels.is_empty() {
1063 return Ok(());
1064 }
1065 let _guard = self.subscription_lock.lock().await;
1066 let topics = channel_topics(&channels);
1067 let params = WsUnsubscribeParams { channels };
1068 let cmd_tx = self.cmd_tx.read().await.clone();
1069 let _: WsUnsubscribeResult = send_request(
1070 &self.rate_limiter,
1071 &cmd_tx,
1072 methods::PUBLIC_UNSUBSCRIBE,
1073 ¶ms,
1074 self.request_timeout,
1075 )
1076 .await?;
1077
1078 for channel in topics {
1079 self.subscriptions.remove(&channel);
1080 }
1081
1082 Ok(())
1083 }
1084
1085 async fn send_subscribe(&self, channel: String, params: &WsSubscribeParams) -> Result<()> {
1086 let _guard = self.subscription_lock.lock().await;
1087 let cmd_tx = self.cmd_tx.read().await.clone();
1088 let result: WsSubscribeResult = send_request(
1089 &self.rate_limiter,
1090 &cmd_tx,
1091 methods::PUBLIC_SUBSCRIBE,
1092 params,
1093 self.request_timeout,
1094 )
1095 .await?;
1096
1097 let (confirmed, failure) =
1098 subscription_outcome(¶ms.channels, &result, SUBSCRIPTION_ACCEPTED_STATUSES);
1099
1100 if confirmed.iter().any(|topic| topic == &channel) {
1101 self.subscriptions.insert(channel, ());
1102 }
1103 failure.map_or(Ok(()), Err)
1104 }
1105
1106 async fn send_unsubscribe(&self, channel: String) -> Result<()> {
1107 let _guard = self.subscription_lock.lock().await;
1108 let params = WsUnsubscribeParams {
1109 channels: vec![DeriveWsChannel::from(channel.clone())],
1110 };
1111 let cmd_tx = self.cmd_tx.read().await.clone();
1112 let _: WsUnsubscribeResult = send_request(
1113 &self.rate_limiter,
1114 &cmd_tx,
1115 methods::PUBLIC_UNSUBSCRIBE,
1116 ¶ms,
1117 self.request_timeout,
1118 )
1119 .await?;
1120
1121 self.subscriptions.remove(&channel);
1122
1123 Ok(())
1124 }
1125}
1126
1127impl DeriveWsExecutionHandle {
1128 #[must_use]
1130 pub fn conn_id(&self) -> String {
1131 self.conn_id.load_full().as_ref().clone()
1132 }
1133
1134 pub async fn submit_order(&self, params: &DeriveOrderParams) -> Result<DeriveOrder> {
1146 let reservation = self
1147 .reserve_matching_request(methods::PRIVATE_ORDER, ¶ms.instrument_name)
1148 .await?;
1149 self.submit_order_after_rate_limit(params, reservation)
1150 .await
1151 }
1152
1153 pub(crate) async fn submit_order_after_rate_limit(
1154 &self,
1155 params: &DeriveOrderParams,
1156 reservation: MatchingRateLimitReservation,
1157 ) -> Result<DeriveOrder> {
1158 self.ensure_authenticated(methods::PRIVATE_ORDER)?;
1159 debug_assert_eq!(reservation.method, methods::PRIVATE_ORDER);
1160 self.refresh_matching_reservation(&reservation).await;
1161 let cmd_tx = self.cmd_tx.read().await.clone();
1162 let result: DeriveOrderResult = send_request_typed_after_rate_limit(
1163 &self.rate_limiter,
1164 &cmd_tx,
1165 methods::PRIVATE_ORDER,
1166 params,
1167 self.request_timeout,
1168 )
1169 .await?;
1170 Ok(result.order)
1171 }
1172
1173 pub async fn submit_trigger_order(
1181 &self,
1182 params: &DeriveTriggerOrderParams,
1183 ) -> Result<DeriveOrder> {
1184 let reservation = self
1185 .reserve_matching_request(
1186 methods::PRIVATE_TRIGGER_ORDER,
1187 ¶ms.order.instrument_name,
1188 )
1189 .await?;
1190 self.submit_trigger_order_after_rate_limit(params, reservation)
1191 .await
1192 }
1193
1194 pub(crate) async fn submit_trigger_order_after_rate_limit(
1195 &self,
1196 params: &DeriveTriggerOrderParams,
1197 reservation: MatchingRateLimitReservation,
1198 ) -> Result<DeriveOrder> {
1199 self.ensure_authenticated(methods::PRIVATE_TRIGGER_ORDER)?;
1200 debug_assert_eq!(reservation.method, methods::PRIVATE_TRIGGER_ORDER);
1201 self.refresh_matching_reservation(&reservation).await;
1202 let cmd_tx = self.cmd_tx.read().await.clone();
1203 let result: DeriveOrderResult = send_request_typed_after_rate_limit(
1204 &self.rate_limiter,
1205 &cmd_tx,
1206 methods::PRIVATE_TRIGGER_ORDER,
1207 params,
1208 self.request_timeout,
1209 )
1210 .await?;
1211 Ok(result.order)
1212 }
1213
1214 pub async fn modify_order(&self, params: &DeriveReplaceParams) -> Result<DeriveReplaceOutcome> {
1223 let reservation = self
1224 .reserve_matching_request(methods::PRIVATE_REPLACE, ¶ms.order.instrument_name)
1225 .await?;
1226 self.modify_order_after_rate_limit(params, reservation)
1227 .await
1228 }
1229
1230 pub(crate) async fn modify_order_after_rate_limit(
1231 &self,
1232 params: &DeriveReplaceParams,
1233 reservation: MatchingRateLimitReservation,
1234 ) -> Result<DeriveReplaceOutcome> {
1235 self.ensure_authenticated(methods::PRIVATE_REPLACE)?;
1236 debug_assert_eq!(reservation.method, methods::PRIVATE_REPLACE);
1237 self.refresh_matching_reservation(&reservation).await;
1238 let cmd_tx = self.cmd_tx.read().await.clone();
1239 let result: DeriveReplaceResult = send_request_typed_after_rate_limit(
1240 &self.rate_limiter,
1241 &cmd_tx,
1242 methods::PRIVATE_REPLACE,
1243 params,
1244 self.request_timeout,
1245 )
1246 .await?;
1247 result
1248 .into_outcome(¶ms.order_id_to_cancel, ¶ms.order.label)
1249 .map_err(|message| {
1250 DeriveWsError::Serde(<serde_json::Error as serde::de::Error>::custom(message))
1251 })
1252 }
1253
1254 pub async fn cancel_order(&self, params: &DeriveCancelParams) -> Result<()> {
1262 self.require_authenticated(methods::PRIVATE_CANCEL).await?;
1263 let cmd_tx = self.cmd_tx.read().await.clone();
1264 let _: DeriveEmptyResult = send_request_for_instrument(
1265 &self.rate_limiter,
1266 &cmd_tx,
1267 methods::PRIVATE_CANCEL,
1268 params,
1269 self.request_timeout,
1270 params.instrument_name,
1271 )
1272 .await?;
1273 Ok(())
1274 }
1275
1276 pub async fn cancel_by_instrument(
1284 &self,
1285 params: &DeriveCancelByInstrumentParams,
1286 ) -> Result<DeriveCancelByInstrumentResult> {
1287 self.require_authenticated(methods::PRIVATE_CANCEL_BY_INSTRUMENT)
1288 .await?;
1289 let cmd_tx = self.cmd_tx.read().await.clone();
1290 send_request_typed_for_instrument(
1291 &self.rate_limiter,
1292 &cmd_tx,
1293 methods::PRIVATE_CANCEL_BY_INSTRUMENT,
1294 params,
1295 self.request_timeout,
1296 params.instrument_name,
1297 )
1298 .await
1299 }
1300
1301 pub async fn cancel_trigger_order(
1309 &self,
1310 params: &DeriveCancelTriggerOrderParams,
1311 ) -> Result<DeriveOrder> {
1312 self.require_authenticated(methods::PRIVATE_CANCEL_TRIGGER_ORDER)
1313 .await?;
1314 let cmd_tx = self.cmd_tx.read().await.clone();
1315 send_request_typed(
1316 &self.rate_limiter,
1317 &cmd_tx,
1318 methods::PRIVATE_CANCEL_TRIGGER_ORDER,
1319 params,
1320 self.request_timeout,
1321 )
1322 .await
1323 }
1324
1325 pub async fn cancel_by_label(
1334 &self,
1335 params: &DeriveCancelByLabelParams,
1336 ) -> Result<DeriveCancelByLabelResult> {
1337 self.require_authenticated(methods::PRIVATE_CANCEL_BY_LABEL)
1338 .await?;
1339 let cmd_tx = self.cmd_tx.read().await.clone();
1340 send_request_typed(
1341 &self.rate_limiter,
1342 &cmd_tx,
1343 methods::PRIVATE_CANCEL_BY_LABEL,
1344 params,
1345 self.request_timeout,
1346 )
1347 .await
1348 }
1349
1350 pub async fn get_trigger_orders(
1359 &self,
1360 params: &DeriveGetTriggerOrdersParams,
1361 ) -> Result<DeriveOpenOrdersResult> {
1362 self.require_authenticated(methods::PRIVATE_GET_TRIGGER_ORDERS)
1363 .await?;
1364 let cmd_tx = self.cmd_tx.read().await.clone();
1365 send_request_typed(
1366 &self.rate_limiter,
1367 &cmd_tx,
1368 methods::PRIVATE_GET_TRIGGER_ORDERS,
1369 params,
1370 self.request_timeout,
1371 )
1372 .await
1373 }
1374
1375 pub async fn cancel_all_orders(&self, params: &DeriveCancelAllParams) -> Result<()> {
1384 self.require_authenticated(methods::PRIVATE_CANCEL_ALL)
1385 .await?;
1386 let cmd_tx = self.cmd_tx.read().await.clone();
1387 let _: DeriveEmptyResult = send_request(
1388 &self.rate_limiter,
1389 &cmd_tx,
1390 methods::PRIVATE_CANCEL_ALL,
1391 params,
1392 self.request_timeout,
1393 )
1394 .await?;
1395 Ok(())
1396 }
1397
1398 pub(crate) async fn reserve_matching_request(
1399 &self,
1400 operation: &'static str,
1401 instrument_name: &Ustr,
1402 ) -> Result<MatchingRateLimitReservation> {
1403 self.require_authenticated(operation).await?;
1404 debug_assert_eq!(rate_class_for_method(operation), RateClass::Matching);
1405 let window = self
1406 .rate_limiter
1407 .await_class_ready(RateClass::Matching, Some(instrument_name))
1408 .await;
1409 self.ensure_authenticated(operation)?;
1410 Ok(MatchingRateLimitReservation {
1411 method: operation,
1412 instrument_name: *instrument_name,
1413 window,
1414 })
1415 }
1416
1417 async fn refresh_matching_reservation(&self, reservation: &MatchingRateLimitReservation) {
1421 self.rate_limiter
1422 .ensure_window_current(
1423 RateClass::Matching,
1424 Some(&reservation.instrument_name),
1425 reservation.window,
1426 )
1427 .await;
1428 }
1429
1430 fn ensure_authenticated(&self, operation: &'static str) -> Result<()> {
1431 if self.auth_tracker.is_authenticated() {
1432 return Ok(());
1433 }
1434
1435 Err(DeriveWsError::Authentication {
1436 operation: operation.to_string(),
1437 reason: "WebSocket session is not authenticated".to_string(),
1438 })
1439 }
1440
1441 async fn require_authenticated(&self, operation: &'static str) -> Result<()> {
1442 if self
1443 .auth_tracker
1444 .wait_for_authenticated(self.request_timeout)
1445 .await
1446 {
1447 return Ok(());
1448 }
1449
1450 Err(DeriveWsError::Authentication {
1451 operation: operation.to_string(),
1452 reason: "WebSocket session is not authenticated".to_string(),
1453 })
1454 }
1455}
1456
1457#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1461enum RequestRateLimit {
1462 Await(Option<Ustr>),
1465 Reserved,
1468}
1469
1470async fn send_raw<P>(
1471 rate_limiter: &WsRateLimiter,
1472 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1473 method: &'static str,
1474 params: &P,
1475 timeout: Duration,
1476) -> Result<Value>
1477where
1478 P: Serialize + ?Sized,
1479{
1480 send_raw_with_rate_limit(
1481 rate_limiter,
1482 cmd_tx,
1483 method,
1484 params,
1485 timeout,
1486 RequestRateLimit::Await(None),
1487 None,
1488 )
1489 .await
1490}
1491
1492async fn send_raw_for_instrument<P>(
1495 rate_limiter: &WsRateLimiter,
1496 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1497 method: &'static str,
1498 params: &P,
1499 timeout: Duration,
1500 instrument_name: Ustr,
1501) -> Result<Value>
1502where
1503 P: Serialize + ?Sized,
1504{
1505 send_raw_with_rate_limit(
1506 rate_limiter,
1507 cmd_tx,
1508 method,
1509 params,
1510 timeout,
1511 RequestRateLimit::Await(Some(instrument_name)),
1512 None,
1513 )
1514 .await
1515}
1516
1517async fn send_raw_after_rate_limit<P>(
1518 rate_limiter: &WsRateLimiter,
1519 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1520 method: &'static str,
1521 params: &P,
1522 timeout: Duration,
1523) -> Result<Value>
1524where
1525 P: Serialize + ?Sized,
1526{
1527 send_raw_with_rate_limit(
1528 rate_limiter,
1529 cmd_tx,
1530 method,
1531 params,
1532 timeout,
1533 RequestRateLimit::Reserved,
1534 None,
1535 )
1536 .await
1537}
1538
1539async fn send_raw_with_rate_limit<P>(
1540 rate_limiter: &WsRateLimiter,
1541 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1542 method: &'static str,
1543 params: &P,
1544 timeout: Duration,
1545 rate_limit: RequestRateLimit,
1546 connection_epoch: Option<u64>,
1547) -> Result<Value>
1548where
1549 P: Serialize + ?Sized,
1550{
1551 let params = SecretString::from(serde_json::to_string(params)?);
1552
1553 if let RequestRateLimit::Await(instrument_name) = rate_limit {
1554 rate_limiter
1555 .await_class_ready(rate_class_for_method(method), instrument_name.as_ref())
1556 .await;
1557 }
1558
1559 let (response_tx, response_rx) = tokio::sync::oneshot::channel();
1560 cmd_tx
1561 .send(HandlerCommand::Request {
1562 method,
1563 params,
1564 connection_epoch,
1565 response_tx,
1566 })
1567 .map_err(|e| DeriveWsError::transport(format!("failed to enqueue `{method}`: {e}")))?;
1568
1569 match tokio::time::timeout(timeout, response_rx).await {
1573 Ok(Ok(outcome)) => outcome,
1574 Ok(Err(_)) => Err(DeriveWsError::RequestCancelled {
1575 method: method.to_owned(),
1576 }),
1577 Err(_) => Err(DeriveWsError::Timeout {
1578 method: method.to_owned(),
1579 }),
1580 }
1581}
1582
1583async fn send_request<P, R>(
1586 rate_limiter: &WsRateLimiter,
1587 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1588 method: &'static str,
1589 params: &P,
1590 timeout: Duration,
1591) -> Result<R>
1592where
1593 P: Serialize + ?Sized,
1594 R: Default + DeserializeOwned,
1595{
1596 let value = send_raw(rate_limiter, cmd_tx, method, params, timeout).await?;
1597 decode_default_result(value)
1598}
1599
1600async fn send_request_for_instrument<P, R>(
1603 rate_limiter: &WsRateLimiter,
1604 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1605 method: &'static str,
1606 params: &P,
1607 timeout: Duration,
1608 instrument_name: Ustr,
1609) -> Result<R>
1610where
1611 P: Serialize + ?Sized,
1612 R: Default + DeserializeOwned,
1613{
1614 let value = send_raw_for_instrument(
1615 rate_limiter,
1616 cmd_tx,
1617 method,
1618 params,
1619 timeout,
1620 instrument_name,
1621 )
1622 .await?;
1623 decode_default_result(value)
1624}
1625
1626async fn send_request_typed_for_instrument<P, R>(
1628 rate_limiter: &WsRateLimiter,
1629 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1630 method: &'static str,
1631 params: &P,
1632 timeout: Duration,
1633 instrument_name: Ustr,
1634) -> Result<R>
1635where
1636 P: Serialize + ?Sized,
1637 R: DeserializeOwned,
1638{
1639 let value = send_raw_for_instrument(
1640 rate_limiter,
1641 cmd_tx,
1642 method,
1643 params,
1644 timeout,
1645 instrument_name,
1646 )
1647 .await?;
1648 Ok(serde_json::from_value(value)?)
1649}
1650
1651fn decode_default_result<R>(value: Value) -> Result<R>
1652where
1653 R: Default + DeserializeOwned,
1654{
1655 if value.is_null() {
1656 Ok(R::default())
1657 } else {
1658 Ok(serde_json::from_value(value)?)
1659 }
1660}
1661
1662async fn send_request_on_connection<P, R>(
1663 rate_limiter: &WsRateLimiter,
1664 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1665 method: &'static str,
1666 params: &P,
1667 timeout: Duration,
1668 connection_epoch: u64,
1669) -> Result<R>
1670where
1671 P: Serialize + ?Sized,
1672 R: Default + DeserializeOwned,
1673{
1674 let value = send_raw_with_rate_limit(
1675 rate_limiter,
1676 cmd_tx,
1677 method,
1678 params,
1679 timeout,
1680 RequestRateLimit::Await(None),
1681 Some(connection_epoch),
1682 )
1683 .await?;
1684
1685 decode_default_result(value)
1686}
1687
1688async fn send_request_typed<P, R>(
1691 rate_limiter: &WsRateLimiter,
1692 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1693 method: &'static str,
1694 params: &P,
1695 timeout: Duration,
1696) -> Result<R>
1697where
1698 P: Serialize + ?Sized,
1699 R: DeserializeOwned,
1700{
1701 let value = send_raw(rate_limiter, cmd_tx, method, params, timeout).await?;
1702 Ok(serde_json::from_value(value)?)
1703}
1704
1705async fn send_request_typed_after_rate_limit<P, R>(
1706 rate_limiter: &WsRateLimiter,
1707 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1708 method: &'static str,
1709 params: &P,
1710 timeout: Duration,
1711) -> Result<R>
1712where
1713 P: Serialize + ?Sized,
1714 R: DeserializeOwned,
1715{
1716 let value = send_raw_after_rate_limit(rate_limiter, cmd_tx, method, params, timeout).await?;
1717 Ok(serde_json::from_value(value)?)
1718}
1719
1720fn channel_topics(channels: &[DeriveWsChannel]) -> Vec<String> {
1721 channels.iter().map(ToString::to_string).collect()
1722}
1723
1724#[expect(
1725 clippy::too_many_arguments,
1726 reason = "authentication state is passed explicitly for epoch fencing"
1727)]
1728async fn login_via_handler(
1729 rate_limiter: &WsRateLimiter,
1730 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1731 auth_tracker: &AuthTracker,
1732 authenticated_epoch: &AtomicU64,
1733 connection_mode: &AtomicU8,
1734 connection_epoch: &AtomicU64,
1735 creds: &DeriveWsCredentials,
1736 timeout: Duration,
1737) -> Result<()> {
1738 let _receiver = auth_tracker.begin();
1739 let expected_epoch = connection_epoch.load(Ordering::Acquire);
1740
1741 match send_login_request(rate_limiter, cmd_tx, creds, timeout, expected_epoch).await {
1742 Ok(())
1743 if complete_session_authentication(
1744 auth_tracker,
1745 authenticated_epoch,
1746 connection_mode,
1747 connection_epoch,
1748 expected_epoch,
1749 ) =>
1750 {
1751 log::debug!("Derive WebSocket authenticated");
1752
1753 Ok(())
1754 }
1755 Ok(()) => {
1756 let e = DeriveWsError::transport(
1757 "connection changed while completing WebSocket authentication",
1758 );
1759 authenticated_epoch.store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
1760 auth_tracker.fail(e.to_string());
1761 Err(e)
1762 }
1763 Err(e) => {
1764 authenticated_epoch.store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
1765 auth_tracker.fail(e.to_string());
1766 Err(e)
1767 }
1768 }
1769}
1770
1771async fn send_login_request(
1772 rate_limiter: &WsRateLimiter,
1773 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1774 creds: &DeriveWsCredentials,
1775 timeout: Duration,
1776 connection_epoch: u64,
1777) -> Result<()> {
1778 let login = build_ws_login(&creds.wallet_address, &creds.signer)?;
1779 let params = Zeroizing::new(WsLoginParams {
1780 wallet: login.wallet,
1781 timestamp: login.timestamp,
1782 signature: login.signature,
1783 });
1784 let result = send_request_on_connection::<_, WsLoginResult>(
1785 rate_limiter,
1786 cmd_tx,
1787 methods::PUBLIC_LOGIN,
1788 &*params,
1789 timeout,
1790 connection_epoch,
1791 )
1792 .await?;
1793
1794 if matches!(result, WsLoginResult::Success { success: false }) {
1795 return Err(DeriveWsError::Authentication {
1796 operation: methods::PUBLIC_LOGIN.to_string(),
1797 reason: "venue returned an unsuccessful login result".to_string(),
1798 });
1799 }
1800
1801 Ok(())
1802}
1803
1804#[expect(
1805 clippy::too_many_arguments,
1806 reason = "recovery state is passed explicitly for epoch fencing"
1807)]
1808async fn recover_session(
1809 rate_limiter: &WsRateLimiter,
1810 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1811 auth_tracker: &AuthTracker,
1812 authenticated_epoch: &AtomicU64,
1813 connection_mode: &AtomicU8,
1814 connection_epoch: &AtomicU64,
1815 creds: Option<&DeriveWsCredentials>,
1816 subscriptions: &DashMap<String, ()>,
1817 subscription_lock: &tokio::sync::Mutex<()>,
1818 timeout: Duration,
1819) -> Result<u64> {
1820 let _guard = subscription_lock.lock().await;
1821 let _receiver = creds.map(|_| auth_tracker.begin());
1822
1823 for attempt in 1..=MAX_SESSION_RECOVERY_ATTEMPTS {
1824 let expected_epoch = wait_for_session_connection(connection_mode, connection_epoch).await?;
1825
1826 let result = async {
1827 if !session_connection_is_active(connection_mode, connection_epoch, expected_epoch) {
1828 return Err(DeriveWsError::transport(
1829 "connection changed before WebSocket session recovery",
1830 ));
1831 }
1832
1833 if let Some(creds) = creds {
1834 send_login_request(rate_limiter, cmd_tx, creds, timeout, expected_epoch).await?;
1835 }
1836 let channels: Vec<String> = subscriptions
1837 .iter()
1838 .map(|entry| entry.key().clone())
1839 .collect();
1840 subscribe_via_handler(rate_limiter, cmd_tx, channels, timeout, expected_epoch).await?;
1841
1842 if !session_connection_is_active(connection_mode, connection_epoch, expected_epoch) {
1843 return Err(DeriveWsError::transport(
1844 "connection changed during WebSocket session recovery",
1845 ));
1846 }
1847
1848 Ok(())
1849 }
1850 .await;
1851
1852 match result {
1853 Ok(()) => {
1854 if creds.is_some()
1855 && !complete_session_authentication(
1856 auth_tracker,
1857 authenticated_epoch,
1858 connection_mode,
1859 connection_epoch,
1860 expected_epoch,
1861 )
1862 {
1863 continue;
1864 }
1865
1866 if creds.is_some() {
1867 log::info!("Derive WebSocket session re-authenticated");
1868 }
1869
1870 return Ok(expected_epoch);
1871 }
1872 Err(e) if attempt < MAX_SESSION_RECOVERY_ATTEMPTS => {
1873 let multiplier = 1_u32 << (attempt - 1);
1874 let delay = RECONNECT_BASE_BACKOFF
1875 .saturating_mul(multiplier)
1876 .min(RECONNECT_MAX_BACKOFF);
1877 log::warn!(
1878 "Derive WebSocket session recovery attempt {attempt}/{MAX_SESSION_RECOVERY_ATTEMPTS} failed: {e}; retrying in {delay:?}",
1879 );
1880 tokio::time::sleep(delay).await;
1881 }
1882 Err(e) => {
1883 authenticated_epoch.store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
1884
1885 if creds.is_some() {
1886 auth_tracker.fail(e.to_string());
1887 }
1888 return Err(e);
1889 }
1890 }
1891 }
1892
1893 let e = DeriveWsError::transport("WebSocket session changed while recovery completed");
1894 authenticated_epoch.store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
1895
1896 if creds.is_some() {
1897 auth_tracker.fail(e.to_string());
1898 }
1899 Err(e)
1900}
1901
1902async fn wait_for_session_connection(
1903 connection_mode: &AtomicU8,
1904 connection_epoch: &AtomicU64,
1905) -> Result<u64> {
1906 loop {
1907 match ConnectionMode::from_atomic(connection_mode) {
1908 ConnectionMode::Active => return Ok(connection_epoch.load(Ordering::Acquire)),
1909 ConnectionMode::Reconnect => tokio::time::sleep(RECONNECT_BASE_BACKOFF).await,
1910 ConnectionMode::Disconnect | ConnectionMode::Closed => {
1911 return Err(DeriveWsError::transport(
1912 "WebSocket closed during session recovery",
1913 ));
1914 }
1915 }
1916 }
1917}
1918
1919fn complete_session_authentication(
1920 auth_tracker: &AuthTracker,
1921 authenticated_epoch: &AtomicU64,
1922 connection_mode: &AtomicU8,
1923 connection_epoch: &AtomicU64,
1924 expected_epoch: u64,
1925) -> bool {
1926 if !session_connection_is_active(connection_mode, connection_epoch, expected_epoch) {
1927 return false;
1928 }
1929
1930 authenticated_epoch.store(expected_epoch, Ordering::Release);
1931 auth_tracker.succeed();
1932
1933 if session_connection_is_active(connection_mode, connection_epoch, expected_epoch) {
1934 true
1935 } else {
1936 authenticated_epoch.store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
1937 auth_tracker.invalidate();
1938 false
1939 }
1940}
1941
1942fn session_connection_is_active(
1943 connection_mode: &AtomicU8,
1944 connection_epoch: &AtomicU64,
1945 expected_epoch: u64,
1946) -> bool {
1947 ConnectionMode::from_atomic(connection_mode).is_active()
1948 && connection_epoch.load(Ordering::Acquire) == expected_epoch
1949}
1950
1951async fn subscribe_via_handler(
1952 rate_limiter: &WsRateLimiter,
1953 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1954 channels: Vec<String>,
1955 timeout: Duration,
1956 connection_epoch: u64,
1957) -> Result<()> {
1958 if channels.is_empty() {
1959 return Ok(());
1960 }
1961
1962 let params = WsSubscribeParams {
1963 channels: channels.into_iter().map(DeriveWsChannel::from).collect(),
1964 };
1965 let result: WsSubscribeResult = send_request_on_connection(
1966 rate_limiter,
1967 cmd_tx,
1968 methods::PUBLIC_SUBSCRIBE,
1969 ¶ms,
1970 timeout,
1971 connection_epoch,
1972 )
1973 .await?;
1974
1975 let (_, failure) = subscription_outcome(
1976 ¶ms.channels,
1977 &result,
1978 SUBSCRIPTION_REPLAY_ACCEPTED_STATUSES,
1979 );
1980 failure.map_or(Ok(()), Err)
1981}
1982
1983fn subscription_outcome(
1984 requested: &[DeriveWsChannel],
1985 result: &WsSubscribeResult,
1986 accepted_statuses: &[&str],
1987) -> (Vec<String>, Option<DeriveWsError>) {
1988 let mut confirmed = Vec::with_capacity(requested.len());
1989 let mut failures = Vec::new();
1990
1991 for channel in requested {
1992 let topic = channel.to_string();
1993 match result.status.get(channel) {
1994 Some(status) if accepted_statuses.contains(&status.as_str()) => confirmed.push(topic),
1995 Some(status) => failures.push(format!("{topic}: {status}")),
1996 None if result.channels.contains(channel) => confirmed.push(topic),
1997 None => failures.push(format!("{topic}: missing channel status")),
1998 }
1999 }
2000
2001 let failure = (!failures.is_empty()).then(|| DeriveWsError::Subscription {
2002 details: failures.join(", "),
2003 });
2004 (confirmed, failure)
2005}
2006
2007#[cfg(test)]
2008mod tests {
2009 use rstest::rstest;
2010
2011 use super::*;
2012 use crate::common::rate_limit::RateBucket;
2013
2014 #[rstest]
2015 fn test_public_client_defaults_to_environment_url() {
2016 let client = DeriveWebSocketClient::new(
2017 None,
2018 DeriveEnvironment::Mainnet,
2019 TransportBackend::default(),
2020 None,
2021 );
2022 assert!(client.url().starts_with("wss://"));
2023 assert!(client.url().contains("api.lyra.finance"));
2024 assert!(!client.is_authenticated());
2025 assert!(!client.is_active());
2026 assert_eq!(client.subscription_count(), 0);
2027 }
2028
2029 #[tokio::test]
2030 async fn test_execution_auth_barrier_waits_for_authentication() {
2031 let client = DeriveWebSocketClient::with_credentials(
2032 None,
2033 DeriveEnvironment::Mainnet,
2034 TransportBackend::default(),
2035 None,
2036 DeriveWsCredentials::new(
2037 "0x000000000000000000000000000000000000aaaa",
2038 "0x2ae8be44db8a590d20bffbe3b6872df9b569147d3bf6801a35a28281a4816bbd",
2039 )
2040 .unwrap(),
2041 None,
2042 None,
2043 );
2044 let execution = client.execution_handle();
2045 let auth_tracker = execution.auth_tracker.clone();
2046 let _receiver = auth_tracker.begin();
2047 let tracker_for_task = auth_tracker.clone();
2048
2049 get_runtime().spawn(async move {
2050 tokio::time::sleep(Duration::from_millis(10)).await;
2051 tracker_for_task.succeed();
2052 });
2053
2054 execution
2055 .require_authenticated(methods::PRIVATE_ORDER)
2056 .await
2057 .expect("barrier should wait for successful authentication");
2058 }
2059
2060 #[tokio::test]
2061 async fn test_execution_auth_barrier_fails_on_terminal_auth_failure() {
2062 let client = DeriveWebSocketClient::with_credentials(
2063 None,
2064 DeriveEnvironment::Mainnet,
2065 TransportBackend::default(),
2066 None,
2067 DeriveWsCredentials::new(
2068 "0x000000000000000000000000000000000000aaaa",
2069 "0x2ae8be44db8a590d20bffbe3b6872df9b569147d3bf6801a35a28281a4816bbd",
2070 )
2071 .unwrap(),
2072 None,
2073 None,
2074 );
2075 let execution = client.execution_handle();
2076 let _receiver = execution.auth_tracker.begin();
2077 execution.auth_tracker.fail("bad signature");
2078
2079 let error = execution
2080 .require_authenticated(methods::PRIVATE_ORDER)
2081 .await
2082 .expect_err("terminal auth failure must reject private operations");
2083
2084 assert!(matches!(error, DeriveWsError::Authentication { .. }));
2085 }
2086
2087 #[tokio::test]
2088 async fn test_session_recovery_waits_for_reconnecting_transport() {
2089 let connection_mode = Arc::new(AtomicU8::new(ConnectionMode::Reconnect as u8));
2090 let connection_epoch = Arc::new(AtomicU64::new(1));
2091 let mode_for_task = Arc::clone(&connection_mode);
2092 let epoch_for_task = Arc::clone(&connection_epoch);
2093
2094 get_runtime().spawn(async move {
2095 tokio::time::sleep(Duration::from_millis(10)).await;
2096 epoch_for_task.store(2, Ordering::Release);
2097 mode_for_task.store(ConnectionMode::Active as u8, Ordering::Release);
2098 });
2099
2100 let epoch = tokio::time::timeout(
2101 Duration::from_secs(1),
2102 wait_for_session_connection(&connection_mode, &connection_epoch),
2103 )
2104 .await
2105 .expect("session recovery should resume after reconnect")
2106 .expect("active replacement connection should be accepted");
2107
2108 assert_eq!(epoch, 2);
2109 }
2110
2111 #[rstest]
2112 fn test_testnet_client_routes_to_demo_url() {
2113 let client = DeriveWebSocketClient::new(
2114 None,
2115 DeriveEnvironment::Testnet,
2116 TransportBackend::default(),
2117 None,
2118 );
2119 assert!(client.url().contains("demo"));
2120 }
2121
2122 #[rstest]
2123 fn test_credentials_constructor_parses_session_key() {
2124 let creds = DeriveWsCredentials::new(
2125 "0x000000000000000000000000000000000000aaaa",
2126 "0x2ae8be44db8a590d20bffbe3b6872df9b569147d3bf6801a35a28281a4816bbd",
2127 )
2128 .unwrap();
2129 assert!(creds.wallet_address.starts_with("0x"));
2130 let client = DeriveWebSocketClient::with_credentials(
2131 None,
2132 DeriveEnvironment::Testnet,
2133 TransportBackend::default(),
2134 None,
2135 creds,
2136 None,
2137 None,
2138 );
2139 assert!(client.url().contains("demo"));
2140 assert!(!client.is_authenticated());
2141 }
2142
2143 #[rstest]
2144 fn test_credentials_debug_redacts_signer() {
2145 let creds = DeriveWsCredentials::new(
2146 "0xWALLET",
2147 "0x2ae8be44db8a590d20bffbe3b6872df9b569147d3bf6801a35a28281a4816bbd",
2148 )
2149 .unwrap();
2150 let debug = format!("{creds:?}");
2151 assert!(debug.contains("redacted"));
2152 assert!(debug.contains("0xWALLET"));
2153 assert!(!debug.contains("2ae8be44"));
2154 }
2155
2156 #[rstest]
2157 fn test_credentials_constructor_rejects_invalid_session_key() {
2158 let err = DeriveWsCredentials::new("0xWALLET", "not-a-hex-key").unwrap_err();
2159 assert!(err.to_string().contains("invalid session key"));
2160 }
2161
2162 #[rstest]
2163 #[tokio::test]
2164 async fn test_send_raw_times_out_when_no_response_arrives() {
2165 let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
2168 let rate_limiter: WsRateLimiter =
2169 FixedWindowLimiter::new(FixedWindowLimits::websocket(None, None), MonotonicClock {});
2170 let err = send_raw(
2171 &rate_limiter,
2172 &cmd_tx,
2173 methods::PRIVATE_ORDER,
2174 &serde_json::json!({}),
2175 Duration::from_millis(50),
2176 )
2177 .await
2178 .expect_err("must time out");
2179
2180 match err {
2181 DeriveWsError::Timeout { method } => assert_eq!(method, methods::PRIVATE_ORDER),
2182 other => panic!("expected Timeout, was {other:?}"),
2183 }
2184 }
2185
2186 #[rstest]
2187 #[tokio::test]
2188 async fn test_send_request_typed_rejects_null_result() {
2189 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
2193 tokio::spawn(async move {
2194 if let Some(HandlerCommand::Request { response_tx, .. }) = cmd_rx.recv().await {
2195 let _ = response_tx.send(Ok(Value::Null));
2196 }
2197 });
2198 let rate_limiter: WsRateLimiter =
2199 FixedWindowLimiter::new(FixedWindowLimits::websocket(None, None), MonotonicClock {});
2200 let result: Result<DeriveOrderResult> = send_request_typed(
2201 &rate_limiter,
2202 &cmd_tx,
2203 methods::PRIVATE_ORDER,
2204 &serde_json::json!({}),
2205 Duration::from_secs(1),
2206 )
2207 .await;
2208 assert!(matches!(result, Err(DeriveWsError::Serde(_))));
2209 }
2210
2211 #[rstest]
2212 #[tokio::test]
2213 async fn test_reserved_send_does_not_wait_for_or_consume_second_quota_cell() {
2214 let rate_limiter: WsRateLimiter =
2215 FixedWindowLimiter::new(FixedWindowLimits::websocket(None, None), MonotonicClock {});
2216
2217 for _ in 0..5 {
2218 rate_limiter
2219 .check_bucket(RateBucket::Matching)
2220 .expect("reservation consumes the matching window");
2221 }
2222 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
2223 tokio::spawn(async move {
2224 if let Some(HandlerCommand::Request { response_tx, .. }) = cmd_rx.recv().await {
2225 let _ = response_tx.send(Ok(serde_json::json!({"accepted": true})));
2226 }
2227 });
2228
2229 let response = tokio::time::timeout(
2230 Duration::from_millis(100),
2231 send_raw_after_rate_limit(
2232 &rate_limiter,
2233 &cmd_tx,
2234 methods::PRIVATE_ORDER,
2235 &serde_json::json!({}),
2236 Duration::from_secs(1),
2237 ),
2238 )
2239 .await
2240 .expect("reserved send must not wait for quota")
2241 .expect("reserved send succeeds");
2242
2243 assert_eq!(response, serde_json::json!({"accepted": true}));
2244 assert!(
2245 rate_limiter.check_bucket(RateBucket::Matching).is_err(),
2246 "reserved send must not consume a second cell",
2247 );
2248 }
2249}