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

acta_maker_sdk/ws/managed/
mod.rs
