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