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