acta_maker_sdk/ws/managed/
mod.rs

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

acta_maker_sdk/ws/managed/
mod.rs
