1mod auth;
4mod protocol;
5mod snapshot_then_stream;
6
7pub use snapshot_then_stream::{SnapshotErrorFn, SnapshotThenStream, SnapshotThenStreamConfig};
8
9use crate::auth::Credentials;
10use crate::errors::{Error, Result};
11use futures_util::{SinkExt, StreamExt};
12use rand_core::{OsRng, RngCore};
13use std::sync::{
14 Arc, Mutex, MutexGuard, RwLock, RwLockReadGuard, RwLockWriteGuard,
15 atomic::{AtomicBool, Ordering},
16};
17use tokio::sync::{mpsc, oneshot, watch};
18use tokio::time::{Duration, timeout};
19use tokio_tungstenite::{
20 connect_async_with_config,
21 tungstenite::{
22 Message,
23 client::IntoClientRequest,
24 http::{HeaderValue, header::SEC_WEBSOCKET_PROTOCOL},
25 protocol::WebSocketConfig,
26 },
27};
28
29const WS_PATH: &str = "/connection/websocket";
30const DEFAULT_QUEUE: usize = 1000;
31const CENTRIFUGO_READ_TIMEOUT: Duration = Duration::from_secs(30);
32const RECONNECT_INITIAL_CAP: Duration = Duration::from_millis(500);
33const RECONNECT_MAX_CAP: Duration = Duration::from_secs(30);
34const CENTRIFUGO_PROTOBUF_SUBPROTOCOL: &str = "centrifuge-protobuf";
35
36type ErrorCallback = Arc<dyn Fn(Error) + Send + Sync>;
37
38pub(crate) fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
39 mutex
40 .lock()
41 .unwrap_or_else(std::sync::PoisonError::into_inner)
42}
43
44pub(crate) fn read_unpoisoned<T>(lock: &RwLock<T>) -> RwLockReadGuard<'_, T> {
45 lock.read()
46 .unwrap_or_else(std::sync::PoisonError::into_inner)
47}
48
49pub(crate) fn write_unpoisoned<T>(lock: &RwLock<T>) -> RwLockWriteGuard<'_, T> {
50 lock.write()
51 .unwrap_or_else(std::sync::PoisonError::into_inner)
52}
53
54#[derive(Default)]
55struct SubscriptionErrorState {
56 last: Option<Error>,
57 callback: Option<ErrorCallback>,
58}
59
60struct SubscriptionAttempt<'a> {
61 stop: &'a mut watch::Receiver<bool>,
62 ready: &'a mut Option<oneshot::Sender<Result<()>>>,
63 gap: &'a ResubscribeGap,
64 error_state: &'a Arc<Mutex<SubscriptionErrorState>>,
65 connected: &'a mut bool,
66}
67
68fn invoke_error_callback(callback: ErrorCallback, err: Error) {
69 let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| callback(err)));
71}
72
73fn record_subscription_error(state: &Arc<Mutex<SubscriptionErrorState>>, err: Error) {
74 let callback = {
75 let mut state = lock_unpoisoned(state);
76 state.last = Some(err.clone());
77 state.callback.clone()
78 };
79 if let Some(callback) = callback {
80 invoke_error_callback(callback, err);
81 }
82}
83
84fn clear_subscription_error(state: &Arc<Mutex<SubscriptionErrorState>>) {
85 lock_unpoisoned(state).last = None;
86}
87
88struct ReconnectBackoff {
89 failures: u32,
90 jitter_state: u64,
91}
92
93impl ReconnectBackoff {
94 fn new() -> Self {
95 let mut rng = OsRng;
96 Self::with_seed(rng.next_u64())
97 }
98
99 fn with_seed(seed: u64) -> Self {
100 Self {
101 failures: 0,
102 jitter_state: if seed == 0 {
103 0x9e37_79b9_7f4a_7c15
104 } else {
105 seed
106 },
107 }
108 }
109
110 fn reset(&mut self) {
111 self.failures = 0;
112 }
113
114 fn next_delay(&mut self) -> Duration {
115 let multiplier = 1u64 << self.failures.min(16);
116 let cap_ms = (RECONNECT_INITIAL_CAP.as_millis() as u64)
117 .saturating_mul(multiplier)
118 .min(RECONNECT_MAX_CAP.as_millis() as u64);
119 self.failures = self.failures.saturating_add(1);
120
121 self.jitter_state ^= self.jitter_state << 13;
124 self.jitter_state ^= self.jitter_state >> 7;
125 self.jitter_state ^= self.jitter_state << 17;
126 let floor_ms = cap_ms / 2;
127 let delay_ms = floor_ms + self.jitter_state % (cap_ms - floor_ms + 1);
128 Duration::from_millis(delay_ms)
129 }
130}
131
132pub const MAX_REALTIME_MESSAGE_BYTES: usize = 8 * 1024 * 1024;
136
137pub fn try_enqueue<T>(
140 tx: &mpsc::Sender<T>,
141 item: T,
142 closed: &AtomicBool,
143 last_error: &std::sync::Mutex<Option<Error>>,
144 message: &str,
145) -> bool {
146 if closed.load(Ordering::SeqCst) {
147 return false;
148 }
149 match tx.try_send(item) {
150 Ok(()) => true,
151 Err(mpsc::error::TrySendError::Full(_)) => {
152 closed.store(true, Ordering::SeqCst);
153 *lock_unpoisoned(last_error) = Some(Error::queue_overflow(message));
154 false
155 }
156 Err(mpsc::error::TrySendError::Closed(_)) => {
157 closed.store(true, Ordering::SeqCst);
158 false
159 }
160 }
161}
162
163fn try_send_direct<T>(tx: &mpsc::Sender<T>, item: T, message: &str) -> Result<bool> {
164 match tx.try_send(item) {
165 Ok(()) => Ok(true),
166 Err(mpsc::error::TrySendError::Full(_)) => Err(Error::queue_overflow(message)),
167 Err(mpsc::error::TrySendError::Closed(_)) => Ok(false),
168 }
169}
170
171#[derive(Clone)]
173pub struct Client {
174 ws_url: String,
175 api_url: String,
176 credentials: Option<Credentials>,
177 max_queue: usize,
178 timeout: Duration,
180}
181
182impl Client {
183 pub fn new(
184 ws_url: impl Into<String>,
185 api_url: impl Into<String>,
186 credentials: Option<Credentials>,
187 max_queue: Option<usize>,
188 ) -> Self {
189 Self::with_timeout(
190 ws_url,
191 api_url,
192 credentials,
193 max_queue,
194 auth::DEFAULT_TOKEN_REQUEST_TIMEOUT,
195 )
196 }
197
198 pub fn with_timeout(
199 ws_url: impl Into<String>,
200 api_url: impl Into<String>,
201 credentials: Option<Credentials>,
202 max_queue: Option<usize>,
203 timeout: Duration,
204 ) -> Self {
205 Self {
206 ws_url: ws_url.into(),
207 api_url: api_url.into(),
208 credentials,
209 max_queue: max_queue.unwrap_or(DEFAULT_QUEUE).max(1),
210 timeout: if timeout.is_zero() {
211 auth::DEFAULT_TOKEN_REQUEST_TIMEOUT
212 } else {
213 timeout
214 },
215 }
216 }
217
218 pub(crate) fn request_timeout(&self) -> Duration {
219 self.timeout
220 }
221
222 fn ws_endpoint(&self) -> String {
223 let base = self.ws_url.trim_end_matches('/');
224 if base.contains(WS_PATH) {
225 base.to_owned()
226 } else {
227 format!("{base}{WS_PATH}")
228 }
229 }
230
231 fn validate_channel(&self, channel: &str) -> Result<()> {
232 if is_private_channel(channel) {
233 if self.credentials.is_none() {
234 return Err(Error::auth(format!(
235 "Cannot subscribe to private channel \"{channel}\" without API-key credentials"
236 )));
237 }
238 if self.api_url.is_empty() {
239 return Err(Error::realtime(
240 "Realtime private channels require api_url".to_owned(),
241 ));
242 }
243 }
244 Ok(())
245 }
246
247 pub async fn subscribe_raw(&self, channel: &str) -> Result<TypedSubscription<Vec<u8>>> {
252 self.subscribe_proto(channel, |bytes| Ok(bytes.to_vec()))
253 .await
254 }
255
256 pub async fn subscribe_proto<T, F>(
261 &self,
262 channel: &str,
263 decode: F,
264 ) -> Result<TypedSubscription<T>>
265 where
266 T: Send + 'static,
267 F: Fn(&[u8]) -> Result<T> + Send + Sync + 'static,
268 {
269 self.subscribe_proto_with_options(channel, decode, true)
270 .await
271 }
272
273 pub async fn subscribe_proto_with_options<T, F>(
278 &self,
279 channel: &str,
280 decode: F,
281 auto_reconnect: bool,
282 ) -> Result<TypedSubscription<T>>
283 where
284 T: Send + 'static,
285 F: Fn(&[u8]) -> Result<T> + Send + Sync + 'static,
286 {
287 self.validate_channel(channel)?;
288
289 let (stop_tx, mut stop_rx) = watch::channel(false);
290 let alive = Arc::new(AtomicBool::new(true));
291 let error_state = Arc::new(Mutex::new(SubscriptionErrorState::default()));
292 let gap = Arc::new(ResubscribeGap::default());
293 let (ready_tx, ready_rx) = oneshot::channel();
294 let (tx, rx) = mpsc::channel::<T>(self.max_queue);
295 let decode = Arc::new(decode);
296
297 let this = self.clone();
298 let channel = channel.to_owned();
299 let alive_task = alive.clone();
300 let error_task = error_state.clone();
301 let gap_task = gap.clone();
302 let task = tokio::spawn(async move {
303 let _guard = AliveGuard(alive_task.clone());
304 let mut ready = Some(ready_tx);
305 let mut backoff = ReconnectBackoff::new();
306 while !*stop_rx.borrow() {
307 let mut connected_this_attempt = false;
308 let mut attempt = SubscriptionAttempt {
309 stop: &mut stop_rx,
310 ready: &mut ready,
311 gap: &gap_task,
312 error_state: &error_task,
313 connected: &mut connected_this_attempt,
314 };
315 match this
316 .run_proto_subscription_once(&channel, decode.as_ref(), &tx, &mut attempt)
317 .await
318 {
319 Ok(()) => break,
320 Err(_) if *stop_rx.borrow() => break,
321 Err(err) if matches!(err, Error::QueueOverflow(_)) => {
322 record_subscription_error(&error_task, err);
323 break;
324 }
325 Err(err) => {
326 if connected_this_attempt {
327 backoff.reset();
328 }
329 record_subscription_error(&error_task, err.clone());
330 if let Some(ready) = ready.take() {
331 let _ = ready.send(Err(err));
332 break;
333 }
334 if *stop_rx.borrow() || !auto_reconnect {
335 break;
336 }
337 let delay = backoff.next_delay();
340 tokio::select! {
341 _ = stop_rx.changed() => {
342 if *stop_rx.borrow() {
343 break;
344 }
345 }
346 _ = tokio::time::sleep(delay) => {}
347 }
348 }
349 }
350 }
351 alive_task.store(false, Ordering::SeqCst);
352 drop(tx);
353 });
354 let mut abort_on_cancel = AbortOnDrop(Some(task));
357 let ready_result = ready_rx
358 .await
359 .map_err(|_| Error::realtime("realtime task ended before handshake".to_owned()))
360 .and_then(|inner| inner);
361 let task = abort_on_cancel
362 .0
363 .take()
364 .expect("subscription task present after ready");
365 ready_result?;
366
367 Ok(TypedSubscription {
368 rx,
369 stop: stop_tx,
370 alive,
371 task,
372 error_state,
373 gap,
374 })
375 }
376
377 async fn handshake_channel<W, R>(
378 &self,
379 write: &mut W,
380 read: &mut R,
381 channel: &str,
382 ) -> Result<()>
383 where
384 W: SinkExt<Message> + Unpin,
385 W::Error: std::fmt::Display,
386 R: StreamExt<Item = std::result::Result<Message, tokio_tungstenite::tungstenite::Error>>
387 + Unpin,
388 {
389 if is_private_channel(channel) {
390 let creds = self
391 .credentials
392 .as_ref()
393 .ok_or_else(|| Error::auth("private channel requires credentials"))?;
394 let connection_token =
395 auth::fetch_connection_token(creds, &self.api_url, self.timeout).await?;
396 centrifugo_connect(write, read, Some(&connection_token)).await?;
397 let subscription_token =
398 auth::fetch_subscription_token(creds, &self.api_url, channel, self.timeout).await?;
399 centrifugo_subscribe(write, read, channel, Some(&subscription_token)).await?;
400 } else {
401 centrifugo_connect(write, read, None).await?;
402 centrifugo_subscribe(write, read, channel, None).await?;
403 }
404 Ok(())
405 }
406
407 async fn run_proto_subscription_once<T, F>(
408 &self,
409 channel: &str,
410 decode: &F,
411 tx: &mpsc::Sender<T>,
412 attempt: &mut SubscriptionAttempt<'_>,
413 ) -> Result<()>
414 where
415 F: Fn(&[u8]) -> Result<T>,
416 {
417 let url = self.ws_endpoint();
418 let mut request = url
419 .into_client_request()
420 .map_err(|e| Error::realtime(format!("ws request: {e}")))?;
421 request.headers_mut().insert(
422 SEC_WEBSOCKET_PROTOCOL,
423 HeaderValue::from_static(CENTRIFUGO_PROTOBUF_SUBPROTOCOL),
424 );
425 let websocket_config = WebSocketConfig::default()
426 .max_message_size(Some(MAX_REALTIME_MESSAGE_BYTES))
427 .max_frame_size(Some(MAX_REALTIME_MESSAGE_BYTES));
428 let (ws, response) = timeout(
429 CENTRIFUGO_READ_TIMEOUT,
430 connect_async_with_config(request, Some(websocket_config), false),
431 )
432 .await
433 .map_err(|_| Error::realtime("websocket connect timed out".to_owned()))?
434 .map_err(|e| Error::realtime(format!("ws connect: {e}")))?;
435 if response
436 .headers()
437 .get(SEC_WEBSOCKET_PROTOCOL)
438 .and_then(|value| value.to_str().ok())
439 != Some(CENTRIFUGO_PROTOBUF_SUBPROTOCOL)
440 {
441 return Err(Error::realtime(
442 "server did not negotiate centrifuge-protobuf websocket subprotocol".to_owned(),
443 ));
444 }
445 let (mut write, mut read) = ws.split();
446 self.handshake_channel(&mut write, &mut read, channel)
447 .await?;
448 *attempt.connected = true;
449 clear_subscription_error(attempt.error_state);
450 if let Some(ready) = attempt.ready.take() {
451 let _ = ready.send(Ok(()));
452 } else {
453 attempt.gap.note_resubscribe();
457 }
458
459 loop {
460 if *attempt.stop.borrow() {
461 return Ok(());
462 }
463 let msg = tokio::select! {
464 changed = attempt.stop.changed() => {
465 if changed.is_err() || *attempt.stop.borrow() {
466 return Ok(());
467 }
468 None
469 }
470 msg = timeout(CENTRIFUGO_READ_TIMEOUT, read.next()) => {
471 Some(match msg {
472 Ok(Some(Ok(msg))) => msg,
473 Ok(Some(Err(e))) => return Err(Error::realtime(e.to_string())),
474 Ok(None) => return Err(Error::realtime("websocket closed".to_owned())),
475 Err(_) => {
477 return Err(Error::realtime("websocket read timeout".to_owned()));
478 }
479 })
480 }
481 };
482 let Some(msg) = msg else {
483 continue;
484 };
485 match msg {
486 Message::Binary(frame) => {
487 for incoming in protocol::decode_replies(&frame)? {
488 match incoming {
489 protocol::Incoming::Ping => {
490 write
491 .send(Message::Binary(protocol::pong_command().into()))
492 .await
493 .map_err(|e| {
494 Error::realtime(format!("protobuf pong send: {e}"))
495 })?;
496 }
497 protocol::Incoming::Publication(bytes) => {
498 let item = decode(&bytes)?;
499 if !try_send_direct(
500 tx,
501 item,
502 "typed realtime subscription queue full; consumer too slow",
503 )? {
504 return Ok(());
505 }
506 }
507 protocol::Incoming::Reply {
508 error: Some(err), ..
509 } => {
510 return Err(centrifugo_protocol_error(err));
511 }
512 protocol::Incoming::Reply { .. } => {}
513 }
514 }
515 }
516 Message::Text(_) => {
517 return Err(Error::realtime(
518 "received JSON text frame on protobuf websocket".to_owned(),
519 ));
520 }
521 Message::Ping(payload) => {
522 write
523 .send(Message::Pong(payload))
524 .await
525 .map_err(|e| Error::realtime(format!("ws pong: {e}")))?;
526 }
527 Message::Close(_) => {
528 return Err(Error::realtime("websocket closed".to_owned()));
529 }
530 _ => {}
531 }
532 }
533 }
534}
535
536#[derive(Default)]
537struct ResubscribeGap {
538 count: std::sync::atomic::AtomicU64,
539 latched: AtomicBool,
540}
541
542impl ResubscribeGap {
543 fn note_resubscribe(&self) {
544 self.count.fetch_add(1, Ordering::SeqCst);
545 self.latched.store(true, Ordering::SeqCst);
546 }
547}
548
549pub struct TypedSubscription<T> {
555 rx: mpsc::Receiver<T>,
556 stop: watch::Sender<bool>,
557 alive: Arc<AtomicBool>,
558 task: tokio::task::JoinHandle<()>,
559 error_state: Arc<Mutex<SubscriptionErrorState>>,
560 gap: Arc<ResubscribeGap>,
561}
562
563impl<T> TypedSubscription<T> {
564 pub async fn recv(&mut self) -> Option<T> {
565 self.rx.recv().await
566 }
567
568 pub async fn recv_result(&mut self) -> Result<Option<T>> {
573 match self.rx.recv().await {
574 Some(item) => Ok(Some(item)),
575 None => match self.take_err() {
576 Some(err) => Err(err),
577 None => Ok(None),
578 },
579 }
580 }
581
582 pub fn set_on_error<F>(&self, callback: F)
587 where
588 F: Fn(Error) + Send + Sync + 'static,
589 {
590 let callback: ErrorCallback = Arc::new(callback);
591 let current = {
592 let mut state = lock_unpoisoned(&self.error_state);
593 state.callback = Some(callback.clone());
594 state.last.clone()
595 };
596 if let Some(err) = current {
597 invoke_error_callback(callback, err);
598 }
599 }
600
601 pub fn is_alive(&self) -> bool {
603 self.alive.load(Ordering::SeqCst) && !self.task.is_finished()
604 }
605
606 pub fn err(&self) -> Option<Error> {
608 lock_unpoisoned(&self.error_state).last.clone()
609 }
610
611 pub fn take_err(&self) -> Option<Error> {
613 lock_unpoisoned(&self.error_state).last.take()
614 }
615
616 pub fn resubscribes(&self) -> u64 {
619 self.gap.count.load(Ordering::SeqCst)
620 }
621
622 pub fn take_resubscribed(&self) -> bool {
625 self.gap.latched.swap(false, Ordering::SeqCst)
626 }
627
628 pub fn close(&self) {
629 let _ = self.stop.send(true);
630 self.task.abort();
631 }
632}
633
634impl<T> Drop for TypedSubscription<T> {
635 fn drop(&mut self) {
636 let _ = self.stop.send(true);
637 self.task.abort();
638 }
639}
640
641struct AliveGuard(Arc<AtomicBool>);
642
643impl Drop for AliveGuard {
644 fn drop(&mut self) {
645 self.0.store(false, Ordering::SeqCst);
646 }
647}
648
649struct AbortOnDrop(Option<tokio::task::JoinHandle<()>>);
651
652impl Drop for AbortOnDrop {
653 fn drop(&mut self) {
654 if let Some(task) = self.0.take() {
655 task.abort();
656 }
657 }
658}
659
660pub fn is_private_channel(channel: &str) -> bool {
661 channel.starts_with("private:")
662}
663
664async fn centrifugo_connect<W, R>(write: &mut W, read: &mut R, token: Option<&str>) -> Result<()>
665where
666 W: SinkExt<Message> + Unpin,
667 W::Error: std::fmt::Display,
668 R: StreamExt<Item = std::result::Result<Message, tokio_tungstenite::tungstenite::Error>>
669 + Unpin,
670{
671 write
672 .send(Message::Binary(protocol::connect_command(1, token).into()))
673 .await
674 .map_err(|e| Error::realtime(format!("connect send: {e}")))?;
675 read_centrifugo_reply(write, read, 1).await
676}
677
678async fn centrifugo_subscribe<W, R>(
679 write: &mut W,
680 read: &mut R,
681 channel: &str,
682 token: Option<&str>,
683) -> Result<()>
684where
685 W: SinkExt<Message> + Unpin,
686 W::Error: std::fmt::Display,
687 R: StreamExt<Item = std::result::Result<Message, tokio_tungstenite::tungstenite::Error>>
688 + Unpin,
689{
690 write
691 .send(Message::Binary(
692 protocol::subscribe_command(2, channel, token).into(),
693 ))
694 .await
695 .map_err(|e| Error::realtime(format!("subscribe send: {e}")))?;
696 read_centrifugo_reply(write, read, 2).await
697}
698
699async fn read_centrifugo_reply<W, R>(write: &mut W, read: &mut R, expected_id: u32) -> Result<()>
700where
701 W: SinkExt<Message> + Unpin,
702 W::Error: std::fmt::Display,
703 R: StreamExt<Item = std::result::Result<Message, tokio_tungstenite::tungstenite::Error>>
704 + Unpin,
705{
706 loop {
707 let msg = timeout(Duration::from_secs(10), read.next())
708 .await
709 .map_err(|_| Error::realtime("centrifugo reply timeout".to_owned()))?
710 .ok_or_else(|| Error::realtime("centrifugo closed before reply".to_owned()))?
711 .map_err(|e| Error::realtime(format!("centrifugo read: {e}")))?;
712 match msg {
713 Message::Binary(frame) => {
714 for incoming in protocol::decode_replies(&frame)? {
715 match incoming {
716 protocol::Incoming::Reply {
717 id,
718 error: Some(err),
719 } if id == expected_id => return Err(centrifugo_protocol_error(err)),
720 protocol::Incoming::Reply { id, error: None } if id == expected_id => {
721 return Ok(());
722 }
723 protocol::Incoming::Ping => {
724 write
725 .send(Message::Binary(protocol::pong_command().into()))
726 .await
727 .map_err(|e| Error::realtime(format!("protobuf pong send: {e}")))?;
728 }
729 _ => {}
730 }
731 }
732 }
733 Message::Text(_) => {
734 return Err(Error::realtime(
735 "received JSON text reply on protobuf websocket".to_owned(),
736 ));
737 }
738 Message::Ping(payload) => {
739 write
740 .send(Message::Pong(payload))
741 .await
742 .map_err(|e| Error::realtime(format!("ws pong: {e}")))?;
743 }
744 Message::Close(_) => {
745 return Err(Error::realtime("centrifugo closed before reply".to_owned()));
746 }
747 _ => {}
748 }
749 }
750}
751
752fn centrifugo_protocol_error(error: protocol::ProtoError) -> Error {
753 Error::realtime(format!(
754 "centrifugo error {}: {}{}",
755 error.code,
756 error.message,
757 if error.temporary { " (temporary)" } else { "" }
758 ))
759}
760
761#[cfg(test)]
762mod tests {
763 use super::*;
764
765 #[test]
766 fn is_private_channel_detects_prefix() {
767 assert!(is_private_channel("private:spot:orders:acct:proto"));
768 assert!(!is_private_channel("public:spot:market:trades:1:proto"));
769 }
770
771 #[test]
772 fn rwlock_helpers_recover_from_poison() {
773 let lock = RwLock::new(7u32);
774 let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
775 let _guard = lock.write().unwrap();
776 panic!("poison rwlock");
777 }));
778 assert!(lock.read().is_err());
779 assert_eq!(*read_unpoisoned(&lock), 7);
780 *write_unpoisoned(&lock) = 9;
781 assert_eq!(*read_unpoisoned(&lock), 9);
782 }
783
784 #[test]
785 fn protobuf_websocket_subprotocol_is_centrifuge_protobuf() {
786 assert_eq!(CENTRIFUGO_PROTOBUF_SUBPROTOCOL, "centrifuge-protobuf");
787 let value = HeaderValue::from_static(CENTRIFUGO_PROTOBUF_SUBPROTOCOL);
788 assert_eq!(value.to_str().unwrap(), "centrifuge-protobuf");
789 }
790
791 #[test]
792 fn try_enqueue_fails_on_overflow_without_silent_drop() {
793 use std::sync::Mutex;
794
795 let (tx, mut rx) = mpsc::channel::<u8>(1);
796 let closed = AtomicBool::new(false);
797 let last_error = Mutex::new(None);
798 assert!(try_enqueue(
799 &tx,
800 1,
801 &closed,
802 &last_error,
803 "orderbook subscription queue full; consumer too slow"
804 ));
805 assert!(!try_enqueue(
806 &tx,
807 2,
808 &closed,
809 &last_error,
810 "orderbook subscription queue full; consumer too slow"
811 ));
812 assert!(closed.load(Ordering::SeqCst));
813 assert!(matches!(
814 last_error.lock().unwrap().as_ref(),
815 Some(Error::QueueOverflow(_))
816 ));
817 assert_eq!(rx.try_recv().unwrap(), 1);
818 assert!(rx.try_recv().is_err());
819 }
820
821 #[test]
822 fn direct_subscription_queue_fails_closed_on_overflow() {
823 let (tx, mut rx) = mpsc::channel::<u8>(1);
824 assert!(try_send_direct(&tx, 1, "full").unwrap());
825 assert!(matches!(
826 try_send_direct(&tx, 2, "full"),
827 Err(Error::QueueOverflow(_))
828 ));
829 assert_eq!(rx.try_recv().unwrap(), 1);
830 assert!(rx.try_recv().is_err());
831 }
832
833 #[test]
834 fn realtime_queue_capacity_is_never_zero() {
835 let client = Client::new(
836 "wss://example.invalid",
837 "https://example.invalid",
838 None,
839 Some(0),
840 );
841 assert_eq!(client.max_queue, 1);
842 }
843
844 #[tokio::test]
845 async fn subscribe_surfaces_initial_handshake_failure() {
846 let client = Client::new("not a websocket URL", "", None, None);
847
848 let result =
849 tokio::time::timeout(Duration::from_secs(2), client.subscribe_raw("public:test"))
850 .await
851 .expect("subscribe should not hang");
852 assert!(result.is_err(), "initial handshake error must be returned");
853 }
854
855 #[test]
856 fn read_timeout_constant_is_positive() {
857 assert!(CENTRIFUGO_READ_TIMEOUT > Duration::from_secs(0));
858 }
859
860 #[tokio::test]
861 async fn dropping_typed_subscription_signals_stop_and_aborts() {
862 let (stop, stop_rx) = watch::channel(false);
863 let (_tx, rx) = mpsc::channel::<u8>(1);
864 let alive = Arc::new(AtomicBool::new(true));
865 let (marker_tx, marker_rx) = tokio::sync::oneshot::channel::<()>();
866 let task = tokio::spawn(async move {
867 let _marker = marker_tx;
868 std::future::pending::<()>().await
869 });
870 let subscription = TypedSubscription {
871 rx,
872 stop,
873 alive,
874 task,
875 error_state: Arc::new(Mutex::new(SubscriptionErrorState::default())),
876 gap: Arc::new(ResubscribeGap::default()),
877 };
878
879 drop(subscription);
880
881 assert!(*stop_rx.borrow(), "close/Drop must signal stop");
882 assert!(
883 tokio::time::timeout(Duration::from_millis(500), marker_rx)
884 .await
885 .expect("abort should drop task locals promptly")
886 .is_err(),
887 "JoinHandle must be aborted on Drop"
888 );
889 }
890
891 #[test]
892 fn reconnect_backoff_is_bounded_exponential_and_jittered() {
893 let mut first = ReconnectBackoff::with_seed(1);
894 let mut second = ReconnectBackoff::with_seed(2);
895 let mut first_delays = Vec::new();
896 let mut second_delays = Vec::new();
897
898 for attempt in 0..12 {
899 let first_delay = first.next_delay();
900 let second_delay = second.next_delay();
901 let cap_ms = (RECONNECT_INITIAL_CAP.as_millis() as u64)
902 .saturating_mul(1u64 << attempt.min(16))
903 .min(RECONNECT_MAX_CAP.as_millis() as u64);
904 let floor = Duration::from_millis(cap_ms / 2);
905 let cap = Duration::from_millis(cap_ms);
906 assert!((floor..=cap).contains(&first_delay));
907 assert!((floor..=cap).contains(&second_delay));
908 first_delays.push(first_delay);
909 second_delays.push(second_delay);
910 }
911
912 assert_ne!(
913 first_delays, second_delays,
914 "independent subscriptions must not share one reconnect schedule"
915 );
916 assert!(first_delays.iter().all(|delay| *delay <= RECONNECT_MAX_CAP));
917
918 first.reset();
919 assert!(first.next_delay() <= RECONNECT_INITIAL_CAP);
920 }
921}