Fix shutdown race-condition by introducing a flush_timeout before dropping data

crates/re_sdk/src/log_sink.rs

@@ -190,12 +190,23 @@ pub struct TcpSink {
 }
 
 impl TcpSink {
+    /// Default timeout to use for the `disconnected_timeout` of [`TcpSink::new`]
+    pub const DEFAULT_TIMEOUT: Option<std::time::Duration> =
+        Some(std::time::Duration::from_secs(2));
+
     /// Connect to the given address in a background thread.
     /// Retries until successful.
-    #[inline]
-    pub fn new(addr: std::net::SocketAddr) -> Self {
+    ///
+    /// `disconnected_timeout` is the minimum time the [`TcpSink`] will wait during a flush
+    /// before potentially dropping data.  Note: Passing `None` here can cause a
+    /// call to `flush` to block indefinitely if a connection cannot be established.
+    #[inline]
+    pub fn new(
+        addr: std::net::SocketAddr,
+        disconnected_timeout: Option<std::time::Duration>,
+    ) -> Self {
         Self {
-            client: re_sdk_comms::Client::new(addr),
+            client: re_sdk_comms::Client::new(addr, disconnected_timeout),
         }
     }
 }

crates/re_sdk/src/recording_stream.rs

@@ -210,20 +210,28 @@ impl RecordingStreamBuilder {
     /// Creates a new [`RecordingStream`] that is pre-configured to stream the data through to a
     /// remote Rerun instance.
     ///
+    /// `disconnected_timeout` is the minimum time the [`TcpSink`][`crate::log_sink::TcpSink`] will
+    /// wait during a flush before potentially dropping data.  Note: Passing `None` here can cause a
+    /// call to `flush` to block indefinitely if a connection cannot be established.
+    ///
     /// ## Example
     ///
     /// ```no_run
     /// let rec_stream = re_sdk::RecordingStreamBuilder::new("my_app")
-    ///     .connect(re_sdk::default_server_addr())?;
+    ///     .connect(re_sdk::default_server_addr(), re_sdk::sink::TcpSink::DEFAULT_TIMEOUT)?;
     /// # Ok::<(), Box<dyn std::error::Error>>(())
     /// ```
-    pub fn connect(self, addr: std::net::SocketAddr) -> RecordingStreamResult<RecordingStream> {
+    pub fn connect(
+        self,
+        addr: std::net::SocketAddr,
+        disconnected_timeout: Option<std::time::Duration>,
+    ) -> RecordingStreamResult<RecordingStream> {
         let (enabled, store_info, batcher_config) = self.into_args();
         if enabled {
             RecordingStream::new(
                 store_info,
                 batcher_config,
-                Box::new(crate::log_sink::TcpSink::new(addr)),
+                Box::new(crate::log_sink::TcpSink::new(addr, disconnected_timeout)),
             )
         } else {
             re_log::debug!("Rerun disabled - call to connect() ignored");
@@ -794,11 +802,22 @@ impl RecordingStream {
     /// Swaps the underlying sink for a [`crate::log_sink::TcpSink`] sink pre-configured to use
     /// the specified address.
     ///
+    /// `disconnected_timeout` is the minimum time the [`TcpSink`][`crate::log_sink::TcpSink`] will
+    /// wait during a flush before potentially dropping data.  Note: Passing `None` here can cause a
+    /// call to `flush` to block indefinitely if a connection cannot be established.
+    ///
     /// This is a convenience wrapper for [`Self::set_sink`] that upholds the same guarantees in
     /// terms of data durability and ordering.
     /// See [`Self::set_sink`] for more information.
-    pub fn connect(&self, addr: std::net::SocketAddr) {
-        self.set_sink(Box::new(crate::log_sink::TcpSink::new(addr)));
+    pub fn connect(
+        &self,
+        addr: std::net::SocketAddr,
+        disconnected_timeout: Option<std::time::Duration>,
+    ) {
+        self.set_sink(Box::new(crate::log_sink::TcpSink::new(
+            addr,
+            disconnected_timeout,
+        )));
     }
 
     /// Swaps the underlying sink for a [`crate::sink::MemorySink`] sink and returns the associated

crates/re_sdk_comms/src/buffered_client.rs

@@ -54,7 +54,11 @@ pub struct Client {
 
 impl Client {
     /// Connect via TCP to this log server.
-    pub fn new(addr: SocketAddr) -> Self {
+    ///
+    /// `disconnected_timeout` is the minimum time the [`TcpSink`][`crate::log_sink::TcpSink`] will
+    /// wait during a flush before potentially dropping data.  Note: Passing `None` here can cause a
+    /// call to `flush` to block indefinitely if a connection cannot be established.
+    pub fn new(addr: SocketAddr, disconnected_timeout: Option<std::time::Duration>) -> Self {
         re_log::debug!("Connecting to remote {addr}…");
 
         // TODO(emilk): keep track of how much memory is in each pipe
@@ -88,7 +92,13 @@ impl Client {
         let send_join = std::thread::Builder::new()
             .name("tcp_sender".into())
             .spawn(move || {
-                tcp_sender(addr, &packet_rx, &send_quit_rx, &flushed_tx);
+                tcp_sender(
+                    addr,
+                    disconnected_timeout,
+                    &packet_rx,
+                    &send_quit_rx,
+                    &flushed_tx,
+                );
             })
             .expect("Failed to spawn thread");
 
@@ -250,11 +260,12 @@ fn msg_encode(
 
 fn tcp_sender(
     addr: SocketAddr,
+    disconnected_timeout: Option<std::time::Duration>,
     packet_rx: &Receiver<PacketMsg>,
     quit_rx: &Receiver<InterruptMsg>,
     flushed_tx: &Sender<FlushedMsg>,
 ) {
-    let mut tcp_client = crate::tcp_client::TcpClient::new(addr);
+    let mut tcp_client = crate::tcp_client::TcpClient::new(addr, disconnected_timeout);
     // Once this flag has been set, we will drop all messages if the tcp_client is
     // no longer connected.
     let mut drop_if_disconnected = false;
@@ -311,14 +322,14 @@ fn send_until_success(
     quit_rx: &Receiver<InterruptMsg>,
 ) -> Option<InterruptMsg> {
     // Early exit if tcp_client is disconnected
-    if drop_if_disconnected && tcp_client.has_disconnected() {
-        re_log::debug_once!("Dropping messages because we're disconnected.");
+    if drop_if_disconnected && tcp_client.has_timed_out() {
+        re_log::warn_once!("Dropping messages because tcp client has timed out.");
         return None;
     }
 
     if let Err(err) = tcp_client.send(packet) {
-        if drop_if_disconnected {
-            re_log::debug_once!("Dropping messages because we're disconnected.");
+        if drop_if_disconnected && tcp_client.has_timed_out() {
+            re_log::warn_once!("Dropping messages because tcp client has timed out.");
             return None;
         }
         // If this is the first time we fail to send the message, produce a warning.
@@ -329,11 +340,17 @@ fn send_until_success(
         loop {
             select! {
                 recv(quit_rx) -> _quit_msg => {
-                    re_log::debug_once!("Dropping messages because we're disconnected or quitting.");
+                    re_log::warn_once!("Dropping messages because tcp client has timed out or quitting.");
                     return Some(_quit_msg.unwrap_or(InterruptMsg::Quit));
                 }
                 default(std::time::Duration::from_millis(sleep_ms)) => {
                     if let Err(new_err) = tcp_client.send(packet) {
+
+                        if drop_if_disconnected && tcp_client.has_timed_out() {
+                            re_log::warn_once!("Dropping messages because tcp client has timed out.");
+                            return None;
+                        }
+
                         const MAX_SLEEP_MS : u64 = 3000;
 
                         sleep_ms = (sleep_ms * 2).min(MAX_SLEEP_MS);

crates/re_sdk_comms/src/tcp_client.rs

@@ -1,6 +1,7 @@
 use std::{
     io::Write,
     net::{SocketAddr, TcpStream},
+    time::{Duration, Instant},
 };
 
 #[derive(thiserror::Error, Debug)]
@@ -29,12 +30,15 @@ pub enum ClientError {
 enum TcpStreamState {
     /// The [`TcpStream`] is yet to be connected.
     ///
+    /// Tracks the duration and connection attempts since the last time the client
+    /// was `Connected.`
+    ///
     /// Behavior: Try to connect on next [`TcpClient::connect`] or [`TcpClient::send`].
     ///
     /// Transitions:
     ///  - Pending -> Connected on successful connection.
-    ///  - Pending -> Disconnected on failed connection.
-    Pending,
+    ///  - Pending -> Pending on failed connection.
+    Pending(Instant, usize),
 
     /// A healthy [`TcpStream`] ready to send packets
     ///
@@ -43,14 +47,12 @@ enum TcpStreamState {
     /// Transitions:
     ///  - Connected -> Disconnected on send error
     Connected(TcpStream),
+}
 
-    /// A broken [`TcpStream`] which experienced a failure to connect or send.
-    ///
-    /// Behavior: Try to re-connect on next [`TcpClient::connect`] or [`TcpClient::send`]
-    ///
-    /// Transitions:
-    ///  - Disconnected -> Connected on successful connection.
-    Disconnected,
+impl TcpStreamState {
+    fn reset() -> Self {
+        Self::Pending(Instant::now(), 0)
+    }
 }
 
 /// Connect to a rerun server and send log messages.
@@ -59,52 +61,49 @@ enum TcpStreamState {
 pub struct TcpClient {
     addr: SocketAddr,
     stream_state: TcpStreamState,
-}
-
-impl Default for TcpClient {
-    fn default() -> Self {
-        Self::new(crate::default_server_addr())
-    }
+    disconnected_timeout: Option<Duration>,
 }
 
 impl TcpClient {
-    pub fn new(addr: SocketAddr) -> Self {
+    pub fn new(addr: SocketAddr, disconnected_timeout: Option<Duration>) -> Self {
         Self {
             addr,
-            stream_state: TcpStreamState::Pending,
+            stream_state: TcpStreamState::reset(),
+            disconnected_timeout,
         }
     }
 
     /// Returns `false` on failure. Does nothing if already connected.
     ///
     /// [`Self::send`] will call this.
     pub fn connect(&mut self) -> Result<(), ClientError> {
-        if let TcpStreamState::Connected(_) = self.stream_state {
-            Ok(())
-        } else {
-            re_log::debug!("Connecting to {:?}…", self.addr);
-            let timeout = std::time::Duration::from_secs(5);
-            match TcpStream::connect_timeout(&self.addr, timeout) {
-                Ok(mut stream) => {
-                    re_log::debug!("Connected to {:?}.", self.addr);
-                    if let Err(err) = stream.write(&crate::PROTOCOL_VERSION.to_le_bytes()) {
-                        self.stream_state = TcpStreamState::Disconnected;
-                        Err(ClientError::Send {
+        match self.stream_state {
+            TcpStreamState::Connected(_) => Ok(()),
+            TcpStreamState::Pending(since, tries) => {
+                re_log::debug!("Connecting to {:?}…", self.addr);
+                let timeout = std::time::Duration::from_secs(5);
+                match TcpStream::connect_timeout(&self.addr, timeout) {
+                    Ok(mut stream) => {
+                        re_log::debug!("Connected to {:?}.", self.addr);
+                        if let Err(err) = stream.write(&crate::PROTOCOL_VERSION.to_le_bytes()) {
+                            self.stream_state = TcpStreamState::Pending(since, tries + 1);
+                            Err(ClientError::Send {
+                                addr: self.addr,
+                                err,
+                            })
+                        } else {
+                            self.stream_state = TcpStreamState::Connected(stream);
+                            Ok(())
+                        }
+                    }
+                    Err(err) => {
+                        self.stream_state = TcpStreamState::Pending(since, tries + 1);
+                        Err(ClientError::Connect {
                             addr: self.addr,
                             err,
                         })
-                    } else {
-                        self.stream_state = TcpStreamState::Connected(stream);
-                        Ok(())
                     }
                 }
-                Err(err) => {
-                    self.stream_state = TcpStreamState::Disconnected;
-                    Err(ClientError::Connect {
-                        addr: self.addr,
-                        err,
-                    })
-                }
             }
         }
     }
@@ -118,15 +117,15 @@ impl TcpClient {
         if let TcpStreamState::Connected(stream) = &mut self.stream_state {
             re_log::trace!("Sending a packet of size {}…", packet.len());
             if let Err(err) = stream.write(&(packet.len() as u32).to_le_bytes()) {
-                self.stream_state = TcpStreamState::Disconnected;
+                self.stream_state = TcpStreamState::reset();
                 return Err(ClientError::Send {
                     addr: self.addr,
                     err,
                 });
             }
 
             if let Err(err) = stream.write(packet) {
-                self.stream_state = TcpStreamState::Disconnected;
+                self.stream_state = TcpStreamState::reset();
                 return Err(ClientError::Send {
                     addr: self.addr,
                     err,
@@ -141,23 +140,36 @@ impl TcpClient {
 
     /// Wait until all logged data have been sent.
     pub fn flush(&mut self) {
-        re_log::trace!("Flushing TCP stream…");
-        if let TcpStreamState::Connected(stream) = &mut self.stream_state {
-            if let Err(err) = stream.flush() {
-                re_log::warn!("Failed to flush: {err}");
-                self.stream_state = TcpStreamState::Disconnected;
+        re_log::debug!("Flushing TCP stream…");
+        match &mut self.stream_state {
+            TcpStreamState::Pending(_, _) => {
+                re_log::warn_once!(
+                    "Tried to flush while TCP stream was still Pending. Data was possibly dropped."
+                );
+            }
+            TcpStreamState::Connected(stream) => {
+                if let Err(err) = stream.flush() {
+                    re_log::warn!("Failed to flush: {err}");
+                    self.stream_state = TcpStreamState::reset();
+                }
             }
         }
-        re_log::trace!("TCP stream flushed.");
+        re_log::debug!("TCP stream flushed.");
     }
 
-    /// Check if the underlying [`TcpStream`] has entered the [`TcpStreamState::Disconnected`] state
+    /// Check if the underlying [`TcpStream`] is in the [`TcpStreamState::Pending`] state
+    /// and has reached the timeout threshold.
     ///
     /// Note that this only occurs after a failure to connect or a failure to send.
-    pub fn has_disconnected(&self) -> bool {
+    pub fn has_timed_out(&self) -> bool {
         match self.stream_state {
-            TcpStreamState::Pending | TcpStreamState::Connected(_) => false,
-            TcpStreamState::Disconnected => true,
+            TcpStreamState::Pending(since, tries) => {
+                // If a timeout wasn't provided, never timeout
+                self.disconnected_timeout.map_or(false, |timeout| {
+                    Instant::now().duration_since(since) > timeout && tries > 0
+                })
+            }
+            TcpStreamState::Connected(_) => false,
         }
     }
 }

crates/rerun/src/clap.rs

@@ -106,7 +106,10 @@ impl RerunArgs {
         }
 
         let sink: Box<dyn re_sdk::sink::LogSink> = match self.to_behavior()? {
-            RerunBehavior::Connect(addr) => Box::new(crate::sink::TcpSink::new(addr)),
+            RerunBehavior::Connect(addr) => Box::new(crate::sink::TcpSink::new(
+                addr,
+                crate::sink::TcpSink::DEFAULT_TIMEOUT,
+            )),
 
             RerunBehavior::Save(path) => Box::new(crate::sink::FileSink::new(path)?),
 

crates/rerun_c/src/lib.rs

@@ -141,7 +141,7 @@ pub unsafe extern "C" fn rr_recording_stream_new(
         .expect("invalid tcp_addr string")
         .parse()
         .expect("invalid tcp_addr");
-    let sink = Box::new(TcpSink::new(tcp_addr));
+    let sink = Box::new(TcpSink::new(tcp_addr, TcpSink::DEFAULT_TIMEOUT));
 
     let rec_stream = RecordingStream::new(store_info, batcher_config, sink).unwrap();