1use std::{
7 collections::BTreeMap,
8 fmt,
9 sync::{
10 Arc, Weak,
11 atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
12 },
13 time::Duration,
14};
15
16use data_encoding::BASE64;
17use futures_util::{SinkExt as _, StreamExt as _};
18use parking_lot::Mutex as ParkingMutex;
19use rand::{TryRng as _, rngs::SysRng};
20use reqwest::{StatusCode, Version, header};
21use serde::{Deserialize, Serialize, de::DeserializeOwned};
22use serde_json::Value;
23use serde_json::value::RawValue;
24use sha1::{Digest as _, Sha1};
25use thiserror::Error;
26use tokio::{
27 sync::{Notify, mpsc, oneshot},
28 task::JoinHandle,
29 time::Instant,
30};
31use tokio_tungstenite::{
32 WebSocketStream,
33 tungstenite::{
34 Error as TungsteniteError, Message,
35 protocol::{Role, WebSocketConfig},
36 },
37};
38use tokio_util::sync::CancellationToken;
39
40use crate::{AccountId, ContractId, Endpoints, Error as ClientError, token::TokenStore};
41
42const SIGNALR_TERMINATOR: char = '\u{001e}';
43const SIGNALR_PING: &str = "{\"type\":6}\u{001e}";
44const WEBSOCKET_READ_BUFFER_SIZE: usize = 64 * 1_024;
46const WEBSOCKET_WRITE_BUFFER_SIZE: usize = 64 * 1_024;
47const EVENT_BASE_WEIGHT: usize = 256;
48const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
49const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
50const CLOSE_TIMEOUT: Duration = Duration::from_secs(5);
51const CLIENT_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15);
52const SOCKET_PROBE_INTERVAL: Duration = Duration::from_secs(15);
55const WATCHDOG_INTERVAL: Duration = Duration::from_secs(5);
56const WEBSOCKET_GUID: &[u8] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
57
58type PendingInvocation = oneshot::Sender<Result<(), ()>>;
59
60struct PendingEntry {
61 generation: u64,
62 reply: PendingInvocation,
63}
64
65type PendingInvocations = Arc<ParkingMutex<BTreeMap<String, PendingEntry>>>;
66
67#[derive(Serialize)]
68#[serde(rename_all = "camelCase")]
69struct OutboundInvocation<'a> {
70 #[serde(rename = "type")]
71 message_type: u8,
72 invocation_id: &'a str,
73 target: &'a str,
74 arguments: &'a [Value],
75}
76
77#[derive(Clone, Copy, Debug, Eq, PartialEq)]
79#[non_exhaustive]
80pub enum Hub {
81 Market,
83 User,
85}
86
87impl Hub {
88 const fn path(self) -> &'static str {
89 match self {
90 Self::Market => "market",
91 Self::User => "user",
92 }
93 }
94}
95
96#[derive(Clone, Debug)]
98pub struct SignalRInvocation {
99 target: String,
100 contract_id: Option<ContractId>,
101 payload: Value,
102 raw_entity: Box<RawValue>,
103}
104
105impl PartialEq for SignalRInvocation {
106 fn eq(&self, other: &Self) -> bool {
107 self.target == other.target
108 && self.contract_id == other.contract_id
109 && self.payload == other.payload
110 && self.raw_entity.get() == other.raw_entity.get()
111 }
112}
113
114impl SignalRInvocation {
115 pub fn from_value(value: Value) -> Result<Option<Self>, RealtimeError> {
124 if value.get("type").and_then(Value::as_i64) != Some(1) {
125 return Ok(None);
126 }
127 let Value::Object(mut object) = value else {
128 return Err(RealtimeError::Protocol(
129 "invocation frame was not an object",
130 ));
131 };
132 let target = object
133 .remove("target")
134 .and_then(|target| target.as_str().map(str::to_owned))
135 .ok_or(RealtimeError::Protocol("invocation target is missing"))?;
136 let Value::Array(arguments) = object
137 .remove("arguments")
138 .ok_or(RealtimeError::Protocol("invocation arguments are missing"))?
139 else {
140 return Err(RealtimeError::Protocol(
141 "invocation arguments were not an array",
142 ));
143 };
144 let mut arguments = arguments.into_iter();
145 let first = arguments
146 .next()
147 .ok_or(RealtimeError::Protocol("invocation payload is missing"))?;
148 let (contract_id, payload) = match arguments.next() {
149 Some(second) => {
150 let contract_id = ContractId::new(first.as_str().ok_or(
151 RealtimeError::Protocol("market contract identifier was not a string"),
152 )?)
153 .map_err(|_| RealtimeError::Protocol("contract identifier is invalid"))?;
154 (Some(contract_id), second)
155 }
156 None => (None, first),
157 };
158 if arguments.next().is_some() {
159 return Err(RealtimeError::Protocol(
160 "invocation contained too many arguments",
161 ));
162 }
163 let entity = payload
164 .get("data")
165 .filter(|data| data.is_object() || data.is_array())
166 .unwrap_or(&payload);
167 let raw_entity =
168 RawValue::from_string(serde_json::to_string(entity).map_err(RealtimeError::Decode)?)
169 .map_err(RealtimeError::Decode)?;
170 Ok(Some(Self {
171 target,
172 contract_id,
173 payload,
174 raw_entity,
175 }))
176 }
177
178 pub fn from_json(json: &str) -> Result<Option<Self>, RealtimeError> {
189 #[derive(Deserialize)]
190 struct InvocationFrame {
191 #[serde(rename = "type")]
192 message_type: u64,
193 target: Option<String>,
194 arguments: Option<Vec<Box<RawValue>>>,
195 }
196
197 let frame: InvocationFrame = serde_json::from_str(json).map_err(RealtimeError::Decode)?;
198 if frame.message_type != 1 {
199 return Ok(None);
200 }
201 let target = frame
202 .target
203 .ok_or(RealtimeError::Protocol("invocation target is missing"))?;
204 let mut arguments = frame
205 .arguments
206 .ok_or(RealtimeError::Protocol("invocation arguments are missing"))?
207 .into_iter();
208 let first = arguments
209 .next()
210 .ok_or(RealtimeError::Protocol("invocation payload is missing"))?;
211 let (contract_id, raw_payload) = match arguments.next() {
212 Some(second) => {
213 let contract = serde_json::from_str::<String>(first.get()).map_err(|_| {
214 RealtimeError::Protocol("market contract identifier was not a string")
215 })?;
216 let contract_id = ContractId::new(contract)
217 .map_err(|_| RealtimeError::Protocol("contract identifier is invalid"))?;
218 (Some(contract_id), second)
219 }
220 None => (None, first),
221 };
222 if arguments.next().is_some() {
223 return Err(RealtimeError::Protocol(
224 "invocation contained too many arguments",
225 ));
226 }
227 let payload: Value =
228 serde_json::from_str(raw_payload.get()).map_err(RealtimeError::Decode)?;
229 let raw_entity = if payload
230 .get("data")
231 .is_some_and(|data| data.is_object() || data.is_array())
232 {
233 let mut object =
234 serde_json::from_str::<BTreeMap<String, Box<RawValue>>>(raw_payload.get())
235 .map_err(RealtimeError::Decode)?;
236 object.remove("data").ok_or(RealtimeError::Protocol(
237 "invocation data envelope is missing",
238 ))?
239 } else {
240 raw_payload
241 };
242 Ok(Some(Self {
243 target,
244 contract_id,
245 payload,
246 raw_entity,
247 }))
248 }
249
250 #[must_use]
252 pub fn target(&self) -> &str {
253 &self.target
254 }
255
256 #[must_use]
258 pub fn contract_id(&self) -> Option<&ContractId> {
259 self.contract_id.as_ref()
260 }
261
262 #[must_use]
264 pub fn payload(&self) -> &Value {
265 &self.payload
266 }
267
268 #[must_use]
271 pub fn entity(&self) -> &Value {
272 self.payload
273 .get("data")
274 .filter(|data| data.is_object() || data.is_array())
275 .unwrap_or(&self.payload)
276 }
277
278 pub fn decode<T>(&self) -> Result<T, RealtimeError>
284 where
285 T: DeserializeOwned,
286 {
287 serde_json::from_str(self.raw_entity.get()).map_err(RealtimeError::Decode)
288 }
289
290 #[must_use]
296 pub fn decode_batch<T>(&self) -> Vec<Result<T, RealtimeError>>
297 where
298 T: DeserializeOwned,
299 {
300 if !self.entity().is_array() {
301 return vec![self.decode()];
302 }
303 let values = match serde_json::from_str::<Vec<Box<RawValue>>>(self.raw_entity.get()) {
304 Ok(values) => values,
305 Err(error) => return vec![Err(RealtimeError::Decode(error))],
306 };
307 values
308 .into_iter()
309 .filter(|value| value.get() != "null")
310 .map(|value| serde_json::from_str(value.get()).map_err(RealtimeError::Decode))
311 .collect()
312 }
313}
314
315#[derive(Clone, Debug, PartialEq)]
317#[non_exhaustive]
318pub enum RealtimeEvent {
319 Connected,
321 Disconnected,
323 Reconnected,
328 TransportGap,
333 Invocation(SignalRInvocation),
335 Message(Value),
340}
341
342mod config;
343mod event_flow;
344pub(crate) use config::RealtimeConfig;
345mod session_handle;
346use event_flow::{EventEnvelope, EventFlow, PublishOutcome};
347pub use event_flow::{RealtimeEventReceiver, RealtimeGeneration, RealtimeMessage};
348pub use session_handle::RealtimeSession;
349
350#[derive(Debug, Error)]
352#[non_exhaustive]
353pub enum RealtimeError {
354 #[error("authentication is required before connecting a real-time hub")]
356 MissingAuthToken,
357 #[error("real-time endpoint configuration is invalid")]
359 Endpoint(#[source] ClientError),
360 #[error("WebSocket transport failed")]
363 Transport,
364 #[error("WebSocket ping/pong check timed out")]
366 PingTimedOut,
367 #[error("SignalR handshake failed: {0}")]
369 Handshake(&'static str),
370 #[error("SignalR protocol error: {0}")]
372 Protocol(&'static str),
373 #[error("SignalR JSON decoding failed")]
375 Decode(#[source] serde_json::Error),
376 #[error("SignalR invocation JSON encoding failed")]
378 Encode(#[source] serde_json::Error),
379 #[error("real-time client is already connected")]
381 AlreadyConnected,
382 #[error("real-time connection setup was cancelled")]
384 ConnectionCancelled,
385 #[error("real-time WebSocket upgrade timed out")]
387 ConnectionTimedOut,
388 #[error("real-time client is not connected")]
390 NotConnected,
391 #[error("real-time socket generation is no longer current")]
393 StaleGeneration,
394 #[error("subscription is not valid for this real-time hub")]
396 WrongHub,
397 #[error("real-time writer queue is full")]
399 SendQueueFull,
400 #[error("real-time writer is closed")]
402 SendClosed,
403 #[error("real-time event queue is full; transport gap latched")]
405 EventQueueFull,
406 #[error("real-time reconnect is fenced by an unacknowledged transport gap")]
408 TransportGapPending,
409 #[error("real-time event receiver is closed")]
411 EventReceiverClosed,
412 #[error("pending SignalR invocation capacity is exhausted")]
414 PendingInvocationCapacity,
415 #[error("real-time transport identifier capacity is exhausted")]
417 IdentifierCapacity,
418 #[error("SignalR invocation `{target}` was rejected")]
420 InvocationRejected {
421 target: String,
423 },
424 #[error("SignalR invocation `{target}` timed out")]
426 InvocationTimedOut {
427 target: String,
429 },
430 #[error("SignalR invocation `{target}` ended with its session")]
432 InvocationSessionEnded {
433 target: String,
435 },
436 #[error("real-time close handshake did not complete")]
438 Close,
439}
440
441pub struct RealtimeClient {
447 inner: Arc<RealtimeInner>,
448}
449
450struct RealtimeInner {
451 hub: Hub,
452 endpoints: Endpoints,
453 http: reqwest::Client,
454 token: Arc<TokenStore>,
455 config: RealtimeConfig,
456 lifecycle: ParkingMutex<Lifecycle>,
457 lifecycle_changed: Notify,
458 generation: AtomicU64,
459 reconnect_enabled: AtomicBool,
460 owner_cancel: CancellationToken,
461 client_handles: AtomicUsize,
462 watchdog_task: ParkingMutex<Option<JoinHandle<()>>>,
463 retirement_task: ParkingMutex<Option<JoinHandle<()>>>,
464 failed_close: AtomicU64,
465 last_activity: ParkingMutex<Instant>,
466 socket_probe: ParkingMutex<Option<(u64, u64)>>,
467 request_counter: AtomicU64,
468 pending: PendingInvocations,
469 event_tx: mpsc::Sender<EventEnvelope>,
470 event_rx: ParkingMutex<Option<RealtimeEventReceiver>>,
471 event_flow: Arc<EventFlow>,
472}
473
474enum Lifecycle {
475 Disconnected,
476 Connecting {
477 generation: u64,
478 cancellation: CancellationToken,
479 },
480 Connected(Session),
481 Closing {
482 generation: u64,
483 was_connected: bool,
484 },
485}
486
487impl Lifecycle {
488 fn generation(&self) -> Option<u64> {
489 match self {
490 Self::Disconnected => None,
491 Self::Connecting { generation, .. } | Self::Closing { generation, .. } => {
492 Some(*generation)
493 }
494 Self::Connected(session) => Some(session.generation),
495 }
496 }
497}
498
499struct Session {
500 runtime: tokio::runtime::Handle,
501 generation: u64,
502 writer: mpsc::Sender<Message>,
503 cancellation: CancellationToken,
504 reader_task: JoinHandle<Result<(), RealtimeError>>,
505 writer_task: JoinHandle<Result<(), RealtimeError>>,
506}
507
508struct ConnectClaim {
509 inner: Weak<RealtimeInner>,
510 generation: u64,
511 cancellation: CancellationToken,
512 complete: bool,
513}
514
515impl ConnectClaim {
516 fn complete(&mut self) {
517 self.complete = true;
518 }
519}
520
521impl Drop for ConnectClaim {
522 fn drop(&mut self) {
523 if !self.complete
524 && let Some(inner) = self.inner.upgrade()
525 {
526 inner.cancel_connect(self.generation);
527 }
528 }
529}
530
531struct PendingGuard {
532 pending: PendingInvocations,
533 invocation_id: String,
534 generation: u64,
535}
536
537impl Drop for PendingGuard {
538 fn drop(&mut self) {
539 {
540 let mut pending = self.pending.lock();
541 if pending
542 .get(&self.invocation_id)
543 .is_some_and(|entry| entry.generation == self.generation)
544 {
545 pending.remove(&self.invocation_id);
546 }
547 }
548 }
549}
550
551enum DisconnectAction {
552 None,
553 Wait(u64),
554 Close(Session),
555}
556
557#[derive(Clone, Copy, Debug, Eq, PartialEq)]
558enum ProcessOutcome {
559 Continue,
560 Close,
561}
562
563impl Drop for Session {
564 fn drop(&mut self) {
565 self.cancellation.cancel();
566 self.reader_task.abort();
567 self.writer_task.abort();
568 }
569}
570
571impl RealtimeClient {
572 pub(crate) fn new(
573 hub: Hub,
574 endpoints: Endpoints,
575 http: reqwest::Client,
576 token: Arc<TokenStore>,
577 config: RealtimeConfig,
578 ) -> Self {
579 let (event_tx, event_rx) = mpsc::channel(config.event_capacity);
580 let event_flow = EventFlow::new();
581 Self {
582 inner: Arc::new(RealtimeInner {
583 hub,
584 endpoints,
585 http,
586 token,
587 config,
588 lifecycle: ParkingMutex::new(Lifecycle::Disconnected),
589 lifecycle_changed: Notify::new(),
590 generation: AtomicU64::new(0),
591 reconnect_enabled: AtomicBool::new(false),
592 owner_cancel: CancellationToken::new(),
593 client_handles: AtomicUsize::new(1),
594 watchdog_task: ParkingMutex::new(None),
595 retirement_task: ParkingMutex::new(None),
596 failed_close: AtomicU64::new(0),
597 last_activity: ParkingMutex::new(Instant::now()),
598 socket_probe: ParkingMutex::new(None),
599 request_counter: AtomicU64::new(1),
600 pending: Arc::default(),
601 event_tx,
602 event_rx: ParkingMutex::new(Some(RealtimeEventReceiver {
603 events: event_rx,
604 flow: Arc::clone(&event_flow),
605 gap_reported: false,
606 })),
607 event_flow,
608 }),
609 }
610 }
611
612 #[must_use]
614 pub fn take_event_receiver(&self) -> Option<RealtimeEventReceiver> {
615 self.inner.event_rx.lock().take()
616 }
617
618 pub async fn connect(&self) -> Result<(), RealtimeError> {
625 self.inner.connect_once(false).await
626 }
627
628 pub async fn disconnect(&self) -> Result<(), RealtimeError> {
635 self.inner.disconnect().await
636 }
637
638 #[must_use]
640 pub fn is_connected(&self) -> bool {
641 self.inner.is_connected()
642 }
643
644 pub fn session(&self) -> Result<RealtimeSession, RealtimeError> {
649 let lifecycle = self.inner.lifecycle.lock();
650 let Lifecycle::Connected(session) = &*lifecycle else {
651 return Err(RealtimeError::NotConnected);
652 };
653 Ok(RealtimeSession {
654 inner: Arc::downgrade(&self.inner),
655 generation: RealtimeGeneration(session.generation),
656 })
657 }
658
659 pub async fn invoke(
668 &self,
669 target: impl Into<String>,
670 arguments: Vec<Value>,
671 ) -> Result<(), RealtimeError> {
672 self.inner
673 .send_invocation(None, target.into(), arguments)
674 .await
675 }
676
677 pub async fn subscribe_contract_trades(
683 &self,
684 contract: &ContractId,
685 ) -> Result<(), RealtimeError> {
686 if self.inner.hub != Hub::Market {
687 return Err(RealtimeError::WrongHub);
688 }
689 self.session()?.subscribe_contract_trades(contract).await
690 }
691
692 pub async fn unsubscribe_contract_trades(
698 &self,
699 contract: &ContractId,
700 ) -> Result<(), RealtimeError> {
701 if self.inner.hub != Hub::Market {
702 return Err(RealtimeError::WrongHub);
703 }
704 self.session()?.unsubscribe_contract_trades(contract).await
705 }
706
707 pub async fn subscribe_contract_quotes(
713 &self,
714 contract: &ContractId,
715 ) -> Result<(), RealtimeError> {
716 if self.inner.hub != Hub::Market {
717 return Err(RealtimeError::WrongHub);
718 }
719 self.session()?.subscribe_contract_quotes(contract).await
720 }
721
722 pub async fn unsubscribe_contract_quotes(
728 &self,
729 contract: &ContractId,
730 ) -> Result<(), RealtimeError> {
731 if self.inner.hub != Hub::Market {
732 return Err(RealtimeError::WrongHub);
733 }
734 self.session()?.unsubscribe_contract_quotes(contract).await
735 }
736
737 pub async fn subscribe_contract_depth(
743 &self,
744 contract: &ContractId,
745 ) -> Result<(), RealtimeError> {
746 if self.inner.hub != Hub::Market {
747 return Err(RealtimeError::WrongHub);
748 }
749 self.session()?.subscribe_contract_depth(contract).await
750 }
751
752 pub async fn unsubscribe_contract_depth(
758 &self,
759 contract: &ContractId,
760 ) -> Result<(), RealtimeError> {
761 if self.inner.hub != Hub::Market {
762 return Err(RealtimeError::WrongHub);
763 }
764 self.session()?.unsubscribe_contract_depth(contract).await
765 }
766
767 pub async fn subscribe_accounts(&self) -> Result<(), RealtimeError> {
773 if self.inner.hub != Hub::User {
774 return Err(RealtimeError::WrongHub);
775 }
776 self.session()?.subscribe_accounts().await
777 }
778
779 pub async fn unsubscribe_accounts(&self) -> Result<(), RealtimeError> {
785 if self.inner.hub != Hub::User {
786 return Err(RealtimeError::WrongHub);
787 }
788 self.session()?.unsubscribe_accounts().await
789 }
790
791 pub async fn subscribe_orders(&self, account: AccountId) -> Result<(), RealtimeError> {
797 if self.inner.hub != Hub::User {
798 return Err(RealtimeError::WrongHub);
799 }
800 self.session()?.subscribe_orders(account).await
801 }
802
803 pub async fn unsubscribe_orders(&self, account: AccountId) -> Result<(), RealtimeError> {
809 if self.inner.hub != Hub::User {
810 return Err(RealtimeError::WrongHub);
811 }
812 self.session()?.unsubscribe_orders(account).await
813 }
814
815 pub async fn subscribe_positions(&self, account: AccountId) -> Result<(), RealtimeError> {
821 if self.inner.hub != Hub::User {
822 return Err(RealtimeError::WrongHub);
823 }
824 self.session()?.subscribe_positions(account).await
825 }
826
827 pub async fn unsubscribe_positions(&self, account: AccountId) -> Result<(), RealtimeError> {
833 if self.inner.hub != Hub::User {
834 return Err(RealtimeError::WrongHub);
835 }
836 self.session()?.unsubscribe_positions(account).await
837 }
838
839 pub async fn subscribe_trades(&self, account: AccountId) -> Result<(), RealtimeError> {
845 if self.inner.hub != Hub::User {
846 return Err(RealtimeError::WrongHub);
847 }
848 self.session()?.subscribe_trades(account).await
849 }
850
851 pub async fn unsubscribe_trades(&self, account: AccountId) -> Result<(), RealtimeError> {
857 if self.inner.hub != Hub::User {
858 return Err(RealtimeError::WrongHub);
859 }
860 self.session()?.unsubscribe_trades(account).await
861 }
862}
863
864impl Clone for RealtimeClient {
865 fn clone(&self) -> Self {
866 self.inner.client_handles.fetch_add(1, Ordering::Relaxed);
867 Self {
868 inner: Arc::clone(&self.inner),
869 }
870 }
871}
872
873impl Drop for RealtimeClient {
874 fn drop(&mut self) {
875 if self.inner.client_handles.fetch_sub(1, Ordering::AcqRel) == 1 {
876 self.inner.shutdown_now();
877 }
878 }
879}
880
881impl fmt::Debug for RealtimeClient {
882 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
883 formatter
884 .debug_struct("RealtimeClient")
885 .field("hub", &self.inner.hub)
886 .field("endpoints", &self.inner.endpoints)
887 .field("connected", &self.is_connected())
888 .finish_non_exhaustive()
889 }
890}
891
892impl RealtimeInner {
893 fn is_connected(&self) -> bool {
894 matches!(*self.lifecycle.lock(), Lifecycle::Connected(_))
895 }
896
897 fn begin_connect(self: &Arc<Self>, reconnecting: bool) -> Result<ConnectClaim, RealtimeError> {
898 self.begin_connect_with_pre_lock(reconnecting, || {})
899 }
900
901 fn begin_connect_with_pre_lock<F>(
902 self: &Arc<Self>,
903 reconnecting: bool,
904 before_lifecycle_lock: F,
905 ) -> Result<ConnectClaim, RealtimeError>
906 where
907 F: FnOnce(),
908 {
909 before_lifecycle_lock();
910 let mut lifecycle = self.lifecycle.lock();
911 if self.event_flow.has_unacknowledged_gap() {
912 return Err(RealtimeError::TransportGapPending);
913 }
914 if !matches!(*lifecycle, Lifecycle::Disconnected) {
915 return Err(RealtimeError::AlreadyConnected);
916 }
917 if !reconnecting {
918 self.reconnect_enabled.store(true, Ordering::Release);
919 } else if !self.reconnect_enabled.load(Ordering::Acquire) {
920 return Err(RealtimeError::ConnectionCancelled);
921 }
922 let previous = self
923 .generation
924 .fetch_update(Ordering::AcqRel, Ordering::Acquire, |value| {
925 value.checked_add(1)
926 })
927 .map_err(|_| RealtimeError::IdentifierCapacity)?;
928 let generation = previous
929 .checked_add(1)
930 .ok_or(RealtimeError::IdentifierCapacity)?;
931 let cancellation = self.owner_cancel.child_token();
932 *lifecycle = Lifecycle::Connecting {
933 generation,
934 cancellation: cancellation.clone(),
935 };
936 drop(lifecycle);
937 self.lifecycle_changed.notify_waiters();
938 Ok(ConnectClaim {
939 inner: Arc::downgrade(self),
940 generation,
941 cancellation,
942 complete: false,
943 })
944 }
945
946 async fn connect_once(self: &Arc<Self>, reconnecting: bool) -> Result<(), RealtimeError> {
947 let mut claim = self.begin_connect(reconnecting)?;
948 let url = self.connection_url()?;
949 let mut stream = tokio::select! {
950 () = claim.cancellation.cancelled() => {
951 return Err(RealtimeError::ConnectionCancelled);
952 }
953 result = tokio::time::timeout(
954 CONNECT_TIMEOUT,
955 upgrade_websocket(&self.http, url),
956 ) => {
957 match result {
958 Ok(result) => result?,
959 Err(_) => return Err(RealtimeError::ConnectionTimedOut),
960 }
961 }
962 };
963 let handshake_tail = tokio::select! {
964 () = claim.cancellation.cancelled() => {
965 return Err(RealtimeError::ConnectionCancelled);
966 }
967 result = tokio::time::timeout(HANDSHAKE_TIMEOUT, negotiate_handshake(&mut stream)) => {
968 match result {
969 Ok(result) => result?,
970 Err(_) => return Err(RealtimeError::Handshake("response timed out")),
971 }
972 }
973 };
974
975 let (write, read) = stream.split();
976 let (writer, writer_rx) = mpsc::channel(self.config.writer_capacity);
977 let start = CancellationToken::new();
978 let session_cancellation = claim.cancellation.child_token();
979 let writer_task = tokio::spawn(run_writer(
980 Arc::downgrade(self),
981 claim.generation,
982 write,
983 writer_rx,
984 session_cancellation.clone(),
985 start.clone(),
986 ));
987 let reader_task = tokio::spawn(run_reader(
988 Arc::downgrade(self),
989 claim.generation,
990 read,
991 handshake_tail,
992 session_cancellation.clone(),
993 start.clone(),
994 ));
995 let session = Session {
996 runtime: tokio::runtime::Handle::current(),
997 generation: claim.generation,
998 writer,
999 cancellation: session_cancellation,
1000 reader_task,
1001 writer_task,
1002 };
1003 let event = if reconnecting {
1004 RealtimeEvent::Reconnected
1005 } else {
1006 RealtimeEvent::Connected
1007 };
1008 self.install_ready(claim.generation, session, event)?;
1009 self.record_activity(claim.generation);
1010 start.cancel();
1011 self.start_watchdog();
1012 claim.complete();
1013 Ok(())
1014 }
1015
1016 fn install_ready(
1017 self: &Arc<Self>,
1018 generation: u64,
1019 session: Session,
1020 event: RealtimeEvent,
1021 ) -> Result<(), RealtimeError> {
1022 let mut lifecycle = self.lifecycle.lock();
1023 let can_install = matches!(
1024 &*lifecycle,
1025 Lifecycle::Connecting {
1026 generation: active,
1027 cancellation,
1028 } if *active == generation && !cancellation.is_cancelled()
1029 );
1030 if !can_install {
1031 return Err(RealtimeError::ConnectionCancelled);
1032 }
1033 self.event_flow.start_generation(generation)?;
1034 *lifecycle = Lifecycle::Connected(session);
1035 let published = self.publish(generation, event, 0);
1036 drop(lifecycle);
1037 self.lifecycle_changed.notify_waiters();
1038 match published {
1039 Ok(PublishOutcome::Published) => Ok(()),
1040 Ok(PublishOutcome::StaleGeneration) => {
1041 self.end_generation(generation);
1042 Err(RealtimeError::ConnectionCancelled)
1043 }
1044 Err(error) => {
1045 self.end_generation(generation);
1046 Err(error)
1047 }
1048 }
1049 }
1050
1051 fn cancel_connect(self: &Arc<Self>, generation: u64) {
1052 self.end_generation(generation);
1053 }
1054
1055 fn begin_disconnect(&self) -> DisconnectAction {
1056 let mut lifecycle = self.lifecycle.lock();
1057 self.reconnect_enabled.store(false, Ordering::Release);
1058 self.stop_watchdog();
1059 let previous = std::mem::replace(&mut *lifecycle, Lifecycle::Disconnected);
1060 let action = match previous {
1061 Lifecycle::Disconnected => DisconnectAction::None,
1062 Lifecycle::Connecting {
1063 generation,
1064 cancellation,
1065 } => {
1066 cancellation.cancel();
1067 *lifecycle = Lifecycle::Closing {
1068 generation,
1069 was_connected: false,
1070 };
1071 DisconnectAction::Wait(generation)
1072 }
1073 Lifecycle::Connected(session) => {
1074 let generation = session.generation;
1075 *lifecycle = Lifecycle::Closing {
1076 generation,
1077 was_connected: true,
1078 };
1079 DisconnectAction::Close(session)
1080 }
1081 closing @ Lifecycle::Closing { generation, .. } => {
1082 *lifecycle = closing;
1083 DisconnectAction::Wait(generation)
1084 }
1085 };
1086 drop(lifecycle);
1087 self.lifecycle_changed.notify_waiters();
1088 action
1089 }
1090
1091 async fn disconnect(self: &Arc<Self>) -> Result<(), RealtimeError> {
1092 match self.begin_disconnect() {
1093 DisconnectAction::None => Ok(()),
1094 DisconnectAction::Wait(generation) => {
1095 let waited = tokio::time::timeout(
1096 CLOSE_TIMEOUT,
1097 self.wait_until_generation_ends(generation),
1098 )
1099 .await;
1100 if waited.is_err() {
1101 Err(RealtimeError::Close)
1102 } else {
1103 Ok(())
1104 }
1105 }
1106 DisconnectAction::Close(session) => {
1107 let generation = session.generation;
1108 self.retire_session(session, true);
1109 self.wait_until_generation_ends(generation).await;
1110 if self.failed_close.load(Ordering::Acquire) == generation {
1111 Err(RealtimeError::Close)
1112 } else {
1113 Ok(())
1114 }
1115 }
1116 }
1117 }
1118
1119 async fn wait_until_generation_ends(&self, generation: u64) {
1120 loop {
1121 let changed = self.lifecycle_changed.notified();
1122 tokio::pin!(changed);
1123 let _ = changed.as_mut().enable();
1124 let is_active = match &*self.lifecycle.lock() {
1125 Lifecycle::Connecting {
1126 generation: active, ..
1127 }
1128 | Lifecycle::Closing {
1129 generation: active, ..
1130 } => *active == generation,
1131 Lifecycle::Connected(session) => session.generation == generation,
1132 Lifecycle::Disconnected => false,
1133 };
1134 if !is_active {
1135 return;
1136 }
1137 changed.await;
1138 }
1139 }
1140
1141 fn finish_closing(&self, generation: u64) {
1142 let mut lifecycle = self.lifecycle.lock();
1143 let should_publish = match &*lifecycle {
1144 Lifecycle::Closing {
1145 generation: active,
1146 was_connected,
1147 } if *active == generation => *was_connected,
1148 _ => return,
1149 };
1150 self.finish_event_generation_locked(generation, should_publish);
1151 *lifecycle = Lifecycle::Disconnected;
1152 drop(lifecycle);
1153 self.fail_pending_generation(generation);
1154 self.lifecycle_changed.notify_waiters();
1155 }
1156
1157 fn end_generation(self: &Arc<Self>, generation: u64) {
1158 let mut lifecycle = self.lifecycle.lock();
1159 let previous = std::mem::replace(&mut *lifecycle, Lifecycle::Disconnected);
1160 match previous {
1161 Lifecycle::Connected(session) if session.generation == generation => {
1162 *lifecycle = Lifecycle::Closing {
1163 generation,
1164 was_connected: true,
1165 };
1166 drop(lifecycle);
1167 self.retire_session(session, false);
1168 }
1169 Lifecycle::Connecting {
1170 generation: active,
1171 cancellation,
1172 } if active == generation => {
1173 cancellation.cancel();
1174 drop(lifecycle);
1175 self.lifecycle_changed.notify_waiters();
1176 }
1177 Lifecycle::Closing {
1178 generation: active,
1179 was_connected: false,
1180 } if active == generation => {
1181 drop(lifecycle);
1182 self.lifecycle_changed.notify_waiters();
1183 }
1184 other => {
1185 *lifecycle = other;
1186 }
1187 }
1188 }
1189
1190 fn retire_session(self: &Arc<Self>, mut session: Session, graceful: bool) {
1193 let generation = session.generation;
1194 if !graceful {
1195 session.cancellation.cancel();
1196 }
1197 self.fail_pending_generation(generation);
1198 let inner = Arc::clone(self);
1199 let mut slot = self.retirement_task.lock();
1200 let runtime = session.runtime.clone();
1203 *slot = Some(runtime.spawn(async move {
1204 if graceful && send_queued(&session.writer, Message::Close(None)).is_err() {
1207 session.cancellation.cancel();
1208 inner.failed_close.store(generation, Ordering::Release);
1209 }
1210 if !join_socket_tasks(&mut session).await {
1211 inner.failed_close.store(generation, Ordering::Release);
1212 }
1213 inner.finish_closing(generation);
1214 }));
1215 }
1216
1217 fn shutdown_now(self: &Arc<Self>) {
1218 self.owner_cancel.cancel();
1219 self.reconnect_enabled.store(false, Ordering::Release);
1220 self.stop_watchdog();
1221 let generation = self.lifecycle.lock().generation();
1222 if let Some(generation) = generation {
1223 self.end_generation(generation);
1224 }
1225 self.pending.lock().clear();
1226 }
1227
1228 fn finish_event_generation_locked(&self, generation: u64, publish_disconnected: bool) {
1231 if publish_disconnected {
1232 let _ = self.publish(generation, RealtimeEvent::Disconnected, 0);
1233 }
1234 self.event_flow.finish_generation(generation);
1235 }
1236
1237 fn start_watchdog(self: &Arc<Self>) {
1238 let mut watchdog = self.watchdog_task.lock();
1239 if !self.reconnect_enabled.load(Ordering::Acquire) || self.owner_cancel.is_cancelled() {
1240 return;
1241 }
1242 if watchdog.as_ref().is_some_and(|task| !task.is_finished()) {
1243 return;
1244 }
1245 if let Some(task) = watchdog.take() {
1246 task.abort();
1247 }
1248 let weak = Arc::downgrade(self);
1249 let cancellation = self.owner_cancel.clone();
1250 *watchdog = Some(tokio::spawn(async move {
1251 let mut ticker = tokio::time::interval(WATCHDOG_INTERVAL);
1252 ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
1253 loop {
1254 tokio::select! {
1255 biased;
1256 () = cancellation.cancelled() => break,
1257 _ = ticker.tick() => {
1258 let Some(inner) = weak.upgrade() else {
1259 break;
1260 };
1261 inner.watchdog_tick().await;
1262 }
1263 }
1264 }
1265 }));
1266 }
1267
1268 fn stop_watchdog(&self) {
1269 if let Some(task) = self.watchdog_task.lock().take() {
1270 task.abort();
1271 }
1272 }
1273
1274 async fn watchdog_tick(self: &Arc<Self>) {
1275 if !self.reconnect_enabled.load(Ordering::Acquire)
1276 || self.event_flow.has_unacknowledged_gap()
1277 || self.owner_cancel.is_cancelled()
1278 {
1279 return;
1280 }
1281 let disconnected = matches!(*self.lifecycle.lock(), Lifecycle::Disconnected);
1282 if disconnected
1283 && let Err(error) = self.connect_once(true).await
1284 && !matches!(
1285 error,
1286 RealtimeError::ConnectionCancelled | RealtimeError::AlreadyConnected
1287 )
1288 {
1289 tracing::warn!(%error, hub = ?self.hub, "ProjectX real-time reconnect failed");
1290 }
1291 }
1292
1293 fn connection_url(&self) -> Result<url::Url, RealtimeError> {
1294 let token = self
1295 .token
1296 .snapshot()
1297 .filter(|token| !token.trim().is_empty())
1298 .ok_or(RealtimeError::MissingAuthToken)?;
1299 let mut url = self
1300 .endpoints
1301 .hub_url(self.hub.path())
1302 .map_err(RealtimeError::Endpoint)?;
1303 url.query_pairs_mut().append_pair("access_token", &token);
1304 Ok(url)
1305 }
1306
1307 async fn send_invocation(
1308 self: &Arc<Self>,
1309 expected: Option<RealtimeGeneration>,
1310 target: String,
1311 arguments: Vec<Value>,
1312 ) -> Result<(), RealtimeError> {
1313 let invocation_id = self
1314 .request_counter
1315 .fetch_update(Ordering::AcqRel, Ordering::Acquire, |value| {
1316 value.checked_add(1)
1317 })
1318 .map_err(|_| RealtimeError::IdentifierCapacity)?
1319 .to_string();
1320 let message = encode_invocation(&invocation_id, &target, &arguments)?;
1321 let (reply_tx, reply_rx) = oneshot::channel();
1322 let (generation, writer) = self.register_invocation(expected, &invocation_id, reply_tx)?;
1323 let _pending = PendingGuard {
1324 pending: Arc::clone(&self.pending),
1325 invocation_id: invocation_id.clone(),
1326 generation,
1327 };
1328 send_queued(&writer, message)?;
1329
1330 match tokio::time::timeout(self.config.invocation_timeout, reply_rx).await {
1331 Ok(Ok(Ok(()))) => Ok(()),
1332 Ok(Ok(Err(()))) => Err(RealtimeError::InvocationRejected { target }),
1333 Ok(Err(_)) => Err(RealtimeError::InvocationSessionEnded { target }),
1334 Err(_) => Err(RealtimeError::InvocationTimedOut { target }),
1335 }
1336 }
1337
1338 fn register_invocation(
1339 &self,
1340 expected: Option<RealtimeGeneration>,
1341 invocation_id: &str,
1342 reply: PendingInvocation,
1343 ) -> Result<(u64, mpsc::Sender<Message>), RealtimeError> {
1344 let lifecycle = self.lifecycle.lock();
1345 let Lifecycle::Connected(session) = &*lifecycle else {
1346 return Err(RealtimeError::NotConnected);
1347 };
1348 if expected.is_some_and(|expected| expected.0 != session.generation) {
1349 return Err(RealtimeError::StaleGeneration);
1350 }
1351 let mut pending = self.pending.lock();
1352 if pending.len() >= self.config.pending_capacity {
1353 return Err(RealtimeError::PendingInvocationCapacity);
1354 }
1355 pending.insert(
1356 invocation_id.to_owned(),
1357 PendingEntry {
1358 generation: session.generation,
1359 reply,
1360 },
1361 );
1362 Ok((session.generation, session.writer.clone()))
1363 }
1364
1365 async fn process_text(
1366 &self,
1367 generation: u64,
1368 text: &str,
1369 ) -> Result<ProcessOutcome, RealtimeError> {
1370 let mut records = text.split(SIGNALR_TERMINATOR).enumerate().peekable();
1371 let quantum = self.config.event_capacity.min(32);
1375 while let Some((index, frame)) = records.next() {
1376 if index != 0 && index % quantum == 0 {
1377 tokio::task::yield_now().await;
1378 }
1379 if records.peek().is_none() {
1380 if !frame.is_empty() {
1381 self.event_flow.mark_gap(generation);
1382 }
1383 break;
1384 }
1385 match self.process_record(generation, frame) {
1386 Ok(ProcessOutcome::Continue) => {}
1387 Ok(ProcessOutcome::Close) => return Ok(ProcessOutcome::Close),
1388 Err(
1389 RealtimeError::Decode(_)
1390 | RealtimeError::Protocol(_)
1391 | RealtimeError::EventQueueFull,
1392 ) => {
1393 self.event_flow.mark_gap(generation);
1394 }
1395 Err(error) => return Err(error),
1396 }
1397 }
1398 Ok(ProcessOutcome::Continue)
1399 }
1400
1401 fn process_record(
1402 &self,
1403 generation: u64,
1404 frame: &str,
1405 ) -> Result<ProcessOutcome, RealtimeError> {
1406 let value: Value = serde_json::from_str(frame).map_err(RealtimeError::Decode)?;
1407 let message_type = value
1408 .get("type")
1409 .ok_or(RealtimeError::Protocol("message type is missing"))?
1410 .as_u64()
1411 .ok_or(RealtimeError::Protocol(
1412 "message type was not an unsigned integer",
1413 ))?;
1414 match message_type {
1415 1 => {
1416 let invocation = SignalRInvocation::from_json(frame)?.ok_or(
1417 RealtimeError::Protocol("type-1 frame was not an invocation"),
1418 )?;
1419 let weight = frame.len();
1420 if self.publish(generation, RealtimeEvent::Invocation(invocation), weight)?
1421 == PublishOutcome::StaleGeneration
1422 {
1423 return Ok(ProcessOutcome::Close);
1424 }
1425 }
1426 3 => {
1427 let (invocation_id, result) = completion(&value)?;
1428 let reply = {
1429 let mut pending = self.pending.lock();
1430 if pending
1431 .get(invocation_id)
1432 .is_some_and(|entry| entry.generation == generation)
1433 {
1434 pending.remove(invocation_id).map(|entry| entry.reply)
1435 } else {
1436 None
1437 }
1438 };
1439 if let Some(reply) = reply {
1440 let _ = reply.send(result);
1441 }
1442 }
1443 6 => {}
1444 7 => {
1445 if value.get("error").is_some_and(|error| !error.is_string()) {
1446 return Err(RealtimeError::Protocol("close error was not a string"));
1447 }
1448 let _provider_hint = match value.get("allowReconnect") {
1449 Some(flag) => flag.as_bool().ok_or(RealtimeError::Protocol(
1450 "close allowReconnect flag was not a boolean",
1451 ))?,
1452 None => false,
1453 };
1454 return Ok(ProcessOutcome::Close);
1455 }
1456 _ => {
1457 let weight = frame.len();
1458 if self.publish(generation, RealtimeEvent::Message(value), weight)?
1459 == PublishOutcome::StaleGeneration
1460 {
1461 return Ok(ProcessOutcome::Close);
1462 }
1463 }
1464 }
1465 Ok(ProcessOutcome::Continue)
1466 }
1467
1468 fn publish(
1469 &self,
1470 generation: u64,
1471 event: RealtimeEvent,
1472 weight: usize,
1473 ) -> Result<PublishOutcome, RealtimeError> {
1474 let weight = weight.max(EVENT_BASE_WEIGHT);
1475 match self
1476 .event_flow
1477 .publish(&self.event_tx, generation, event, weight)
1478 {
1479 Err(RealtimeError::EventReceiverClosed) => {
1480 self.reconnect_enabled.store(false, Ordering::Release);
1481 self.stop_watchdog();
1482 Err(RealtimeError::EventReceiverClosed)
1483 }
1484 result => result,
1485 }
1486 }
1487
1488 fn fail_pending_generation(&self, generation: u64) {
1489 self.pending
1490 .lock()
1491 .retain(|_, entry| entry.generation != generation);
1492 }
1493
1494 fn acknowledge_probe(&self, payload: &[u8]) {
1495 let Ok(payload) = <[u8; 16]>::try_from(payload) else {
1496 return;
1497 };
1498 let (left, right) = payload.split_at(8);
1499 let (Ok(left), Ok(right)) = (left.try_into(), right.try_into()) else {
1500 return;
1501 };
1502 let answer = (u64::from_be_bytes(left), u64::from_be_bytes(right));
1503 let mut probe = self.socket_probe.lock();
1504 if *probe == Some(answer) {
1505 *probe = None;
1506 }
1507 }
1508
1509 fn record_activity(&self, generation: u64) {
1510 let lifecycle = self.lifecycle.lock();
1511 if matches!(
1512 &*lifecycle,
1513 Lifecycle::Connected(session) if session.generation == generation
1514 ) {
1515 *self.last_activity.lock() = Instant::now();
1516 }
1517 }
1518}
1519
1520impl Drop for RealtimeInner {
1521 fn drop(&mut self) {
1522 self.owner_cancel.cancel();
1523 if let Some(task) = self.watchdog_task.get_mut().take() {
1524 task.abort();
1525 }
1526 match std::mem::replace(self.lifecycle.get_mut(), Lifecycle::Disconnected) {
1527 Lifecycle::Connecting { cancellation, .. } => cancellation.cancel(),
1528 Lifecycle::Connected(session) => drop(session),
1529 Lifecycle::Disconnected | Lifecycle::Closing { .. } => {}
1530 }
1531 }
1532}
1533
1534async fn upgrade_websocket(
1535 http: &reqwest::Client,
1536 mut url: url::Url,
1537) -> Result<WebSocketStream<reqwest::Upgraded>, RealtimeError> {
1538 let http_scheme = match url.scheme() {
1539 "wss" => "https",
1540 "ws" => "http",
1541 _ => return Err(RealtimeError::Transport),
1542 };
1543 url.set_scheme(http_scheme)
1544 .map_err(|()| RealtimeError::Transport)?;
1545
1546 let key = websocket_key()?;
1547 let expected_accept = websocket_accept(&key);
1548 let response = http
1549 .get(url)
1550 .version(Version::HTTP_11)
1551 .header(header::CONNECTION, "Upgrade")
1552 .header(header::UPGRADE, "websocket")
1553 .header(header::SEC_WEBSOCKET_VERSION, "13")
1554 .header(header::SEC_WEBSOCKET_KEY, key)
1555 .send()
1556 .await
1557 .map_err(|_| RealtimeError::Transport)?;
1558 validate_websocket_upgrade(
1559 response.status(),
1560 response.version(),
1561 response.headers(),
1562 &expected_accept,
1563 )?;
1564 let upgraded = response
1565 .upgrade()
1566 .await
1567 .map_err(|_| RealtimeError::Transport)?;
1568 Ok(WebSocketStream::from_raw_socket(upgraded, Role::Client, Some(websocket_config())).await)
1569}
1570
1571fn websocket_key() -> Result<String, RealtimeError> {
1572 let mut nonce = [0_u8; 16];
1573 SysRng
1574 .try_fill_bytes(&mut nonce)
1575 .map_err(|_| RealtimeError::Transport)?;
1576 Ok(BASE64.encode(&nonce))
1577}
1578
1579fn websocket_accept(key: &str) -> String {
1580 let mut digest = Sha1::new();
1581 digest.update(key.as_bytes());
1582 digest.update(WEBSOCKET_GUID);
1583 BASE64.encode(&digest.finalize())
1584}
1585
1586fn validate_websocket_upgrade(
1587 status: StatusCode,
1588 version: Version,
1589 headers: &header::HeaderMap,
1590 expected_accept: &str,
1591) -> Result<(), RealtimeError> {
1592 let mut accept_values = headers.get_all(header::SEC_WEBSOCKET_ACCEPT).iter();
1593 let accept_matches = accept_values
1594 .next()
1595 .is_some_and(|value| value.as_bytes() == expected_accept.as_bytes())
1596 && accept_values.next().is_none();
1597 if status != StatusCode::SWITCHING_PROTOCOLS
1598 || version != Version::HTTP_11
1599 || !header_contains_token(headers, &header::CONNECTION, "upgrade")
1600 || !header_contains_token(headers, &header::UPGRADE, "websocket")
1601 || !accept_matches
1602 || headers.contains_key(header::SEC_WEBSOCKET_EXTENSIONS)
1603 || headers.contains_key(header::SEC_WEBSOCKET_PROTOCOL)
1604 {
1605 return Err(RealtimeError::Transport);
1606 }
1607 Ok(())
1608}
1609
1610fn header_contains_token(
1611 headers: &header::HeaderMap,
1612 name: &header::HeaderName,
1613 expected: &str,
1614) -> bool {
1615 headers.get_all(name).iter().any(|value| {
1616 value.to_str().is_ok_and(|value| {
1617 value
1618 .split(',')
1619 .any(|token| token.trim().eq_ignore_ascii_case(expected))
1620 })
1621 })
1622}
1623
1624fn websocket_config() -> WebSocketConfig {
1625 WebSocketConfig::default()
1626 .read_buffer_size(WEBSOCKET_READ_BUFFER_SIZE)
1627 .write_buffer_size(WEBSOCKET_WRITE_BUFFER_SIZE)
1628 .max_message_size(None)
1631 .max_frame_size(None)
1632}
1633
1634fn encode_invocation(
1635 invocation_id: &str,
1636 target: &str,
1637 arguments: &[Value],
1638) -> Result<Message, RealtimeError> {
1639 let payload = OutboundInvocation {
1640 message_type: 1,
1641 invocation_id,
1642 target,
1643 arguments,
1644 };
1645 let mut text = serde_json::to_string(&payload).map_err(RealtimeError::Encode)?;
1646 text.push(SIGNALR_TERMINATOR);
1647 Ok(Message::Text(text.into()))
1648}
1649
1650fn send_queued(writer: &mpsc::Sender<Message>, message: Message) -> Result<(), RealtimeError> {
1651 writer.try_send(message).map_err(|error| match error {
1652 mpsc::error::TrySendError::Full(_) => RealtimeError::SendQueueFull,
1653 mpsc::error::TrySendError::Closed(_) => RealtimeError::SendClosed,
1654 })
1655}
1656
1657async fn join_socket_tasks(session: &mut Session) -> bool {
1658 let deadline = tokio::time::sleep(CLOSE_TIMEOUT);
1659 tokio::pin!(deadline);
1660 let mut reader_done = false;
1661 let mut writer_done = false;
1662 loop {
1663 if reader_done && writer_done {
1664 return true;
1665 }
1666 tokio::select! {
1667 _ = &mut session.reader_task, if !reader_done => reader_done = true,
1668 _ = &mut session.writer_task, if !writer_done => writer_done = true,
1669 () = &mut deadline => break,
1670 }
1671 }
1672 if !reader_done {
1673 session.reader_task.abort();
1674 let _ = (&mut session.reader_task).await;
1675 }
1676 if !writer_done {
1677 session.writer_task.abort();
1678 let _ = (&mut session.writer_task).await;
1679 }
1680 false
1681}
1682
1683async fn negotiate_handshake<S>(stream: &mut S) -> Result<Option<String>, RealtimeError>
1684where
1685 S: futures_util::Stream<Item = Result<Message, TungsteniteError>>
1686 + futures_util::Sink<Message, Error = TungsteniteError>
1687 + Unpin,
1688{
1689 let payload = format!(
1690 "{}{}",
1691 serde_json::json!({"protocol":"json","version":1}),
1692 SIGNALR_TERMINATOR
1693 );
1694 stream
1695 .send(Message::Text(payload.into()))
1696 .await
1697 .map_err(|_| RealtimeError::Transport)?;
1698 loop {
1699 match stream.next().await {
1700 Some(Ok(Message::Text(text))) => return validate_handshake(text.as_ref()),
1701 Some(Ok(Message::Binary(bytes))) => {
1702 let text = std::str::from_utf8(bytes.as_ref())
1703 .map_err(|_| RealtimeError::Handshake("response was not UTF-8"))?;
1704 return validate_handshake(text);
1705 }
1706 Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => {}
1709 Some(Ok(Message::Close(_))) | None => {
1710 return Err(RealtimeError::Handshake(
1711 "connection closed before the response",
1712 ));
1713 }
1714 Some(Err(_)) => return Err(RealtimeError::Transport),
1715 }
1716 }
1717}
1718
1719async fn run_writer<S>(
1720 inner: Weak<RealtimeInner>,
1721 generation: u64,
1722 mut write: S,
1723 mut messages: mpsc::Receiver<Message>,
1724 cancellation: CancellationToken,
1725 start: CancellationToken,
1726) -> Result<(), RealtimeError>
1727where
1728 S: futures_util::Sink<Message, Error = TungsteniteError> + Unpin,
1729{
1730 tokio::select! {
1731 biased;
1732 () = cancellation.cancelled() => {
1733 let _ = tokio::time::timeout(
1734 Duration::from_secs(1),
1735 write.send(Message::Close(None)),
1736 )
1737 .await;
1738 return Ok(());
1739 }
1740 () = start.cancelled() => {}
1741 }
1742 let keepalive = tokio::time::sleep(CLIENT_KEEPALIVE_INTERVAL);
1743 tokio::pin!(keepalive);
1744 let probe_tick = tokio::time::sleep(SOCKET_PROBE_INTERVAL);
1745 tokio::pin!(probe_tick);
1746 let mut probe_id = 0_u64;
1747 let result = loop {
1748 let message = tokio::select! {
1749 biased;
1750 () = cancellation.cancelled() => {
1751 let _ = tokio::time::timeout(
1752 Duration::from_secs(1),
1753 write.send(Message::Close(None)),
1754 )
1755 .await;
1756 break Ok(());
1757 },
1758 () = &mut probe_tick => {
1759 let Some(inner) = inner.upgrade() else { break Ok(()); };
1760 let mut probe = inner.socket_probe.lock();
1761 if probe_id != 0 && *probe == Some((generation, probe_id)) {
1762 break Err(RealtimeError::PingTimedOut);
1763 }
1764 let Some(next) = probe_id.checked_add(1) else { break Err(RealtimeError::IdentifierCapacity); };
1765 probe_id = next;
1766 *probe = Some((generation, probe_id));
1767 let mut payload = Vec::with_capacity(16);
1768 payload.extend_from_slice(&generation.to_be_bytes());
1769 payload.extend_from_slice(&probe_id.to_be_bytes());
1770 probe_tick.as_mut().reset(Instant::now() + SOCKET_PROBE_INTERVAL);
1771 Message::Ping(payload.into())
1772 }
1773 message = messages.recv() => {
1774 match message {
1775 Some(message) => message,
1776 None => break Err(RealtimeError::SendClosed),
1777 }
1778 }
1779 () = &mut keepalive => Message::Text(SIGNALR_PING.into()),
1780 };
1781 let is_close = matches!(message, Message::Close(_));
1782 let is_signalr = matches!(message, Message::Text(_) | Message::Binary(_));
1783 if write.send(message).await.is_err() {
1784 break Err(RealtimeError::Transport);
1785 }
1786 if is_close {
1787 break Ok(());
1788 }
1789 if is_signalr {
1790 keepalive
1791 .as_mut()
1792 .reset(Instant::now() + CLIENT_KEEPALIVE_INTERVAL);
1793 }
1794 };
1795 if result.is_err() && !cancellation.is_cancelled() {
1796 cancellation.cancel();
1797 if let Some(inner) = inner.upgrade() {
1798 inner.end_generation(generation);
1799 }
1800 }
1801 result
1802}
1803
1804async fn run_reader<S>(
1805 inner: Weak<RealtimeInner>,
1806 generation: u64,
1807 mut read: S,
1808 handshake_tail: Option<String>,
1809 cancellation: CancellationToken,
1810 start: CancellationToken,
1811) -> Result<(), RealtimeError>
1812where
1813 S: futures_util::Stream<Item = Result<Message, TungsteniteError>> + Unpin,
1814{
1815 tokio::select! {
1816 biased;
1817 () = cancellation.cancelled() => return Ok(()),
1818 () = start.cancelled() => {}
1819 }
1820 let result = async {
1821 if let Some(tail) = handshake_tail {
1822 let Some(inner) = inner.upgrade() else {
1823 return Ok(());
1824 };
1825 if inner.process_text(generation, &tail).await? == ProcessOutcome::Close {
1826 return Ok(());
1827 }
1828 }
1829 loop {
1830 let message = tokio::select! {
1831 biased;
1832 () = cancellation.cancelled() => return Ok(()),
1833 message = read.next() => message,
1834 };
1835 match message {
1836 Some(Ok(Message::Text(text))) => {
1837 let Some(inner) = inner.upgrade() else {
1838 return Ok(());
1839 };
1840 let outcome = inner.process_text(generation, text.as_ref()).await?;
1841 inner.record_activity(generation);
1842 if outcome == ProcessOutcome::Close {
1843 return Ok(());
1844 }
1845 }
1846 Some(Ok(Message::Binary(bytes))) => {
1847 let Some(inner) = inner.upgrade() else {
1848 return Ok(());
1849 };
1850 let Ok(text) = std::str::from_utf8(bytes.as_ref()) else {
1851 inner.event_flow.mark_gap(generation);
1852 continue;
1853 };
1854 let outcome = inner.process_text(generation, text).await?;
1855 inner.record_activity(generation);
1856 if outcome == ProcessOutcome::Close {
1857 return Ok(());
1858 }
1859 }
1860 Some(Ok(Message::Ping(_))) => {
1861 let Some(inner) = inner.upgrade() else {
1862 return Ok(());
1863 };
1864 inner.record_activity(generation);
1865 }
1867 Some(Ok(Message::Close(_))) | None => return Ok(()),
1868 Some(Err(_)) => return Err(RealtimeError::Transport),
1869 Some(Ok(Message::Pong(payload))) => {
1870 let Some(inner) = inner.upgrade() else {
1871 return Ok(());
1872 };
1873 inner.acknowledge_probe(&payload);
1874 inner.record_activity(generation);
1875 }
1876 Some(Ok(Message::Frame(_))) => {}
1877 }
1878 }
1879 }
1880 .await;
1881 if !cancellation.is_cancelled() {
1882 cancellation.cancel();
1883 if let Some(inner) = inner.upgrade() {
1884 inner.end_generation(generation);
1885 }
1886 }
1887 result
1888}
1889
1890fn validate_handshake(text: &str) -> Result<Option<String>, RealtimeError> {
1891 let (payload, tail) = text
1892 .split_once(SIGNALR_TERMINATOR)
1893 .ok_or(RealtimeError::Handshake("response frame was incomplete"))?;
1894 let response: Value = serde_json::from_str(payload)
1895 .map_err(|_| RealtimeError::Handshake("response was invalid JSON"))?;
1896 if response
1897 .get("error")
1898 .and_then(Value::as_str)
1899 .is_some_and(|detail| !detail.trim().is_empty())
1900 {
1901 return Err(RealtimeError::Handshake("provider rejected the handshake"));
1902 }
1903 if response.as_object().is_some_and(serde_json::Map::is_empty) {
1904 Ok((!tail.is_empty()).then(|| tail.to_owned()))
1905 } else {
1906 Err(RealtimeError::Handshake(
1907 "provider returned an unexpected response shape",
1908 ))
1909 }
1910}
1911
1912fn completion(value: &Value) -> Result<(&str, Result<(), ()>), RealtimeError> {
1913 let object = value
1914 .as_object()
1915 .ok_or(RealtimeError::Protocol("completion was not an object"))?;
1916 let invocation_id =
1917 object
1918 .get("invocationId")
1919 .and_then(Value::as_str)
1920 .ok_or(RealtimeError::Protocol(
1921 "completion invocationId was not a string",
1922 ))?;
1923 if object.contains_key("error") && object.contains_key("result") {
1924 return Err(RealtimeError::Protocol(
1925 "completion contained both error and result",
1926 ));
1927 }
1928 let result = match object.get("error") {
1929 Some(Value::String(_)) => Err(()),
1930 Some(_) => {
1931 return Err(RealtimeError::Protocol("completion error was not a string"));
1932 }
1933 None => Ok(()),
1934 };
1935 Ok((invocation_id, result))
1936}
1937
1938#[cfg(test)]
1939mod tests {
1940 use serde_json::json;
1941
1942 use super::*;
1943
1944 fn fixture_realtime() -> RealtimeClient {
1945 let credentials = crate::Credentials::new("user", "key")
1946 .unwrap_or_else(|error| panic!("fixture credentials must be valid: {error}"));
1947 let client = crate::Client::builder(credentials)
1948 .build()
1949 .unwrap_or_else(|error| panic!("fixture client must build: {error}"));
1950 client.realtime(Hub::Market)
1951 }
1952
1953 fn install_connected_generation(inner: &Arc<RealtimeInner>, generation: u64) {
1954 let (writer, _writer_rx) = mpsc::channel(1);
1955 let session = Session {
1956 runtime: tokio::runtime::Handle::current(),
1957 generation,
1958 writer,
1959 cancellation: CancellationToken::new(),
1960 reader_task: tokio::spawn(std::future::pending::<Result<(), RealtimeError>>()),
1961 writer_task: tokio::spawn(std::future::pending::<Result<(), RealtimeError>>()),
1962 };
1963 let mut lifecycle = inner.lifecycle.lock();
1964 inner
1965 .event_flow
1966 .start_generation(generation)
1967 .unwrap_or_else(|error| panic!("fixture generation must start: {error}"));
1968 inner.generation.store(generation, Ordering::Release);
1969 *lifecycle = Lifecycle::Connected(session);
1970 }
1971
1972 fn install_paused_connected_generation(
1973 inner: &Arc<RealtimeInner>,
1974 generation: u64,
1975 ) -> mpsc::Receiver<Message> {
1976 let (writer, writer_rx) = mpsc::channel(1);
1977 let session = Session {
1978 runtime: tokio::runtime::Handle::current(),
1979 generation,
1980 writer,
1981 cancellation: CancellationToken::new(),
1982 reader_task: tokio::spawn(std::future::pending::<Result<(), RealtimeError>>()),
1983 writer_task: tokio::spawn(std::future::pending::<Result<(), RealtimeError>>()),
1984 };
1985 let mut lifecycle = inner.lifecycle.lock();
1986 inner
1987 .event_flow
1988 .start_generation(generation)
1989 .unwrap_or_else(|error| panic!("fixture generation must start: {error}"));
1990 inner.generation.store(generation, Ordering::Release);
1991 *lifecycle = Lifecycle::Connected(session);
1992 writer_rx
1993 }
1994
1995 #[test]
1996 fn invocation_extracts_market_contract_and_payload() {
1997 let value = json!({
1998 "type": 1,
1999 "target": "GatewayTrade",
2000 "arguments": ["CON.F.US.MNQ.M26", {"price": 1.25}],
2001 });
2002 let invocation = SignalRInvocation::from_value(value)
2003 .and_then(|value| value.ok_or(RealtimeError::Protocol("missing invocation")))
2004 .unwrap_or_else(|error| panic!("fixture invocation must decode: {error}"));
2005 assert_eq!(invocation.target(), "GatewayTrade");
2006 assert_eq!(
2007 invocation.contract_id().map(ContractId::as_str),
2008 Some("CON.F.US.MNQ.M26")
2009 );
2010 assert_eq!(invocation.payload(), &json!({"price": 1.25}));
2011 }
2012
2013 #[test]
2014 fn invocation_rejects_malformed_market_arguments() {
2015 for value in [
2016 json!({
2017 "type": 1,
2018 "target": "GatewayTrade",
2019 "arguments": [42, {"price": 1.25}],
2020 }),
2021 json!({
2022 "type": 1,
2023 "target": "GatewayTrade",
2024 "arguments": ["CON.F.US.MNQ.M26", {"price": 1.25}, "extra"],
2025 }),
2026 ] {
2027 assert!(matches!(
2028 SignalRInvocation::from_value(value),
2029 Err(RealtimeError::Protocol(_))
2030 ));
2031 }
2032 }
2033
2034 #[test]
2035 fn handshake_accepts_coalesced_tail() {
2036 let tail = validate_handshake("{}\u{001e}{\"type\":6}\u{001e}")
2037 .unwrap_or_else(|error| panic!("fixture handshake must decode: {error}"));
2038 assert_eq!(tail.as_deref(), Some("{\"type\":6}\u{001e}"));
2039 }
2040
2041 #[test]
2042 fn handshake_rejects_provider_error_without_retaining_detail() {
2043 let result = validate_handshake("{\"error\":\"secret provider detail\"}\u{001e}");
2044 assert!(matches!(result, Err(RealtimeError::Handshake(_))));
2045 assert!(!format!("{:?}", result.err()).contains("secret provider detail"));
2046 }
2047
2048 #[test]
2049 fn completion_requires_an_unambiguous_protocol_shape() {
2050 for value in [
2051 json!({"type": 3}),
2052 json!({"type": 3, "invocationId": 1}),
2053 json!({"type": 3, "invocationId": "1", "error": false}),
2054 json!({"type": 3, "invocationId": "1", "error": "rejected", "result": null}),
2055 ] {
2056 assert!(matches!(
2057 completion(&value),
2058 Err(RealtimeError::Protocol(_))
2059 ));
2060 }
2061
2062 assert_eq!(
2063 completion(&json!({"type": 3, "invocationId": "1"}))
2064 .unwrap_or_else(|error| panic!("void completion must decode: {error}")),
2065 ("1", Ok(()))
2066 );
2067 assert_eq!(
2068 completion(&json!({"type": 3, "invocationId": "1", "result": 42}))
2069 .unwrap_or_else(|error| panic!("result completion must decode: {error}")),
2070 ("1", Ok(()))
2071 );
2072 assert_eq!(
2073 completion(&json!({"type": 3, "invocationId": "1", "error": "rejected"}))
2074 .unwrap_or_else(|error| panic!("error completion must decode: {error}")),
2075 ("1", Err(()))
2076 );
2077 }
2078
2079 #[test]
2080 fn malformed_transport_frames_are_not_published_as_messages() {
2081 let credentials = crate::Credentials::new("user", "key")
2082 .unwrap_or_else(|error| panic!("fixture credentials must be valid: {error}"));
2083 let client = crate::Client::builder(credentials)
2084 .build()
2085 .unwrap_or_else(|error| panic!("fixture client must build: {error}"));
2086 let realtime = client.realtime(Hub::Market);
2087 let mut events = realtime
2088 .take_event_receiver()
2089 .unwrap_or_else(|| panic!("event receiver must be available"));
2090
2091 for frame in [
2092 "{}\u{001e}",
2093 "{\"type\":\"1\"}\u{001e}",
2094 "{\"type\":3,\"invocationId\":1}\u{001e}",
2095 "{\"type\":3,\"invocationId\":\"1\",\"error\":null}\u{001e}",
2096 "{\"type\":3,\"invocationId\":\"1\",\"error\":\"x\",\"result\":null}\u{001e}",
2097 "{\"type\":7,\"error\":null}\u{001e}",
2098 ] {
2099 assert!(matches!(
2100 realtime
2101 .inner
2102 .process_record(1, frame.trim_end_matches(SIGNALR_TERMINATOR)),
2103 Err(RealtimeError::Protocol(_))
2104 ));
2105 }
2106 assert!(matches!(
2107 events.events.try_recv(),
2108 Err(mpsc::error::TryRecvError::Empty)
2109 ));
2110 }
2111
2112 #[tokio::test]
2113 async fn malformed_record_framing_is_rejected_before_publication() {
2114 let realtime = fixture_realtime();
2115 let mut events = realtime
2116 .take_event_receiver()
2117 .unwrap_or_else(|| panic!("event receiver must be available"));
2118 install_connected_generation(&realtime.inner, 1);
2119
2120 for batch in [
2121 "{\"type\":1}",
2122 "{\"type\":1}\u{001e}\u{001e}{\"type\":1}\u{001e}",
2123 "\u{001e}{\"type\":1}\u{001e}",
2124 "\u{001e}",
2125 ] {
2126 assert!(matches!(
2127 realtime.inner.process_text(1, batch).await,
2128 Ok(ProcessOutcome::Continue)
2129 ));
2130 }
2131 assert!(matches!(
2132 events.events.try_recv(),
2133 Err(mpsc::error::TryRecvError::Empty)
2134 ));
2135 }
2136
2137 #[test]
2138 fn connect_claim_rechecks_the_gap_after_waiting_for_lifecycle() {
2139 let realtime = fixture_realtime();
2140 let before_lock = Arc::new(std::sync::Barrier::new(2));
2141 let release_lock = Arc::new(std::sync::Barrier::new(2));
2142 let worker_inner = Arc::clone(&realtime.inner);
2143 let worker_before = Arc::clone(&before_lock);
2144 let worker_release = Arc::clone(&release_lock);
2145 let claim = std::thread::spawn(move || {
2146 worker_inner.begin_connect_with_pre_lock(false, || {
2147 worker_before.wait();
2148 worker_release.wait();
2149 })
2150 });
2151
2152 before_lock.wait();
2153 realtime
2154 .inner
2155 .event_flow
2156 .start_generation(1)
2157 .unwrap_or_else(|error| panic!("fixture generation must start: {error}"));
2158 realtime.inner.event_flow.mark_gap(1);
2159 realtime.inner.event_flow.finish_generation(1);
2160 release_lock.wait();
2161
2162 let result = claim
2163 .join()
2164 .unwrap_or_else(|payload| std::panic::resume_unwind(payload));
2165 assert!(matches!(result, Err(RealtimeError::TransportGapPending)));
2166 assert!(matches!(
2167 *realtime.inner.lifecycle.lock(),
2168 Lifecycle::Disconnected
2169 ));
2170 }
2171
2172 #[tokio::test(start_paused = true)]
2173 async fn disconnect_waiter_timeout_keeps_the_owner_session_fenced() {
2174 let realtime = fixture_realtime();
2175 let _writer_rx = install_paused_connected_generation(&realtime.inner, 1);
2176 let DisconnectAction::Close(owner_session) = realtime.inner.begin_disconnect() else {
2177 panic!("fixture owner must claim the connected session");
2178 };
2179 let waiter_client = realtime.clone();
2180 let waiter = tokio::spawn(async move { waiter_client.disconnect().await });
2181 tokio::task::yield_now().await;
2182
2183 tokio::time::advance(CLOSE_TIMEOUT).await;
2184 let result = waiter
2185 .await
2186 .unwrap_or_else(|error| panic!("waiter must join: {error}"));
2187 assert!(matches!(result, Err(RealtimeError::Close)));
2188 assert!(matches!(
2189 *realtime.inner.lifecycle.lock(),
2190 Lifecycle::Closing {
2191 generation: 1,
2192 was_connected: true,
2193 }
2194 ));
2195 assert!(matches!(
2196 realtime.connect().await,
2197 Err(RealtimeError::AlreadyConnected)
2198 ));
2199
2200 drop(owner_session);
2201 realtime.inner.finish_closing(1);
2202 assert!(matches!(
2203 *realtime.inner.lifecycle.lock(),
2204 Lifecycle::Disconnected
2205 ));
2206 }
2207
2208 #[tokio::test(start_paused = true)]
2209 async fn disconnect_waiter_timeout_keeps_the_connect_claim_fenced() {
2210 let realtime = fixture_realtime();
2211 let cancellation = realtime.inner.owner_cancel.child_token();
2212 *realtime.inner.lifecycle.lock() = Lifecycle::Connecting {
2213 generation: 1,
2214 cancellation: cancellation.clone(),
2215 };
2216 realtime.inner.generation.store(1, Ordering::Release);
2217 let owner_claim = ConnectClaim {
2218 inner: Arc::downgrade(&realtime.inner),
2219 generation: 1,
2220 cancellation,
2221 complete: false,
2222 };
2223 assert!(matches!(
2224 realtime.inner.begin_disconnect(),
2225 DisconnectAction::Wait(1)
2226 ));
2227 let waiter_client = realtime.clone();
2228 let waiter = tokio::spawn(async move { waiter_client.disconnect().await });
2229 tokio::task::yield_now().await;
2230
2231 tokio::time::advance(CLOSE_TIMEOUT).await;
2232 let result = waiter
2233 .await
2234 .unwrap_or_else(|error| panic!("waiter must join: {error}"));
2235 assert!(matches!(result, Err(RealtimeError::Close)));
2236 assert!(matches!(
2237 *realtime.inner.lifecycle.lock(),
2238 Lifecycle::Closing {
2239 generation: 1,
2240 was_connected: false,
2241 }
2242 ));
2243 assert!(matches!(
2244 realtime.connect().await,
2245 Err(RealtimeError::AlreadyConnected)
2246 ));
2247
2248 drop(owner_claim);
2249 assert!(matches!(
2250 *realtime.inner.lifecycle.lock(),
2251 Lifecycle::Disconnected
2252 ));
2253 }
2254
2255 #[tokio::test]
2256 async fn ended_boundary_waits_for_both_socket_producers_to_drop() {
2257 struct Probe(Arc<AtomicUsize>);
2258 impl Drop for Probe {
2259 fn drop(&mut self) {
2260 self.0.fetch_add(1, Ordering::SeqCst);
2261 }
2262 }
2263 let realtime = fixture_realtime();
2264 let stopped = Arc::new(AtomicUsize::new(0));
2265 let (release, wait) = oneshot::channel::<()>();
2266 let (writer, _writes) = mpsc::channel(1);
2267 let cancel = CancellationToken::new();
2268 let reader_probe = Probe(Arc::clone(&stopped));
2269 let writer_probe = Probe(Arc::clone(&stopped));
2270 let writer_cancel = cancel.clone();
2271 let session = Session {
2272 runtime: tokio::runtime::Handle::current(),
2273 generation: 1,
2274 writer,
2275 cancellation: cancel,
2276 reader_task: tokio::spawn(async move {
2277 let _probe = reader_probe;
2278 let _ = wait.await;
2279 Ok(())
2280 }),
2281 writer_task: tokio::spawn(async move {
2282 let _probe = writer_probe;
2283 writer_cancel.cancelled().await;
2284 Ok(())
2285 }),
2286 };
2287 realtime
2288 .inner
2289 .event_flow
2290 .start_generation(1)
2291 .unwrap_or_else(|e| panic!("start: {e}"));
2292 *realtime.inner.lifecycle.lock() = Lifecycle::Connected(session);
2293 let mut events = realtime
2294 .take_event_receiver()
2295 .unwrap_or_else(|| panic!("receiver"));
2296 realtime.inner.end_generation(1);
2297 assert!(
2298 tokio::time::timeout(Duration::from_millis(10), events.recv_message())
2299 .await
2300 .is_err()
2301 );
2302 assert!(stopped.load(Ordering::SeqCst) < 2);
2303 assert!(release.send(()).is_ok());
2304 let ended = events
2305 .recv_message()
2306 .await
2307 .unwrap_or_else(|| panic!("ended boundary"));
2308 assert_eq!(ended.generation, RealtimeGeneration(1));
2309 assert_eq!(ended.event, RealtimeEvent::Disconnected);
2310 assert_eq!(stopped.load(Ordering::SeqCst), 2);
2311 }
2312
2313 #[tokio::test]
2314 async fn gap_follows_accepted_events_without_waiting_for_generation_end() {
2315 let flow = EventFlow::new();
2316 let (events_tx, events_rx) = mpsc::channel(2);
2317 let mut receiver = RealtimeEventReceiver {
2318 events: events_rx,
2319 flow: Arc::clone(&flow),
2320 gap_reported: false,
2321 };
2322 let generation = 7;
2323 flow.start_generation(generation)
2324 .unwrap_or_else(|error| panic!("generation must start: {error}"));
2325 flow.publish(
2326 &events_tx,
2327 generation,
2328 RealtimeEvent::Message(json!({"sequence": 1})),
2329 EVENT_BASE_WEIGHT,
2330 )
2331 .unwrap_or_else(|error| panic!("first event must enter the queue: {error}"));
2332
2333 let (terminal_staged_tx, terminal_staged_rx) = oneshot::channel();
2334 let (finish_tx, finish_rx) = oneshot::channel();
2335 let producer_flow = Arc::clone(&flow);
2336 let producer = tokio::spawn(async move {
2337 producer_flow.mark_gap(generation);
2338 producer_flow
2339 .publish(
2340 &events_tx,
2341 generation,
2342 RealtimeEvent::Disconnected,
2343 EVENT_BASE_WEIGHT,
2344 )
2345 .unwrap_or_else(|error| panic!("terminal event must be staged: {error}"));
2346 terminal_staged_tx
2347 .send(())
2348 .unwrap_or_else(|()| panic!("terminal-staged signal must send"));
2349 finish_rx
2350 .await
2351 .unwrap_or_else(|error| panic!("finish signal must arrive: {error}"));
2352 producer_flow.finish_generation(generation);
2353 });
2354
2355 assert_eq!(
2356 receiver.recv().await,
2357 Some(RealtimeEvent::Message(json!({"sequence": 1})))
2358 );
2359 terminal_staged_rx
2360 .await
2361 .unwrap_or_else(|error| panic!("terminal-staged signal must arrive: {error}"));
2362 assert_eq!(receiver.recv().await, Some(RealtimeEvent::TransportGap));
2363 receiver.acknowledge_transport_gap();
2364 finish_tx
2365 .send(())
2366 .unwrap_or_else(|()| panic!("finish signal must send"));
2367 producer
2368 .await
2369 .unwrap_or_else(|error| panic!("producer must join: {error}"));
2370
2371 assert_eq!(receiver.recv().await, Some(RealtimeEvent::Disconnected));
2372 flow.start_generation(generation + 1)
2373 .unwrap_or_else(|error| {
2374 panic!("acknowledgement must allow the next generation: {error}")
2375 });
2376 assert_eq!(receiver.recv().await, None);
2377 assert_eq!(flow.queued_weight.load(Ordering::Acquire), 0);
2378 }
2379
2380 #[tokio::test]
2381 async fn stale_generation_cannot_publish_or_latch_a_gap_after_replacement() {
2382 let flow = EventFlow::new();
2383 let (events_tx, events_rx) = mpsc::channel(8);
2384 let mut receiver = RealtimeEventReceiver {
2385 events: events_rx,
2386 flow: Arc::clone(&flow),
2387 gap_reported: false,
2388 };
2389
2390 flow.start_generation(1)
2391 .unwrap_or_else(|error| panic!("first generation must start: {error}"));
2392 assert!(matches!(
2393 flow.start_generation(2),
2394 Err(RealtimeError::ConnectionCancelled)
2395 ));
2396 assert_eq!(
2397 flow.publish(
2398 &events_tx,
2399 1,
2400 RealtimeEvent::Disconnected,
2401 EVENT_BASE_WEIGHT,
2402 )
2403 .unwrap_or_else(|error| panic!("disconnect must publish: {error}")),
2404 PublishOutcome::Published
2405 );
2406 flow.finish_generation(1);
2407 flow.start_generation(2)
2408 .unwrap_or_else(|error| panic!("replacement generation must start: {error}"));
2409 assert_eq!(
2410 flow.publish(&events_tx, 2, RealtimeEvent::Reconnected, EVENT_BASE_WEIGHT,)
2411 .unwrap_or_else(|error| panic!("reconnect must publish: {error}")),
2412 PublishOutcome::Published
2413 );
2414
2415 for weight in [EVENT_BASE_WEIGHT, usize::MAX] {
2416 assert_eq!(
2417 flow.publish(
2418 &events_tx,
2419 1,
2420 RealtimeEvent::Message(json!({"generation": 1})),
2421 weight,
2422 )
2423 .unwrap_or_else(|error| panic!("stale publication must be ignored: {error}")),
2424 PublishOutcome::StaleGeneration
2425 );
2426 }
2427 assert!(!flow.has_unacknowledged_gap());
2428 assert_eq!(
2429 flow.publish(
2430 &events_tx,
2431 2,
2432 RealtimeEvent::Message(json!({"generation": 2})),
2433 EVENT_BASE_WEIGHT,
2434 )
2435 .unwrap_or_else(|error| panic!("replacement message must publish: {error}")),
2436 PublishOutcome::Published
2437 );
2438 flow.finish_generation(2);
2439 drop(events_tx);
2440
2441 assert_eq!(receiver.recv().await, Some(RealtimeEvent::Disconnected));
2442 assert_eq!(receiver.recv().await, Some(RealtimeEvent::Reconnected));
2443 assert_eq!(
2444 receiver.recv().await,
2445 Some(RealtimeEvent::Message(json!({"generation": 2})))
2446 );
2447 assert_eq!(receiver.recv().await, None);
2448 }
2449
2450 #[tokio::test]
2451 async fn stale_close_frame_cannot_disable_the_replacement_watchdog() {
2452 let realtime = fixture_realtime();
2453 install_connected_generation(&realtime.inner, 2);
2454 realtime
2455 .inner
2456 .reconnect_enabled
2457 .store(true, Ordering::Release);
2458 realtime.inner.start_watchdog();
2459 assert!(
2460 realtime
2461 .inner
2462 .watchdog_task
2463 .lock()
2464 .as_ref()
2465 .is_some_and(|task| !task.is_finished())
2466 );
2467
2468 assert_eq!(
2469 realtime
2470 .inner
2471 .process_text(1, "{\"type\":7,\"allowReconnect\":false}\u{001e}")
2472 .await
2473 .unwrap_or_else(|error| panic!("close frame must decode: {error}")),
2474 ProcessOutcome::Close
2475 );
2476 assert!(realtime.inner.reconnect_enabled.load(Ordering::Acquire));
2477 assert!(
2478 realtime
2479 .inner
2480 .watchdog_task
2481 .lock()
2482 .as_ref()
2483 .is_some_and(|task| !task.is_finished())
2484 );
2485
2486 assert_eq!(
2487 realtime
2488 .inner
2489 .process_text(2, "{\"type\":7,\"allowReconnect\":false}\u{001e}")
2490 .await
2491 .unwrap_or_else(|error| panic!("close frame must decode: {error}")),
2492 ProcessOutcome::Close
2493 );
2494 assert!(realtime.inner.reconnect_enabled.load(Ordering::Acquire));
2495 assert!(realtime.inner.watchdog_task.lock().is_some());
2496 }
2497
2498 #[tokio::test]
2499 async fn stale_generation_cannot_refresh_replacement_liveness() {
2500 let realtime = fixture_realtime();
2501 install_connected_generation(&realtime.inner, 2);
2502 let original = Instant::now() - Duration::from_secs(10);
2503 *realtime.inner.last_activity.lock() = original;
2504
2505 realtime.inner.record_activity(1);
2506 assert_eq!(*realtime.inner.last_activity.lock(), original);
2507
2508 realtime.inner.record_activity(2);
2509 assert!(*realtime.inner.last_activity.lock() > original);
2510 }
2511
2512 #[tokio::test]
2513 async fn typed_subscriptions_reject_the_wrong_hub_before_network_io() {
2514 let credentials = crate::Credentials::new("user", "key")
2515 .unwrap_or_else(|error| panic!("fixture credentials must be valid: {error}"));
2516 let client = crate::Client::builder(credentials)
2517 .build()
2518 .unwrap_or_else(|error| panic!("fixture client must build: {error}"));
2519 let market = client.realtime(Hub::Market);
2520 let user = client.realtime(Hub::User);
2521 let contract = ContractId::new("CON.F.US.MNQ.M26")
2522 .unwrap_or_else(|error| panic!("fixture contract must be valid: {error}"));
2523
2524 assert!(matches!(
2525 market.subscribe_accounts().await,
2526 Err(RealtimeError::WrongHub)
2527 ));
2528 assert!(matches!(
2529 user.subscribe_contract_trades(&contract).await,
2530 Err(RealtimeError::WrongHub)
2531 ));
2532 }
2533
2534 #[tokio::test]
2535 async fn transport_gap_fences_connect_until_acknowledged() {
2536 let credentials = crate::Credentials::new("user", "key")
2537 .unwrap_or_else(|error| panic!("fixture credentials must be valid: {error}"));
2538 let client = crate::Client::builder(credentials)
2539 .build()
2540 .unwrap_or_else(|error| panic!("fixture client must build: {error}"));
2541 let realtime = client.realtime(Hub::Market);
2542 let mut events = realtime
2543 .take_event_receiver()
2544 .unwrap_or_else(|| panic!("event receiver must be available"));
2545 let generation = 1;
2546 realtime
2547 .inner
2548 .event_flow
2549 .start_generation(generation)
2550 .unwrap_or_else(|error| panic!("generation must start: {error}"));
2551 realtime.inner.event_flow.mark_gap(generation);
2552 realtime
2553 .inner
2554 .event_flow
2555 .publish(
2556 &realtime.inner.event_tx,
2557 generation,
2558 RealtimeEvent::Disconnected,
2559 EVENT_BASE_WEIGHT,
2560 )
2561 .unwrap_or_else(|error| panic!("terminal event must be staged: {error}"));
2562 realtime.inner.event_flow.finish_generation(generation);
2563
2564 assert!(matches!(
2565 realtime.connect().await,
2566 Err(RealtimeError::TransportGapPending)
2567 ));
2568 assert!(matches!(
2569 events.recv().await,
2570 Some(RealtimeEvent::TransportGap)
2571 ));
2572 assert!(matches!(
2573 events.recv().await,
2574 Some(RealtimeEvent::Disconnected)
2575 ));
2576 assert!(matches!(
2577 realtime.connect().await,
2578 Err(RealtimeError::TransportGapPending)
2579 ));
2580
2581 events.acknowledge_transport_gap();
2582 assert!(matches!(
2583 realtime.connect().await,
2584 Err(RealtimeError::MissingAuthToken)
2585 ));
2586 }
2587
2588 #[tokio::test]
2589 async fn disabled_reconnect_does_not_spawn_a_watchdog() {
2590 let credentials = crate::Credentials::new("user", "key")
2591 .unwrap_or_else(|error| panic!("fixture credentials must be valid: {error}"));
2592 let client = crate::Client::builder(credentials)
2593 .build()
2594 .unwrap_or_else(|error| panic!("fixture client must build: {error}"));
2595 let realtime = client.realtime(Hub::Market);
2596
2597 realtime.inner.start_watchdog();
2598
2599 assert!(realtime.inner.watchdog_task.lock().is_none());
2600 }
2601
2602 #[test]
2603 fn websocket_configuration_accepts_large_provider_batches() {
2604 let config = websocket_config();
2605 assert_eq!(config.read_buffer_size, WEBSOCKET_READ_BUFFER_SIZE);
2606 assert_eq!(config.write_buffer_size, WEBSOCKET_WRITE_BUFFER_SIZE);
2607 assert_eq!(config.max_message_size, None);
2608 assert_eq!(config.max_frame_size, None);
2609 assert!(config.max_write_buffer_size > config.write_buffer_size);
2610 }
2611
2612 #[test]
2613 fn websocket_accept_matches_the_rfc_6455_vector() {
2614 assert_eq!(
2615 websocket_accept("dGhlIHNhbXBsZSBub25jZQ=="),
2616 "s3pPLMBiTxaQ9kYGzzhZRbK+xOo="
2617 );
2618 let key = websocket_key()
2619 .unwrap_or_else(|error| panic!("fixture key generation must succeed: {error}"));
2620 let decoded = BASE64
2621 .decode(key.as_bytes())
2622 .unwrap_or_else(|error| panic!("generated key must be base64: {error}"));
2623 assert_eq!(decoded.len(), 16);
2624 }
2625
2626 #[test]
2627 fn websocket_upgrade_validation_rejects_untrusted_responses() {
2628 let expected_accept = websocket_accept("dGhlIHNhbXBsZSBub25jZQ==");
2629 let mut valid = header::HeaderMap::new();
2630 valid.insert(
2631 header::CONNECTION,
2632 header::HeaderValue::from_static("keep-alive, Upgrade"),
2633 );
2634 valid.insert(
2635 header::UPGRADE,
2636 header::HeaderValue::from_static("WebSocket"),
2637 );
2638 valid.insert(
2639 header::SEC_WEBSOCKET_ACCEPT,
2640 header::HeaderValue::from_str(&expected_accept)
2641 .unwrap_or_else(|error| panic!("fixture accept header must be valid: {error}")),
2642 );
2643
2644 assert!(
2645 validate_websocket_upgrade(
2646 StatusCode::SWITCHING_PROTOCOLS,
2647 Version::HTTP_11,
2648 &valid,
2649 &expected_accept,
2650 )
2651 .is_ok()
2652 );
2653
2654 let mut cases = Vec::new();
2655 cases.push((StatusCode::OK, Version::HTTP_11, valid.clone()));
2656 cases.push((
2657 StatusCode::SWITCHING_PROTOCOLS,
2658 Version::HTTP_2,
2659 valid.clone(),
2660 ));
2661 for missing in [
2662 header::CONNECTION,
2663 header::UPGRADE,
2664 header::SEC_WEBSOCKET_ACCEPT,
2665 ] {
2666 let mut headers = valid.clone();
2667 headers.remove(missing);
2668 cases.push((StatusCode::SWITCHING_PROTOCOLS, Version::HTTP_11, headers));
2669 }
2670 for unsolicited in [
2671 header::SEC_WEBSOCKET_EXTENSIONS,
2672 header::SEC_WEBSOCKET_PROTOCOL,
2673 ] {
2674 let mut headers = valid.clone();
2675 headers.insert(unsolicited, header::HeaderValue::from_static("unsupported"));
2676 cases.push((StatusCode::SWITCHING_PROTOCOLS, Version::HTTP_11, headers));
2677 }
2678 let mut duplicate_accept = valid.clone();
2679 duplicate_accept.append(
2680 header::SEC_WEBSOCKET_ACCEPT,
2681 header::HeaderValue::from_str(&expected_accept)
2682 .unwrap_or_else(|error| panic!("fixture accept header must be valid: {error}")),
2683 );
2684 cases.push((
2685 StatusCode::SWITCHING_PROTOCOLS,
2686 Version::HTTP_11,
2687 duplicate_accept,
2688 ));
2689
2690 for (status, version, headers) in cases {
2691 assert!(matches!(
2692 validate_websocket_upgrade(status, version, &headers, &expected_accept),
2693 Err(RealtimeError::Transport)
2694 ));
2695 }
2696 }
2697
2698 #[test]
2699 fn invocation_encoder_accepts_small_terminated_messages() {
2700 let message = encode_invocation("1", "SubscribeAccounts", &[])
2701 .unwrap_or_else(|error| panic!("small invocation must encode: {error}"));
2702 let Message::Text(text) = message else {
2703 panic!("invocation must encode as text");
2704 };
2705 assert!(text.ends_with(SIGNALR_TERMINATOR));
2706 }
2707
2708 #[test]
2709 fn invocation_encoder_preserves_large_argument_batches() {
2710 let payload = "x".repeat(2 * 1_024 * 1_024);
2711 let message = encode_invocation("1", "Batch", &[Value::String(payload.clone())])
2712 .unwrap_or_else(|error| panic!("large invocation: {error}"));
2713 let Message::Text(text) = message else {
2714 panic!("text invocation");
2715 };
2716 let value: Value = serde_json::from_str(text.trim_end_matches(SIGNALR_TERMINATOR))
2717 .unwrap_or_else(|error| panic!("decode: {error}"));
2718 assert_eq!(value["arguments"][0], payload);
2719 }
2720 #[test]
2721 fn dropping_ready_client_outside_a_runtime_joins_on_its_original_runtime() {
2722 let runtime = tokio::runtime::Builder::new_current_thread()
2723 .enable_all()
2724 .build()
2725 .unwrap_or_else(|error| panic!("fixture runtime: {error}"));
2726 let (realtime, mut events) = runtime.block_on(async {
2727 let realtime = fixture_realtime();
2728 install_connected_generation(&realtime.inner, 1);
2729 let events = realtime
2730 .take_event_receiver()
2731 .unwrap_or_else(|| panic!("receiver"));
2732 (realtime, events)
2733 });
2734 drop(realtime);
2736 runtime.block_on(async {
2737 tokio::time::pause();
2738 assert_eq!(events.recv().await, Some(RealtimeEvent::Disconnected));
2739 assert_eq!(events.recv().await, None);
2740 });
2741 }
2742}