Skip to main content

nautilus_network/socket/
client.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Raw TCP client with optional TLS, suffix framing, heartbeats, and automatic reconnection.
17//!
18//! # State management
19//!
20//! The client tracks active, reconnecting, disconnecting, and closed states. State changes notify
21//! waiting tasks immediately instead of relying only on periodic checks.
22//!
23//! # Connection ownership
24//!
25//! A controller owns the connection lifecycle. A dedicated reader task passes complete messages to
26//! the configured callback, while a dedicated writer task serializes concurrent sends received over
27//! a channel.
28//!
29//! # Framing and heartbeats
30//!
31//! The configured suffix frames the byte stream in both directions. The writer appends it to sent
32//! messages and heartbeats, and the reader uses it to split incoming data into complete messages.
33//! Heartbeats are optional and run in a separate task.
34//!
35//! # Reconnection
36//!
37//! The writer buffers messages while reconnecting. A successful reconnect installs the replacement
38//! writer, drains that buffer, restarts the reader, and then invokes the configured
39//! post-reconnection callback.
40
41use 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
77// Connection timing constants
78const 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
83// Maximum buffer size for read operations (10 MB)
84const 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
92/// Produces protocol messages that must precede buffered application writes after reconnect.
93pub 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    /// Connects to a URL with the specified configuration.
115    ///
116    /// # Errors
117    ///
118    /// Returns an error if connection fails or configuration is invalid.
119    async fn connect_url(
120        config: SocketConfig,
121        state_sink: Option<SocketStateSink>,
122    ) -> anyhow::Result<Self> {
123        install_cryptographic_provider();
124
125        // Validate suffix is non-empty to prevent panic in read loop (windows(0) panics)
126        if config.suffix.is_empty() {
127            anyhow::bail!("Socket suffix cannot be empty: suffix is required for message framing");
128        }
129
130        // Adapters build this config by struct literal, bypassing the builder, so this is the only
131        // place the field invariants are enforced for them.
132        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, // immediate-first
142        )?;
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        // Retry initial connection with exponential backoff to handle transient DNS/network issues
152        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    /// Parses a URL into its socket address and request URL.
276    ///
277    /// Accepts either:
278    /// - Raw socket address: "host:port" → returns ("host:port", "scheme://host:port")
279    /// - Full URL: "scheme://host:port/path" → returns ("host:port", original URL)
280    ///
281    /// # Errors
282    ///
283    /// Returns an error if the URL is invalid or missing required components.
284    fn parse_socket_url(url: &str, mode: Mode) -> Result<(String, String), Error> {
285        if url.contains("://") {
286            // URL with scheme (e.g., "wss://host:port/path")
287            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            // Raw socket address (e.g., "host:port")
315            // Construct a proper URL for the request based on mode
316            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    /// Establish a TLS or plain TCP connection with the server.
326    ///
327    /// Accepts either a raw socket address (e.g., "host:port") or a full URL with scheme
328    /// (e.g., "wss://host:port"). For FIX/raw socket connections, use the host:port format.
329    /// For WebSocket-style connections, include the scheme.
330    ///
331    /// # Errors
332    ///
333    /// Returns an error if the connection cannot be established.
334    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    /// Reconnects to the server.
363    ///
364    /// Makes a new connection with server, uses the new read and write halves
365    /// to update the reader and writer.
366    ///
367    /// The reconnect timeout bounds only connection establishment. Once the
368    /// new writer is handed to the writer task the swap runs to completion,
369    /// so buffered messages can never drain into a connection that lost its
370    /// reader to a timeout; the writer task bounds both the old-writer
371    /// shutdown and the buffer drain with its graceful-shutdown timeout.
372    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        // Bound only connection establishment; the swap below must run to completion
384        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        // Use a oneshot channel to synchronize with the writer task.
411        // We must verify that the buffer was successfully drained before transitioning to ACTIVE
412        // to prevent silent message loss if the new connection drops immediately.
413        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        // Wait for writer to confirm it has drained the buffer
429        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 error to trigger retry logic in controller
434                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        // Delay before closing connection
448        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        // Atomically transition from Reconnect to Active
463        // This prevents race condition where disconnect could be requested between check and store
464        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        // Spawn new read task
474        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    /// Returns whether the read and write tasks are still running.
489    ///
490    /// Returns `true` if both the read and write tasks are still running.
491    /// There may be some delay between the connection closing and the
492    /// client detecting it.
493    #[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        // Interval between checking the connection mode
514        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                // Connection has been terminated or vector buffer is complete
558                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                // Received bytes of data
567                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    /// Drains buffered messages after reconnection completes.
632    ///
633    /// Attempts to send all buffered messages that were queued during reconnection.
634    /// Uses a peek-and-pop pattern to preserve messages if sending fails midway through the buffer.
635    ///
636    /// # Returns
637    ///
638    /// Returns `true` if a send error occurred (buffer may still contain unsent messages),
639    /// `false` if all messages were sent successfully (buffer is empty).
640    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        // Interval between checking the connection mode
755        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                        // Re-check connection mode after receiving a message
805                        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                                    // CAS: a disconnect landing mid-write must not be overwritten
916                                    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                        // Channel closed - writer task should terminate
929                        log::debug!("Writer channel closed, terminating writer task");
930                        break;
931                    }
932                    Err(_) => {
933                        // Timeout - just continue the loop
934                    }
935                }
936            }
937
938            // Attempt to shutdown the writer gracefully before exiting,
939            // we ignore any error as the writer may already be closed.
940            _ = 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
1024/// A suffix-framed TCP client with optional TLS and automatic reconnection.
1025///
1026/// The internal writer task serializes concurrent calls to [`Self::send_bytes`]. The configured
1027/// suffix frames all sent and received messages, and an optional heartbeat task sends its payload
1028/// at the configured interval. See [`SocketConfig`] for framing and reconnect policy.
1029pub 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/// Cloneable controller handle for requesting one raw socket reconnect.
1041#[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    /// Requests that the controller replace the active transport.
1062    ///
1063    /// An accepted request reports the transport unavailable and wakes the controller. Rejected
1064    /// requests leave the transport state unchanged.
1065    #[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    /// Returns a builder for a new [`SocketClient`] connection.
1093    ///
1094    /// After a successful reconnect, `post_reconnection` runs after the replacement writer is
1095    /// installed, buffered messages are drained, and the replacement reader is started. The
1096    /// callback does not run after the initial connection.
1097    ///
1098    /// `state_sink` receives connection state changes. `reconnect_replay` produces protocol setup
1099    /// messages that the client sends before buffered application messages after a reconnect.
1100    ///
1101    /// # Errors
1102    ///
1103    /// Returns an error if connecting to the server fails.
1104    #[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    /// Returns a cloneable handle to this client's reconnect controller.
1144    #[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    /// Requests that the controller replace the active transport.
1155    ///
1156    /// Returns `true` only when this call transitions the client from active to reconnecting.
1157    /// Duplicate, disconnecting, and closed requests return `false`.
1158    #[must_use]
1159    pub fn request_reconnect(&self) -> bool {
1160        self.reconnect_handle().request_reconnect() == ReconnectRequestOutcome::Accepted
1161    }
1162
1163    /// Returns the current connection mode.
1164    #[must_use]
1165    pub fn connection_mode(&self) -> ConnectionMode {
1166        ConnectionMode::from_atomic(&self.connection_mode)
1167    }
1168
1169    /// Returns whether the client connection is active.
1170    ///
1171    /// Returns `true` if the client is connected and has not been signaled to disconnect.
1172    /// The client will automatically retry connection based on its configuration.
1173    #[inline]
1174    #[must_use]
1175    pub fn is_active(&self) -> bool {
1176        self.connection_mode().is_active()
1177    }
1178
1179    /// Returns whether the client is reconnecting.
1180    ///
1181    /// Returns `true` if the client lost connection and is attempting to reestablish it.
1182    /// The client will automatically retry connection based on its configuration.
1183    #[inline]
1184    #[must_use]
1185    pub fn is_reconnecting(&self) -> bool {
1186        self.connection_mode().is_reconnect()
1187    }
1188
1189    /// Returns whether the client is disconnecting.
1190    ///
1191    /// Returns `true` if the client is in disconnect mode.
1192    #[inline]
1193    #[must_use]
1194    pub fn is_disconnecting(&self) -> bool {
1195        self.connection_mode().is_disconnect()
1196    }
1197
1198    /// Returns whether the client is closed.
1199    ///
1200    /// Returns `true` if the client has been explicitly disconnected or reached
1201    /// maximum reconnection attempts. In this state, the client cannot be reused
1202    /// and a new client must be created for further connections.
1203    #[inline]
1204    #[must_use]
1205    pub fn is_closed(&self) -> bool {
1206        self.connection_mode().is_closed()
1207    }
1208
1209    /// Close the client.
1210    ///
1211    /// Controller task will periodically check the disconnect mode
1212    /// and shutdown the client if it is not alive.
1213    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    /// Requests shutdown without waiting for the controller task.
1220    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    /// Checks whether the connection is in a terminal state (disconnecting or closed).
1250    ///
1251    /// Single atomic load to fail fast before waiting.
1252    #[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    /// Waits for the client to become active before sending.
1261    ///
1262    /// Uses `state_notify` for event-driven wakeup so sends resume immediately
1263    /// after reconnection completes. A fallback interval guards against missed
1264    /// notifications.
1265    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                // Enable before the state check: an unpolled Notified is unregistered and misses notifies
1284                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    /// Sends a message of the given `data`.
1309    ///
1310    /// Returns `Ok(())` when the message is enqueued to the writer channel. This does not
1311    /// guarantee delivery: if a disconnect occurs concurrently, the writer task may drop the
1312    /// message. During reconnection, messages are buffered and replayed on the new connection.
1313    ///
1314    /// # Errors
1315    ///
1316    /// Returns an error if the writer capacity is exhausted, or if the client cannot send.
1317    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                        // Delay awaiting graceful shutdown
1359                        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; // Controller finished
1384                }
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                    // Check max reconnection attempts before attempting reconnect
1434                    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                    // Race reconnect against disconnect notification
1484                    let reconnect_result = tokio::select! {
1485                        biased;
1486                        result = inner.reconnect(reconnect_replay.as_ref()) => Some(result),
1487                        () = async {
1488                            loop {
1489                                // Enable before the check so a disconnect notify between iterations is not missed
1490                                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                            // The outcome records a completed reconnection; emit recovery
1512                            // callbacks only while the replacement is still `Active`, not
1513                            // after a teardown or another drop.
1514                            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
1558// Dropping cancels background work without reporting a terminal socket state transition.
1559impl 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)))] // transport-layer I/O not simulated
1573#[cfg(target_os = "linux")] // Only run network tests on Linux (CI stability)
1574mod 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                        // Remove trailing \r\n
1609                        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        // Ready sends do not yield on this current-thread runtime, keeping the queue occupied
1650        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        // Wait a bit for the server to echo them back
1756        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); // We drop it immediately -> no server is listening
1767
1768        // Wait until port is truly unavailable (OS has released it)
1769        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); // We drop it immediately -> no server is listening
1812
1813        // Wait until port is truly unavailable (OS has released it)
1814        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        // Asserted positively so the blanket secret assertion below cannot pass vacuously by
1866        // capturing no warning at all.
1867        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            // Accept the first connection then drop it; never accept again so
1940            // the client exhausts its reconnect attempts and transitions to CLOSED
1941            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        // Closing an already CLOSED client must return promptly (no 5s spin
1974        // waiting for a controller that has already exited) and must not
1975        // regress the terminal state to DISCONNECT
1976        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        // Wait ~3 seconds to collect some heartbeats
2087        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        // Spawn a server task that:
2110        // 1. Accepts the first connection and then drops it after a short delay (simulate disconnect)
2111        // 2. Waits a bit and then accepts a new connection and runs the echo server
2112        let server_task = task::spawn(async move {
2113            // Accept first connection
2114            let (mut socket, _) = listener.accept().await.expect("First accept failed");
2115
2116            // Wait briefly and then force-close the connection
2117            sleep(Duration::from_millis(500)).await;
2118            let _ = socket.shutdown().await;
2119
2120            // Wait for the client's reconnect attempt
2121            sleep(Duration::from_millis(500)).await;
2122
2123            // Run the echo server on the new connection
2124            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        // Initially, the client should be active
2153        assert!(client.is_active(), "Client should start as active");
2154
2155        // Wait until the client loses connection (i.e. not active),
2156        // then wait until it reconnects (active again).
2157        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)))] // transport-layer I/O not simulated
2281mod 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        // Bind an ephemeral port
3708        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3709        let port = listener.local_addr().unwrap().port();
3710
3711        // Server task: accept one connection and then drop it
3712        let server = task::spawn(async move {
3713            if let Ok((mut sock, _)) = listener.accept().await {
3714                drop(sock.shutdown());
3715            }
3716
3717            // Keep listener alive briefly to avoid premature exit
3718            sleep(Duration::from_secs(1)).await;
3719        });
3720
3721        // Configure client with a short reconnect backoff
3722        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        // Connect client (handler=None)
3741        let client = SocketClient::builder()
3742            .config(config.clone())
3743            .connect()
3744            .await
3745            .unwrap();
3746
3747        // Wait for client to detect dropped connection and enter reconnect state
3748        wait_until_async(
3749            || async { client.is_reconnecting() },
3750            Duration::from_secs(2),
3751        )
3752        .await;
3753
3754        // Now close the client
3755        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        // Bind an ephemeral port and accept a single connection which we immediately close.
3764        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            // Give the client a moment to observe the closed connection.
3772            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        // Raw socket address with TLS mode
3812        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        // Raw socket address with Plain mode
3818        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        // Full URL with wss scheme
3827        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        // Full URL with ws scheme
3833        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        // wss without explicit port defaults to 443
3842        let (socket_addr, _) =
3843            SocketClientInner::parse_socket_url("wss://example.com", Mode::Tls).unwrap();
3844        assert_eq!(socket_addr, "example.com:443");
3845
3846        // ws without explicit port defaults to 80
3847        let (socket_addr, _) =
3848            SocketClientInner::parse_socket_url("ws://example.com", Mode::Plain).unwrap();
3849        assert_eq!(socket_addr, "example.com:80");
3850
3851        // https defaults to 443
3852        let (socket_addr, _) =
3853            SocketClientInner::parse_socket_url("https://example.com", Mode::Tls).unwrap();
3854        assert_eq!(socket_addr, "example.com:443");
3855
3856        // http defaults to 80
3857        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        // Unknown scheme defaults to mode-based port
3865        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        // IPv6 address with port
3877        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        // IPv6 in URL
3883        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        // Test that raw socket addresses (host:port) work correctly
3892        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}"), // Raw socket address format
3904            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        // Should successfully connect with raw socket address
3921        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        // Test that URLs with schemes also work
3938        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}"), // URL with scheme
3950            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        // Should successfully connect with URL format
3967        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        // IPv6 with zone ID (link-local address)
3983        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        // Verify zone is preserved in URL format too
3989        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        // Test IPv6 loopback address connection
3999        // Skip if IPv6 is not available on the system
4000        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                    // Echo back
4011                    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}"), // IPv6 loopback
4020            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        // Test that send operations wait for reconnection to complete (up to configured timeout)
4053        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            // First connection - accept and immediately close
4060            if let Ok((sock, _)) = listener.accept().await {
4061                drop(sock);
4062            }
4063
4064            // Wait before accepting second connection
4065            sleep(Duration::from_millis(500)).await;
4066
4067            // Second connection - accept and keep alive
4068            if let Ok((mut sock, _)) = listener.accept().await {
4069                // Echo messages
4070                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), // 5s timeout - enough for reconnect
4090            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 for reconnection to trigger
4108        wait_until_async(
4109            || async { client.is_reconnecting() },
4110            Duration::from_secs(2),
4111        )
4112        .await;
4113
4114        // Try to send while reconnecting - should wait and succeed after reconnect
4115        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        // Test that send_bytes operations respect the configured connect_timeout.
4134        // When a client is stuck in RECONNECT longer than the timeout, sends should fail with Timeout.
4135        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            // Accept first connection and immediately close it
4142            if let Ok((sock, _)) = listener.accept().await {
4143                drop(sock);
4144            }
4145
4146            // Drop listener entirely so reconnection fails completely
4147            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), // 1s timeout for faster test
4158            reconnect_delay_initial_ms: Some(200), // Short backoff (but > timeout) to keep client in RECONNECT
4159            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 for client to enter RECONNECT state
4176        wait_until_async(
4177            || async { client.is_reconnecting() },
4178            Duration::from_secs(3),
4179        )
4180        .await;
4181
4182        // Attempt send while stuck in RECONNECT - should timeout after 1s (configured timeout)
4183        // The client will try to reconnect for 1s, fail, then wait 5s backoff before next attempt
4184        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        // Verify timeout respects configured value (1s), but don't check upper bound
4197        // as CI scheduler jitter can cause legitimate delays beyond the timeout
4198        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        // Server accepts connection but sends nothing (simulates silent death)
4214        let server = task::spawn(async move {
4215            let (_sock1, _) = listener.accept().await.unwrap();
4216            // Hold connection open but send nothing, wait for reconnect attempt
4217            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 for the dead-peer timeout to fire and the client to enter reconnect
4247        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        // Server sends data every 200ms (well within the 1s idle timeout)
4269        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        // Wait 1.5s - data arrives every 200ms so idle timeout (1s) should NOT fire
4308        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        // Verify that close() interrupts backoff sleep (Finding 1).
4323        // Server accepts then drops, no second listener -> reconnect fails -> enters backoff.
4324        // We close while backing off and assert the client shuts down quickly.
4325        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            // Accept first connection, close immediately
4330            if let Ok((mut sock, _)) = listener.accept().await {
4331                drop(sock.shutdown());
4332            }
4333
4334            // Don't accept again so reconnect fails and enters backoff
4335            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), // 10s backoff to ensure we're sleeping
4346            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 for client to enter reconnect
4363        wait_until_async(
4364            || async { client.is_reconnecting() },
4365            Duration::from_secs(3),
4366        )
4367        .await;
4368
4369        // Wait for the reconnect attempt to fail and enter backoff sleep
4370        sleep(Duration::from_millis(1_500)).await;
4371
4372        // Close while backing off
4373        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        // Should exit well before the 10s backoff sleep completes
4379        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}