1mod config;
2mod desired_subscriptions;
3mod endpoint;
4mod inbound;
5mod outbound;
6mod reconnect_window;
7mod session;
8mod signer;
9mod state;
10#[cfg(feature = "test-helpers")]
11mod test_peer;
12mod tracker;
13mod writer;
14
15use std::sync::Arc;
16use std::sync::Mutex as StdMutex;
17use std::sync::atomic::{AtomicU64, Ordering};
18use std::time::Duration;
19
20use tokio::sync::{broadcast, mpsc, oneshot, watch};
21
22use crate::ws::client::{PreparedClientMessage, ServerClose};
23use crate::ws::error::{UnreadableFrame, WsClientError};
24use crate::ws::types::{ClientMessage, ServerMessage, VaultAuthSuccessData};
25
26pub use config::{GapPolicy, ManagedWsConfig, ManagedWsConfigError};
27#[cfg(all(feature = "vault", feature = "http"))]
28pub(crate) use endpoint::ws_base_url;
29pub use endpoint::{
30 MakerWsEndpoint, normalize_maker_data_ws_url, normalize_maker_ws_url,
31 normalize_maker_ws_url_for_endpoint,
32};
33pub(crate) use inbound::InboundItem;
34pub use inbound::{
35 ManagedInbound, ManagedMessageReceiver, ManagedReceiveError, ReferenceChangeReceiver,
36};
37pub use outbound::{OutboundMessageError, SendTicket};
38#[cfg(feature = "vault")]
39pub(crate) use signer::ChallengeSigning;
40pub use state::{ManagedWsState, ManagedWsTerminationReason, WaitUntilReadyError};
41#[cfg(feature = "test-helpers")]
42pub use test_peer::ManagedWsTestPeer;
43use tracker::{AwaitRegistration, AwaitTracker};
44
45fn send_event(tx: &broadcast::Sender<ManagedWsEvent>, event: ManagedWsEvent) {
46 if tx.send(event).is_err() {
47 tracing::trace!("no event receivers");
48 }
49}
50
51fn transition_state(
52 state_tx: &watch::Sender<ManagedWsState>,
53 events_tx: &broadcast::Sender<ManagedWsEvent>,
54 state: ManagedWsState,
55 event: ManagedWsEvent,
56) {
57 state::set_state(state_tx, state);
58 send_event(events_tx, event);
59}
60
61#[derive(Debug, Clone)]
62#[non_exhaustive]
63pub enum ManagedWsEvent {
64 Connected,
65 Authenticated,
66 Reconciled,
67 Ready,
68 Reconnecting {
69 attempt: u64,
70 delay_ms: u64,
71 },
72 Disconnected {
74 close: Option<ServerClose>,
75 },
76 Error(ManagedWsFailure),
77 UnreadableFrame(UnreadableFrame),
80 Terminated(ManagedWsTerminationReason),
81}
82
83#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85pub enum ManagedWsFailureKind {
86 Network,
88 Auth,
90 Readiness,
92 Server,
94}
95
96#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
97#[error("{detail}")]
98pub struct ManagedWsFailure {
99 pub kind: ManagedWsFailureKind,
100 pub detail: String,
101}
102
103impl ManagedWsFailure {
104 pub(super) fn event(kind: ManagedWsFailureKind, detail: impl ToString) -> ManagedWsEvent {
105 ManagedWsEvent::Error(Self {
106 kind,
107 detail: detail.to_string(),
108 })
109 }
110}
111
112#[derive(Debug, thiserror::Error)]
113#[non_exhaustive]
114pub enum ManagedWsError {
115 #[error("managed ws connection is closed")]
116 Closed,
117 #[error("managed ws send queue is full")]
120 QueueFull,
121 #[error("a subscription change is already awaiting acknowledgement")]
122 SubscriptionPending,
123 #[error("managed ws session is not ready")]
124 NotReady,
125 #[error("managed ws is disconnected")]
126 Disconnected,
127 #[error("websocket write timed out")]
128 WriteTimeout,
129 #[error("websocket write failed")]
130 Write(#[source] WsClientError),
131 #[error("managed websocket task failed")]
132 TaskJoin(#[source] tokio::task::JoinError),
133 #[error(transparent)]
134 InvalidMessage(#[from] OutboundMessageError),
135}
136
137#[derive(Debug, thiserror::Error)]
138#[non_exhaustive]
139pub enum SendAwaitError {
140 #[error("a subscription change is already awaiting acknowledgement")]
141 SubscriptionPending,
142 #[error("managed ws session is not ready")]
143 NotReady,
144 #[error("connection closed")]
145 Disconnected,
146 #[error("request timed out")]
147 Timeout,
148 #[error("response could not be decoded: {0}")]
149 UnreadableResponse(Box<UnreadableFrame>),
150 #[error("message has no stable correlation key")]
151 NoCorrelationKey,
152 #[error("an identical request is already awaiting a response")]
153 DuplicateInFlight,
154 #[error("too many pending requests (limit {limit})")]
155 TooManyPending { limit: usize },
156 #[error("managed ws send queue is full")]
157 QueueFull,
158 #[error(transparent)]
159 InvalidMessage(#[from] OutboundMessageError),
160}
161
162pub(crate) enum ManagedCommand {
163 Send {
164 message: PreparedClientMessage,
165 connection_epoch: Option<u64>,
166 tx: oneshot::Sender<Result<(), ManagedWsError>>,
167 },
168 SendAwait {
169 await_id: u64,
170 message: PreparedClientMessage,
171 connection_epoch: Option<u64>,
172 registration: AwaitRegistration,
173 tx: oneshot::Sender<Result<Arc<ServerMessage>, SendAwaitError>>,
174 },
175}
176
177#[derive(Clone)]
178pub struct ManagedWsHandle {
179 cmd_tx: mpsc::Sender<ManagedCommand>,
180 cancel_tx: mpsc::UnboundedSender<u64>,
181 gap_tx: Option<mpsc::UnboundedSender<()>>,
182 messages_tx: broadcast::Sender<InboundItem>,
183 events_tx: broadcast::Sender<ManagedWsEvent>,
184 initial_messages_rx: Arc<StdMutex<Option<broadcast::Receiver<InboundItem>>>>,
185 initial_events_rx: Arc<StdMutex<Option<broadcast::Receiver<ManagedWsEvent>>>>,
186 state_rx: watch::Receiver<ManagedWsState>,
187 vault_auth_rx: watch::Receiver<Option<VaultAuthSuccessData>>,
188 shutdown: Arc<ShutdownOnLastHandle>,
189 next_await_id: Arc<AtomicU64>,
190 task: Arc<tokio::sync::Mutex<Option<tokio::task::JoinHandle<()>>>>,
191 max_batch_quotes: usize,
192 max_outbound_message_size: usize,
193 #[cfg(any(test, feature = "test-helpers"))]
194 next_inbound_sequence: Arc<AtomicU64>,
195 #[cfg(any(test, feature = "test-helpers"))]
196 state_tx: watch::Sender<ManagedWsState>,
197}
198
199struct ShutdownOnLastHandle {
200 tx: watch::Sender<bool>,
201}
202
203impl ShutdownOnLastHandle {
204 fn request(&self) {
205 self.tx.send_replace(true);
206 }
207}
208
209impl Drop for ShutdownOnLastHandle {
210 fn drop(&mut self) {
211 self.request();
212 }
213}
214
215impl ManagedWsHandle {
216 pub fn subscribe_messages(&self) -> ManagedMessageReceiver {
220 let initial = self
221 .initial_messages_rx
222 .lock()
223 .unwrap_or_else(std::sync::PoisonError::into_inner)
224 .take();
225 ManagedMessageReceiver {
226 inner: initial.unwrap_or_else(|| self.messages_tx.subscribe()),
227 gap_tx: self.gap_tx.clone(),
228 }
229 }
230
231 pub fn subscribe_reference_changes(&self) -> ReferenceChangeReceiver {
232 ReferenceChangeReceiver::new(self.subscribe_messages())
233 }
234
235 pub fn subscribe_events(&self) -> broadcast::Receiver<ManagedWsEvent> {
237 self.initial_events_rx
238 .lock()
239 .unwrap_or_else(std::sync::PoisonError::into_inner)
240 .take()
241 .unwrap_or_else(|| self.events_tx.subscribe())
242 }
243
244 pub fn subscribe_state(&self) -> watch::Receiver<ManagedWsState> {
246 self.state_rx.clone()
247 }
248
249 #[must_use]
250 pub fn state(&self) -> ManagedWsState {
251 self.state_rx.borrow().clone()
252 }
253
254 #[must_use]
256 pub fn vault_auth(&self) -> Option<VaultAuthSuccessData> {
257 self.vault_auth_rx.borrow().clone()
258 }
259
260 #[must_use]
262 pub fn subscribe_vault_auth(&self) -> watch::Receiver<Option<VaultAuthSuccessData>> {
263 self.vault_auth_rx.clone()
264 }
265
266 #[must_use]
267 pub fn is_ready(&self) -> bool {
268 matches!(*self.state_rx.borrow(), ManagedWsState::Ready { .. })
269 }
270
271 pub fn ensure_ready(&self) -> Result<(), ManagedWsError> {
272 self.ready_epoch().map(|_| ())
273 }
274
275 pub async fn wait_until_ready(&self) -> Result<u64, WaitUntilReadyError> {
277 let mut state_rx = self.state_rx.clone();
278 loop {
279 let state = state_rx.borrow().clone();
280 match state {
281 ManagedWsState::Ready { connection_epoch } => return Ok(connection_epoch),
282 ManagedWsState::Closed { reason } => {
283 return Err(WaitUntilReadyError::Terminated(reason));
284 }
285 _ => {}
286 }
287 state_rx
288 .changed()
289 .await
290 .map_err(|_| WaitUntilReadyError::StateChannelClosed)?;
291 }
292 }
293
294 fn command_epoch(&self, message: &ClientMessage) -> Result<Option<u64>, ManagedWsError> {
295 match self.state() {
296 ManagedWsState::Ready { connection_epoch } => Ok(Some(connection_epoch)),
297 _ if message.requires_ready() => Err(ManagedWsError::NotReady),
298 _ => Ok(None),
299 }
300 }
301
302 fn ready_epoch(&self) -> Result<u64, ManagedWsError> {
303 match self.state() {
304 ManagedWsState::Ready { connection_epoch } => Ok(connection_epoch),
305 _ => Err(ManagedWsError::NotReady),
306 }
307 }
308
309 pub async fn send(&self, message: ClientMessage) -> Result<(), ManagedWsError> {
310 let connection_epoch = self.command_epoch(&message)?;
311 self.send_with_epoch(message, connection_epoch).await
312 }
313
314 pub async fn send_in_epoch(
317 &self,
318 message: ClientMessage,
319 connection_epoch: u64,
320 ) -> Result<(), ManagedWsError> {
321 if self.ready_epoch()? != connection_epoch {
322 return Err(ManagedWsError::NotReady);
323 }
324 self.send_with_epoch(message, Some(connection_epoch)).await
325 }
326
327 async fn send_with_epoch(
328 &self,
329 message: ClientMessage,
330 connection_epoch: Option<u64>,
331 ) -> Result<(), ManagedWsError> {
332 let message = self.prepare_outbound(&message)?;
333 let (tx, rx) = oneshot::channel();
334 self.cmd_tx
335 .send(ManagedCommand::Send {
336 message,
337 connection_epoch,
338 tx,
339 })
340 .await
341 .map_err(|_| ManagedWsError::Closed)?;
342 SendTicket { rx }.wait().await
343 }
344
345 pub fn try_send(&self, message: ClientMessage) -> Result<SendTicket, ManagedWsError> {
346 let connection_epoch = self.command_epoch(&message)?;
347 self.try_send_with_epoch(message, connection_epoch)
348 }
349
350 pub fn try_send_in_epoch(
352 &self,
353 message: ClientMessage,
354 connection_epoch: u64,
355 ) -> Result<SendTicket, ManagedWsError> {
356 if self.ready_epoch()? != connection_epoch {
357 return Err(ManagedWsError::NotReady);
358 }
359 self.try_send_with_epoch(message, Some(connection_epoch))
360 }
361
362 fn try_send_with_epoch(
363 &self,
364 message: ClientMessage,
365 connection_epoch: Option<u64>,
366 ) -> Result<SendTicket, ManagedWsError> {
367 let message = self.prepare_outbound(&message)?;
368 let (tx, rx) = oneshot::channel();
369 match self.cmd_tx.try_send(ManagedCommand::Send {
370 message,
371 connection_epoch,
372 tx,
373 }) {
374 Ok(()) => Ok(SendTicket { rx }),
375 Err(mpsc::error::TrySendError::Closed(_)) => Err(ManagedWsError::Closed),
376 Err(mpsc::error::TrySendError::Full(_)) => Err(ManagedWsError::QueueFull),
377 }
378 }
379
380 pub async fn send_await(
381 &self,
382 message: ClientMessage,
383 timeout_duration: Duration,
384 ) -> Result<Arc<ServerMessage>, SendAwaitError> {
385 let connection_epoch = self
386 .command_epoch(&message)
387 .map_err(|_| SendAwaitError::NotReady)?;
388 self.send_await_with_epoch(message, timeout_duration, connection_epoch)
389 .await
390 }
391
392 pub(crate) async fn send_await_ready(
393 &self,
394 message: ClientMessage,
395 timeout_duration: Duration,
396 ) -> Result<Arc<ServerMessage>, SendAwaitError> {
397 let connection_epoch = self.ready_epoch().map_err(|_| SendAwaitError::NotReady)?;
398 self.send_await_with_epoch(message, timeout_duration, Some(connection_epoch))
399 .await
400 }
401
402 async fn send_await_with_epoch(
403 &self,
404 message: ClientMessage,
405 timeout_duration: Duration,
406 connection_epoch: Option<u64>,
407 ) -> Result<Arc<ServerMessage>, SendAwaitError> {
408 let deadline = tokio::time::Instant::now() + timeout_duration;
409 let (command, response) = self.await_command(message, connection_epoch, deadline)?;
410 if tokio::time::timeout_at(deadline, self.cmd_tx.send(command))
411 .await
412 .map_err(|_| SendAwaitError::Timeout)?
413 .is_err()
414 {
415 return Err(SendAwaitError::Disconnected);
416 }
417 response.await
418 }
419
420 #[cfg(feature = "vault")]
423 pub(crate) fn try_send_await_in_epoch(
424 &self,
425 message: ClientMessage,
426 timeout_duration: Duration,
427 connection_epoch: u64,
428 ) -> Result<
429 impl Future<Output = Result<Arc<ServerMessage>, SendAwaitError>> + use<'_>,
430 SendAwaitError,
431 > {
432 if self.ready_epoch().ok() != Some(connection_epoch) {
433 return Err(SendAwaitError::NotReady);
434 }
435 let deadline = tokio::time::Instant::now() + timeout_duration;
436 let (command, response) = self.await_command(message, Some(connection_epoch), deadline)?;
437 match self.cmd_tx.try_send(command) {
438 Ok(()) => Ok(response),
439 Err(mpsc::error::TrySendError::Full(_)) => Err(SendAwaitError::QueueFull),
440 Err(mpsc::error::TrySendError::Closed(_)) => Err(SendAwaitError::Disconnected),
441 }
442 }
443
444 fn await_command(
445 &self,
446 message: ClientMessage,
447 connection_epoch: Option<u64>,
448 deadline: tokio::time::Instant,
449 ) -> Result<
450 (
451 ManagedCommand,
452 impl Future<Output = Result<Arc<ServerMessage>, SendAwaitError>> + use<'_>,
453 ),
454 SendAwaitError,
455 > {
456 let prepared = self.prepare_outbound(&message)?;
457 let registration = AwaitTracker::prepare(&message)?;
458 let await_id = self.next_await_id.fetch_add(1, Ordering::Relaxed);
459 let (tx, rx) = oneshot::channel();
460 let command = ManagedCommand::SendAwait {
461 await_id,
462 message: prepared,
463 connection_epoch,
464 registration,
465 tx,
466 };
467 let response = async move {
468 match tokio::time::timeout_at(deadline, rx).await {
469 Ok(Ok(result)) => result,
470 Ok(Err(_)) => Err(SendAwaitError::Disconnected),
471 Err(_) => {
472 let _ = self.cancel_tx.send(await_id);
473 Err(SendAwaitError::Timeout)
474 }
475 }
476 };
477 Ok((command, response))
478 }
479
480 fn prepare_outbound(
481 &self,
482 message: &ClientMessage,
483 ) -> Result<PreparedClientMessage, OutboundMessageError> {
484 outbound::prepare_outbound(
485 message,
486 self.max_batch_quotes,
487 self.max_outbound_message_size,
488 )
489 }
490
491 pub async fn close(&self) -> Result<(), ManagedWsError> {
492 self.shutdown.request();
493 let mut task = self.task.lock().await;
494 if let Some(handle) = task.as_mut() {
495 let result = handle.await;
496 task.take();
497 result.map_err(ManagedWsError::TaskJoin)?;
498 }
499 Ok(())
500 }
501
502 #[cfg(any(test, feature = "test-helpers"))]
503 fn make_test_handle(
504 cmd_buffer: usize,
505 broadcast_buffer: usize,
506 ) -> (Self, mpsc::Receiver<ManagedCommand>) {
507 let (cmd_tx, cmd_rx) = mpsc::channel(cmd_buffer);
508 let (cancel_tx, _cancel_rx) = mpsc::unbounded_channel();
509 let (messages_tx, initial_messages_rx) =
510 broadcast::channel::<InboundItem>(broadcast_buffer);
511 let (events_tx, initial_events_rx) = broadcast::channel(broadcast_buffer);
512 let (state_tx, state_rx) = watch::channel(ManagedWsState::Connecting);
513 let (shutdown_tx, _shutdown_rx) = watch::channel(false);
514 let (_vault_auth_tx, vault_auth_rx) = watch::channel(None);
515 (
516 Self {
517 cmd_tx,
518 cancel_tx,
519 gap_tx: None,
520 messages_tx,
521 events_tx,
522 initial_messages_rx: Arc::new(StdMutex::new(Some(initial_messages_rx))),
523 initial_events_rx: Arc::new(StdMutex::new(Some(initial_events_rx))),
524 state_rx,
525 vault_auth_rx,
526 shutdown: Arc::new(ShutdownOnLastHandle { tx: shutdown_tx }),
527 next_await_id: Arc::new(AtomicU64::new(1)),
528 task: Arc::new(tokio::sync::Mutex::new(None)),
529 max_batch_quotes: 50,
530 max_outbound_message_size: 1024 * 1024,
531 #[cfg(any(test, feature = "test-helpers"))]
532 next_inbound_sequence: Arc::new(AtomicU64::new(1)),
533 #[cfg(any(test, feature = "test-helpers"))]
534 state_tx,
535 },
536 cmd_rx,
537 )
538 }
539
540 #[cfg(test)]
541 pub(crate) fn test_handle(
542 cmd_buffer: usize,
543 broadcast_buffer: usize,
544 ) -> (Self, mpsc::Receiver<ManagedCommand>) {
545 Self::make_test_handle(cmd_buffer, broadcast_buffer)
546 }
547
548 #[cfg(all(feature = "test-helpers", not(test)))]
549 pub fn test_handle(cmd_buffer: usize, broadcast_buffer: usize) -> (Self, ManagedWsTestPeer) {
550 let (handle, commands) = Self::make_test_handle(cmd_buffer, broadcast_buffer);
551 (handle, ManagedWsTestPeer::new(commands))
552 }
553
554 #[cfg(any(test, feature = "test-helpers"))]
555 pub fn inject_message(&self, msg: ServerMessage) {
556 let sequence = self.next_inbound_sequence.fetch_add(1, Ordering::Relaxed);
557 let _ = self.messages_tx.send(InboundItem::Message(ManagedInbound {
558 connection_epoch: 0,
559 sequence,
560 received_at: std::time::Instant::now(),
561 message: Arc::new(msg),
562 }));
563 }
564
565 #[cfg(any(test, feature = "test-helpers"))]
566 pub fn inject_event(&self, event: ManagedWsEvent) {
567 let _ = self.events_tx.send(event);
568 }
569
570 #[cfg(any(test, feature = "test-helpers"))]
571 pub fn inject_state(&self, state: ManagedWsState) {
572 self.state_tx.send_replace(state);
573 }
574}
575
576pub fn spawn_managed_ws(config: ManagedWsConfig) -> Result<ManagedWsHandle, ManagedWsConfigError> {
577 config.validate()?;
578 let command_buffer = config.command_buffer;
579 let broadcast_buffer = config.broadcast_buffer;
580 let max_batch_quotes = config.max_batch_quotes;
581 let max_outbound_message_size = config.max_outbound_message_size;
582 let (cmd_tx, cmd_rx) = mpsc::channel(command_buffer);
583 let (cancel_tx, cancel_rx) = mpsc::unbounded_channel();
584 let (gap_tx, gap_rx) = mpsc::unbounded_channel();
585 let gap_tx = config.reconnect_on_gap().then_some(gap_tx);
586 let (messages_tx, initial_messages_rx) = broadcast::channel::<InboundItem>(broadcast_buffer);
587 let (events_tx, initial_events_rx) = broadcast::channel(broadcast_buffer);
588 let (state_tx, state_rx) = watch::channel(ManagedWsState::Connecting);
589 let (shutdown_tx, shutdown_rx) = watch::channel(false);
590 let (vault_auth_tx, vault_auth_rx) = watch::channel(None);
591 let connection_open = Arc::new(std::sync::atomic::AtomicBool::new(false));
592
593 let session = session::run_managed_ws(
594 config,
595 session::HandleInputs {
596 cmd_rx,
597 cancel_rx,
598 gap_rx,
599 connection_open: Arc::clone(&connection_open),
600 },
601 messages_tx.clone(),
602 events_tx.clone(),
603 state_tx.clone(),
604 shutdown_rx,
605 vault_auth_tx.clone(),
606 );
607 let panic_events_tx = events_tx.clone();
608 let panic_state_tx = state_tx.clone();
609 let task = tokio::spawn(async move {
610 use futures_util::FutureExt;
611 if let Err(payload) = std::panic::AssertUnwindSafe(session).catch_unwind().await {
612 vault_auth_tx.send_if_modified(|grant| grant.take().is_some());
613 session::connection_ended(&connection_open, &panic_events_tx, None);
614 session::terminate(
615 &panic_events_tx,
616 &panic_state_tx,
617 ManagedWsTerminationReason::SessionPanicked,
618 );
619 std::panic::resume_unwind(payload);
620 }
621 });
622
623 Ok(ManagedWsHandle {
624 cmd_tx,
625 cancel_tx,
626 gap_tx,
627 messages_tx,
628 events_tx,
629 initial_messages_rx: Arc::new(StdMutex::new(Some(initial_messages_rx))),
630 initial_events_rx: Arc::new(StdMutex::new(Some(initial_events_rx))),
631 state_rx,
632 vault_auth_rx,
633 shutdown: Arc::new(ShutdownOnLastHandle { tx: shutdown_tx }),
634 next_await_id: Arc::new(AtomicU64::new(1)),
635 task: Arc::new(tokio::sync::Mutex::new(Some(task))),
636 max_batch_quotes,
637 max_outbound_message_size,
638 #[cfg(any(test, feature = "test-helpers"))]
639 next_inbound_sequence: Arc::new(AtomicU64::new(1)),
640 #[cfg(any(test, feature = "test-helpers"))]
641 state_tx,
642 })
643}
644
645#[cfg(test)]
646#[path = "../managed_tests.rs"]
647mod tests;