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

acta_maker_sdk/ws/managed/
mod.rs
