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
215fn take_initial<T>(slot: &StdMutex<Option<T>>) -> Option<T> {
216 match slot.lock() {
217 Ok(mut guard) => guard.take(),
218 Err(poisoned) => poisoned.into_inner().take(),
219 }
220}
221
222impl ManagedWsHandle {
223 pub fn subscribe_messages(&self) -> ManagedMessageReceiver {
227 let inner = match take_initial(&self.initial_messages_rx) {
228 Some(initial) => initial,
229 None => self.messages_tx.subscribe(),
230 };
231 ManagedMessageReceiver {
232 inner,
233 gap_tx: self.gap_tx.clone(),
234 }
235 }
236
237 pub fn subscribe_reference_changes(&self) -> ReferenceChangeReceiver {
238 ReferenceChangeReceiver::new(self.subscribe_messages())
239 }
240
241 pub fn subscribe_events(&self) -> broadcast::Receiver<ManagedWsEvent> {
243 match take_initial(&self.initial_events_rx) {
244 Some(initial) => initial,
245 None => self.events_tx.subscribe(),
246 }
247 }
248
249 pub fn subscribe_state(&self) -> watch::Receiver<ManagedWsState> {
251 self.state_rx.clone()
252 }
253
254 #[must_use]
255 pub fn state(&self) -> ManagedWsState {
256 self.state_rx.borrow().clone()
257 }
258
259 #[must_use]
261 pub fn vault_auth(&self) -> Option<VaultAuthSuccessData> {
262 self.vault_auth_rx.borrow().clone()
263 }
264
265 #[must_use]
267 pub fn subscribe_vault_auth(&self) -> watch::Receiver<Option<VaultAuthSuccessData>> {
268 self.vault_auth_rx.clone()
269 }
270
271 #[must_use]
272 pub fn is_ready(&self) -> bool {
273 matches!(*self.state_rx.borrow(), ManagedWsState::Ready { .. })
274 }
275
276 pub fn ensure_ready(&self) -> Result<(), ManagedWsError> {
277 self.ready_epoch().map(|_| ())
278 }
279
280 pub async fn wait_until_ready(&self) -> Result<u64, WaitUntilReadyError> {
282 let mut state_rx = self.state_rx.clone();
283 loop {
284 let state = state_rx.borrow().clone();
285 match state {
286 ManagedWsState::Ready { connection_epoch } => return Ok(connection_epoch),
287 ManagedWsState::Closed { reason } => {
288 return Err(WaitUntilReadyError::Terminated(reason));
289 }
290 _ => {}
291 }
292 state_rx
293 .changed()
294 .await
295 .map_err(|_| WaitUntilReadyError::StateChannelClosed)?;
296 }
297 }
298
299 fn command_epoch(&self, message: &ClientMessage) -> Result<Option<u64>, ManagedWsError> {
300 match self.state() {
301 ManagedWsState::Ready { connection_epoch } => Ok(Some(connection_epoch)),
302 _ if message.requires_ready() => Err(ManagedWsError::NotReady),
303 _ => Ok(None),
304 }
305 }
306
307 fn ready_epoch(&self) -> Result<u64, ManagedWsError> {
308 match self.state() {
309 ManagedWsState::Ready { connection_epoch } => Ok(connection_epoch),
310 _ => Err(ManagedWsError::NotReady),
311 }
312 }
313
314 pub async fn send(&self, message: ClientMessage) -> Result<(), ManagedWsError> {
315 let connection_epoch = self.command_epoch(&message)?;
316 self.send_with_epoch(message, connection_epoch).await
317 }
318
319 pub async fn send_in_epoch(
322 &self,
323 message: ClientMessage,
324 connection_epoch: u64,
325 ) -> Result<(), ManagedWsError> {
326 if self.ready_epoch()? != connection_epoch {
327 return Err(ManagedWsError::NotReady);
328 }
329 self.send_with_epoch(message, Some(connection_epoch)).await
330 }
331
332 async fn send_with_epoch(
333 &self,
334 message: ClientMessage,
335 connection_epoch: Option<u64>,
336 ) -> Result<(), ManagedWsError> {
337 let message = self.prepare_outbound(&message)?;
338 let (tx, rx) = oneshot::channel();
339 self.cmd_tx
340 .send(ManagedCommand::Send {
341 message,
342 connection_epoch,
343 tx,
344 })
345 .await
346 .map_err(|_| ManagedWsError::Closed)?;
347 SendTicket { rx }.wait().await
348 }
349
350 pub fn try_send(&self, message: ClientMessage) -> Result<SendTicket, ManagedWsError> {
351 let connection_epoch = self.command_epoch(&message)?;
352 self.try_send_with_epoch(message, connection_epoch)
353 }
354
355 pub fn try_send_in_epoch(
357 &self,
358 message: ClientMessage,
359 connection_epoch: u64,
360 ) -> Result<SendTicket, ManagedWsError> {
361 if self.ready_epoch()? != connection_epoch {
362 return Err(ManagedWsError::NotReady);
363 }
364 self.try_send_with_epoch(message, Some(connection_epoch))
365 }
366
367 fn try_send_with_epoch(
368 &self,
369 message: ClientMessage,
370 connection_epoch: Option<u64>,
371 ) -> Result<SendTicket, ManagedWsError> {
372 let message = self.prepare_outbound(&message)?;
373 let (tx, rx) = oneshot::channel();
374 match self.cmd_tx.try_send(ManagedCommand::Send {
375 message,
376 connection_epoch,
377 tx,
378 }) {
379 Ok(()) => Ok(SendTicket { rx }),
380 Err(mpsc::error::TrySendError::Closed(_)) => Err(ManagedWsError::Closed),
381 Err(mpsc::error::TrySendError::Full(_)) => Err(ManagedWsError::QueueFull),
382 }
383 }
384
385 pub async fn send_await(
386 &self,
387 message: ClientMessage,
388 timeout_duration: Duration,
389 ) -> Result<Arc<ServerMessage>, SendAwaitError> {
390 let connection_epoch = self
391 .command_epoch(&message)
392 .map_err(|_| SendAwaitError::NotReady)?;
393 self.send_await_with_epoch(message, timeout_duration, connection_epoch)
394 .await
395 }
396
397 pub(crate) async fn send_await_ready(
398 &self,
399 message: ClientMessage,
400 timeout_duration: Duration,
401 ) -> Result<Arc<ServerMessage>, SendAwaitError> {
402 let connection_epoch = self.ready_epoch().map_err(|_| SendAwaitError::NotReady)?;
403 self.send_await_with_epoch(message, timeout_duration, Some(connection_epoch))
404 .await
405 }
406
407 async fn send_await_with_epoch(
408 &self,
409 message: ClientMessage,
410 timeout_duration: Duration,
411 connection_epoch: Option<u64>,
412 ) -> Result<Arc<ServerMessage>, SendAwaitError> {
413 let deadline = tokio::time::Instant::now() + timeout_duration;
414 let (command, response) = self.await_command(message, connection_epoch, deadline)?;
415 if tokio::time::timeout_at(deadline, self.cmd_tx.send(command))
416 .await
417 .map_err(|_| SendAwaitError::Timeout)?
418 .is_err()
419 {
420 return Err(SendAwaitError::Disconnected);
421 }
422 response.await
423 }
424
425 #[cfg(feature = "vault")]
428 pub(crate) fn try_send_await_in_epoch(
429 &self,
430 message: ClientMessage,
431 timeout_duration: Duration,
432 connection_epoch: u64,
433 ) -> Result<
434 impl Future<Output = Result<Arc<ServerMessage>, SendAwaitError>> + use<'_>,
435 SendAwaitError,
436 > {
437 if self.ready_epoch().ok() != Some(connection_epoch) {
438 return Err(SendAwaitError::NotReady);
439 }
440 let deadline = tokio::time::Instant::now() + timeout_duration;
441 let (command, response) = self.await_command(message, Some(connection_epoch), deadline)?;
442 match self.cmd_tx.try_send(command) {
443 Ok(()) => Ok(response),
444 Err(mpsc::error::TrySendError::Full(_)) => Err(SendAwaitError::QueueFull),
445 Err(mpsc::error::TrySendError::Closed(_)) => Err(SendAwaitError::Disconnected),
446 }
447 }
448
449 fn await_command(
450 &self,
451 message: ClientMessage,
452 connection_epoch: Option<u64>,
453 deadline: tokio::time::Instant,
454 ) -> Result<
455 (
456 ManagedCommand,
457 impl Future<Output = Result<Arc<ServerMessage>, SendAwaitError>> + use<'_>,
458 ),
459 SendAwaitError,
460 > {
461 let prepared = self.prepare_outbound(&message)?;
462 let registration = AwaitTracker::prepare(&message)?;
463 let await_id = self.next_await_id.fetch_add(1, Ordering::Relaxed);
464 let (tx, rx) = oneshot::channel();
465 let command = ManagedCommand::SendAwait {
466 await_id,
467 message: prepared,
468 connection_epoch,
469 registration,
470 tx,
471 };
472 let response = async move {
473 match tokio::time::timeout_at(deadline, rx).await {
474 Ok(Ok(result)) => result,
475 Ok(Err(_)) => Err(SendAwaitError::Disconnected),
476 Err(_) => {
477 let _ = self.cancel_tx.send(await_id);
478 Err(SendAwaitError::Timeout)
479 }
480 }
481 };
482 Ok((command, response))
483 }
484
485 fn prepare_outbound(
486 &self,
487 message: &ClientMessage,
488 ) -> Result<PreparedClientMessage, OutboundMessageError> {
489 outbound::prepare_outbound(
490 message,
491 self.max_batch_quotes,
492 self.max_outbound_message_size,
493 )
494 }
495
496 pub async fn close(&self) -> Result<(), ManagedWsError> {
497 self.shutdown.request();
498 let mut task = self.task.lock().await;
499 if let Some(handle) = task.as_mut() {
500 let result = handle.await;
501 task.take();
502 result.map_err(ManagedWsError::TaskJoin)?;
503 }
504 Ok(())
505 }
506
507 #[cfg(any(test, feature = "test-helpers"))]
508 fn make_test_handle(
509 cmd_buffer: usize,
510 broadcast_buffer: usize,
511 ) -> (Self, mpsc::Receiver<ManagedCommand>) {
512 let (cmd_tx, cmd_rx) = mpsc::channel(cmd_buffer);
513 let (cancel_tx, _cancel_rx) = mpsc::unbounded_channel();
514 let (messages_tx, initial_messages_rx) =
515 broadcast::channel::<InboundItem>(broadcast_buffer);
516 let (events_tx, initial_events_rx) = broadcast::channel(broadcast_buffer);
517 let (state_tx, state_rx) = watch::channel(ManagedWsState::Connecting);
518 let (shutdown_tx, _shutdown_rx) = watch::channel(false);
519 let (_vault_auth_tx, vault_auth_rx) = watch::channel(None);
520 (
521 Self {
522 cmd_tx,
523 cancel_tx,
524 gap_tx: None,
525 messages_tx,
526 events_tx,
527 initial_messages_rx: Arc::new(StdMutex::new(Some(initial_messages_rx))),
528 initial_events_rx: Arc::new(StdMutex::new(Some(initial_events_rx))),
529 state_rx,
530 vault_auth_rx,
531 shutdown: Arc::new(ShutdownOnLastHandle { tx: shutdown_tx }),
532 next_await_id: Arc::new(AtomicU64::new(1)),
533 task: Arc::new(tokio::sync::Mutex::new(None)),
534 max_batch_quotes: 50,
535 max_outbound_message_size: 1024 * 1024,
536 #[cfg(any(test, feature = "test-helpers"))]
537 next_inbound_sequence: Arc::new(AtomicU64::new(1)),
538 #[cfg(any(test, feature = "test-helpers"))]
539 state_tx,
540 },
541 cmd_rx,
542 )
543 }
544
545 #[cfg(test)]
546 pub(crate) fn test_handle(
547 cmd_buffer: usize,
548 broadcast_buffer: usize,
549 ) -> (Self, mpsc::Receiver<ManagedCommand>) {
550 Self::make_test_handle(cmd_buffer, broadcast_buffer)
551 }
552
553 #[cfg(all(feature = "test-helpers", not(test)))]
554 pub fn test_handle(cmd_buffer: usize, broadcast_buffer: usize) -> (Self, ManagedWsTestPeer) {
555 let (handle, commands) = Self::make_test_handle(cmd_buffer, broadcast_buffer);
556 (handle, ManagedWsTestPeer::new(commands))
557 }
558
559 #[cfg(any(test, feature = "test-helpers"))]
560 pub fn inject_message(&self, msg: ServerMessage) {
561 let sequence = self.next_inbound_sequence.fetch_add(1, Ordering::Relaxed);
562 let _ = self.messages_tx.send(InboundItem::Message(ManagedInbound {
563 connection_epoch: 0,
564 sequence,
565 received_at: std::time::Instant::now(),
566 message: Arc::new(msg),
567 }));
568 }
569
570 #[cfg(any(test, feature = "test-helpers"))]
571 pub fn inject_event(&self, event: ManagedWsEvent) {
572 let _ = self.events_tx.send(event);
573 }
574
575 #[cfg(any(test, feature = "test-helpers"))]
576 pub fn inject_state(&self, state: ManagedWsState) {
577 self.state_tx.send_replace(state);
578 }
579}
580
581pub fn spawn_managed_ws(config: ManagedWsConfig) -> Result<ManagedWsHandle, ManagedWsConfigError> {
582 config.validate()?;
583 let command_buffer = config.command_buffer;
584 let broadcast_buffer = config.broadcast_buffer;
585 let max_batch_quotes = config.max_batch_quotes;
586 let max_outbound_message_size = config.max_outbound_message_size;
587 let (cmd_tx, cmd_rx) = mpsc::channel(command_buffer);
588 let (cancel_tx, cancel_rx) = mpsc::unbounded_channel();
589 let (gap_tx, gap_rx) = mpsc::unbounded_channel();
590 let gap_tx = config.reconnect_on_gap().then_some(gap_tx);
591 let (messages_tx, initial_messages_rx) = broadcast::channel::<InboundItem>(broadcast_buffer);
592 let (events_tx, initial_events_rx) = broadcast::channel(broadcast_buffer);
593 let (state_tx, state_rx) = watch::channel(ManagedWsState::Connecting);
594 let (shutdown_tx, shutdown_rx) = watch::channel(false);
595 let (vault_auth_tx, vault_auth_rx) = watch::channel(None);
596 let connection_open = Arc::new(std::sync::atomic::AtomicBool::new(false));
597
598 let session = session::run_managed_ws(
599 config,
600 session::HandleInputs {
601 cmd_rx,
602 cancel_rx,
603 gap_rx,
604 connection_open: Arc::clone(&connection_open),
605 },
606 messages_tx.clone(),
607 events_tx.clone(),
608 state_tx.clone(),
609 shutdown_rx,
610 vault_auth_tx.clone(),
611 );
612 let panic_events_tx = events_tx.clone();
613 let panic_state_tx = state_tx.clone();
614 let task = tokio::spawn(async move {
615 use futures_util::FutureExt;
616 if let Err(payload) = std::panic::AssertUnwindSafe(session).catch_unwind().await {
617 vault_auth_tx.send_if_modified(|grant| grant.take().is_some());
618 session::connection_ended(&connection_open, &panic_events_tx, None);
619 session::terminate(
620 &panic_events_tx,
621 &panic_state_tx,
622 ManagedWsTerminationReason::SessionPanicked,
623 );
624 std::panic::resume_unwind(payload);
625 }
626 });
627
628 Ok(ManagedWsHandle {
629 cmd_tx,
630 cancel_tx,
631 gap_tx,
632 messages_tx,
633 events_tx,
634 initial_messages_rx: Arc::new(StdMutex::new(Some(initial_messages_rx))),
635 initial_events_rx: Arc::new(StdMutex::new(Some(initial_events_rx))),
636 state_rx,
637 vault_auth_rx,
638 shutdown: Arc::new(ShutdownOnLastHandle { tx: shutdown_tx }),
639 next_await_id: Arc::new(AtomicU64::new(1)),
640 task: Arc::new(tokio::sync::Mutex::new(Some(task))),
641 max_batch_quotes,
642 max_outbound_message_size,
643 #[cfg(any(test, feature = "test-helpers"))]
644 next_inbound_sequence: Arc::new(AtomicU64::new(1)),
645 #[cfg(any(test, feature = "test-helpers"))]
646 state_tx,
647 })
648}
649
650#[cfg(test)]
651#[path = "../managed_tests.rs"]
652mod tests;