simulator_client/managed/
mod.rs

1//! Connection-managed wrappers over the backtest WebSocket protocol.
2//!
3//! Most callers should use [`ManagedBacktestSession`]. It owns the control and
4//! subscription manager tasks, waits for session creation, gates `Continue`
5//! sends on live connections, and provides one shutdown path.
6//!
7//! The control and subscription WebSockets are each owned by a dedicated task
8//! that handles its own lifecycle: connect, handshake, keepalive, reconnect.
9//! Workload code interacts via channels and a status watcher, never with the
10//! WebSocket directly.
11
12use std::{
13    future::Future,
14    sync::{
15        Arc,
16        atomic::{AtomicUsize, Ordering},
17    },
18    time::{Duration, Instant},
19};
20
21use futures::SinkExt;
22use rand::Rng;
23use simulator_api::{BacktestError, BacktestRequest};
24use tokio::{
25    net::TcpStream,
26    sync::{Notify, OwnedSemaphorePermit, Semaphore, watch},
27};
28use tokio_tungstenite::{
29    MaybeTlsStream, WebSocketStream, connect_async,
30    tungstenite::{Error as WsError, Message, client::IntoClientRequest, http::HeaderValue},
31};
32use tokio_util::sync::CancellationToken;
33use tracing::warn;
34
35use crate::error::err_chain;
36
37mod control;
38mod parallel;
39mod session;
40mod subscription;
41
42pub use control::{ControlEvent, ControlHandle, spawn_control_manager};
43pub use parallel::{ManagedParallelSession, ParallelSubSession};
44pub use session::{ManagedBacktestSession, ManagedEvent, ManagedSessionError};
45pub use subscription::{
46    SubscriptionHandle, SubscriptionNotification, spawn_account_diff_subscription_manager,
47    spawn_action_subscription_manager, spawn_transaction_subscription_manager,
48};
49
50/// Timeout for the initial WebSocket connect (TCP + TLS + HTTP upgrade).
51pub const CONNECT_TIMEOUT: Duration = Duration::from_secs(30);
52
53/// Timeout for any single response during a handshake (Create, Attach, Resume,
54/// subscribe ack). Does not apply to per-message reads during an established
55/// session — those are bounded by ping/pong liveness instead.
56///
57/// Set generously because the management service can park subscribe upgrades
58/// for the duration of session startup (observed ~100s on cold runtimes); a
59/// shorter timeout fires before management can forward the subscribe message
60/// and gets reported as "subscribe ack timeout" even though the session is
61/// fine. The proper fix is server-side (reject parked upgrades cleanly), but
62/// this keeps the client from giving up early in the meantime.
63pub const HANDSHAKE_RESPONSE_TIMEOUT: Duration = Duration::from_secs(120);
64
65/// How often to send a WebSocket ping during an established connection.
66pub const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15);
67
68/// How long without inbound traffic before declaring the connection dead.
69/// Roughly three missed pings.
70pub const KEEPALIVE_MISS_DEADLINE: Duration = Duration::from_secs(45);
71
72/// Total budget for the graceful close handshake (close request + close frame).
73/// A supervisor's SIGTERM→SIGKILL grace must exceed this or the close is cut short.
74pub const GRACEFUL_CLOSE_TIMEOUT: Duration = Duration::from_secs(5);
75
76pub const RECONNECT_INITIAL_BACKOFF: Duration = Duration::from_secs(1);
77pub const RECONNECT_MAX_BACKOFF: Duration = Duration::from_secs(30);
78pub const RECONNECT_BACKOFF_MULTIPLIER: f64 = 2.0;
79pub const RECONNECT_JITTER: f64 = 0.2;
80pub const RECONNECT_MAX_TOTAL: Duration = Duration::from_secs(5 * 60);
81pub const RECONNECT_MAX_ATTEMPTS: u32 = 20;
82
83/// Reconnect attempts that fire immediately, before the cross-session reconnect
84/// gate engages. A transient drop recovers fast; only once these rapid retries
85/// are exhausted does a session fall back to waiting for siblings to drain.
86pub const RECONNECT_UNGATED_ATTEMPTS: u32 = 5;
87
88/// A connection that stays `Up` for this long resets the reconnect counter.
89pub const RECONNECT_UPTIME_RESET: Duration = Duration::from_secs(30);
90
91/// Connection state as observed from outside the manager task.
92///
93/// `Down` is a transient "currently not connected, manager is retrying"
94/// state. `Failed` is terminal; the retry budget was exhausted.
95#[derive(Clone, Debug, PartialEq, Eq)]
96pub enum ConnectionStatus {
97    Up,
98    Down,
99    Failed(String),
100}
101
102/// Identity of a created backtest session.
103#[derive(Clone, Debug)]
104pub struct SessionInfo {
105    pub session_id: String,
106    pub rpc_endpoint: String,
107    /// Opaque `task_id` reported by the server for this session, if any.
108    pub task_id: Option<String>,
109}
110
111/// Tracks reconnect attempts and enforces the reconnect policy.
112pub(crate) struct ReconnectBudget {
113    attempts: u32,
114    started_at: std::time::Instant,
115    current_backoff: Duration,
116}
117
118impl ReconnectBudget {
119    pub fn new() -> Self {
120        Self {
121            attempts: 0,
122            started_at: std::time::Instant::now(),
123            current_backoff: RECONNECT_INITIAL_BACKOFF,
124        }
125    }
126
127    pub fn reset(&mut self) {
128        self.attempts = 0;
129        self.started_at = std::time::Instant::now();
130        self.current_backoff = RECONNECT_INITIAL_BACKOFF;
131    }
132
133    pub fn attempt(&self) -> u32 {
134        self.attempts
135    }
136
137    /// Discount time the subscription spent parked behind the reconnect
138    /// coordinator. Parking is deliberate waiting for siblings to free the
139    /// link, not retry thrashing, so it must not count against the wall-clock
140    /// budget — otherwise a session that correctly steps aside for minutes is
141    /// failed the moment it reconnects. The attempt-count ceiling still bounds
142    /// total retries.
143    pub fn discount_parked(&mut self, parked: Duration) {
144        self.started_at += parked;
145    }
146
147    /// Record an attempt and return its backoff, or `None` if the budget is
148    /// exhausted.
149    pub fn next_backoff(&mut self) -> Option<Duration> {
150        if self.attempts >= RECONNECT_MAX_ATTEMPTS
151            || self.started_at.elapsed() >= RECONNECT_MAX_TOTAL
152        {
153            return None;
154        }
155        self.attempts += 1;
156        let backoff = with_jitter(self.current_backoff);
157        self.current_backoff = std::cmp::min(
158            RECONNECT_MAX_BACKOFF,
159            Duration::from_secs_f64(
160                self.current_backoff.as_secs_f64() * RECONNECT_BACKOFF_MULTIPLIER,
161            ),
162        );
163        Some(backoff)
164    }
165}
166
167/// Coordinates subscription reconnects across a parallel batch sharing one
168/// bandwidth-constrained link.
169///
170/// A dropped subscription competing with its still-streaming siblings just
171/// starves — the handshake can't get through a saturated pipe. So once a session
172/// has burned its fast ungated attempts, it stops fighting and **parks until no
173/// sibling is streaming** (the connected sessions have finished and freed their
174/// bandwidth), then reconnects into the quiet link and resumes via
175/// `replayFromSlot`. Stepping aside also speeds the streaming siblings up, so the
176/// batch drains and the parked sessions recover — serially, one at a time, since
177/// at a tight rate only one stream fits.
178pub struct ReconnectCoordinator {
179    /// Subscriptions currently `Up` and streaming. A parked session waits for
180    /// this to reach zero before reconnecting.
181    streaming: AtomicUsize,
182    /// Woken whenever `streaming` decreases, so parked sessions re-check.
183    drained: Notify,
184    /// One in-flight reconnect handshake at a time, so the all-parked case
185    /// recovers serially rather than storming the link the moment it goes quiet.
186    handshake: Arc<Semaphore>,
187}
188
189impl Default for ReconnectCoordinator {
190    fn default() -> Self {
191        Self::new()
192    }
193}
194
195impl ReconnectCoordinator {
196    pub fn new() -> Self {
197        Self {
198            streaming: AtomicUsize::new(0),
199            drained: Notify::new(),
200            handshake: Arc::new(Semaphore::new(1)),
201        }
202    }
203
204    /// Mark this subscription as streaming until the returned guard drops.
205    /// Call once a (re)subscribe succeeds; the guard's `Drop` decrements the
206    /// count and wakes any sibling parked in [`Self::reconnect_slot`]. RAII so
207    /// a panicking or aborted task can't leak the count and wedge the batch.
208    pub fn enter(self: &Arc<Self>) -> StreamingGuard {
209        self.streaming.fetch_add(1, Ordering::SeqCst);
210        StreamingGuard(self.clone())
211    }
212
213    /// Park until no sibling is streaming, then take the single reconnect slot.
214    /// The returned permit must be held across connect + subscribe and dropped
215    /// once [`Self::enter`] has been called (or the attempt failed), so the next
216    /// parked session proceeds only after this one is streaming again. Returns
217    /// `None` if cancelled.
218    pub async fn reconnect_slot(&self, cancel: &CancellationToken) -> Option<OwnedSemaphorePermit> {
219        loop {
220            // Wait for the link to go quiet. Register interest before reading so
221            // a `leave()` between the check and the wait isn't lost.
222            loop {
223                let drained = self.drained.notified();
224                if self.streaming.load(Ordering::SeqCst) == 0 {
225                    break;
226                }
227                tokio::select! {
228                    biased;
229                    _ = cancel.cancelled() => return None,
230                    _ = drained => {}
231                }
232            }
233            let permit = tokio::select! {
234                biased;
235                _ = cancel.cancelled() => return None,
236                p = self.handshake.clone().acquire_owned() => p.ok()?,
237            };
238            // A sibling may have re-entered while we queued for the permit; if so,
239            // release it and wait for quiet again rather than reconnect into
240            // contention.
241            if self.streaming.load(Ordering::SeqCst) == 0 {
242                return Some(permit);
243            }
244        }
245    }
246}
247
248/// Holds a [`ReconnectCoordinator`]'s streaming count for the lifetime of one
249/// connection; decrements and wakes parked siblings on drop.
250pub struct StreamingGuard(Arc<ReconnectCoordinator>);
251
252impl Drop for StreamingGuard {
253    fn drop(&mut self) {
254        self.0.streaming.fetch_sub(1, Ordering::SeqCst);
255        self.0.drained.notify_waiters();
256    }
257}
258
259fn with_jitter(d: Duration) -> Duration {
260    let jitter = rand::rng().random_range(-RECONNECT_JITTER..RECONNECT_JITTER);
261    let secs = (d.as_secs_f64() * (1.0 + jitter)).max(0.0);
262    Duration::from_secs_f64(secs)
263}
264
265/// Sleep for `delay`, returning early on cancellation. Returns `true` if the
266/// sleep completed, `false` if cancelled.
267pub(crate) async fn cancellable_sleep(delay: Duration, cancel: &CancellationToken) -> bool {
268    tokio::select! {
269        _ = tokio::time::sleep(delay) => true,
270        _ = cancel.cancelled() => false,
271    }
272}
273
274/// A managed control WebSocket stream.
275pub(super) type Ws = WebSocketStream<MaybeTlsStream<TcpStream>>;
276
277/// Open a control WebSocket to `url`, attaching the API-key header, bounded by
278/// [`CONNECT_TIMEOUT`]. Shared by the single-session and parallel managers.
279pub(super) async fn connect_ws(url: &str, api_key: &str) -> Result<Ws, String> {
280    let mut request = url
281        .into_client_request()
282        .map_err(|e| format!("build request: {}", err_chain(&e)))?;
283    request.headers_mut().insert(
284        "X-API-Key",
285        HeaderValue::from_str(api_key).map_err(|e| format!("api key header: {}", err_chain(&e)))?,
286    );
287
288    let connect = tokio::time::timeout(CONNECT_TIMEOUT, connect_async(request))
289        .await
290        .map_err(|_| format!("connect timeout after {CONNECT_TIMEOUT:?}"))?
291        .map_err(|e| format!("connect: {}", err_chain(&e)))?;
292    Ok(connect.0)
293}
294
295/// Serialize and send a `BacktestRequest` over a control WebSocket.
296pub(super) async fn send_request(ws: &mut Ws, req: &BacktestRequest) -> Result<(), String> {
297    let text = serde_json::to_string(req).map_err(|e| format!("serialize: {}", err_chain(&e)))?;
298    ws.send(Message::Text(text))
299        .await
300        .map_err(|e| format!("send: {}", err_chain(&e)))
301}
302
303/// Resolve a server-reported `rpc_endpoint` to an absolute URL: pass through an
304/// already-absolute endpoint, otherwise join it onto the control `base`.
305pub(super) fn resolve_rpc_url(base: &str, endpoint: &str) -> String {
306    if endpoint.starts_with("http://") || endpoint.starts_with("https://") {
307        endpoint.to_string()
308    } else {
309        format!("{}/{}", base, endpoint.trim_start_matches('/'))
310    }
311}
312
313/// Outcome of a control-WS handshake (create/attach/resume).
314pub(super) enum HandshakeError {
315    /// Recoverable (network blip, or the session is momentarily contended).
316    Transient(String),
317    /// Unrecoverable — the server rejected the session outright.
318    Fatal(String),
319}
320
321/// Classify a `BacktestError` returned during a handshake `stage`: a contended
322/// session is expected to become claimable shortly (transient); anything else is
323/// fatal.
324pub(super) fn handshake_error_for_response(
325    stage: &'static str,
326    err: BacktestError,
327) -> HandshakeError {
328    match err {
329        BacktestError::SessionOwnershipBusy { .. } => {
330            HandshakeError::Transient(format!("{stage} contended: {}", err_chain(&err)))
331        }
332        _ => HandshakeError::Fatal(format!("{stage} rejected: {}", err_chain(&err))),
333    }
334}
335
336/// Why a control task's message loop exited. Shared by the single-session and
337/// parallel control managers.
338pub(super) enum MessageLoopExit {
339    /// Server reported the session(s) done, or the driver(s) dropped the request
340    /// channel. Calls for a graceful WS close.
341    SessionEnded,
342    /// Cancellation token fired. Still graceful-closes so the server reaps the
343    /// runtime now instead of after the disconnect timeout.
344    Cancelled,
345    /// Connection lost — attempt to reconnect.
346    ConnectionLost(String),
347    /// Protocol or application error with no sensible recovery.
348    Terminal(String),
349}
350
351/// Publish a `ConnectionStatus` transition, suppressing no-op updates.
352pub(super) fn publish_status(
353    status_tx: &watch::Sender<ConnectionStatus>,
354    status: ConnectionStatus,
355) {
356    status_tx.send_if_modified(|current| {
357        if *current == status {
358            false
359        } else {
360            *current = status;
361            true
362        }
363    });
364}
365
366/// Send `CloseBacktestSession` then the WS close frame, bounded by a single
367/// [`GRACEFUL_CLOSE_TIMEOUT`] budget. The close request makes the server reap
368/// the runtime(s) immediately instead of waiting out the disconnect timeout; the
369/// close frame is best-effort cleanup.
370pub(super) async fn graceful_close(ws: &mut Ws) {
371    let _ = tokio::time::timeout(GRACEFUL_CLOSE_TIMEOUT, async {
372        let _ = send_request(ws, &BacktestRequest::CloseBacktestSession).await;
373        let _ = ws.close(None).await;
374    })
375    .await;
376}
377
378/// An inbound control frame, classified for the message loop.
379pub(super) enum InboundFrame {
380    /// A text (or UTF-8 binary) payload to parse.
381    Text(String),
382    /// A frame with no application payload (ping/pong/raw/non-UTF-8 binary).
383    Ignore,
384    /// The connection ended or errored; carries the loss reason.
385    Lost(String),
386}
387
388/// Classify the result of `ws.next()` for a control message loop.
389pub(super) fn classify_inbound(msg: Option<Result<Message, WsError>>) -> InboundFrame {
390    match msg {
391        Some(Ok(Message::Text(t))) => InboundFrame::Text(t),
392        Some(Ok(Message::Binary(b))) => match String::from_utf8(b) {
393            Ok(t) => InboundFrame::Text(t),
394            Err(_) => InboundFrame::Ignore,
395        },
396        Some(Ok(Message::Pong(_) | Message::Ping(_) | Message::Frame(_))) => InboundFrame::Ignore,
397        Some(Ok(Message::Close(frame))) => InboundFrame::Lost(format!("remote close: {frame:?}")),
398        Some(Err(e)) => InboundFrame::Lost(format!("ws read: {}", err_chain(&e))),
399        None => InboundFrame::Lost("ws stream ended".into()),
400    }
401}
402
403/// Send a keepalive ping unless inbound traffic has been silent past the
404/// deadline. Returns the loss reason if the connection should be considered
405/// dead, else `None`.
406pub(super) async fn send_keepalive_ping(ws: &mut Ws, last_inbound: Instant) -> Option<String> {
407    if last_inbound.elapsed() > KEEPALIVE_MISS_DEADLINE {
408        return Some(format!("no traffic for {:?}", last_inbound.elapsed()));
409    }
410    if let Err(e) = ws.send(Message::Ping(vec![])).await {
411        return Some(format!("ping send: {}", err_chain(&e)));
412    }
413    None
414}
415
416/// A control WebSocket task: connect, handshake, run a message loop, reconnect.
417/// [`run_control_loop`] drives the shared connect/backoff/exit state machine;
418/// implementors supply the protocol-specific handshake and message loop.
419pub(super) trait ControlConnection: Send + 'static {
420    fn url(&self) -> &str;
421    fn api_key(&self) -> &str;
422    fn cancel(&self) -> &CancellationToken;
423    /// Prefix for reconnect log lines, e.g. `"control"` / `"parallel control"`.
424    fn label(&self) -> &'static str;
425    fn status_tx(&self) -> &watch::Sender<ConnectionStatus>;
426    /// Fail the pending create one-shot with `reason` (no-op once fired).
427    fn fail_pending(&mut self, reason: String);
428    /// Create on first connect, attach/resume on reconnect.
429    fn handshake(&mut self, ws: Ws) -> impl Future<Output = Result<Ws, HandshakeError>> + Send;
430    /// Drive the established connection until it exits.
431    fn message_loop(&mut self, ws: Ws) -> impl Future<Output = MessageLoopExit> + Send;
432
433    fn publish(&self, status: ConnectionStatus) {
434        publish_status(self.status_tx(), status);
435    }
436
437    fn finish_failed(&mut self, reason: String) {
438        self.fail_pending(reason.clone());
439        self.publish(ConnectionStatus::Failed(reason));
440    }
441}
442
443/// Drive a [`ControlConnection`] through its connect → handshake → run →
444/// reconnect lifecycle until the session ends, is cancelled, or the reconnect
445/// budget is exhausted.
446pub(super) async fn run_control_loop<T: ControlConnection>(mut task: T) {
447    let mut budget = ReconnectBudget::new();
448
449    loop {
450        if task.cancel().is_cancelled() {
451            task.fail_pending("cancelled before session created".to_string());
452            return;
453        }
454        task.publish(ConnectionStatus::Down);
455
456        let ws = match connect_ws(task.url(), task.api_key()).await {
457            Ok(ws) => ws,
458            Err(why) => {
459                if let Some(delay) = budget.next_backoff() {
460                    warn!(attempt = budget.attempt(), error = %why, ?delay, "{} connect failed, retrying", task.label());
461                    if !cancellable_sleep(delay, task.cancel()).await {
462                        return;
463                    }
464                    continue;
465                }
466                task.finish_failed(format!("connect: {why}"));
467                return;
468            }
469        };
470
471        let ws = match task.handshake(ws).await {
472            Ok(ws) => ws,
473            Err(HandshakeError::Fatal(why)) => {
474                task.finish_failed(format!("handshake: {why}"));
475                return;
476            }
477            Err(HandshakeError::Transient(why)) => {
478                if let Some(delay) = budget.next_backoff() {
479                    warn!(attempt = budget.attempt(), error = %why, ?delay, "{} handshake failed, retrying", task.label());
480                    if !cancellable_sleep(delay, task.cancel()).await {
481                        return;
482                    }
483                    continue;
484                }
485                task.finish_failed(format!("handshake: {why}"));
486                return;
487            }
488        };
489
490        task.publish(ConnectionStatus::Up);
491        let connected_at = Instant::now();
492
493        match task.message_loop(ws).await {
494            MessageLoopExit::SessionEnded | MessageLoopExit::Cancelled => return,
495            MessageLoopExit::ConnectionLost(why) => {
496                if connected_at.elapsed() >= RECONNECT_UPTIME_RESET {
497                    budget.reset();
498                }
499                if let Some(delay) = budget.next_backoff() {
500                    warn!(attempt = budget.attempt(), reason = %why, ?delay, "{} connection lost, reconnecting", task.label());
501                    if !cancellable_sleep(delay, task.cancel()).await {
502                        return;
503                    }
504                    continue;
505                }
506                task.finish_failed(format!("connection lost: {why}"));
507                return;
508            }
509            MessageLoopExit::Terminal(why) => {
510                task.finish_failed(why);
511                return;
512            }
513        }
514    }
515}
516
517#[cfg(test)]
518mod tests {
519    use super::*;
520
521    #[test]
522    fn budget_exhausts_after_max_attempts() {
523        let mut b = ReconnectBudget::new();
524        for _ in 0..RECONNECT_MAX_ATTEMPTS {
525            assert!(b.next_backoff().is_some());
526        }
527        assert!(b.next_backoff().is_none());
528    }
529
530    #[test]
531    fn budget_reset_restores_full_budget() {
532        let mut b = ReconnectBudget::new();
533        b.next_backoff();
534        b.next_backoff();
535        b.reset();
536        assert_eq!(b.attempt(), 0);
537    }
538
539    #[test]
540    fn streaming_guard_balances_the_count() {
541        let coord = Arc::new(ReconnectCoordinator::new());
542        assert_eq!(coord.streaming.load(Ordering::SeqCst), 0);
543        let g = coord.enter();
544        assert_eq!(coord.streaming.load(Ordering::SeqCst), 1);
545        drop(g);
546        assert_eq!(coord.streaming.load(Ordering::SeqCst), 0);
547    }
548
549    #[tokio::test]
550    async fn reconnect_slot_available_when_link_is_quiet() {
551        let coord = Arc::new(ReconnectCoordinator::new());
552        let cancel = CancellationToken::new();
553        assert!(coord.reconnect_slot(&cancel).await.is_some());
554    }
555
556    #[tokio::test]
557    async fn reconnect_slot_unparks_when_last_sibling_leaves() {
558        let coord = Arc::new(ReconnectCoordinator::new());
559        let cancel = CancellationToken::new();
560        let guard = coord.enter(); // streaming == 1, so a slot request must park
561
562        let waiter = tokio::spawn({
563            let coord = coord.clone();
564            let cancel = cancel.clone();
565            async move { coord.reconnect_slot(&cancel).await.is_some() }
566        });
567
568        // While a sibling streams, the parked request stays pending.
569        tokio::task::yield_now().await;
570        assert!(!waiter.is_finished());
571
572        drop(guard); // streaming -> 0, wakes the parked request
573        assert!(waiter.await.unwrap());
574    }
575
576    #[tokio::test]
577    async fn reconnect_slot_returns_none_on_cancel_while_parked() {
578        let coord = Arc::new(ReconnectCoordinator::new());
579        let _guard = coord.enter(); // keep the link busy so the request parks
580        let cancel = CancellationToken::new();
581        cancel.cancel();
582        assert!(coord.reconnect_slot(&cancel).await.is_none());
583    }
584
585    #[test]
586    fn discount_parked_does_not_consume_the_budget() {
587        let mut b = ReconnectBudget::new();
588        // A long park must leave the full attempt budget intact afterwards.
589        b.discount_parked(2 * RECONNECT_MAX_TOTAL);
590        for _ in 0..RECONNECT_MAX_ATTEMPTS {
591            assert!(b.next_backoff().is_some());
592        }
593    }
594}
simulator_client/managed/mod.rs

simulator_client/managed/
mod.rs
