1use std::{
42 collections::VecDeque,
43 fmt::Debug,
44 path::Path,
45 pin::pin,
46 sync::{
47 Arc,
48 atomic::{AtomicU8, Ordering},
49 },
50 time::Duration,
51};
52
53use bytes::Bytes;
54use nautilus_core::string::secret::REDACTED;
55use nautilus_cryptography::providers::install_cryptographic_provider;
56use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
57use tokio_tungstenite::tungstenite::{Error, client::IntoClientRequest, stream::Mode};
58
59use super::{SocketConfig, TcpMessageHandler, TcpReader, TcpWriter, WriterCommand};
60use crate::{
61 SocketStateSink,
62 backoff::{
63 ExponentialBackoff, RECONNECT_STABILITY_THRESHOLD, ReconnectThrottle, wait_reconnect_delay,
64 },
65 dst,
66 error::{SendError, is_connection_drop_io_error},
67 logging::{log_task_aborted, log_task_started, log_task_stopped},
68 mode::{
69 ConnectionMode, ControllerLifecycle, ReadSessionFence, ReconnectOutcome,
70 ReconnectRequestOutcome,
71 },
72 net::TcpStream,
73 tls::{create_tls_config_from_certs_dir, tcp_tls},
74 writer::{self, DEFAULT_WRITER_CAPACITY, WriterReceiver, WriterSender},
75};
76
77const CONNECTION_STATE_CHECK_INTERVAL_MS: u64 = 10;
79const GRACEFUL_SHUTDOWN_DELAY_MS: u64 = 100;
80const GRACEFUL_SHUTDOWN_TIMEOUT_SECS: u64 = 5;
81const WRITE_TIMEOUT_SECS: u64 = 5;
82
83const MAX_READ_BUFFER_BYTES: usize = 10 * 1024 * 1024;
85
86struct BufferedWrite {
87 data: Bytes,
88 replay_key: Option<u64>,
89 _permit: tokio::sync::OwnedSemaphorePermit,
90}
91
92pub type SocketReconnectReplay = Arc<dyn Fn() -> Vec<Bytes> + Send + Sync>;
94
95struct SocketClientInner {
96 config: SocketConfig,
97 connector: Option<Arc<rustls::ClientConfig>>,
98 read_task: tokio::task::JoinHandle<()>,
99 read_fence: ReadSessionFence,
100 write_task: tokio::task::JoinHandle<()>,
101 writer_tx: WriterSender<WriterCommand>,
102 heartbeat_task: Option<tokio::task::JoinHandle<()>>,
103 connection_mode: Arc<AtomicU8>,
104 state_notify: Arc<tokio::sync::Notify>,
105 connect_timeout: Duration,
106 backoff: ExponentialBackoff,
107 reconnect_throttle: ReconnectThrottle,
108 reconnect_max_attempts: Option<u32>,
109 reconnect_attempt_count: u32,
110 state_sink: Option<SocketStateSink>,
111}
112
113impl SocketClientInner {
114 async fn connect_url(
120 config: SocketConfig,
121 state_sink: Option<SocketStateSink>,
122 ) -> anyhow::Result<Self> {
123 install_cryptographic_provider();
124
125 if config.suffix.is_empty() {
127 anyhow::bail!("Socket suffix cannot be empty: suffix is required for message framing");
128 }
129
130 config.validate()?;
133
134 let connect_timeout = Duration::from_millis(config.connect_timeout_ms.unwrap_or(10_000));
135 let reconnect_backoff = ExponentialBackoff::new(
136 Duration::from_millis(config.reconnect_delay_initial_ms.unwrap_or(2_000)),
137 Duration::from_millis(config.reconnect_delay_max_ms.unwrap_or(30_000)),
138 config.reconnect_backoff_factor.unwrap_or(1.5),
139 config.reconnect_jitter_ms.unwrap_or(100),
140 true, )?;
142 let connector = if let Some(dir) = &config.certs_dir {
143 let config = create_tls_config_from_certs_dir(Path::new(dir), false)?;
144 Some(Arc::new(config))
145 } else {
146 None
147 };
148
149 let max_retries = config.connection_max_retries.unwrap_or(5);
151
152 let mut backoff = ExponentialBackoff::new(
153 Duration::from_millis(500),
154 Duration::from_secs(5),
155 2.0,
156 250,
157 false,
158 )?;
159
160 let mut attempt = 0;
161 let (reader, writer) = loop {
162 attempt += 1;
163
164 let last_error = match dst::time::timeout(
165 connect_timeout,
166 Self::tls_connect_with_server(&config.url, config.mode, connector.clone()),
167 )
168 .await
169 {
170 Ok(Ok(result)) => {
171 if attempt > 1 {
172 log::info!("Socket connection established after {attempt} attempts");
173 }
174 break result;
175 }
176 Ok(Err(e)) => {
177 let error = e.to_string();
178 log::warn!(
179 "Socket connection attempt {attempt}/{max_retries} to {REDACTED} failed: {error}"
180 );
181 error
182 }
183 Err(_) => {
184 let error = format!(
185 "Connection timeout after {:.1}s (possible DNS resolution failure)",
186 connect_timeout.as_secs_f64()
187 );
188 log::warn!(
189 "Socket connection attempt {attempt}/{max_retries} to {REDACTED} timed out"
190 );
191 error
192 }
193 };
194
195 if attempt >= max_retries {
196 anyhow::bail!(
197 "Failed to connect to {REDACTED} after {max_retries} attempts: {last_error}. \
198 If this is a DNS error, check your network configuration and DNS settings."
199 );
200 }
201
202 let delay = backoff.next_duration();
203 log::debug!(
204 "Retrying in {delay:?} (attempt {}/{})",
205 attempt + 1,
206 max_retries
207 );
208 dst::time::sleep(delay).await;
209 };
210
211 log::debug!("Connected");
212
213 let connection_mode = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
214 let state_notify = Arc::new(tokio::sync::Notify::new());
215 let outcome =
216 ConnectionMode::complete_reconnect_with_sink(&connection_mode, state_sink.as_ref());
217 debug_assert_eq!(outcome, ReconnectOutcome::Reconnected);
218 let read_fence = ReadSessionFence::new();
219
220 let read_task = Self::spawn_read_task(
221 Arc::clone(&connection_mode),
222 read_fence.clone(),
223 reader,
224 config.message_handler.clone(),
225 config.suffix.clone(),
226 config.resolved_heartbeat_timeout(),
227 );
228
229 let (writer_tx, writer_rx) = writer::channel::<WriterCommand>(
230 config.writer_capacity.unwrap_or(DEFAULT_WRITER_CAPACITY),
231 );
232
233 let write_task = Self::spawn_write_task(
234 Arc::clone(&connection_mode),
235 Arc::clone(&state_notify),
236 writer,
237 writer_rx,
238 config.suffix.clone(),
239 state_sink.clone(),
240 );
241
242 let heartbeat_task = config.heartbeat.as_ref().map(|heartbeat| {
243 Self::spawn_heartbeat_task(
244 Arc::clone(&connection_mode),
245 heartbeat.interval_secs,
246 heartbeat.payload.clone(),
247 writer_tx.clone(),
248 )
249 });
250 let reconnect_max_attempts = config.reconnect_max_attempts;
251
252 Ok(Self {
253 config,
254 connector,
255 read_task,
256 read_fence,
257 write_task,
258 writer_tx,
259 heartbeat_task,
260 connection_mode,
261 state_notify,
262 connect_timeout,
263 backoff: reconnect_backoff,
264 reconnect_throttle: ReconnectThrottle::default(),
265 reconnect_max_attempts,
266 reconnect_attempt_count: 0,
267 state_sink,
268 })
269 }
270
271 fn parse_socket_url(url: &str, mode: Mode) -> Result<(String, String), Error> {
281 if url.contains("://") {
282 let parsed = url.parse::<http::Uri>().map_err(|e| {
284 Error::Io(std::io::Error::new(
285 std::io::ErrorKind::InvalidInput,
286 format!("Invalid URL: {e}"),
287 ))
288 })?;
289
290 let host = parsed.host().ok_or_else(|| {
291 Error::Io(std::io::Error::new(
292 std::io::ErrorKind::InvalidInput,
293 "URL missing host",
294 ))
295 })?;
296
297 let port = parsed
298 .port_u16()
299 .unwrap_or_else(|| match parsed.scheme_str() {
300 Some("wss" | "https") => 443,
301 Some("ws" | "http") => 80,
302 _ => match mode {
303 Mode::Tls => 443,
304 Mode::Plain => 80,
305 },
306 });
307
308 Ok((format!("{host}:{port}"), url.to_string()))
309 } else {
310 let scheme = match mode {
313 Mode::Tls => "wss",
314 Mode::Plain => "ws",
315 };
316 Ok((url.to_string(), format!("{scheme}://{url}")))
317 }
318 }
319
320 pub(crate) async fn tls_connect_with_server(
330 url: &str,
331 mode: Mode,
332 connector: Option<Arc<rustls::ClientConfig>>,
333 ) -> Result<(TcpReader, TcpWriter), Error> {
334 log::debug!("Connecting to {REDACTED}");
335
336 let (socket_addr, request_url) = Self::parse_socket_url(url, mode)?;
337 let tcp_result = TcpStream::connect(&socket_addr).await;
338
339 match tcp_result {
340 Ok(stream) => {
341 log::debug!("TCP connection established to {socket_addr}, proceeding with TLS");
342
343 crate::net::apply_socket_options(&stream);
344
345 let request = request_url.into_client_request()?;
346 tcp_tls(&request, mode, stream, connector)
347 .await
348 .map(tokio::io::split)
349 }
350 Err(e) => {
351 log::warn!("TCP connection failed to {socket_addr}: {e:?}");
352 Err(Error::Io(e))
353 }
354 }
355 }
356
357 async fn reconnect(
368 &mut self,
369 reconnect_replay: Option<&SocketReconnectReplay>,
370 ) -> Result<ReconnectOutcome, Error> {
371 log::info!("Reconnecting");
372
373 if ConnectionMode::from_atomic(&self.connection_mode).is_disconnect() {
374 log::debug!("Reconnect aborted due to disconnect state");
375 return Ok(ReconnectOutcome::Aborted);
376 }
377
378 let (reader, new_writer) = dst::time::timeout(
380 self.connect_timeout,
381 Self::tls_connect_with_server(
382 &self.config.url,
383 self.config.mode,
384 self.connector.clone(),
385 ),
386 )
387 .await
388 .map_err(|_| {
389 Error::Io(std::io::Error::new(
390 std::io::ErrorKind::TimedOut,
391 format!(
392 "reconnection timed out after {}s",
393 self.connect_timeout.as_secs_f64()
394 ),
395 ))
396 })??;
397
398 if ConnectionMode::from_atomic(&self.connection_mode).is_disconnect() {
399 log::debug!("Reconnect aborted mid-flight (after connect)");
400 return Ok(ReconnectOutcome::Aborted);
401 }
402 log::debug!("Connected");
403
404 let (tx, rx) = tokio::sync::oneshot::channel();
408 let command = if let Some(reconnect_replay) = reconnect_replay {
409 WriterCommand::UpdateWithReplay(new_writer, reconnect_replay(), tx)
410 } else {
411 WriterCommand::Update(new_writer, tx)
412 };
413
414 if let Err(e) = self.writer_tx.send_update(command) {
415 log::error!("{e}");
416 return Err(Error::Io(std::io::Error::new(
417 std::io::ErrorKind::BrokenPipe,
418 format!("Failed to send update command: {e}"),
419 )));
420 }
421
422 match rx.await {
424 Ok(true) => log::debug!("Writer confirmed buffer drain success"),
425 Ok(false) => {
426 log::warn!("Writer failed to drain buffer, aborting reconnect");
427 return Err(Error::Io(std::io::Error::other(
429 "Failed to drain reconnection buffer",
430 )));
431 }
432 Err(e) => {
433 log::error!("Writer dropped update channel: {e}");
434 return Err(Error::Io(std::io::Error::new(
435 std::io::ErrorKind::BrokenPipe,
436 "Writer task dropped response channel",
437 )));
438 }
439 }
440
441 dst::time::sleep(Duration::from_millis(GRACEFUL_SHUTDOWN_DELAY_MS)).await;
443
444 if ConnectionMode::from_atomic(&self.connection_mode).is_disconnect() {
445 log::debug!("Reconnect aborted mid-flight (after delay)");
446 return Ok(ReconnectOutcome::Aborted);
447 }
448
449 self.read_fence.invalidate();
450
451 if !self.read_task.is_finished() {
452 self.read_task.abort();
453 log_task_aborted("read");
454 }
455
456 if ConnectionMode::complete_reconnect_with_sink(
459 &self.connection_mode,
460 self.state_sink.as_ref(),
461 ) == ReconnectOutcome::Aborted
462 {
463 log::debug!("Reconnect aborted (state changed during reconnect)");
464 return Ok(ReconnectOutcome::Aborted);
465 }
466
467 self.read_fence = ReadSessionFence::new();
469 self.read_task = Self::spawn_read_task(
470 Arc::clone(&self.connection_mode),
471 self.read_fence.clone(),
472 reader,
473 self.config.message_handler.clone(),
474 self.config.suffix.clone(),
475 self.config.resolved_heartbeat_timeout(),
476 );
477
478 log::info!("Reconnect succeeded");
479 Ok(ReconnectOutcome::Reconnected)
480 }
481
482 #[inline]
488 #[must_use]
489 pub(crate) fn is_alive(&self) -> bool {
490 !self.read_task.is_finished() && !self.write_task.is_finished()
491 }
492
493 #[must_use]
494 fn spawn_read_task<R>(
495 connection_state: Arc<AtomicU8>,
496 read_fence: ReadSessionFence,
497 reader: R,
498 handler: Option<TcpMessageHandler>,
499 suffix: Vec<u8>,
500 heartbeat_timeout_secs: Option<u64>,
501 ) -> tokio::task::JoinHandle<()>
502 where
503 R: AsyncRead + Unpin + Send + 'static,
504 {
505 log_task_started("read");
506
507 let check_interval = Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS);
509 let heartbeat_timeout = heartbeat_timeout_secs.map(Duration::from_secs);
510
511 tokio::task::spawn(Self::run_read_loop(
512 connection_state,
513 read_fence,
514 reader,
515 handler,
516 suffix,
517 heartbeat_timeout,
518 check_interval,
519 ))
520 }
521
522 async fn run_read_loop<R>(
523 connection_state: Arc<AtomicU8>,
524 read_fence: ReadSessionFence,
525 mut reader: R,
526 handler: Option<TcpMessageHandler>,
527 suffix: Vec<u8>,
528 heartbeat_timeout: Option<Duration>,
529 check_interval: Duration,
530 ) where
531 R: AsyncRead + Unpin,
532 {
533 let mut buf = Vec::new();
534 let mut last_data_time = dst::time::Instant::now();
535
536 'read: loop {
537 if !ConnectionMode::from_atomic(&connection_state).is_active() || !read_fence.is_valid()
538 {
539 if !buf.is_empty() {
540 log::debug!(
541 "Dropping {} buffered bytes after socket session ended",
542 buf.len()
543 );
544 buf.clear();
545 }
546 break;
547 }
548
549 match dst::time::timeout(check_interval, reader.read_buf(&mut buf)).await {
550 Ok(Ok(0)) => {
552 log::debug!("Connection closed by server");
553 break;
554 }
555 Ok(Err(e)) => {
556 log::debug!("Connection ended: {e}");
557 break;
558 }
559 Ok(Ok(bytes)) => {
561 log::trace!("Received <binary> {bytes} bytes");
562 last_data_time = dst::time::Instant::now();
563
564 if !ConnectionMode::from_atomic(&connection_state).is_active()
565 || !read_fence.is_valid()
566 {
567 log::debug!(
568 "Dropping {} buffered bytes after socket session ended",
569 buf.len()
570 );
571 buf.clear();
572 break;
573 }
574
575 while let Some((i, _)) = &buf
576 .windows(suffix.len())
577 .enumerate()
578 .find(|(_, pair)| pair.eq(&suffix))
579 {
580 let mut data: Vec<u8> = buf.drain(0..i + suffix.len()).collect();
581 data.truncate(data.len() - suffix.len());
582
583 if let Some(ref handler) = handler {
584 if !ConnectionMode::from_atomic(&connection_state).is_active()
585 || !read_fence.is_valid()
586 {
587 log::debug!(
588 "Dropping {} buffered bytes after socket session ended",
589 data.len() + buf.len()
590 );
591 buf.clear();
592 break 'read;
593 }
594 handler(&data);
595 }
596 }
597
598 if buf.len() > MAX_READ_BUFFER_BYTES {
599 log::error!(
600 "Read buffer exceeded maximum size ({MAX_READ_BUFFER_BYTES} bytes), closing connection"
601 );
602 break;
603 }
604 }
605 Err(_) => {
606 if let Some(timeout) = heartbeat_timeout {
607 let silent_for = last_data_time.elapsed();
608 if silent_for >= timeout {
609 log::warn!(
610 "Heartbeat timeout: no bytes received for {:.1}s",
611 silent_for.as_secs_f64()
612 );
613 break;
614 }
615 }
616 }
617 }
618 }
619
620 log_task_stopped("read");
621 }
622
623 async fn drain_reconnect_buffer<W>(
633 buffer: &mut VecDeque<BufferedWrite>,
634 writer: &mut W,
635 suffix: &[u8],
636 replay: &[Bytes],
637 ) -> bool
638 where
639 W: AsyncWrite + Unpin,
640 {
641 if buffer.is_empty() {
642 return false;
643 }
644
645 let initial_buffer_len = buffer.len();
646 log::info!("Sending {initial_buffer_len} buffered messages after reconnection");
647
648 while let Some(buffered) = buffer.front() {
649 if buffered.replay_key.is_some()
650 && replay.iter().any(|replayed| replayed == &buffered.data)
651 {
652 buffer.pop_front();
653 continue;
654 }
655
656 let mut combined_msg = Vec::with_capacity(buffered.data.len() + suffix.len());
657 combined_msg.extend_from_slice(&buffered.data);
658 combined_msg.extend_from_slice(suffix);
659
660 if let Err(e) = writer.write_all(&combined_msg).await {
661 if is_connection_drop_io_error(&e) {
662 log::warn!(
663 "Failed to send buffered message with suffix after reconnection: {e}, {} messages remain in buffer",
664 buffer.len()
665 );
666 } else {
667 log::error!(
668 "Failed to send buffered message with suffix after reconnection: {e}, {} messages remain in buffer",
669 buffer.len()
670 );
671 }
672 return true;
673 }
674
675 buffer.pop_front();
676 }
677
678 log::info!("Successfully sent all {initial_buffer_len} buffered messages");
679
680 false
681 }
682
683 async fn replace_writer<W>(
684 active_writer: &mut W,
685 new_writer: W,
686 replay: Vec<Bytes>,
687 reconnect_buffer: &mut VecDeque<BufferedWrite>,
688 suffix: &[u8],
689 ) -> bool
690 where
691 W: AsyncWrite + Unpin,
692 {
693 log::debug!("Received new writer");
694 dst::time::sleep(Duration::from_millis(GRACEFUL_SHUTDOWN_DELAY_MS)).await;
695
696 _ = dst::time::timeout(
697 Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS),
698 active_writer.shutdown(),
699 )
700 .await;
701
702 *active_writer = new_writer;
703 log::debug!("Updated writer");
704
705 let drain_result =
706 dst::time::timeout(Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS), async {
707 for replay_msg in &replay {
708 let mut framed = Vec::with_capacity(replay_msg.len() + suffix.len());
709 framed.extend_from_slice(replay_msg);
710 framed.extend_from_slice(suffix);
711 if let Err(e) = active_writer.write_all(&framed).await {
712 log::warn!("Failed to send reconnect replay: {e}");
713 return true;
714 }
715 }
716
717 Self::drain_reconnect_buffer(reconnect_buffer, active_writer, suffix, &replay).await
718 })
719 .await;
720
721 let send_error = drain_result.unwrap_or_else(|_| {
722 log::warn!(
723 "Timed out sending reconnect replay and buffered messages, {} buffered messages remain",
724 reconnect_buffer.len()
725 );
726 true
727 });
728 !send_error
729 }
730
731 fn spawn_write_task<W>(
732 connection_state: Arc<AtomicU8>,
733 state_notify: Arc<tokio::sync::Notify>,
734 writer: W,
735 mut writer_rx: WriterReceiver<WriterCommand<W>>,
736 suffix: Vec<u8>,
737 state_sink: Option<SocketStateSink>,
738 ) -> tokio::task::JoinHandle<()>
739 where
740 W: AsyncWrite + Unpin + Send + 'static,
741 {
742 log_task_started("write");
743
744 let check_interval = Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS);
746
747 tokio::task::spawn(async move {
748 let mut active_writer = writer;
749 let mut reconnect_buffer: VecDeque<BufferedWrite> = VecDeque::new();
750 let mut write_buf: Vec<u8> = Vec::new();
751
752 loop {
753 let mode = ConnectionMode::from_atomic(&connection_state);
754 if matches!(mode, ConnectionMode::Disconnect | ConnectionMode::Closed) {
755 break;
756 }
757
758 if mode.is_active() && !reconnect_buffer.is_empty() {
759 let drain_result = dst::time::timeout(
760 Duration::from_secs(WRITE_TIMEOUT_SECS),
761 Self::drain_reconnect_buffer(
762 &mut reconnect_buffer,
763 &mut active_writer,
764 &suffix,
765 &[],
766 ),
767 )
768 .await;
769 let send_error = drain_result.unwrap_or_else(|_| {
770 log::warn!(
771 "Timed out draining reconnect buffer after {WRITE_TIMEOUT_SECS}s, {} messages remain",
772 reconnect_buffer.len()
773 );
774 true
775 });
776
777 if send_error
778 && ConnectionMode::request_reconnect_with_sink(
779 &connection_state,
780 state_sink.as_ref(),
781 )
782 {
783 log::warn!("Writer triggering reconnect");
784 state_notify.notify_one();
785 }
786 continue;
787 }
788
789 match dst::time::timeout(check_interval, writer_rx.recv()).await {
790 Ok(Some((msg, permit))) => {
791 let mode = ConnectionMode::from_atomic(&connection_state);
793 if matches!(mode, ConnectionMode::Disconnect | ConnectionMode::Closed) {
794 break;
795 }
796
797 match msg {
798 WriterCommand::Update(new_writer, tx) => {
799 let sent = Self::replace_writer(
800 &mut active_writer,
801 new_writer,
802 Vec::new(),
803 &mut reconnect_buffer,
804 &suffix,
805 )
806 .await;
807
808 if let Err(e) = tx.send(sent) {
809 log::error!(
810 "Failed to report drain status to controller: {e:?}"
811 );
812 }
813 }
814 WriterCommand::UpdateWithReplay(new_writer, replay, tx) => {
815 let sent = Self::replace_writer(
816 &mut active_writer,
817 new_writer,
818 replay,
819 &mut reconnect_buffer,
820 &suffix,
821 )
822 .await;
823
824 if let Err(e) = tx.send(sent) {
825 log::error!(
826 "Failed to report drain status to controller: {e:?}"
827 );
828 }
829 }
830 WriterCommand::Send(data)
831 if mode.is_reconnect() || !reconnect_buffer.is_empty() =>
832 {
833 log::debug!(
834 "Buffering message until reconnect drain completes ({} bytes)",
835 data.len()
836 );
837 Self::buffer_reconnect_write(
838 &mut reconnect_buffer,
839 data,
840 None,
841 permit,
842 );
843 }
844 WriterCommand::SendOrReplay { key, data } if mode.is_reconnect() => {
845 log::debug!(
846 "Buffering replayable message while reconnecting ({} bytes)",
847 data.len()
848 );
849 Self::buffer_reconnect_write(
850 &mut reconnect_buffer,
851 data,
852 Some(key),
853 permit,
854 );
855 }
856 command @ (WriterCommand::Send(_)
857 | WriterCommand::SendOrReplay { .. }) => {
858 let (msg, replay_key) = match command {
859 WriterCommand::Send(data) => (data, None),
860 WriterCommand::SendOrReplay { key, data } => (data, Some(key)),
861 _ => unreachable!(),
862 };
863 write_buf.clear();
864 write_buf.extend_from_slice(&msg);
865 write_buf.extend_from_slice(&suffix);
866
867 let write_result = dst::time::timeout(
868 Duration::from_secs(WRITE_TIMEOUT_SECS),
869 active_writer.write_all(&write_buf),
870 )
871 .await;
872 let write_failed = match write_result {
873 Ok(Ok(())) => false,
874 Ok(Err(e)) => {
875 if is_connection_drop_io_error(&e) {
876 log::warn!("Failed to send message: {e}");
877 } else {
878 log::error!("Failed to send message: {e}");
879 }
880 true
881 }
882 Err(_) => {
883 log::warn!(
884 "Timed out sending message after {WRITE_TIMEOUT_SECS}s"
885 );
886 true
887 }
888 };
889
890 if write_failed {
891 Self::buffer_reconnect_write(
892 &mut reconnect_buffer,
893 msg,
894 replay_key,
895 permit,
896 );
897
898 if ConnectionMode::request_reconnect_with_sink(
900 &connection_state,
901 state_sink.as_ref(),
902 ) {
903 log::warn!("Writer triggering reconnect");
904 state_notify.notify_one();
905 }
906 }
907 }
908 }
909 }
910 Ok(None) => {
911 log::debug!("Writer channel closed, terminating writer task");
913 break;
914 }
915 Err(_) => {
916 }
918 }
919 }
920
921 _ = dst::time::timeout(
924 Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS),
925 active_writer.shutdown(),
926 )
927 .await;
928
929 log_task_stopped("write");
930 })
931 }
932
933 fn buffer_reconnect_write(
934 buffer: &mut VecDeque<BufferedWrite>,
935 data: Bytes,
936 replay_key: Option<u64>,
937 permit: tokio::sync::OwnedSemaphorePermit,
938 ) {
939 if let Some(key) = replay_key {
940 buffer.retain(|buffered| buffered.replay_key != Some(key));
941 }
942
943 buffer.push_back(BufferedWrite {
944 data,
945 replay_key,
946 _permit: permit,
947 });
948 }
949
950 fn spawn_heartbeat_task(
951 connection_state: Arc<AtomicU8>,
952 interval_secs: u64,
953 message: Vec<u8>,
954 writer_tx: WriterSender<WriterCommand>,
955 ) -> tokio::task::JoinHandle<()> {
956 log_task_started("heartbeat");
957
958 tokio::task::spawn(async move {
959 let interval = Duration::from_secs(interval_secs);
960
961 loop {
962 dst::time::sleep(interval).await;
963
964 match ConnectionMode::from_u8(connection_state.load(Ordering::SeqCst)) {
965 ConnectionMode::Active => {
966 let msg = WriterCommand::Send(message.clone().into());
967
968 match writer_tx.send(msg) {
969 Ok(()) => log::trace!("Sent heartbeat to writer task"),
970 Err(e) => {
971 log::warn!("Failed to send heartbeat to writer task: {e}");
972 }
973 }
974 }
975 ConnectionMode::Reconnect => {}
976 ConnectionMode::Disconnect | ConnectionMode::Closed => break,
977 }
978 }
979
980 log_task_stopped("heartbeat");
981 })
982 }
983}
984
985impl Drop for SocketClientInner {
986 fn drop(&mut self) {
987 self.read_fence.invalidate();
988
989 if !self.read_task.is_finished() {
990 self.read_task.abort();
991 log_task_aborted("read");
992 }
993
994 if !self.write_task.is_finished() {
995 self.write_task.abort();
996 log_task_aborted("write");
997 }
998
999 if let Some(ref handle) = self.heartbeat_task.take()
1000 && !handle.is_finished()
1001 {
1002 handle.abort();
1003 log_task_aborted("heartbeat");
1004 }
1005 }
1006}
1007
1008pub struct SocketClient {
1014 pub(crate) controller_task: tokio::task::JoinHandle<()>,
1015 pub(crate) connection_mode: Arc<AtomicU8>,
1016 pub(crate) state_notify: Arc<tokio::sync::Notify>,
1017 pub(crate) connect_timeout: Duration,
1018 pub writer_tx: WriterSender<WriterCommand>,
1019 state_sink: Option<SocketStateSink>,
1020 controller_lifecycle: Arc<ControllerLifecycle>,
1021 controller_notify: Arc<tokio::sync::Notify>,
1022}
1023
1024#[derive(Clone)]
1026pub struct SocketReconnectHandle {
1027 connection_mode: Arc<AtomicU8>,
1028 state_sink: Option<SocketStateSink>,
1029 controller_lifecycle: Arc<ControllerLifecycle>,
1030 controller_notify: Arc<tokio::sync::Notify>,
1031}
1032
1033impl Debug for SocketReconnectHandle {
1034 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1035 f.debug_struct(stringify!(SocketReconnectHandle))
1036 .field(
1037 "connection_mode",
1038 &ConnectionMode::from_atomic(&self.connection_mode),
1039 )
1040 .finish_non_exhaustive()
1041 }
1042}
1043
1044impl SocketReconnectHandle {
1045 #[must_use]
1050 pub fn request_reconnect(&self) -> ReconnectRequestOutcome {
1051 let Some(_request) = self.controller_lifecycle.enter_request() else {
1052 return ReconnectRequestOutcome::Closed;
1053 };
1054
1055 let outcome = ConnectionMode::request_reconnect_outcome_with_sink(
1056 &self.connection_mode,
1057 self.state_sink.as_ref(),
1058 );
1059
1060 if outcome == ReconnectRequestOutcome::Accepted {
1061 self.controller_notify.notify_one();
1062 }
1063 outcome
1064 }
1065}
1066
1067impl Debug for SocketClient {
1068 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1069 f.debug_struct(stringify!(SocketClient)).finish()
1070 }
1071}
1072
1073#[bon::bon]
1074impl SocketClient {
1075 #[builder(finish_fn = connect)]
1088 pub async fn builder(
1089 config: SocketConfig,
1090 post_reconnection: Option<Arc<dyn Fn() + Send + Sync>>,
1091 state_sink: Option<SocketStateSink>,
1092 reconnect_replay: Option<SocketReconnectReplay>,
1093 ) -> anyhow::Result<Self> {
1094 let inner = SocketClientInner::connect_url(config, state_sink).await?;
1095 let writer_tx = inner.writer_tx.clone();
1096 let connection_mode = Arc::clone(&inner.connection_mode);
1097 let state_notify = Arc::clone(&inner.state_notify);
1098 let connect_timeout = inner.connect_timeout;
1099 let state_sink = inner.state_sink.clone();
1100 let controller_lifecycle = Arc::new(ControllerLifecycle::new());
1101 let controller_notify = Arc::new(tokio::sync::Notify::new());
1102 let controller_task = Self::spawn_controller_task(
1103 inner,
1104 Arc::clone(&connection_mode),
1105 Arc::clone(&state_notify),
1106 Arc::clone(&controller_lifecycle),
1107 Arc::clone(&controller_notify),
1108 post_reconnection,
1109 reconnect_replay,
1110 );
1111 let abort_handle = controller_task.abort_handle();
1112 controller_lifecycle.set_abort(move || abort_handle.abort());
1113
1114 Ok(Self {
1115 controller_task,
1116 connection_mode,
1117 state_notify,
1118 connect_timeout,
1119 writer_tx,
1120 state_sink,
1121 controller_lifecycle,
1122 controller_notify,
1123 })
1124 }
1125
1126 #[must_use]
1128 pub fn reconnect_handle(&self) -> SocketReconnectHandle {
1129 SocketReconnectHandle {
1130 connection_mode: Arc::clone(&self.connection_mode),
1131 state_sink: self.state_sink.clone(),
1132 controller_lifecycle: Arc::clone(&self.controller_lifecycle),
1133 controller_notify: Arc::clone(&self.controller_notify),
1134 }
1135 }
1136
1137 #[must_use]
1142 pub fn request_reconnect(&self) -> bool {
1143 self.reconnect_handle().request_reconnect() == ReconnectRequestOutcome::Accepted
1144 }
1145
1146 #[must_use]
1148 pub fn connection_mode(&self) -> ConnectionMode {
1149 ConnectionMode::from_atomic(&self.connection_mode)
1150 }
1151
1152 #[inline]
1157 #[must_use]
1158 pub fn is_active(&self) -> bool {
1159 self.connection_mode().is_active()
1160 }
1161
1162 #[inline]
1167 #[must_use]
1168 pub fn is_reconnecting(&self) -> bool {
1169 self.connection_mode().is_reconnect()
1170 }
1171
1172 #[inline]
1176 #[must_use]
1177 pub fn is_disconnecting(&self) -> bool {
1178 self.connection_mode().is_disconnect()
1179 }
1180
1181 #[inline]
1187 #[must_use]
1188 pub fn is_closed(&self) -> bool {
1189 self.connection_mode().is_closed()
1190 }
1191
1192 pub async fn close(&self) {
1197 self.begin_shutdown();
1198 self.close_with_timeout(Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS))
1199 .await;
1200 }
1201
1202 pub fn begin_shutdown(&self) {
1204 ConnectionMode::request_disconnect(&self.connection_mode);
1205 self.state_notify.notify_waiters();
1206 }
1207
1208 async fn close_with_timeout(&self, shutdown_timeout: Duration) {
1209 self.begin_shutdown();
1210
1211 if dst::time::timeout(shutdown_timeout, async {
1212 while !self.controller_task.is_finished() {
1213 dst::time::sleep(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
1214 }
1215 })
1216 .await
1217 .is_err()
1218 {
1219 log::warn!("Timeout waiting for controller task to finish");
1220 }
1221
1222 if !self.controller_task.is_finished() {
1223 self.controller_task.abort();
1224 log_task_aborted("controller");
1225 }
1226
1227 self.connection_mode
1228 .store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
1229 self.state_notify.notify_waiters();
1230 }
1231
1232 #[inline]
1236 fn check_not_terminal(&self) -> Result<(), SendError> {
1237 match self.connection_mode() {
1238 ConnectionMode::Disconnect | ConnectionMode::Closed => Err(SendError::Closed),
1239 _ => Ok(()),
1240 }
1241 }
1242
1243 async fn wait_for_active(&self) -> Result<(), SendError> {
1249 const FALLBACK_INTERVAL_MS: u64 = 100;
1250
1251 let mode = self.connection_mode();
1252 if mode.is_active() {
1253 return Ok(());
1254 }
1255
1256 if matches!(mode, ConnectionMode::Disconnect | ConnectionMode::Closed) {
1257 return Err(SendError::Closed);
1258 }
1259
1260 log::debug!("Waiting for client to become ACTIVE before sending...");
1261
1262 let fallback_interval = Duration::from_millis(FALLBACK_INTERVAL_MS);
1263
1264 dst::time::timeout(self.connect_timeout, async {
1265 loop {
1266 let mut notified = pin!(self.state_notify.notified());
1268 notified.as_mut().enable();
1269
1270 let mode = self.connection_mode();
1271 if mode.is_active() {
1272 return Ok(());
1273 }
1274
1275 if matches!(mode, ConnectionMode::Disconnect | ConnectionMode::Closed) {
1276 return Err(());
1277 }
1278
1279 tokio::select! {
1280 biased;
1281 () = notified => {}
1282 () = dst::time::sleep(fallback_interval) => {}
1283 }
1284 }
1285 })
1286 .await
1287 .map_err(|_| SendError::Timeout)?
1288 .map_err(|()| SendError::Closed)
1289 }
1290
1291 pub async fn send_bytes(&self, data: Vec<u8>) -> Result<(), SendError> {
1301 self.check_not_terminal()?;
1302 self.wait_for_active().await?;
1303
1304 let msg = WriterCommand::Send(data.into());
1305 self.writer_tx.send(msg)
1306 }
1307
1308 fn spawn_controller_task(
1309 mut inner: SocketClientInner,
1310 connection_mode: Arc<AtomicU8>,
1311 state_notify: Arc<tokio::sync::Notify>,
1312 controller_lifecycle: Arc<ControllerLifecycle>,
1313 controller_notify: Arc<tokio::sync::Notify>,
1314 post_reconnection: Option<Arc<dyn Fn() + Send + Sync>>,
1315 reconnect_replay: Option<SocketReconnectReplay>,
1316 ) -> tokio::task::JoinHandle<()> {
1317 const CONTROLLER_FALLBACK_INTERVAL_MS: u64 = 100;
1318
1319 tokio::task::spawn(async move {
1320 let _activity = controller_lifecycle.activity();
1321 log_task_started("controller");
1322
1323 let fallback_interval = Duration::from_millis(CONTROLLER_FALLBACK_INTERVAL_MS);
1324 let mut reconnected_at = None;
1325
1326 loop {
1327 tokio::select! {
1328 biased;
1329 () = controller_notify.notified() => {}
1330 () = state_notify.notified() => {}
1331 () = dst::time::sleep(fallback_interval) => {}
1332 }
1333
1334 let mut mode = ConnectionMode::from_atomic(&connection_mode);
1335
1336 if mode.is_disconnect() {
1337 log::debug!("Disconnecting");
1338
1339 let timeout = Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS);
1340 if dst::time::timeout(timeout, async {
1341 dst::time::sleep(Duration::from_millis(GRACEFUL_SHUTDOWN_DELAY_MS)).await;
1343
1344 inner.read_fence.invalidate();
1345 if !inner.read_task.is_finished() {
1346 inner.read_task.abort();
1347 log_task_aborted("read");
1348 }
1349
1350 if let Some(task) = &inner.heartbeat_task
1351 && !task.is_finished()
1352 {
1353 task.abort();
1354 log_task_aborted("heartbeat");
1355 }
1356 })
1357 .await
1358 .is_err()
1359 {
1360 log::warn!("Shutdown timed out after {}s", timeout.as_secs());
1361 }
1362
1363 log::debug!("Closed");
1364 connection_mode.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
1365 state_notify.notify_waiters();
1366 break; }
1368
1369 if mode.is_closed() {
1370 log::debug!("Connection closed");
1371
1372 inner.read_fence.invalidate();
1373 if !inner.read_task.is_finished() {
1374 inner.read_task.abort();
1375 log_task_aborted("read");
1376 }
1377
1378 if let Some(task) = &inner.heartbeat_task
1379 && !task.is_finished()
1380 {
1381 task.abort();
1382 log_task_aborted("heartbeat");
1383 }
1384
1385 state_notify.notify_waiters();
1386 break;
1387 }
1388
1389 if mode.is_active() && !inner.is_alive() {
1390 if ConnectionMode::request_reconnect_with_sink(
1391 &connection_mode,
1392 inner.state_sink.as_ref(),
1393 ) {
1394 log::info!("Detected dead read task, transitioning to RECONNECT");
1395 }
1396 mode = ConnectionMode::from_atomic(&connection_mode);
1397 }
1398
1399 if mode.is_reconnect() {
1400 let reconnect_uptime = reconnected_at
1401 .take()
1402 .map(|started: dst::time::Instant| started.elapsed());
1403 let previous_reconnect_stable = reconnect_uptime
1404 .is_some_and(|uptime| uptime >= RECONNECT_STABILITY_THRESHOLD);
1405
1406 if previous_reconnect_stable {
1407 inner.backoff.reset();
1408 inner.reconnect_attempt_count = 0;
1409 log::debug!(
1410 "Socket remained active for at least {}s, resetting reconnect cycle",
1411 RECONNECT_STABILITY_THRESHOLD.as_secs()
1412 );
1413 }
1414
1415 if let Some(max_attempts) = inner.reconnect_max_attempts
1417 && inner.reconnect_attempt_count >= max_attempts
1418 {
1419 log::error!(
1420 "Max reconnection attempts ({max_attempts}) exceeded, transitioning to CLOSED"
1421 );
1422
1423 if connection_mode
1424 .compare_exchange(
1425 ConnectionMode::Reconnect.as_u8(),
1426 ConnectionMode::Closed.as_u8(),
1427 Ordering::SeqCst,
1428 Ordering::SeqCst,
1429 )
1430 .is_ok()
1431 {
1432 state_notify.notify_waiters();
1433 break;
1434 }
1435 continue;
1436 }
1437
1438 let backoff_delay = if reconnect_uptime.is_some() && !previous_reconnect_stable
1439 {
1440 inner.backoff.next_duration()
1441 } else {
1442 Duration::ZERO
1443 };
1444
1445 let duration = inner.reconnect_throttle.gated_delay(backoff_delay);
1446 if !duration.is_zero() {
1447 log::warn!("Backing off for {}s...", duration.as_secs_f64());
1448
1449 if !wait_reconnect_delay(
1450 duration,
1451 connection_mode.as_ref(),
1452 state_notify.as_ref(),
1453 )
1454 .await
1455 {
1456 log::debug!("Backoff interrupted by terminal state");
1457 continue;
1458 }
1459 }
1460
1461 inner.reconnect_attempt_count += 1;
1462 inner.reconnect_throttle.record_attempt();
1463
1464 let reconnect_result = tokio::select! {
1466 biased;
1467 result = inner.reconnect(reconnect_replay.as_ref()) => Some(result),
1468 () = async {
1469 loop {
1470 let mut notified = pin!(state_notify.notified());
1472 notified.as_mut().enable();
1473
1474 if ConnectionMode::from_atomic(&connection_mode).is_disconnect() {
1475 break;
1476 }
1477 notified.await;
1478 }
1479 } => None,
1480 };
1481
1482 match reconnect_result {
1483 None => {
1484 log::debug!("Reconnect interrupted by disconnect");
1485 }
1486 Some(Ok(ReconnectOutcome::Reconnected)) => {
1487 log::debug!("Reconnected successfully");
1488 reconnected_at = Some(dst::time::Instant::now());
1489
1490 state_notify.notify_waiters();
1491
1492 if ConnectionMode::from_atomic(&connection_mode).is_active() {
1496 if let Some(ref handler) = post_reconnection {
1497 handler();
1498 log::debug!("Called `post_reconnection` handler");
1499 }
1500 } else {
1501 log::debug!(
1502 "Skipping post_reconnection handlers due to disconnect state"
1503 );
1504 }
1505 }
1506 Some(Ok(ReconnectOutcome::Aborted)) => {
1507 log::debug!("Reconnect aborted");
1508 }
1509 Some(Err(e)) => {
1510 let duration = inner.backoff.next_duration();
1511 log::warn!(
1512 "Reconnect attempt {} failed: {e}",
1513 inner.reconnect_attempt_count
1514 );
1515
1516 if !duration.is_zero() {
1517 log::warn!("Backing off for {}s...", duration.as_secs_f64());
1518 if !wait_reconnect_delay(
1519 duration,
1520 connection_mode.as_ref(),
1521 state_notify.as_ref(),
1522 )
1523 .await
1524 {
1525 log::debug!("Backoff interrupted by terminal state");
1526 }
1527 }
1528 }
1529 }
1530 }
1531 }
1532 log_task_stopped("controller");
1533 })
1534 }
1535}
1536
1537impl Drop for SocketClient {
1539 fn drop(&mut self) {
1540 let controller_running = !self.controller_task.is_finished();
1541 self.controller_lifecycle.close_and_abort();
1542
1543 if controller_running {
1544 log_task_aborted("controller");
1545 }
1546 }
1547}
1548
1549#[cfg(test)]
1550#[cfg(not(feature = "turmoil"))]
1551#[cfg(not(all(feature = "simulation", madsim)))] #[cfg(target_os = "linux")] mod tests {
1554 use log::Level;
1555 use nautilus_common::testing::wait_until_async;
1556 use parking_lot::Mutex as BlockingMutex;
1557 use rstest::rstest;
1558 use tokio::{
1559 io::{AsyncReadExt, AsyncWriteExt},
1560 net::{TcpListener, TcpStream},
1561 sync::Mutex,
1562 task,
1563 time::{Duration, sleep},
1564 };
1565
1566 use super::*;
1567 use crate::{SocketState, logging::tests::capture_logs_for, socket::SocketHeartbeat};
1568
1569 async fn bind_test_server() -> (u16, TcpListener) {
1570 let listener = TcpListener::bind("127.0.0.1:0")
1571 .await
1572 .expect("Failed to bind ephemeral port");
1573 let port = listener.local_addr().unwrap().port();
1574 (port, listener)
1575 }
1576
1577 async fn run_echo_server(mut socket: TcpStream) {
1578 let mut buf = Vec::new();
1579 loop {
1580 match socket.read_buf(&mut buf).await {
1581 Ok(0) => {
1582 break;
1583 }
1584 Ok(_n) => {
1585 while let Some(idx) = buf.array_windows().position(|w| w == b"\r\n") {
1586 let mut line = buf.drain(..idx + 2).collect::<Vec<u8>>();
1587 line.truncate(line.len() - 2);
1589
1590 if line == b"close" {
1591 let _ = socket.shutdown().await;
1592 return;
1593 }
1594
1595 let mut echo_data = line;
1596 echo_data.extend_from_slice(b"\r\n");
1597 if socket.write_all(&echo_data).await.is_err() {
1598 break;
1599 }
1600 }
1601 }
1602 Err(e) => {
1603 eprintln!("Server read error: {e}");
1604 break;
1605 }
1606 }
1607 }
1608 }
1609
1610 #[rstest]
1611 #[tokio::test]
1612 async fn test_writer_capacity_from_config() {
1613 let (port, listener) = bind_test_server().await;
1614 let config = SocketConfig::builder()
1615 .url(format!("127.0.0.1:{port}"))
1616 .mode(Mode::Plain)
1617 .suffix(b"\r\n".to_vec())
1618 .writer_capacity(2)
1619 .build()
1620 .unwrap();
1621 let client = SocketClient::builder()
1622 .config(config)
1623 .connect()
1624 .await
1625 .unwrap();
1626 let (_peer, _) = listener.accept().await.unwrap();
1627
1628 client.send_bytes(vec![1]).await.unwrap();
1630 client.send_bytes(vec![2]).await.unwrap();
1631 let overflow = client.send_bytes(vec![3]).await;
1632
1633 assert_eq!(client.connection_mode(), ConnectionMode::Active);
1634 assert!(matches!(overflow, Err(SendError::BufferFull)));
1635 client.close().await;
1636 }
1637
1638 #[rstest]
1639 #[case::drain_failed(
1640 Some(false),
1641 std::io::ErrorKind::Other,
1642 "Failed to drain reconnection buffer"
1643 )]
1644 #[case::reply_dropped(
1645 None,
1646 std::io::ErrorKind::BrokenPipe,
1647 "Writer task dropped response channel"
1648 )]
1649 #[tokio::test]
1650 async fn test_reconnect_writer_failure_preserves_reconnect_state(
1651 #[case] reply: Option<bool>,
1652 #[case] kind: std::io::ErrorKind,
1653 #[case] message: &str,
1654 ) {
1655 let (port, listener) = bind_test_server().await;
1656 let config = SocketConfig::builder()
1657 .url(format!("127.0.0.1:{port}"))
1658 .mode(Mode::Plain)
1659 .suffix(b"\r\n".to_vec())
1660 .build()
1661 .unwrap();
1662 let mut client = SocketClientInner::connect_url(config, None).await.unwrap();
1663 let (_peer, _) = listener.accept().await.unwrap();
1664 client.read_task.abort();
1665 client.write_task.abort();
1666 let _ = (&mut client.read_task).await;
1667 let _ = (&mut client.write_task).await;
1668 let (writer_tx, mut writer_rx) = writer::channel(DEFAULT_WRITER_CAPACITY);
1669 client.writer_tx = writer_tx;
1670 client
1671 .connection_mode
1672 .store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
1673
1674 let writer = tokio::spawn(async move {
1675 let WriterCommand::Update(_, sender) = writer_rx.recv().await.unwrap().0 else {
1676 panic!("expected writer update");
1677 };
1678
1679 if let Some(reply) = reply {
1680 sender.send(reply).unwrap();
1681 }
1682 });
1683
1684 let error = client.reconnect(None).await.unwrap_err();
1685 writer.await.unwrap();
1686
1687 let Error::Io(error) = error else {
1688 panic!("expected I/O error, was {error}");
1689 };
1690
1691 assert_eq!(error.kind(), kind);
1692 assert_eq!(error.to_string(), message);
1693 assert_eq!(
1694 ConnectionMode::from_atomic(&client.connection_mode),
1695 ConnectionMode::Reconnect
1696 );
1697 }
1698
1699 #[tokio::test]
1700 async fn test_basic_send_receive() {
1701 let (port, listener) = bind_test_server().await;
1702 let server_task = task::spawn(async move {
1703 let (socket, _) = listener.accept().await.unwrap();
1704 run_echo_server(socket).await;
1705 });
1706
1707 let config = SocketConfig {
1708 url: format!("127.0.0.1:{port}"),
1709 mode: Mode::Plain,
1710 suffix: b"\r\n".to_vec(),
1711 message_handler: None,
1712 heartbeat: None,
1713 connect_timeout_ms: None,
1714 reconnect_delay_initial_ms: None,
1715 reconnect_backoff_factor: None,
1716 reconnect_delay_max_ms: None,
1717 reconnect_jitter_ms: None,
1718 reconnect_max_attempts: None,
1719 connection_max_retries: None,
1720 heartbeat_timeout_secs: None,
1721 writer_capacity: None,
1722 certs_dir: None,
1723 };
1724
1725 let client = SocketClient::builder()
1726 .config(config)
1727 .connect()
1728 .await
1729 .expect("Client connect failed unexpectedly");
1730
1731 client.send_bytes(b"Hello".into()).await.unwrap();
1732 client.send_bytes(b"World".into()).await.unwrap();
1733
1734 sleep(Duration::from_millis(100)).await;
1736
1737 client.send_bytes(b"close".into()).await.unwrap();
1738 server_task.await.unwrap();
1739 assert!(!client.is_closed());
1740 }
1741
1742 #[tokio::test]
1743 async fn test_reconnect_fail_exhausted() {
1744 let (port, listener) = bind_test_server().await;
1745 drop(listener); wait_until_async(
1749 || async {
1750 TcpStream::connect(format!("127.0.0.1:{port}"))
1751 .await
1752 .is_err()
1753 },
1754 Duration::from_secs(2),
1755 )
1756 .await;
1757
1758 let config = SocketConfig {
1759 url: format!("127.0.0.1:{port}"),
1760 mode: Mode::Plain,
1761 suffix: b"\r\n".to_vec(),
1762 message_handler: None,
1763 heartbeat: None,
1764 connect_timeout_ms: Some(100),
1765 reconnect_delay_initial_ms: Some(50),
1766 reconnect_backoff_factor: Some(1.0),
1767 reconnect_delay_max_ms: Some(50),
1768 reconnect_jitter_ms: Some(0),
1769 connection_max_retries: Some(1),
1770 reconnect_max_attempts: None,
1771 heartbeat_timeout_secs: None,
1772 writer_capacity: None,
1773 certs_dir: None,
1774 };
1775
1776 let client_res = SocketClient::builder().config(config).connect().await;
1777 assert!(
1778 client_res.is_err(),
1779 "Should fail quickly with no server listening"
1780 );
1781 }
1782
1783 #[rstest]
1784 #[tokio::test]
1785 async fn test_connect_failure_redacts_url_from_logs_and_error() {
1786 const ENDPOINT_PATH_SECRET: &str = "unique-endpoint-path-secret";
1787 const ENDPOINT_QUERY_SECRET: &str = "unique-endpoint-query-secret";
1788
1789 let (port, listener) = bind_test_server().await;
1790 drop(listener); wait_until_async(
1794 || async {
1795 TcpStream::connect(format!("127.0.0.1:{port}"))
1796 .await
1797 .is_err()
1798 },
1799 Duration::from_secs(2),
1800 )
1801 .await;
1802
1803 let capture = capture_logs_for(&["nautilus_network::socket::client"]).await;
1804 let config = SocketConfig {
1805 url: format!(
1806 "ws://127.0.0.1:{port}/{ENDPOINT_PATH_SECRET}?api_key={ENDPOINT_QUERY_SECRET}"
1807 ),
1808 mode: Mode::Plain,
1809 suffix: b"\r\n".to_vec(),
1810 message_handler: None,
1811 heartbeat: None,
1812 connect_timeout_ms: Some(1_000),
1813 reconnect_delay_initial_ms: Some(50),
1814 reconnect_delay_max_ms: Some(50),
1815 reconnect_backoff_factor: Some(1.0),
1816 reconnect_jitter_ms: Some(0),
1817 connection_max_retries: Some(1),
1818 reconnect_max_attempts: None,
1819 heartbeat_timeout_secs: None,
1820 writer_capacity: None,
1821 certs_dir: None,
1822 };
1823
1824 let error = SocketClient::builder()
1825 .config(config)
1826 .connect()
1827 .await
1828 .expect_err("connection should fail with no server listening");
1829
1830 let error_message = error.to_string();
1831 assert!(
1832 error_message.starts_with(&format!("Failed to connect to {REDACTED}")),
1833 "connect error omitted the redaction placeholder: {error_message}"
1834 );
1835 assert!(
1836 !error_message.contains(ENDPOINT_PATH_SECRET)
1837 && !error_message.contains(ENDPOINT_QUERY_SECRET),
1838 "connect error exposed the secret marker: {error_message}"
1839 );
1840
1841 let messages = capture.messages();
1842 let (_, attempt_warning) = messages
1845 .iter()
1846 .find(|(level, message)| {
1847 *level == Level::Warn && message.starts_with("Socket connection attempt 1/1 ")
1848 })
1849 .expect("connect-failure warning was not captured");
1850 assert!(
1851 attempt_warning.contains(REDACTED),
1852 "connect-failure warning omitted the redaction placeholder: {attempt_warning}"
1853 );
1854 assert!(
1855 messages.iter().any(|(level, message)| {
1856 *level == Level::Debug && *message == format!("Connecting to {REDACTED}")
1857 }),
1858 "connect debug log should redact the endpoint, was {messages:?}"
1859 );
1860 assert!(
1861 messages.iter().all(|(_, message)| {
1862 !message.contains(ENDPOINT_PATH_SECRET) && !message.contains(ENDPOINT_QUERY_SECRET)
1863 }),
1864 "socket client logs exposed the secret marker: {messages:?}"
1865 );
1866 }
1867
1868 #[tokio::test]
1869 async fn test_user_disconnect() {
1870 let (port, listener) = bind_test_server().await;
1871 let server_task = task::spawn(async move {
1872 let (socket, _) = listener.accept().await.unwrap();
1873 let mut buf = [0u8; 1024];
1874 let _ = socket.try_read(&mut buf);
1875
1876 loop {
1877 sleep(Duration::from_secs(1)).await;
1878 }
1879 });
1880
1881 let config = SocketConfig {
1882 url: format!("127.0.0.1:{port}"),
1883 mode: Mode::Plain,
1884 suffix: b"\r\n".to_vec(),
1885 message_handler: None,
1886 heartbeat: None,
1887 connect_timeout_ms: None,
1888 reconnect_delay_initial_ms: None,
1889 reconnect_backoff_factor: None,
1890 reconnect_delay_max_ms: None,
1891 reconnect_jitter_ms: None,
1892 reconnect_max_attempts: None,
1893 connection_max_retries: None,
1894 heartbeat_timeout_secs: None,
1895 writer_capacity: None,
1896 certs_dir: None,
1897 };
1898
1899 let client = SocketClient::builder()
1900 .config(config)
1901 .connect()
1902 .await
1903 .unwrap();
1904
1905 client.close().await;
1906 assert!(client.is_closed());
1907 server_task.abort();
1908 }
1909
1910 #[tokio::test]
1911 async fn test_close_after_closed_returns_fast_and_preserves_state() {
1912 let (port, listener) = bind_test_server().await;
1913
1914 let server_task = task::spawn(async move {
1915 let (socket, _) = listener.accept().await.unwrap();
1918 drop(socket);
1919 drop(listener);
1920 sleep(Duration::from_secs(5)).await;
1921 });
1922
1923 let config = SocketConfig {
1924 url: format!("127.0.0.1:{port}"),
1925 mode: Mode::Plain,
1926 suffix: b"\r\n".to_vec(),
1927 message_handler: None,
1928 heartbeat: None,
1929 connect_timeout_ms: Some(200),
1930 reconnect_delay_initial_ms: Some(50),
1931 reconnect_backoff_factor: Some(1.0),
1932 reconnect_delay_max_ms: Some(50),
1933 reconnect_jitter_ms: Some(0),
1934 connection_max_retries: None,
1935 reconnect_max_attempts: Some(1),
1936 heartbeat_timeout_secs: None,
1937 writer_capacity: None,
1938 certs_dir: None,
1939 };
1940
1941 let client = SocketClient::builder()
1942 .config(config)
1943 .connect()
1944 .await
1945 .unwrap();
1946
1947 wait_until_async(|| async { client.is_closed() }, Duration::from_secs(5)).await;
1948
1949 let start = std::time::Instant::now();
1953 client.close().await;
1954 let elapsed = start.elapsed();
1955
1956 assert!(client.is_closed(), "Client should remain CLOSED");
1957 assert!(
1958 !client.is_disconnecting(),
1959 "Closed client should not report DISCONNECT after close()"
1960 );
1961 assert!(
1962 elapsed < Duration::from_secs(2),
1963 "close() on a closed client should return fast, took {elapsed:?}"
1964 );
1965
1966 server_task.abort();
1967 }
1968
1969 #[rstest]
1970 #[tokio::test]
1971 async fn test_heartbeat_task_warns_when_writer_channel_closed() {
1972 let capture = capture_logs_for(&["nautilus_network::socket::client"]).await;
1973 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
1974 let (writer_tx, mut writer_rx) = writer::channel(DEFAULT_WRITER_CAPACITY);
1975 let task = SocketClientInner::spawn_heartbeat_task(
1976 Arc::clone(&connection_state),
1977 1,
1978 b"ping".to_vec(),
1979 writer_tx,
1980 );
1981
1982 let _ = tokio::time::timeout(Duration::from_secs(2), writer_rx.recv())
1983 .await
1984 .expect("timed out waiting for the first heartbeat")
1985 .expect("heartbeat channel closed before the first heartbeat");
1986 drop(writer_rx);
1987 tokio::time::sleep(Duration::from_millis(1_200)).await;
1988
1989 let messages = capture.messages();
1990 assert!(
1991 messages.iter().any(|(level, message)| {
1992 *level == Level::Warn && message.contains("Failed to send heartbeat to writer task")
1993 }),
1994 "expected WARN when the writer channel is closed, was {messages:?}"
1995 );
1996 assert!(
1997 !messages.iter().any(|(level, _)| *level == Level::Error),
1998 "a closed writer channel must not log ERROR, was {messages:?}"
1999 );
2000
2001 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
2002 tokio::time::timeout(Duration::from_secs(2), task)
2003 .await
2004 .expect("heartbeat task should stop after close")
2005 .unwrap();
2006 }
2007
2008 #[tokio::test]
2009 async fn test_heartbeat() {
2010 let (port, listener) = bind_test_server().await;
2011 let received = Arc::new(Mutex::new(Vec::new()));
2012 let received2 = Arc::clone(&received);
2013
2014 let server_task = task::spawn(async move {
2015 let (socket, _) = listener.accept().await.unwrap();
2016
2017 let mut buf = Vec::new();
2018 loop {
2019 match socket.try_read_buf(&mut buf) {
2020 Ok(0) => break,
2021 Ok(_) => {
2022 while let Some(idx) = buf.array_windows().position(|w| w == b"\r\n") {
2023 let mut line = buf.drain(..idx + 2).collect::<Vec<u8>>();
2024 line.truncate(line.len() - 2);
2025 received2.lock().await.push(line);
2026 }
2027 }
2028 Err(_) => {
2029 tokio::time::sleep(Duration::from_millis(10)).await;
2030 }
2031 }
2032 }
2033 });
2034
2035 let config = SocketConfig {
2036 url: format!("127.0.0.1:{port}"),
2037 mode: Mode::Plain,
2038 suffix: b"\r\n".to_vec(),
2039 message_handler: None,
2040 heartbeat: Some(SocketHeartbeat {
2041 interval_secs: 1,
2042 payload: b"ping".to_vec(),
2043 }),
2044 connect_timeout_ms: None,
2045 reconnect_delay_initial_ms: None,
2046 reconnect_backoff_factor: None,
2047 reconnect_delay_max_ms: None,
2048 reconnect_jitter_ms: None,
2049 reconnect_max_attempts: None,
2050 connection_max_retries: None,
2051 heartbeat_timeout_secs: None,
2052 writer_capacity: None,
2053 certs_dir: None,
2054 };
2055
2056 let client = SocketClient::builder()
2057 .config(config)
2058 .connect()
2059 .await
2060 .unwrap();
2061
2062 sleep(Duration::from_secs(3)).await;
2064
2065 {
2066 let lock = received.lock().await;
2067 let pings = lock
2068 .iter()
2069 .filter(|line| line == &&b"ping".to_vec())
2070 .count();
2071 assert!(
2072 pings >= 2,
2073 "Expected at least 2 heartbeat pings; got {pings}"
2074 );
2075 }
2076
2077 client.close().await;
2078 server_task.abort();
2079 }
2080
2081 #[tokio::test]
2082 async fn test_reconnect_success() {
2083 let (port, listener) = bind_test_server().await;
2084
2085 let server_task = task::spawn(async move {
2089 let (mut socket, _) = listener.accept().await.expect("First accept failed");
2091
2092 sleep(Duration::from_millis(500)).await;
2094 let _ = socket.shutdown().await;
2095
2096 sleep(Duration::from_millis(500)).await;
2098
2099 let (socket, _) = listener.accept().await.expect("Second accept failed");
2101 run_echo_server(socket).await;
2102 });
2103
2104 let config = SocketConfig {
2105 url: format!("127.0.0.1:{port}"),
2106 mode: Mode::Plain,
2107 suffix: b"\r\n".to_vec(),
2108 message_handler: None,
2109 heartbeat: None,
2110 connect_timeout_ms: Some(5_000),
2111 reconnect_delay_initial_ms: Some(500),
2112 reconnect_delay_max_ms: Some(5_000),
2113 reconnect_backoff_factor: Some(2.0),
2114 reconnect_jitter_ms: Some(50),
2115 reconnect_max_attempts: None,
2116 connection_max_retries: None,
2117 heartbeat_timeout_secs: None,
2118 writer_capacity: None,
2119 certs_dir: None,
2120 };
2121
2122 let client = SocketClient::builder()
2123 .config(config)
2124 .connect()
2125 .await
2126 .expect("Client connect failed unexpectedly");
2127
2128 assert!(client.is_active(), "Client should start as active");
2130
2131 wait_until_async(|| async { client.is_active() }, Duration::from_secs(10)).await;
2134
2135 client
2136 .send_bytes(b"TestReconnect".into())
2137 .await
2138 .expect("Send failed");
2139
2140 client.close().await;
2141 server_task.abort();
2142 }
2143
2144 #[rstest]
2145 #[tokio::test]
2146 async fn test_state_sink_reports_initial_loss_and_recovery() {
2147 let (port, listener) = bind_test_server().await;
2148 let server_task = task::spawn(async move {
2149 let (mut socket, _) = listener.accept().await.unwrap();
2150 sleep(Duration::from_millis(100)).await;
2151 socket.shutdown().await.unwrap();
2152
2153 let (socket, _) = listener.accept().await.unwrap();
2154 run_echo_server(socket).await;
2155 });
2156 let config = SocketConfig {
2157 url: format!("127.0.0.1:{port}"),
2158 mode: Mode::Plain,
2159 suffix: b"\r\n".to_vec(),
2160 message_handler: None,
2161 heartbeat: None,
2162 connect_timeout_ms: Some(1_000),
2163 reconnect_delay_initial_ms: Some(10),
2164 reconnect_backoff_factor: Some(1.0),
2165 reconnect_delay_max_ms: Some(10),
2166 reconnect_jitter_ms: Some(0),
2167 reconnect_max_attempts: Some(3),
2168 connection_max_retries: Some(1),
2169 heartbeat_timeout_secs: None,
2170 writer_capacity: None,
2171 certs_dir: None,
2172 };
2173 let states = Arc::new(BlockingMutex::new(Vec::new()));
2174 let states_callback = Arc::clone(&states);
2175 let sink = SocketStateSink::new(move |state| {
2176 states_callback.lock().push(state);
2177 });
2178
2179 let client = SocketClient::builder()
2180 .config(config)
2181 .state_sink(sink)
2182 .connect()
2183 .await
2184 .unwrap();
2185
2186 assert_eq!(*states.lock(), vec![SocketState::Connected]);
2187
2188 wait_until_async(
2189 || {
2190 let states = Arc::clone(&states);
2191 async move { states.lock().len() == 3 }
2192 },
2193 Duration::from_secs(5),
2194 )
2195 .await;
2196 assert_eq!(
2197 *states.lock(),
2198 vec![
2199 SocketState::Connected,
2200 SocketState::Disconnected,
2201 SocketState::Connected,
2202 ]
2203 );
2204
2205 client.close().await;
2206 assert_eq!(states.lock().len(), 3);
2207 server_task.abort();
2208 }
2209
2210 #[rstest]
2211 #[tokio::test]
2212 async fn test_state_sink_ignores_initial_connection_failure() {
2213 let (port, listener) = bind_test_server().await;
2214 drop(listener);
2215 let config = SocketConfig {
2216 url: format!("127.0.0.1:{port}"),
2217 mode: Mode::Plain,
2218 suffix: b"\r\n".to_vec(),
2219 message_handler: None,
2220 heartbeat: None,
2221 connect_timeout_ms: Some(100),
2222 reconnect_delay_initial_ms: Some(1),
2223 reconnect_backoff_factor: Some(1.0),
2224 reconnect_delay_max_ms: Some(1),
2225 reconnect_jitter_ms: Some(0),
2226 reconnect_max_attempts: Some(1),
2227 connection_max_retries: Some(1),
2228 heartbeat_timeout_secs: None,
2229 writer_capacity: None,
2230 certs_dir: None,
2231 };
2232 let states = Arc::new(BlockingMutex::new(Vec::new()));
2233 let states_callback = Arc::clone(&states);
2234 let sink = SocketStateSink::new(move |state| {
2235 states_callback.lock().push(state);
2236 });
2237
2238 let result = SocketClient::builder()
2239 .config(config)
2240 .state_sink(sink)
2241 .connect()
2242 .await;
2243
2244 assert!(result.is_err());
2245 assert_eq!(*states.lock(), Vec::new());
2246 }
2247}
2248
2249#[cfg(test)]
2250#[cfg(not(feature = "turmoil"))]
2251#[cfg(not(all(feature = "simulation", madsim)))] mod rust_tests {
2253 use std::{
2254 pin::Pin,
2255 sync::{
2256 Arc,
2257 atomic::{AtomicBool, AtomicUsize},
2258 },
2259 task::{Context, Poll, Waker},
2260 };
2261
2262 use nautilus_common::testing::wait_until_async;
2263 use parking_lot::{Condvar, Mutex};
2264 use rstest::rstest;
2265 use tokio::{
2266 io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, DuplexStream, ReadBuf},
2267 net::TcpListener,
2268 sync::oneshot,
2269 task::{self, yield_now},
2270 time::{Duration, sleep},
2271 };
2272
2273 use super::*;
2274 use crate::{SocketState, socket::SocketHeartbeat};
2275
2276 const TEST_TIMEOUT: Duration = Duration::from_secs(10);
2277
2278 struct CondvarReleaseGuard {
2279 release: Arc<(Mutex<bool>, Condvar)>,
2280 }
2281
2282 impl CondvarReleaseGuard {
2283 fn new(release: Arc<(Mutex<bool>, Condvar)>) -> Self {
2284 Self { release }
2285 }
2286
2287 fn release(&self) {
2288 let (lock, condvar) = self.release.as_ref();
2289 let mut released = lock.lock();
2290 *released = true;
2291 condvar.notify_all();
2292 }
2293 }
2294
2295 impl Drop for CondvarReleaseGuard {
2296 fn drop(&mut self) {
2297 self.release();
2298 }
2299 }
2300
2301 async fn recv_rendezvous<T: Send + 'static>(
2302 receiver: std::sync::mpsc::Receiver<T>,
2303 name: &'static str,
2304 ) -> T {
2305 let receive_task = tokio::task::spawn_blocking(move || receiver.recv_timeout(TEST_TIMEOUT));
2306
2307 match tokio::time::timeout(TEST_TIMEOUT * 2, receive_task).await {
2308 Ok(Ok(Ok(value))) => value,
2309 Ok(Ok(Err(e))) => {
2310 panic!("{name} did not arrive within the test timeout: {e}")
2311 }
2312 Ok(Err(e)) => panic!("{name} receive task failed: {e}"),
2313 Err(e) => panic!("{name} receive task did not finish: {e}"),
2314 }
2315 }
2316
2317 async fn await_task_termination(
2318 task: tokio::task::JoinHandle<()>,
2319 name: &'static str,
2320 cancellation_ok: bool,
2321 ) {
2322 match tokio::time::timeout(TEST_TIMEOUT, task).await {
2323 Ok(Ok(())) => {}
2324 Ok(Err(e)) if cancellation_ok && e.is_cancelled() => {}
2325 Ok(Err(e)) => panic!("{name} failed: {e}"),
2326 Err(e) => panic!("{name} did not terminate within the test timeout: {e}"),
2327 }
2328 }
2329
2330 fn reconnect_test_config(port: u16) -> SocketConfig {
2331 SocketConfig {
2332 url: format!("127.0.0.1:{port}"),
2333 mode: Mode::Plain,
2334 suffix: b"\r\n".to_vec(),
2335 message_handler: None,
2336 heartbeat: None,
2337 connect_timeout_ms: Some(1_000),
2338 reconnect_delay_initial_ms: None,
2339 reconnect_backoff_factor: None,
2340 reconnect_delay_max_ms: None,
2341 reconnect_jitter_ms: None,
2342 connection_max_retries: Some(1),
2343 reconnect_max_attempts: None,
2344 heartbeat_timeout_secs: None,
2345 writer_capacity: None,
2346 certs_dir: None,
2347 }
2348 }
2349
2350 #[rstest]
2351 #[tokio::test]
2352 async fn test_reconnect_outcome_is_aborted_before_connect() {
2353 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
2354 let port = listener.local_addr().unwrap().port();
2355 let server = tokio::spawn(async move {
2356 let (_socket, _) = listener.accept().await.unwrap();
2357 std::future::pending::<()>().await;
2358 });
2359 let mut inner = SocketClientInner::connect_url(reconnect_test_config(port), None)
2360 .await
2361 .unwrap();
2362 inner
2363 .connection_mode
2364 .store(ConnectionMode::Disconnect.as_u8(), Ordering::SeqCst);
2365
2366 let outcome = inner.reconnect(None).await.unwrap();
2367
2368 assert_eq!(outcome, ReconnectOutcome::Aborted);
2369 server.abort();
2370 }
2371
2372 #[rstest]
2373 #[tokio::test]
2374 async fn test_reconnect_outcome_is_reconnected() {
2375 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
2376 let port = listener.local_addr().unwrap().port();
2377 let server = tokio::spawn(async move {
2378 let (_first, _) = listener.accept().await.unwrap();
2379 let (_second, _) = listener.accept().await.unwrap();
2380 std::future::pending::<()>().await;
2381 });
2382 let mut inner = SocketClientInner::connect_url(reconnect_test_config(port), None)
2383 .await
2384 .unwrap();
2385 inner
2386 .connection_mode
2387 .store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
2388
2389 let outcome = inner.reconnect(None).await.unwrap();
2390
2391 assert_eq!(outcome, ReconnectOutcome::Reconnected);
2392 assert_eq!(
2393 ConnectionMode::from_atomic(&inner.connection_mode),
2394 ConnectionMode::Active
2395 );
2396 server.abort();
2397 }
2398
2399 #[rstest]
2400 #[tokio::test]
2401 async fn test_reconnect_aborts_retired_reader() {
2402 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
2403 let port = listener.local_addr().unwrap().port();
2404 let mut inner = SocketClientInner::connect_url(reconnect_test_config(port), None)
2405 .await
2406 .unwrap();
2407 let (_peer, _) = listener.accept().await.unwrap();
2408 let (dropped_tx, dropped_rx) = tokio::sync::oneshot::channel::<()>();
2409
2410 let pending_reader = tokio::spawn(async move {
2411 let _sender = dropped_tx;
2412 std::future::pending::<()>().await;
2413 });
2414
2415 let previous_reader = std::mem::replace(&mut inner.read_task, pending_reader);
2416 previous_reader.abort();
2417 let _ = previous_reader.await;
2418 let old_fence = inner.read_fence.clone();
2419 inner
2420 .connection_mode
2421 .store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
2422
2423 let outcome = inner.reconnect(None).await.unwrap();
2424 let reader_result = tokio::time::timeout(Duration::from_secs(1), dropped_rx)
2425 .await
2426 .expect("retired reader task must be canceled");
2427
2428 assert_eq!(outcome, ReconnectOutcome::Reconnected);
2429 assert!(
2430 reader_result.is_err(),
2431 "retired reader must drop its sender without sending"
2432 );
2433 assert!(!old_fence.is_valid());
2434 assert!(inner.read_fence.is_valid());
2435 assert_eq!(
2436 ConnectionMode::from_atomic(&inner.connection_mode),
2437 ConnectionMode::Active
2438 );
2439 }
2440
2441 #[rstest]
2442 #[tokio::test]
2443 async fn test_inner_drop_invalidates_read_fence_and_aborts_tasks() {
2444 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
2445 let port = listener.local_addr().unwrap().port();
2446 let server = tokio::spawn(async move {
2447 let (_socket, _) = listener.accept().await.unwrap();
2448 std::future::pending::<()>().await;
2449 });
2450 let mut config = reconnect_test_config(port);
2451 config.heartbeat = Some(SocketHeartbeat {
2452 interval_secs: 60,
2453 payload: b"ping\r\n".to_vec(),
2454 });
2455 let inner = SocketClientInner::connect_url(config, None).await.unwrap();
2456 let read_fence = inner.read_fence.clone();
2457 let read_abort = inner.read_task.abort_handle();
2458 let write_abort = inner.write_task.abort_handle();
2459 let heartbeat_abort = inner
2460 .heartbeat_task
2461 .as_ref()
2462 .expect("heartbeat task should be spawned for a configured heartbeat")
2463 .abort_handle();
2464
2465 assert!(read_fence.is_valid(), "read fence should start valid");
2466 assert!(
2467 !read_abort.is_finished(),
2468 "read task should be running before drop"
2469 );
2470 assert!(
2471 !write_abort.is_finished(),
2472 "write task should be running before drop"
2473 );
2474 assert!(
2475 !heartbeat_abort.is_finished(),
2476 "heartbeat task should be running before drop"
2477 );
2478
2479 drop(inner);
2480 wait_until_async(
2481 || async {
2482 read_abort.is_finished()
2483 && write_abort.is_finished()
2484 && heartbeat_abort.is_finished()
2485 },
2486 TEST_TIMEOUT,
2487 )
2488 .await;
2489
2490 assert!(!read_fence.is_valid(), "read fence was not invalidated");
2491 assert!(read_abort.is_finished(), "read task was not aborted");
2492 assert!(write_abort.is_finished(), "write task was not aborted");
2493 assert!(
2494 heartbeat_abort.is_finished(),
2495 "heartbeat task was not aborted"
2496 );
2497 server.abort();
2498 }
2499
2500 struct ScriptedReader {
2501 first: Option<Vec<u8>>,
2502 remainder: Arc<Mutex<Option<Vec<u8>>>>,
2503 pending_tx: Option<oneshot::Sender<()>>,
2504 waker: Arc<Mutex<Option<Waker>>>,
2505 }
2506
2507 impl AsyncRead for ScriptedReader {
2508 fn poll_read(
2509 mut self: Pin<&mut Self>,
2510 cx: &mut Context<'_>,
2511 buf: &mut ReadBuf<'_>,
2512 ) -> Poll<std::io::Result<()>> {
2513 if let Some(first) = self.first.take() {
2514 buf.put_slice(&first);
2515 return Poll::Ready(Ok(()));
2516 }
2517
2518 if let Some(remainder) = self.remainder.lock().take() {
2519 buf.put_slice(&remainder);
2520 return Poll::Ready(Ok(()));
2521 }
2522
2523 if let Some(tx) = self.pending_tx.take() {
2524 let _ = tx.send(());
2525 }
2526 *self.waker.lock() = Some(cx.waker().clone());
2527 Poll::Pending
2528 }
2529 }
2530
2531 struct LimitProbeReader {
2532 remaining: usize,
2533 read_after_limit: Arc<AtomicBool>,
2534 }
2535
2536 impl AsyncRead for LimitProbeReader {
2537 fn poll_read(
2538 mut self: Pin<&mut Self>,
2539 _cx: &mut Context<'_>,
2540 buf: &mut ReadBuf<'_>,
2541 ) -> Poll<std::io::Result<()>> {
2542 if self.remaining == 0 {
2543 self.read_after_limit.store(true, Ordering::SeqCst);
2544 return Poll::Ready(Ok(()));
2545 }
2546
2547 let read_len = self.remaining.min(buf.remaining());
2548 buf.initialize_unfilled_to(read_len).fill(b'x');
2549 buf.advance(read_len);
2550 self.remaining -= read_len;
2551 Poll::Ready(Ok(()))
2552 }
2553 }
2554
2555 struct BackpressuredWriter {
2556 stream: DuplexStream,
2557 pending_tx: Option<oneshot::Sender<()>>,
2558 }
2559
2560 impl AsyncWrite for BackpressuredWriter {
2561 fn poll_write(
2562 mut self: Pin<&mut Self>,
2563 cx: &mut Context<'_>,
2564 buf: &[u8],
2565 ) -> Poll<std::io::Result<usize>> {
2566 let result = Pin::new(&mut self.stream).poll_write(cx, buf);
2567 if result.is_pending()
2568 && let Some(tx) = self.pending_tx.take()
2569 {
2570 let _ = tx.send(());
2571 }
2572 result
2573 }
2574
2575 fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
2576 Pin::new(&mut self.stream).poll_flush(cx)
2577 }
2578
2579 fn poll_shutdown(
2580 mut self: Pin<&mut Self>,
2581 cx: &mut Context<'_>,
2582 ) -> Poll<std::io::Result<()>> {
2583 Pin::new(&mut self.stream).poll_shutdown(cx)
2584 }
2585 }
2586
2587 struct FailingWriter;
2588
2589 impl AsyncWrite for FailingWriter {
2590 fn poll_write(
2591 self: Pin<&mut Self>,
2592 _cx: &mut Context<'_>,
2593 _buf: &[u8],
2594 ) -> Poll<std::io::Result<usize>> {
2595 Poll::Ready(Err(std::io::Error::new(
2596 std::io::ErrorKind::BrokenPipe,
2597 "test writer failure",
2598 )))
2599 }
2600
2601 fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
2602 Poll::Ready(Ok(()))
2603 }
2604
2605 fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
2606 Poll::Ready(Ok(()))
2607 }
2608 }
2609
2610 struct RecordingWriter {
2611 bytes: Arc<Mutex<Vec<u8>>>,
2612 }
2613
2614 impl AsyncWrite for RecordingWriter {
2615 fn poll_write(
2616 self: Pin<&mut Self>,
2617 _cx: &mut Context<'_>,
2618 buf: &[u8],
2619 ) -> Poll<std::io::Result<usize>> {
2620 self.bytes.lock().extend_from_slice(buf);
2621 Poll::Ready(Ok(buf.len()))
2622 }
2623
2624 fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
2625 Poll::Ready(Ok(()))
2626 }
2627
2628 fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
2629 Poll::Ready(Ok(()))
2630 }
2631 }
2632
2633 fn test_socket_client(
2634 connection_state: Arc<AtomicU8>,
2635 state_notify: Arc<tokio::sync::Notify>,
2636 controller_task: tokio::task::JoinHandle<()>,
2637 ) -> SocketClient {
2638 let (writer_tx, _writer_rx) = writer::channel(DEFAULT_WRITER_CAPACITY);
2639 let controller_lifecycle = Arc::new(ControllerLifecycle::new());
2640 let abort_handle = controller_task.abort_handle();
2641 controller_lifecycle.set_abort(move || abort_handle.abort());
2642
2643 SocketClient {
2644 controller_task,
2645 connection_mode: connection_state,
2646 state_notify,
2647 connect_timeout: Duration::from_secs(1),
2648 writer_tx,
2649 controller_lifecycle,
2650 controller_notify: Arc::new(tokio::sync::Notify::new()),
2651 state_sink: None,
2652 }
2653 }
2654
2655 #[rstest]
2656 #[tokio::test]
2657 async fn test_send_reports_writer_buffer_full() {
2658 let state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
2659 let notify = Arc::new(tokio::sync::Notify::new());
2660
2661 let task = tokio::spawn(std::future::pending::<()>());
2662 let mut client = test_socket_client(state, notify, task);
2663 let (writer_tx, mut writer_rx) = writer::channel(1);
2664 client.writer_tx = writer_tx;
2665
2666 client.send_bytes(vec![1]).await.unwrap();
2667 assert!(matches!(
2668 client.send_bytes(vec![2]).await,
2669 Err(SendError::BufferFull)
2670 ));
2671 let (command, permit) = writer_rx.recv().await.unwrap();
2672 assert!(matches!(command, WriterCommand::Send(data) if data == vec![1]));
2673 assert!(matches!(
2674 client.send_bytes(vec![2]).await,
2675 Err(SendError::BufferFull)
2676 ));
2677 drop(permit);
2678 client.send_bytes(vec![3]).await.unwrap();
2679 let (command, _) = writer_rx.recv().await.unwrap();
2680 assert!(matches!(command, WriterCommand::Send(data) if data == vec![3]));
2681 }
2682
2683 #[rstest]
2684 #[tokio::test]
2685 async fn test_reconnect_handle_is_closed_after_client_drop() {
2686 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
2687 let state_notify = Arc::new(tokio::sync::Notify::new());
2688 let controller_task = tokio::spawn(std::future::pending::<()>());
2689 let client = test_socket_client(connection_state, state_notify, controller_task);
2690 let handle = client.reconnect_handle();
2691
2692 drop(client);
2693
2694 assert_eq!(handle.request_reconnect(), ReconnectRequestOutcome::Closed);
2695 }
2696
2697 #[rstest]
2698 #[tokio::test]
2699 async fn test_concurrent_drop_defers_controller_abort_until_request_completes() {
2700 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
2701 let state_notify = Arc::new(tokio::sync::Notify::new());
2702
2703 let controller_task = tokio::spawn(std::future::pending::<()>());
2704 let controller_abort = controller_task.abort_handle();
2705 let mut client = test_socket_client(connection_state, state_notify, controller_task);
2706 let callback_release = Arc::new((Mutex::new(false), Condvar::new()));
2707 let callback_release_guard = CondvarReleaseGuard::new(Arc::clone(&callback_release));
2708 let callback_release_clone = Arc::clone(&callback_release);
2709 let (callback_entered_tx, callback_entered_rx) = std::sync::mpsc::channel();
2710 let states = Arc::new(Mutex::new(Vec::new()));
2711 let states_clone = Arc::clone(&states);
2712 client.state_sink = Some(SocketStateSink::new(move |state| {
2713 states_clone.lock().push(state);
2714 callback_entered_tx.send(()).unwrap();
2715 let (lock, condvar) = callback_release_clone.as_ref();
2716 let mut released = lock.lock();
2717
2718 while !*released {
2719 condvar.wait(&mut released);
2720 }
2721 }));
2722 let handle = client.reconnect_handle();
2723 let surviving_handle = handle.clone();
2724 let controller_notify = Arc::clone(&handle.controller_notify);
2725 let request_thread = std::thread::spawn(move || handle.request_reconnect());
2726
2727 recv_rendezvous(callback_entered_rx, "socket state callback entry").await;
2728 let (drop_finished_tx, drop_finished_rx) = std::sync::mpsc::channel();
2729 let drop_thread = std::thread::spawn(move || {
2730 drop(client);
2731 drop_finished_tx.send(()).unwrap();
2732 });
2733
2734 recv_rendezvous(drop_finished_rx, "client drop").await;
2735 drop_thread.join().unwrap();
2736 assert!(!controller_abort.is_finished());
2737
2738 callback_release_guard.release();
2739 assert_eq!(
2740 request_thread.join().unwrap(),
2741 ReconnectRequestOutcome::Accepted
2742 );
2743 tokio::time::timeout(Duration::from_millis(10), controller_notify.notified())
2744 .await
2745 .expect("accepted request should notify before deferred controller abort");
2746 wait_until_async(|| async { controller_abort.is_finished() }, TEST_TIMEOUT).await;
2747
2748 assert_eq!(
2749 surviving_handle.request_reconnect(),
2750 ReconnectRequestOutcome::Closed
2751 );
2752 assert_eq!(*states.lock(), vec![SocketState::Disconnected]);
2753 assert!(
2754 tokio::time::timeout(Duration::from_millis(10), controller_notify.notified())
2755 .await
2756 .is_err(),
2757 "closed request should not notify controller",
2758 );
2759 }
2760
2761 #[rstest]
2762 #[tokio::test]
2763 async fn test_reconnect_state_callback_can_drop_client() {
2764 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
2765 let state_notify = Arc::new(tokio::sync::Notify::new());
2766
2767 let controller_task = tokio::spawn(std::future::pending::<()>());
2768 let controller_abort = controller_task.abort_handle();
2769 let mut client = test_socket_client(connection_state, state_notify, controller_task);
2770 let client_slot = Arc::new(Mutex::new(None));
2771 let client_slot_callback = Arc::clone(&client_slot);
2772 let states = Arc::new(Mutex::new(Vec::new()));
2773 let states_callback = Arc::clone(&states);
2774 client.state_sink = Some(SocketStateSink::new(move |state| {
2775 states_callback.lock().push(state);
2776 drop(client_slot_callback.lock().take());
2777 }));
2778 let handle = client.reconnect_handle();
2779 let controller_notify = Arc::clone(&handle.controller_notify);
2780 *client_slot.lock() = Some(client);
2781 let (result_tx, result_rx) = std::sync::mpsc::channel();
2782 std::thread::spawn(move || result_tx.send(handle.request_reconnect()).unwrap());
2783
2784 assert_eq!(
2785 recv_rendezvous(result_rx, "reconnect request").await,
2786 ReconnectRequestOutcome::Accepted
2787 );
2788 tokio::time::timeout(Duration::from_millis(10), controller_notify.notified())
2789 .await
2790 .expect("accepted request should notify before deferred controller abort");
2791 wait_until_async(|| async { controller_abort.is_finished() }, TEST_TIMEOUT).await;
2792
2793 assert!(client_slot.lock().is_none());
2794 assert_eq!(*states.lock(), vec![SocketState::Disconnected]);
2795 }
2796
2797 #[rstest]
2798 #[tokio::test]
2799 async fn test_manual_reconnect_uses_controller_path() {
2800 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
2801 let port = listener.local_addr().unwrap().port();
2802 let (first_accepted_tx, first_accepted_rx) = oneshot::channel();
2803 let (second_accepted_tx, second_accepted_rx) = oneshot::channel();
2804 let (payload_tx, payload_rx) = oneshot::channel();
2805 let (release_tx, release_rx) = oneshot::channel();
2806
2807 let server = task::spawn(async move {
2808 let (_first, _) = listener.accept().await.unwrap();
2809 first_accepted_tx.send(()).unwrap();
2810
2811 let (mut second, _) = listener.accept().await.unwrap();
2812 second_accepted_tx.send(()).unwrap();
2813 let mut payload = [0_u8; 8];
2814 second.read_exact(&mut payload).await.unwrap();
2815 payload_tx.send(payload).unwrap();
2816 let _ = release_rx.await;
2817 });
2818 let states = Arc::new(Mutex::new(Vec::new()));
2819 let states_callback = Arc::clone(&states);
2820 let sink = SocketStateSink::new(move |state| {
2821 states_callback.lock().push(state);
2822 });
2823 let callback_count = Arc::new(AtomicUsize::new(0));
2824 let callback_count_clone = Arc::clone(&callback_count);
2825 let post_reconnection = Arc::new(move || {
2826 callback_count_clone.fetch_add(1, Ordering::SeqCst);
2827 });
2828 let client = SocketClient::builder()
2829 .config(reconnect_test_config(port))
2830 .post_reconnection(post_reconnection)
2831 .state_sink(sink)
2832 .connect()
2833 .await
2834 .unwrap();
2835 first_accepted_rx.await.unwrap();
2836 let handle = client.reconnect_handle();
2837
2838 assert!(client.request_reconnect());
2839 assert_eq!(
2840 handle.request_reconnect(),
2841 ReconnectRequestOutcome::AlreadyReconnecting
2842 );
2843 assert!(!client.request_reconnect());
2844 assert_eq!(
2845 *states.lock(),
2846 vec![SocketState::Connected, SocketState::Disconnected]
2847 );
2848
2849 tokio::time::timeout(TEST_TIMEOUT, second_accepted_rx)
2850 .await
2851 .expect("controller should establish a replacement connection")
2852 .unwrap();
2853 wait_until_async(
2854 || async { client.is_active() && callback_count.load(Ordering::SeqCst) == 1 },
2855 TEST_TIMEOUT,
2856 )
2857 .await;
2858 client.send_bytes(b"manual".to_vec()).await.unwrap();
2859
2860 assert_eq!(
2861 tokio::time::timeout(TEST_TIMEOUT, payload_rx)
2862 .await
2863 .expect("replacement connection should receive the framed payload")
2864 .unwrap(),
2865 *b"manual\r\n"
2866 );
2867 assert_eq!(callback_count.load(Ordering::SeqCst), 1);
2868 assert_eq!(
2869 *states.lock(),
2870 vec![
2871 SocketState::Connected,
2872 SocketState::Disconnected,
2873 SocketState::Connected,
2874 ]
2875 );
2876
2877 client.close().await;
2878 release_tx.send(()).unwrap();
2879 server.await.unwrap();
2880 assert_eq!(callback_count.load(Ordering::SeqCst), 1);
2881 assert_eq!(states.lock().len(), 3);
2882 }
2883
2884 #[rstest]
2885 #[case(ConnectionMode::Disconnect)]
2886 #[case(ConnectionMode::Closed)]
2887 #[tokio::test]
2888 async fn test_send_bytes_rejects_terminal_state(#[case] mode: ConnectionMode) {
2889 let connection_state = Arc::new(AtomicU8::new(mode.as_u8()));
2890 let state_notify = Arc::new(tokio::sync::Notify::new());
2891 let controller_task = tokio::spawn(std::future::pending::<()>());
2892 let client = test_socket_client(connection_state, state_notify, controller_task);
2893
2894 let result = client.send_bytes(b"terminal".to_vec()).await;
2895
2896 assert!(matches!(result, Err(SendError::Closed)));
2897 }
2898
2899 #[rstest]
2900 #[tokio::test(start_paused = true)]
2901 async fn test_close_sets_closed_after_controller_was_aborted() {
2902 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
2903 let state_notify = Arc::new(tokio::sync::Notify::new());
2904
2905 let controller_task = tokio::spawn(std::future::pending::<()>());
2906 controller_task.abort();
2907 let client = test_socket_client(connection_state, state_notify, controller_task);
2908
2909 client.close_with_timeout(Duration::from_millis(1)).await;
2910
2911 assert!(client.is_closed());
2912 }
2913
2914 #[rstest]
2915 #[tokio::test]
2916 async fn test_reconnect_exhaustion_closes_client() {
2917 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
2918 let port = listener.local_addr().unwrap().port();
2919 let (accepted_tx, accepted_rx) = oneshot::channel();
2920 let (release_tx, release_rx) = oneshot::channel();
2921
2922 let server = task::spawn(async move {
2923 let (socket, _) = listener.accept().await.unwrap();
2924 let _ = accepted_tx.send(());
2925 let _ = release_rx.await;
2926 drop(socket);
2927 drop(listener);
2928 });
2929
2930 let config = SocketConfig {
2931 url: format!("127.0.0.1:{port}"),
2932 mode: Mode::Plain,
2933 suffix: b"\r\n".to_vec(),
2934 message_handler: None,
2935 heartbeat: None,
2936 connect_timeout_ms: Some(100),
2937 reconnect_delay_initial_ms: Some(1),
2938 reconnect_delay_max_ms: Some(1),
2939 reconnect_backoff_factor: Some(1.0),
2940 reconnect_jitter_ms: Some(0),
2941 connection_max_retries: Some(1),
2942 reconnect_max_attempts: Some(1),
2943 heartbeat_timeout_secs: None,
2944 writer_capacity: None,
2945 certs_dir: None,
2946 };
2947
2948 let states = Arc::new(Mutex::new(Vec::new()));
2949 let states_callback = Arc::clone(&states);
2950 let sink = SocketStateSink::new(move |state| {
2951 states_callback.lock().push(state);
2952 });
2953 let client = SocketClient::builder()
2954 .config(config)
2955 .state_sink(sink)
2956 .connect()
2957 .await
2958 .unwrap();
2959 accepted_rx.await.unwrap();
2960 release_tx.send(()).unwrap();
2961
2962 wait_until_async(|| async { client.is_closed() }, Duration::from_secs(5)).await;
2963 assert!(client.is_closed());
2964 assert_eq!(
2965 *states.lock(),
2966 vec![SocketState::Connected, SocketState::Disconnected]
2967 );
2968
2969 client.close().await;
2970 assert!(client.is_closed());
2971 assert_eq!(states.lock().len(), 2);
2972 server.await.unwrap();
2973 }
2974
2975 #[rstest]
2976 #[tokio::test]
2977 async fn test_drop_suppresses_socket_state_event() {
2978 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
2979 let port = listener.local_addr().unwrap().port();
2980 let (accepted_tx, accepted_rx) = oneshot::channel();
2981
2982 let server = task::spawn(async move {
2983 let (_socket, _) = listener.accept().await.unwrap();
2984 let _ = accepted_tx.send(());
2985 std::future::pending::<()>().await;
2986 });
2987
2988 let config = SocketConfig {
2989 url: format!("127.0.0.1:{port}"),
2990 mode: Mode::Plain,
2991 suffix: b"\r\n".to_vec(),
2992 message_handler: None,
2993 heartbeat: None,
2994 connect_timeout_ms: Some(100),
2995 reconnect_delay_initial_ms: Some(1),
2996 reconnect_delay_max_ms: Some(1),
2997 reconnect_backoff_factor: Some(1.0),
2998 reconnect_jitter_ms: Some(0),
2999 connection_max_retries: Some(1),
3000 reconnect_max_attempts: Some(1),
3001 heartbeat_timeout_secs: None,
3002 writer_capacity: None,
3003 certs_dir: None,
3004 };
3005 let states = Arc::new(Mutex::new(Vec::new()));
3006 let states_callback = Arc::clone(&states);
3007 let sink = SocketStateSink::new(move |state| {
3008 states_callback.lock().push(state);
3009 });
3010 let client = SocketClient::builder()
3011 .config(config)
3012 .state_sink(sink)
3013 .connect()
3014 .await
3015 .unwrap();
3016 accepted_rx.await.unwrap();
3017
3018 drop(client);
3019 sleep(Duration::from_millis(25)).await;
3020
3021 assert_eq!(*states.lock(), vec![SocketState::Connected]);
3022 server.abort();
3023 }
3024
3025 #[rstest]
3026 #[tokio::test]
3027 async fn test_graceful_close_is_idempotent() {
3028 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3029 let port = listener.local_addr().unwrap().port();
3030 let (accepted_tx, accepted_rx) = oneshot::channel();
3031
3032 let server = task::spawn(async move {
3033 let (_socket, _) = listener.accept().await.unwrap();
3034 let _ = accepted_tx.send(());
3035 std::future::pending::<()>().await;
3036 });
3037 let config = SocketConfig {
3038 url: format!("127.0.0.1:{port}"),
3039 mode: Mode::Plain,
3040 suffix: b"\r\n".to_vec(),
3041 message_handler: None,
3042 heartbeat: None,
3043 connect_timeout_ms: None,
3044 reconnect_delay_initial_ms: None,
3045 reconnect_delay_max_ms: None,
3046 reconnect_backoff_factor: None,
3047 reconnect_jitter_ms: None,
3048 connection_max_retries: None,
3049 reconnect_max_attempts: None,
3050 heartbeat_timeout_secs: None,
3051 writer_capacity: None,
3052 certs_dir: None,
3053 };
3054 let client = SocketClient::builder()
3055 .config(config)
3056 .connect()
3057 .await
3058 .unwrap();
3059 accepted_rx.await.unwrap();
3060
3061 client.close().await;
3062 client.close().await;
3063
3064 assert!(client.is_closed());
3065 server.abort();
3066 }
3067
3068 #[rstest]
3069 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3070 async fn test_read_loop_drops_remaining_frames_after_session_replaced() {
3071 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
3072 let read_fence = ReadSessionFence::new();
3073 let received = Arc::new(Mutex::new(Vec::<Vec<u8>>::new()));
3074 let received_clone = Arc::clone(&received);
3075 let handler_release = Arc::new((Mutex::new(false), Condvar::new()));
3076 let handler_release_guard = CondvarReleaseGuard::new(Arc::clone(&handler_release));
3077 let handler_release_clone = Arc::clone(&handler_release);
3078 let (handler_entered_tx, handler_entered_rx) = std::sync::mpsc::channel();
3079 let handler: TcpMessageHandler = Arc::new(move |data| {
3080 received_clone.lock().push(data.to_vec());
3081 handler_entered_tx.send(()).unwrap();
3082 let (lock, condvar) = handler_release_clone.as_ref();
3083 let mut released = lock.lock();
3084
3085 while !*released {
3086 condvar.wait(&mut released);
3087 }
3088 });
3089 let reader = ScriptedReader {
3090 first: Some(b"first\r\nsecond\r\n".to_vec()),
3091 remainder: Arc::new(Mutex::new(None)),
3092 pending_tx: None,
3093 waker: Arc::new(Mutex::new(None)),
3094 };
3095 let read_task = SocketClientInner::spawn_read_task(
3096 Arc::clone(&connection_state),
3097 read_fence.clone(),
3098 reader,
3099 Some(handler),
3100 b"\r\n".to_vec(),
3101 None,
3102 );
3103
3104 recv_rendezvous(handler_entered_rx, "socket first-handler entry").await;
3105 connection_state.store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
3106 read_fence.invalidate();
3107 read_task.abort();
3108 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
3109
3110 handler_release_guard.release();
3111 await_task_termination(read_task, "old socket read task", true).await;
3112
3113 assert_eq!(received.lock().as_slice(), &[b"first".to_vec()]);
3114 }
3115
3116 #[rstest]
3117 #[tokio::test(start_paused = true)]
3118 async fn test_read_loop_drops_partial_old_session_frame() {
3119 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
3120 let old_read_fence = ReadSessionFence::new();
3121 let received = Arc::new(Mutex::new(Vec::<Vec<u8>>::new()));
3122 let received_clone = Arc::clone(&received);
3123 let handler: TcpMessageHandler =
3124 Arc::new(move |data| received_clone.lock().push(data.to_vec()));
3125 let (pending_tx, pending_rx) = oneshot::channel();
3126 let remainder = Arc::new(Mutex::new(None));
3127 let waker = Arc::new(Mutex::new(None));
3128 let reader = ScriptedReader {
3129 first: Some(b"old".to_vec()),
3130 remainder: Arc::clone(&remainder),
3131 pending_tx: Some(pending_tx),
3132 waker: Arc::clone(&waker),
3133 };
3134
3135 let read_task = SocketClientInner::spawn_read_task(
3136 Arc::clone(&connection_state),
3137 old_read_fence.clone(),
3138 reader,
3139 Some(handler),
3140 b"\r\n".to_vec(),
3141 None,
3142 );
3143
3144 pending_rx.await.unwrap();
3145 connection_state.store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
3146 old_read_fence.invalidate();
3147 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
3148 *remainder.lock() = Some(b"\r\nnew\r\n".to_vec());
3149 waker.lock().take().unwrap().wake();
3150 read_task.await.unwrap();
3151
3152 let fresh_reader = ScriptedReader {
3153 first: Some(b"new\r\n".to_vec()),
3154 remainder: Arc::new(Mutex::new(None)),
3155 pending_tx: None,
3156 waker: Arc::new(Mutex::new(None)),
3157 };
3158 let fresh_connection_state = Arc::clone(&connection_state);
3159 let fresh_received = Arc::clone(&received);
3160 let fresh_handler: TcpMessageHandler = Arc::new(move |data| {
3161 fresh_received.lock().push(data.to_vec());
3162 fresh_connection_state.store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
3163 });
3164 SocketClientInner::spawn_read_task(
3165 Arc::clone(&connection_state),
3166 ReadSessionFence::new(),
3167 fresh_reader,
3168 Some(fresh_handler),
3169 b"\r\n".to_vec(),
3170 None,
3171 )
3172 .await
3173 .unwrap();
3174
3175 assert_eq!(received.lock().as_slice(), &[b"new".to_vec()]);
3176 }
3177
3178 #[rstest]
3179 #[tokio::test(start_paused = true)]
3180 async fn test_read_loop_stops_when_first_handler_ends_session() {
3181 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
3182 let handler_state = Arc::clone(&connection_state);
3183 let received = Arc::new(Mutex::new(Vec::<Vec<u8>>::new()));
3184 let received_clone = Arc::clone(&received);
3185 let handler: TcpMessageHandler = Arc::new(move |data| {
3186 received_clone.lock().push(data.to_vec());
3187 handler_state.store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
3188 });
3189 let (pending_tx, _pending_rx) = oneshot::channel();
3190 let reader = ScriptedReader {
3191 first: Some(b"first\r\nsecond\r\n".to_vec()),
3192 remainder: Arc::new(Mutex::new(None)),
3193 pending_tx: Some(pending_tx),
3194 waker: Arc::new(Mutex::new(None)),
3195 };
3196
3197 SocketClientInner::spawn_read_task(
3198 connection_state,
3199 ReadSessionFence::new(),
3200 reader,
3201 Some(handler),
3202 b"\r\n".to_vec(),
3203 None,
3204 )
3205 .await
3206 .unwrap();
3207
3208 assert_eq!(received.lock().as_slice(), &[b"first".to_vec()]);
3209 }
3210
3211 #[tokio::test]
3212 async fn test_read_loop_closes_when_unframed_buffer_exceeds_limit() {
3213 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
3214 let read_after_limit = Arc::new(AtomicBool::new(false));
3215 let reader = LimitProbeReader {
3216 remaining: MAX_READ_BUFFER_BYTES + 1,
3217 read_after_limit: Arc::clone(&read_after_limit),
3218 };
3219
3220 SocketClientInner::run_read_loop(
3221 connection_state,
3222 ReadSessionFence::new(),
3223 reader,
3224 None,
3225 b"\r\n".to_vec(),
3226 None,
3227 Duration::from_secs(1),
3228 )
3229 .await;
3230
3231 assert!(!read_after_limit.load(Ordering::SeqCst));
3232 }
3233
3234 #[rstest]
3235 #[tokio::test(start_paused = true)]
3236 async fn test_stalled_socket_write_sends_reconnect_replay_before_buffer() {
3237 type TestWriter = Pin<Box<dyn AsyncWrite + Send>>;
3238
3239 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
3240 let state_notify = Arc::new(tokio::sync::Notify::new());
3241 let (stream, _non_reading_peer) = tokio::io::duplex(1);
3242 let (pending_tx, pending_rx) = oneshot::channel();
3243 let writer: TestWriter = Box::pin(BackpressuredWriter {
3244 stream,
3245 pending_tx: Some(pending_tx),
3246 });
3247
3248 let (writer_tx, writer_rx) = writer::channel::<WriterCommand<TestWriter>>(2);
3249 let states = Arc::new(Mutex::new(Vec::new()));
3250 let states_callback = Arc::clone(&states);
3251 let sink = SocketStateSink::new(move |state| {
3252 states_callback.lock().push(state);
3253 });
3254 let write_task = SocketClientInner::spawn_write_task(
3255 Arc::clone(&connection_state),
3256 Arc::clone(&state_notify),
3257 writer,
3258 writer_rx,
3259 b"\r\n".to_vec(),
3260 Some(sink),
3261 );
3262
3263 writer_tx
3264 .send(WriterCommand::Send(Bytes::from_static(b"complete-message")))
3265 .unwrap();
3266 pending_rx.await.unwrap();
3267 writer_tx
3268 .send(WriterCommand::Send(Bytes::from_static(b"second-message")))
3269 .unwrap();
3270 assert_eq!(
3271 ConnectionMode::from_atomic(&connection_state),
3272 ConnectionMode::Active
3273 );
3274
3275 for _ in 0..100 {
3276 assert!(matches!(
3277 writer_tx.send(WriterCommand::Send(Bytes::from_static(b"overflow"))),
3278 Err(SendError::BufferFull)
3279 ));
3280 }
3281
3282 tokio::time::advance(Duration::from_secs(WRITE_TIMEOUT_SECS)).await;
3283 tokio::task::yield_now().await;
3284 assert_eq!(
3285 ConnectionMode::from_atomic(&connection_state),
3286 ConnectionMode::Reconnect
3287 );
3288 assert!(matches!(
3289 writer_tx.send(WriterCommand::Send(Bytes::from_static(b"overflow"))),
3290 Err(SendError::BufferFull)
3291 ));
3292
3293 let recorded = Arc::new(Mutex::new(Vec::new()));
3294 let new_writer: TestWriter = Box::pin(RecordingWriter {
3295 bytes: Arc::clone(&recorded),
3296 });
3297 let (update_tx, update_rx) = oneshot::channel();
3298 writer_tx
3299 .send_update(WriterCommand::UpdateWithReplay(
3300 new_writer,
3301 vec![Bytes::from_static(b"authentication")],
3302 update_tx,
3303 ))
3304 .unwrap();
3305
3306 tokio::time::advance(Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS)).await;
3307 yield_now().await;
3308
3309 assert!(
3310 tokio::time::timeout(Duration::from_secs(10), update_rx)
3311 .await
3312 .expect("writer update should not remain queued behind a stalled write")
3313 .unwrap(),
3314 "writer should report successful replay"
3315 );
3316 assert_eq!(
3317 ConnectionMode::from_atomic(&connection_state),
3318 ConnectionMode::Reconnect
3319 );
3320 assert_eq!(
3321 recorded.lock().as_slice(),
3322 b"authentication\r\ncomplete-message\r\nsecond-message\r\n"
3323 );
3324
3325 writer_tx
3326 .send(WriterCommand::Send(Bytes::from_static(b"second-message")))
3327 .unwrap();
3328
3329 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
3330 state_notify.notify_waiters();
3331 drop(writer_tx);
3332 write_task.await.unwrap();
3333
3334 assert_eq!(*states.lock(), vec![SocketState::Disconnected]);
3335 }
3336
3337 #[rstest]
3338 #[tokio::test(start_paused = true)]
3339 async fn test_send_after_writer_update_drains_when_reconnect_completes() {
3340 type TestWriter = Pin<Box<dyn AsyncWrite + Send>>;
3341
3342 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
3343 let state_notify = Arc::new(tokio::sync::Notify::new());
3344 let initial_writer: TestWriter = Box::pin(RecordingWriter {
3345 bytes: Arc::new(Mutex::new(Vec::new())),
3346 });
3347 let (writer_tx, writer_rx) =
3348 writer::channel::<WriterCommand<TestWriter>>(DEFAULT_WRITER_CAPACITY);
3349 let write_task = SocketClientInner::spawn_write_task(
3350 Arc::clone(&connection_state),
3351 Arc::clone(&state_notify),
3352 initial_writer,
3353 writer_rx,
3354 b"\r\n".to_vec(),
3355 None,
3356 );
3357
3358 let recorded = Arc::new(Mutex::new(Vec::new()));
3359 let new_writer: TestWriter = Box::pin(RecordingWriter {
3360 bytes: Arc::clone(&recorded),
3361 });
3362 let (update_tx, update_rx) = oneshot::channel();
3363 writer_tx
3364 .send(WriterCommand::Update(new_writer, update_tx))
3365 .unwrap();
3366 assert!(update_rx.await.unwrap());
3367
3368 writer_tx
3369 .send(WriterCommand::Send(Bytes::from_static(b"late")))
3370 .unwrap();
3371 yield_now().await;
3372
3373 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
3374 state_notify.notify_waiters();
3375 writer_tx
3376 .send(WriterCommand::Send(Bytes::from_static(b"new")))
3377 .unwrap();
3378 tokio::time::advance(Duration::from_millis(
3379 CONNECTION_STATE_CHECK_INTERVAL_MS * 2,
3380 ))
3381 .await;
3382 yield_now().await;
3383
3384 let actual = recorded.lock().clone();
3385 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
3386 state_notify.notify_waiters();
3387 drop(writer_tx);
3388 write_task.await.unwrap();
3389
3390 assert_eq!(actual, b"late\r\nnew\r\n");
3391 }
3392
3393 #[rstest]
3394 #[case::latest_wins(
3395 &["subscription-a", "subscription-b"],
3396 &["subscription-b"],
3397 b"subscription-b\r\n",
3398 )]
3399 #[case::different_replay_drains(
3400 &["subscription-b"],
3401 &["subscription-a"],
3402 b"subscription-a\r\nsubscription-b\r\n",
3403 )]
3404 #[tokio::test]
3405 async fn test_send_or_replay_buffering(
3406 #[case] buffered: &[&str],
3407 #[case] replay: &[&str],
3408 #[case] expected: &[u8],
3409 ) {
3410 type TestWriter = Pin<Box<dyn AsyncWrite + Send>>;
3411
3412 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
3413 let state_notify = Arc::new(tokio::sync::Notify::new());
3414 let initial_writer: TestWriter = Box::pin(RecordingWriter {
3415 bytes: Arc::new(Mutex::new(Vec::new())),
3416 });
3417 let (writer_tx, writer_rx) =
3418 writer::channel::<WriterCommand<TestWriter>>(DEFAULT_WRITER_CAPACITY);
3419 let write_task = SocketClientInner::spawn_write_task(
3420 Arc::clone(&connection_state),
3421 Arc::clone(&state_notify),
3422 initial_writer,
3423 writer_rx,
3424 b"\r\n".to_vec(),
3425 None,
3426 );
3427
3428 for data in buffered {
3429 writer_tx
3430 .send(WriterCommand::SendOrReplay {
3431 key: 7,
3432 data: Bytes::copy_from_slice(data.as_bytes()),
3433 })
3434 .unwrap();
3435 }
3436 let recorded = Arc::new(Mutex::new(Vec::new()));
3437 let new_writer: TestWriter = Box::pin(RecordingWriter {
3438 bytes: Arc::clone(&recorded),
3439 });
3440 let (update_tx, update_rx) = oneshot::channel();
3441 writer_tx
3442 .send(WriterCommand::UpdateWithReplay(
3443 new_writer,
3444 replay
3445 .iter()
3446 .map(|data| Bytes::copy_from_slice(data.as_bytes()))
3447 .collect(),
3448 update_tx,
3449 ))
3450 .unwrap();
3451
3452 assert!(update_rx.await.unwrap());
3453 assert_eq!(recorded.lock().as_slice(), expected);
3454 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
3455 state_notify.notify_waiters();
3456 drop(writer_tx);
3457 write_task.await.unwrap();
3458 }
3459
3460 #[rstest]
3461 #[tokio::test(start_paused = true)]
3462 async fn test_active_reconnect_buffer_write_failure_reconnects_and_retries() {
3463 type TestWriter = Pin<Box<dyn AsyncWrite + Send>>;
3464
3465 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
3466 let state_notify = Arc::new(tokio::sync::Notify::new());
3467 let initial_writer: TestWriter = Box::pin(RecordingWriter {
3468 bytes: Arc::new(Mutex::new(Vec::new())),
3469 });
3470 let (writer_tx, writer_rx) =
3471 writer::channel::<WriterCommand<TestWriter>>(DEFAULT_WRITER_CAPACITY);
3472 let states = Arc::new(Mutex::new(Vec::new()));
3473 let states_callback = Arc::clone(&states);
3474 let sink = SocketStateSink::new(move |state| {
3475 states_callback.lock().push(state);
3476 });
3477 let write_task = SocketClientInner::spawn_write_task(
3478 Arc::clone(&connection_state),
3479 Arc::clone(&state_notify),
3480 initial_writer,
3481 writer_rx,
3482 b"\r\n".to_vec(),
3483 Some(sink),
3484 );
3485
3486 let (update_tx, update_rx) = oneshot::channel();
3487 let failing_writer: TestWriter = Box::pin(FailingWriter);
3488 writer_tx
3489 .send(WriterCommand::Update(failing_writer, update_tx))
3490 .unwrap();
3491 assert!(update_rx.await.unwrap());
3492
3493 writer_tx
3494 .send(WriterCommand::Send(Bytes::from_static(b"late")))
3495 .unwrap();
3496 yield_now().await;
3497 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
3498 state_notify.notify_waiters();
3499 tokio::time::advance(Duration::from_millis(
3500 CONNECTION_STATE_CHECK_INTERVAL_MS * 2,
3501 ))
3502 .await;
3503 yield_now().await;
3504
3505 assert_eq!(
3506 ConnectionMode::from_atomic(&connection_state),
3507 ConnectionMode::Reconnect
3508 );
3509 assert_eq!(*states.lock(), vec![SocketState::Disconnected]);
3510
3511 let recorded = Arc::new(Mutex::new(Vec::new()));
3512 let replacement: TestWriter = Box::pin(RecordingWriter {
3513 bytes: Arc::clone(&recorded),
3514 });
3515 let (retry_tx, retry_rx) = oneshot::channel();
3516 writer_tx
3517 .send(WriterCommand::Update(replacement, retry_tx))
3518 .unwrap();
3519 assert!(retry_rx.await.unwrap());
3520
3521 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
3522 state_notify.notify_waiters();
3523 drop(writer_tx);
3524 write_task.await.unwrap();
3525
3526 assert_eq!(recorded.lock().as_slice(), b"late\r\n");
3527 }
3528
3529 #[rstest]
3530 #[tokio::test(start_paused = true)]
3531 async fn test_failed_send_or_replay_is_not_duplicated_after_replay() {
3532 type TestWriter = Pin<Box<dyn AsyncWrite + Send>>;
3533
3534 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
3535 let state_notify = Arc::new(tokio::sync::Notify::new());
3536 let (stream, _non_reading_peer) = tokio::io::duplex(1);
3537 let (pending_tx, pending_rx) = oneshot::channel();
3538 let writer: TestWriter = Box::pin(BackpressuredWriter {
3539 stream,
3540 pending_tx: Some(pending_tx),
3541 });
3542 let (writer_tx, writer_rx) =
3543 writer::channel::<WriterCommand<TestWriter>>(DEFAULT_WRITER_CAPACITY);
3544 let write_task = SocketClientInner::spawn_write_task(
3545 Arc::clone(&connection_state),
3546 Arc::clone(&state_notify),
3547 writer,
3548 writer_rx,
3549 b"\r\n".to_vec(),
3550 None,
3551 );
3552 let subscription = Bytes::from_static(b"subscription");
3553
3554 writer_tx
3555 .send(WriterCommand::SendOrReplay {
3556 key: 7,
3557 data: subscription.clone(),
3558 })
3559 .unwrap();
3560 pending_rx.await.unwrap();
3561
3562 let recorded = Arc::new(Mutex::new(Vec::new()));
3563 let new_writer: TestWriter = Box::pin(RecordingWriter {
3564 bytes: Arc::clone(&recorded),
3565 });
3566 let (update_tx, update_rx) = oneshot::channel();
3567 writer_tx
3568 .send(WriterCommand::UpdateWithReplay(
3569 new_writer,
3570 vec![subscription],
3571 update_tx,
3572 ))
3573 .unwrap();
3574
3575 tokio::time::advance(Duration::from_secs(WRITE_TIMEOUT_SECS)).await;
3576 yield_now().await;
3577 assert!(update_rx.await.unwrap());
3578 assert_eq!(recorded.lock().as_slice(), b"subscription\r\n");
3579
3580 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
3581 state_notify.notify_waiters();
3582 drop(writer_tx);
3583 write_task.await.unwrap();
3584 }
3585
3586 #[rstest]
3587 #[tokio::test]
3588 async fn test_connect_url_rejects_invalid_reconnect_backoff_before_connect() {
3589 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
3590 listener.set_nonblocking(true).unwrap();
3591 let port = listener.local_addr().unwrap().port();
3592 let config = SocketConfig {
3593 url: format!("127.0.0.1:{port}"),
3594 mode: Mode::Plain,
3595 suffix: b"\r\n".to_vec(),
3596 message_handler: None,
3597 heartbeat: None,
3598 connect_timeout_ms: Some(1_000),
3599 reconnect_delay_initial_ms: Some(50),
3600 reconnect_delay_max_ms: Some(100),
3601 reconnect_backoff_factor: Some(100.1),
3602 reconnect_jitter_ms: Some(0),
3603 connection_max_retries: Some(1),
3604 reconnect_max_attempts: None,
3605 heartbeat_timeout_secs: None,
3606 writer_capacity: None,
3607 certs_dir: None,
3608 };
3609
3610 let error = match SocketClientInner::connect_url(config, None).await {
3611 Ok(_) => panic!("invalid reconnect backoff should be rejected"),
3612 Err(e) => e,
3613 };
3614
3615 assert!(
3616 error.to_string().contains("factor"),
3617 "error should mention the invalid factor, was: {error}"
3618 );
3619 assert_eq!(
3620 listener.accept().unwrap_err().kind(),
3621 std::io::ErrorKind::WouldBlock,
3622 "invalid reconnect backoff must be rejected before connecting"
3623 );
3624 }
3625
3626 #[rstest]
3627 #[tokio::test]
3628 async fn test_reconnect_then_close() {
3629 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3631 let port = listener.local_addr().unwrap().port();
3632
3633 let server = task::spawn(async move {
3635 if let Ok((mut sock, _)) = listener.accept().await {
3636 drop(sock.shutdown());
3637 }
3638 sleep(Duration::from_secs(1)).await;
3640 });
3641
3642 let config = SocketConfig {
3644 url: format!("127.0.0.1:{port}"),
3645 mode: Mode::Plain,
3646 suffix: b"\r\n".to_vec(),
3647 message_handler: None,
3648 heartbeat: None,
3649 connect_timeout_ms: Some(1_000),
3650 reconnect_delay_initial_ms: Some(50),
3651 reconnect_delay_max_ms: Some(100),
3652 reconnect_backoff_factor: Some(1.0),
3653 reconnect_jitter_ms: Some(0),
3654 connection_max_retries: Some(1),
3655 reconnect_max_attempts: None,
3656 heartbeat_timeout_secs: None,
3657 writer_capacity: None,
3658 certs_dir: None,
3659 };
3660
3661 let client = SocketClient::builder()
3663 .config(config.clone())
3664 .connect()
3665 .await
3666 .unwrap();
3667
3668 wait_until_async(
3670 || async { client.is_reconnecting() },
3671 Duration::from_secs(2),
3672 )
3673 .await;
3674
3675 client.close().await;
3677 assert!(client.is_closed());
3678 server.abort();
3679 }
3680
3681 #[rstest]
3682 #[tokio::test]
3683 async fn test_reconnect_state_flips_when_reader_stops() {
3684 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3686 let port = listener.local_addr().unwrap().port();
3687
3688 let server = task::spawn(async move {
3689 if let Ok((sock, _)) = listener.accept().await {
3690 drop(sock);
3691 }
3692 sleep(Duration::from_millis(50)).await;
3694 });
3695
3696 let config = SocketConfig {
3697 url: format!("127.0.0.1:{port}"),
3698 mode: Mode::Plain,
3699 suffix: b"\r\n".to_vec(),
3700 message_handler: None,
3701 heartbeat: None,
3702 connect_timeout_ms: Some(1_000),
3703 reconnect_delay_initial_ms: Some(50),
3704 reconnect_delay_max_ms: Some(100),
3705 reconnect_backoff_factor: Some(1.0),
3706 reconnect_jitter_ms: Some(0),
3707 connection_max_retries: Some(1),
3708 reconnect_max_attempts: None,
3709 heartbeat_timeout_secs: None,
3710 writer_capacity: None,
3711 certs_dir: None,
3712 };
3713
3714 let client = SocketClient::builder()
3715 .config(config)
3716 .connect()
3717 .await
3718 .unwrap();
3719
3720 wait_until_async(
3721 || async { client.is_reconnecting() },
3722 Duration::from_secs(2),
3723 )
3724 .await;
3725
3726 client.close().await;
3727 server.abort();
3728 }
3729
3730 #[rstest]
3731 fn test_parse_socket_url_raw_address() {
3732 let (socket_addr, request_url) =
3734 SocketClientInner::parse_socket_url("example.com:6130", Mode::Tls).unwrap();
3735 assert_eq!(socket_addr, "example.com:6130");
3736 assert_eq!(request_url, "wss://example.com:6130");
3737
3738 let (socket_addr, request_url) =
3740 SocketClientInner::parse_socket_url("localhost:8080", Mode::Plain).unwrap();
3741 assert_eq!(socket_addr, "localhost:8080");
3742 assert_eq!(request_url, "ws://localhost:8080");
3743 }
3744
3745 #[rstest]
3746 fn test_parse_socket_url_with_scheme() {
3747 let (socket_addr, request_url) =
3749 SocketClientInner::parse_socket_url("wss://example.com:443/path", Mode::Tls).unwrap();
3750 assert_eq!(socket_addr, "example.com:443");
3751 assert_eq!(request_url, "wss://example.com:443/path");
3752
3753 let (socket_addr, request_url) =
3755 SocketClientInner::parse_socket_url("ws://localhost:8080", Mode::Plain).unwrap();
3756 assert_eq!(socket_addr, "localhost:8080");
3757 assert_eq!(request_url, "ws://localhost:8080");
3758 }
3759
3760 #[rstest]
3761 fn test_parse_socket_url_default_ports() {
3762 let (socket_addr, _) =
3764 SocketClientInner::parse_socket_url("wss://example.com", Mode::Tls).unwrap();
3765 assert_eq!(socket_addr, "example.com:443");
3766
3767 let (socket_addr, _) =
3769 SocketClientInner::parse_socket_url("ws://example.com", Mode::Plain).unwrap();
3770 assert_eq!(socket_addr, "example.com:80");
3771
3772 let (socket_addr, _) =
3774 SocketClientInner::parse_socket_url("https://example.com", Mode::Tls).unwrap();
3775 assert_eq!(socket_addr, "example.com:443");
3776
3777 let (socket_addr, _) =
3779 SocketClientInner::parse_socket_url("http://example.com", Mode::Plain).unwrap();
3780 assert_eq!(socket_addr, "example.com:80");
3781 }
3782
3783 #[rstest]
3784 fn test_parse_socket_url_unknown_scheme_uses_mode() {
3785 let (socket_addr, _) =
3787 SocketClientInner::parse_socket_url("custom://example.com", Mode::Tls).unwrap();
3788 assert_eq!(socket_addr, "example.com:443");
3789
3790 let (socket_addr, _) =
3791 SocketClientInner::parse_socket_url("custom://example.com", Mode::Plain).unwrap();
3792 assert_eq!(socket_addr, "example.com:80");
3793 }
3794
3795 #[rstest]
3796 fn test_parse_socket_url_ipv6() {
3797 let (socket_addr, request_url) =
3799 SocketClientInner::parse_socket_url("[::1]:8080", Mode::Plain).unwrap();
3800 assert_eq!(socket_addr, "[::1]:8080");
3801 assert_eq!(request_url, "ws://[::1]:8080");
3802
3803 let (socket_addr, _) =
3805 SocketClientInner::parse_socket_url("ws://[::1]:8080", Mode::Plain).unwrap();
3806 assert_eq!(socket_addr, "[::1]:8080");
3807 }
3808
3809 #[rstest]
3810 #[tokio::test]
3811 async fn test_url_parsing_raw_socket_address() {
3812 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3814 let port = listener.local_addr().unwrap().port();
3815
3816 let server = task::spawn(async move {
3817 if let Ok((sock, _)) = listener.accept().await {
3818 drop(sock);
3819 }
3820 sleep(Duration::from_millis(50)).await;
3821 });
3822
3823 let config = SocketConfig {
3824 url: format!("127.0.0.1:{port}"), mode: Mode::Plain,
3826 suffix: b"\r\n".to_vec(),
3827 message_handler: None,
3828 heartbeat: None,
3829 connect_timeout_ms: Some(1_000),
3830 reconnect_delay_initial_ms: Some(50),
3831 reconnect_delay_max_ms: Some(100),
3832 reconnect_backoff_factor: Some(1.0),
3833 reconnect_jitter_ms: Some(0),
3834 connection_max_retries: Some(1),
3835 reconnect_max_attempts: None,
3836 heartbeat_timeout_secs: None,
3837 writer_capacity: None,
3838 certs_dir: None,
3839 };
3840
3841 let client = SocketClient::builder().config(config).connect().await;
3843 assert!(
3844 client.is_ok(),
3845 "Client should connect with raw socket address format"
3846 );
3847
3848 if let Ok(client) = client {
3849 client.close().await;
3850 }
3851 server.abort();
3852 }
3853
3854 #[rstest]
3855 #[tokio::test]
3856 async fn test_url_parsing_with_scheme() {
3857 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3859 let port = listener.local_addr().unwrap().port();
3860
3861 let server = task::spawn(async move {
3862 if let Ok((sock, _)) = listener.accept().await {
3863 drop(sock);
3864 }
3865 sleep(Duration::from_millis(50)).await;
3866 });
3867
3868 let config = SocketConfig {
3869 url: format!("ws://127.0.0.1:{port}"), mode: Mode::Plain,
3871 suffix: b"\r\n".to_vec(),
3872 message_handler: None,
3873 heartbeat: None,
3874 connect_timeout_ms: Some(1_000),
3875 reconnect_delay_initial_ms: Some(50),
3876 reconnect_delay_max_ms: Some(100),
3877 reconnect_backoff_factor: Some(1.0),
3878 reconnect_jitter_ms: Some(0),
3879 connection_max_retries: Some(1),
3880 reconnect_max_attempts: None,
3881 heartbeat_timeout_secs: None,
3882 writer_capacity: None,
3883 certs_dir: None,
3884 };
3885
3886 let client = SocketClient::builder().config(config).connect().await;
3888 assert!(
3889 client.is_ok(),
3890 "Client should connect with URL scheme format"
3891 );
3892
3893 if let Ok(client) = client {
3894 client.close().await;
3895 }
3896 server.abort();
3897 }
3898
3899 #[rstest]
3900 fn test_parse_socket_url_ipv6_with_zone() {
3901 let (socket_addr, request_url) =
3903 SocketClientInner::parse_socket_url("[fe80::1%eth0]:8080", Mode::Plain).unwrap();
3904 assert_eq!(socket_addr, "[fe80::1%eth0]:8080");
3905 assert_eq!(request_url, "ws://[fe80::1%eth0]:8080");
3906
3907 let (socket_addr, request_url) =
3909 SocketClientInner::parse_socket_url("ws://[fe80::1%lo]:9090", Mode::Plain).unwrap();
3910 assert_eq!(socket_addr, "[fe80::1%lo]:9090");
3911 assert_eq!(request_url, "ws://[fe80::1%lo]:9090");
3912 }
3913
3914 #[rstest]
3915 #[tokio::test]
3916 async fn test_ipv6_loopback_connection() {
3917 if TcpListener::bind("[::1]:0").await.is_err() {
3920 return;
3921 }
3922
3923 let listener = TcpListener::bind("[::1]:0").await.unwrap();
3924 let port = listener.local_addr().unwrap().port();
3925
3926 let server = task::spawn(async move {
3927 if let Ok((mut sock, _)) = listener.accept().await {
3928 let mut buf = vec![0u8; 1024];
3929 if let Ok(n) = sock.read(&mut buf).await {
3930 let _ = sock.write_all(&buf[..n]).await;
3932 }
3933 }
3934 sleep(Duration::from_millis(50)).await;
3935 });
3936
3937 let config = SocketConfig {
3938 url: format!("[::1]:{port}"), mode: Mode::Plain,
3940 suffix: b"\r\n".to_vec(),
3941 message_handler: None,
3942 heartbeat: None,
3943 connect_timeout_ms: Some(1_000),
3944 reconnect_delay_initial_ms: Some(50),
3945 reconnect_delay_max_ms: Some(100),
3946 reconnect_backoff_factor: Some(1.0),
3947 reconnect_jitter_ms: Some(0),
3948 connection_max_retries: Some(1),
3949 reconnect_max_attempts: None,
3950 heartbeat_timeout_secs: None,
3951 writer_capacity: None,
3952 certs_dir: None,
3953 };
3954
3955 let client = SocketClient::builder().config(config).connect().await;
3956 assert!(
3957 client.is_ok(),
3958 "Client should connect to IPv6 loopback address"
3959 );
3960
3961 if let Ok(client) = client {
3962 client.close().await;
3963 }
3964 server.abort();
3965 }
3966
3967 #[rstest]
3968 #[tokio::test]
3969 async fn test_send_waits_during_reconnection() {
3970 use nautilus_common::testing::wait_until_async;
3972
3973 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3974 let port = listener.local_addr().unwrap().port();
3975
3976 let server = task::spawn(async move {
3977 if let Ok((sock, _)) = listener.accept().await {
3979 drop(sock);
3980 }
3981
3982 sleep(Duration::from_millis(500)).await;
3984
3985 if let Ok((mut sock, _)) = listener.accept().await {
3987 let mut buf = vec![0u8; 1024];
3989 while let Ok(n) = sock.read(&mut buf).await {
3990 if n == 0 {
3991 break;
3992 }
3993
3994 if sock.write_all(&buf[..n]).await.is_err() {
3995 break;
3996 }
3997 }
3998 }
3999 });
4000
4001 let config = SocketConfig {
4002 url: format!("127.0.0.1:{port}"),
4003 mode: Mode::Plain,
4004 suffix: b"\r\n".to_vec(),
4005 message_handler: None,
4006 heartbeat: None,
4007 connect_timeout_ms: Some(5_000), reconnect_delay_initial_ms: Some(100),
4009 reconnect_delay_max_ms: Some(200),
4010 reconnect_backoff_factor: Some(1.0),
4011 reconnect_jitter_ms: Some(0),
4012 connection_max_retries: Some(1),
4013 reconnect_max_attempts: None,
4014 heartbeat_timeout_secs: None,
4015 writer_capacity: None,
4016 certs_dir: None,
4017 };
4018
4019 let client = SocketClient::builder()
4020 .config(config)
4021 .connect()
4022 .await
4023 .unwrap();
4024
4025 wait_until_async(
4027 || async { client.is_reconnecting() },
4028 Duration::from_secs(2),
4029 )
4030 .await;
4031
4032 let send_result = tokio::time::timeout(
4034 Duration::from_secs(3),
4035 client.send_bytes(b"test_message".to_vec()),
4036 )
4037 .await;
4038
4039 assert!(
4040 send_result.is_ok() && send_result.unwrap().is_ok(),
4041 "Send should succeed after waiting for reconnection"
4042 );
4043
4044 client.close().await;
4045 server.abort();
4046 }
4047
4048 #[rstest]
4049 #[tokio::test]
4050 async fn test_send_bytes_timeout_uses_configured_connect_timeout() {
4051 use nautilus_common::testing::wait_until_async;
4054
4055 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4056 let port = listener.local_addr().unwrap().port();
4057
4058 let server = task::spawn(async move {
4059 if let Ok((sock, _)) = listener.accept().await {
4061 drop(sock);
4062 }
4063 drop(listener);
4065 sleep(Duration::from_mins(1)).await;
4066 });
4067
4068 let config = SocketConfig {
4069 url: format!("127.0.0.1:{port}"),
4070 mode: Mode::Plain,
4071 suffix: b"\r\n".to_vec(),
4072 message_handler: None,
4073 heartbeat: None,
4074 connect_timeout_ms: Some(1_000), reconnect_delay_initial_ms: Some(200), reconnect_delay_max_ms: Some(200),
4077 reconnect_backoff_factor: Some(1.0),
4078 reconnect_jitter_ms: Some(0),
4079 connection_max_retries: Some(1),
4080 reconnect_max_attempts: None,
4081 heartbeat_timeout_secs: None,
4082 writer_capacity: None,
4083 certs_dir: None,
4084 };
4085
4086 let client = SocketClient::builder()
4087 .config(config)
4088 .connect()
4089 .await
4090 .unwrap();
4091
4092 wait_until_async(
4094 || async { client.is_reconnecting() },
4095 Duration::from_secs(3),
4096 )
4097 .await;
4098
4099 let start = std::time::Instant::now();
4102 let send_result = client.send_bytes(b"test".to_vec()).await;
4103 let elapsed = start.elapsed();
4104
4105 assert!(
4106 send_result.is_err(),
4107 "Send should fail when client stuck in RECONNECT, was: {send_result:?}"
4108 );
4109 assert!(
4110 matches!(send_result, Err(crate::error::SendError::Timeout)),
4111 "Send should return Timeout error, was: {send_result:?}"
4112 );
4113 assert!(
4116 elapsed >= Duration::from_millis(900),
4117 "Send should timeout after at least 1s (configured timeout), took {elapsed:?}"
4118 );
4119
4120 client.close().await;
4121 server.abort();
4122 }
4123
4124 #[rstest]
4125 #[tokio::test]
4126 async fn test_heartbeat_timeout_triggers_reconnect() {
4127 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4128 let port = listener.local_addr().unwrap().port();
4129
4130 let server = task::spawn(async move {
4132 let (_sock1, _) = listener.accept().await.unwrap();
4133 sleep(Duration::from_secs(5)).await;
4135 });
4136
4137 let config = SocketConfig {
4138 url: format!("127.0.0.1:{port}"),
4139 mode: Mode::Plain,
4140 suffix: b"\r\n".to_vec(),
4141 message_handler: None,
4142 heartbeat: None,
4143 connect_timeout_ms: Some(2_000),
4144 reconnect_delay_initial_ms: Some(50),
4145 reconnect_delay_max_ms: Some(100),
4146 reconnect_backoff_factor: Some(1.0),
4147 reconnect_jitter_ms: Some(0),
4148 connection_max_retries: Some(1),
4149 reconnect_max_attempts: Some(1),
4150 heartbeat_timeout_secs: Some(1),
4151 writer_capacity: None,
4152 certs_dir: None,
4153 };
4154
4155 let client = SocketClient::builder()
4156 .config(config)
4157 .connect()
4158 .await
4159 .unwrap();
4160
4161 assert!(client.is_active());
4162
4163 wait_until_async(
4165 || async { client.is_reconnecting() || client.is_closed() },
4166 Duration::from_secs(4),
4167 )
4168 .await;
4169
4170 assert!(
4171 !client.is_active(),
4172 "Client should not be active after the dead-peer timeout"
4173 );
4174
4175 client.close().await;
4176 server.abort();
4177 }
4178
4179 #[rstest]
4180 #[tokio::test]
4181 async fn test_heartbeat_timeout_resets_on_inbound_bytes() {
4182 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4183 let port = listener.local_addr().unwrap().port();
4184
4185 let server = task::spawn(async move {
4187 let (mut sock, _) = listener.accept().await.unwrap();
4188 for _ in 0..10 {
4189 sleep(Duration::from_millis(200)).await;
4190
4191 if sock.write_all(b"ping\r\n").await.is_err() {
4192 break;
4193 }
4194 }
4195 });
4196
4197 let config = SocketConfig {
4198 url: format!("127.0.0.1:{port}"),
4199 mode: Mode::Plain,
4200 suffix: b"\r\n".to_vec(),
4201 message_handler: None,
4202 heartbeat: None,
4203 connect_timeout_ms: Some(2_000),
4204 reconnect_delay_initial_ms: Some(50),
4205 reconnect_delay_max_ms: Some(100),
4206 reconnect_backoff_factor: Some(1.0),
4207 reconnect_jitter_ms: Some(0),
4208 connection_max_retries: Some(1),
4209 reconnect_max_attempts: Some(1),
4210 heartbeat_timeout_secs: Some(1),
4211 writer_capacity: None,
4212 certs_dir: None,
4213 };
4214
4215 let client = SocketClient::builder()
4216 .config(config)
4217 .connect()
4218 .await
4219 .unwrap();
4220
4221 assert!(client.is_active());
4222
4223 sleep(Duration::from_millis(1_500)).await;
4225
4226 assert!(
4227 client.is_active(),
4228 "Client should remain active when data is flowing"
4229 );
4230
4231 client.close().await;
4232 server.abort();
4233 }
4234
4235 #[rstest]
4236 #[tokio::test]
4237 async fn test_close_during_backoff_exits_promptly() {
4238 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4242 let port = listener.local_addr().unwrap().port();
4243
4244 let server = task::spawn(async move {
4245 if let Ok((mut sock, _)) = listener.accept().await {
4247 drop(sock.shutdown());
4248 }
4249 sleep(Duration::from_mins(1)).await;
4251 });
4252
4253 let config = SocketConfig {
4254 url: format!("127.0.0.1:{port}"),
4255 mode: Mode::Plain,
4256 suffix: b"\r\n".to_vec(),
4257 message_handler: None,
4258 heartbeat: None,
4259 connect_timeout_ms: Some(1_000),
4260 reconnect_delay_initial_ms: Some(10_000), reconnect_delay_max_ms: Some(10_000),
4262 reconnect_backoff_factor: Some(1.0),
4263 reconnect_jitter_ms: Some(0),
4264 connection_max_retries: None,
4265 reconnect_max_attempts: None,
4266 heartbeat_timeout_secs: None,
4267 writer_capacity: None,
4268 certs_dir: None,
4269 };
4270
4271 let client = SocketClient::builder()
4272 .config(config)
4273 .connect()
4274 .await
4275 .unwrap();
4276
4277 wait_until_async(
4279 || async { client.is_reconnecting() },
4280 Duration::from_secs(3),
4281 )
4282 .await;
4283
4284 sleep(Duration::from_millis(1_500)).await;
4286
4287 let start = std::time::Instant::now();
4289 client.close().await;
4290 let elapsed = start.elapsed();
4291
4292 assert!(client.is_closed(), "Client should be closed");
4293 assert!(
4295 elapsed < Duration::from_secs(2),
4296 "Close should interrupt backoff sleep, took {elapsed:?}"
4297 );
4298
4299 server.abort();
4300 }
4301
4302 #[rstest]
4303 #[tokio::test]
4304 async fn test_zero_heartbeat_timeout_rejected() {
4305 let config = SocketConfig {
4306 url: "127.0.0.1:9999".to_string(),
4307 mode: Mode::Plain,
4308 suffix: b"\r\n".to_vec(),
4309 message_handler: None,
4310 heartbeat: None,
4311 connect_timeout_ms: None,
4312 reconnect_delay_initial_ms: None,
4313 reconnect_delay_max_ms: None,
4314 reconnect_backoff_factor: None,
4315 reconnect_jitter_ms: None,
4316 reconnect_max_attempts: None,
4317 connection_max_retries: Some(1),
4318 heartbeat_timeout_secs: Some(0),
4319 writer_capacity: None,
4320 certs_dir: None,
4321 };
4322
4323 let result = SocketClient::builder().config(config).connect().await;
4324
4325 assert!(result.is_err(), "Zero heartbeat timeout should be rejected");
4326 let err_msg = result.unwrap_err().to_string();
4327 assert!(
4328 err_msg.contains("heartbeat_timeout_secs"),
4329 "Error should name the offending field, was: {err_msg}"
4330 );
4331 }
4332
4333 #[rstest]
4334 #[tokio::test]
4335 async fn test_empty_suffix_rejected() {
4336 let config = SocketConfig {
4337 url: "127.0.0.1:9999".to_string(),
4338 mode: Mode::Plain,
4339 suffix: vec![],
4340 message_handler: None,
4341 heartbeat: None,
4342 connect_timeout_ms: None,
4343 reconnect_delay_initial_ms: None,
4344 reconnect_delay_max_ms: None,
4345 reconnect_backoff_factor: None,
4346 reconnect_jitter_ms: None,
4347 reconnect_max_attempts: None,
4348 connection_max_retries: Some(1),
4349 heartbeat_timeout_secs: None,
4350 writer_capacity: None,
4351 certs_dir: None,
4352 };
4353
4354 let result = SocketClient::builder().config(config).connect().await;
4355
4356 assert!(
4357 result.is_err(),
4358 "Empty suffix should cause connection to fail"
4359 );
4360 let err_msg = result.unwrap_err().to_string();
4361 assert!(
4362 err_msg.contains("suffix cannot be empty"),
4363 "Error should mention empty suffix, was: {err_msg}"
4364 );
4365 }
4366}
4367
4368#[cfg(test)]
4369mod reconnect_request_tests {
4370 use std::sync::{Arc, atomic::AtomicU8};
4371
4372 use parking_lot::Mutex;
4373 use rstest::rstest;
4374
4375 use super::*;
4376 use crate::SocketState;
4377
4378 fn handle(
4379 mode: ConnectionMode,
4380 ) -> (
4381 SocketReconnectHandle,
4382 Arc<tokio::sync::Notify>,
4383 Arc<Mutex<Vec<SocketState>>>,
4384 ) {
4385 let controller_notify = Arc::new(tokio::sync::Notify::new());
4386 let states = Arc::new(Mutex::new(Vec::new()));
4387 let states_callback = Arc::clone(&states);
4388 let state_sink = SocketStateSink::new(move |state| {
4389 states_callback.lock().push(state);
4390 });
4391 let handle = SocketReconnectHandle {
4392 connection_mode: Arc::new(AtomicU8::new(mode.as_u8())),
4393 state_sink: Some(state_sink),
4394 controller_lifecycle: Arc::new(ControllerLifecycle::new()),
4395 controller_notify: Arc::clone(&controller_notify),
4396 };
4397 (handle, controller_notify, states)
4398 }
4399
4400 #[rstest]
4401 #[tokio::test]
4402 async fn accepted_request_reports_loss_and_wakes_controller_once() {
4403 let (handle, controller_notify, states) = handle(ConnectionMode::Active);
4404
4405 assert_eq!(
4406 handle.request_reconnect(),
4407 ReconnectRequestOutcome::Accepted
4408 );
4409 assert_eq!(*states.lock(), vec![SocketState::Disconnected]);
4410 tokio::time::timeout(Duration::from_millis(10), controller_notify.notified())
4411 .await
4412 .expect("accepted request should notify controller");
4413
4414 assert_eq!(
4415 handle.request_reconnect(),
4416 ReconnectRequestOutcome::AlreadyReconnecting
4417 );
4418 assert_eq!(*states.lock(), vec![SocketState::Disconnected]);
4419 assert!(
4420 tokio::time::timeout(Duration::from_millis(10), controller_notify.notified())
4421 .await
4422 .is_err(),
4423 "duplicate request should not notify controller",
4424 );
4425 }
4426
4427 #[rstest]
4428 #[case(
4429 ConnectionMode::Reconnect,
4430 ReconnectRequestOutcome::AlreadyReconnecting
4431 )]
4432 #[case(ConnectionMode::Disconnect, ReconnectRequestOutcome::Disconnected)]
4433 #[case(ConnectionMode::Closed, ReconnectRequestOutcome::Closed)]
4434 #[tokio::test]
4435 async fn rejected_request_preserves_state_and_does_not_wake_controller(
4436 #[case] mode: ConnectionMode,
4437 #[case] expected: ReconnectRequestOutcome,
4438 ) {
4439 let (handle, controller_notify, states) = handle(mode);
4440
4441 assert_eq!(handle.request_reconnect(), expected);
4442 assert!(states.lock().is_empty());
4443 assert!(
4444 tokio::time::timeout(Duration::from_millis(10), controller_notify.notified())
4445 .await
4446 .is_err(),
4447 "rejected request should not notify controller",
4448 );
4449 }
4450}