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