Improve robustness of rrd loader (#11302)

Author: emilk

.vscode/launch.json

@@ -93,7 +93,7 @@
             }
         },
         {
-            "name": "Debug 'rerun' colmap.rrd from url",
+            "name": "Debug 'rerun' bar_chart.rrd from url",
             "type": "lldb",
             "request": "launch",
             "cargo": {
@@ -110,9 +110,12 @@
                 }
             },
             "args": [
-                "https://demo.rerun.io/commit/0f89b62/examples/colmap/data.rrd"
+                "https://app.rerun.io/version/0.25.0/examples/snippets/bar_chart.rrd"
             ],
-            "cwd": "${workspaceFolder}"
+            "cwd": "${workspaceFolder}",
+            "env": {
+                "RUST_LOG": "trace"
+            }
         },
         {
             "name": "Debug 'minimal' example",

crates/store/re_log_encoding/src/codec/file/decoder.rs

@@ -19,9 +19,8 @@ pub(crate) fn decode_to_app(
     app_id_injector: &mut impl ApplicationIdInjector,
     data: &mut impl std::io::Read,
 ) -> Result<(u64, Option<LogMsg>), DecodeError> {
-    let mut read_bytes = 0u64;
     let header = MessageHeader::decode(data)?;
-    read_bytes += std::mem::size_of::<MessageHeader>() as u64 + header.len;
+    let read_bytes = std::mem::size_of::<MessageHeader>() as u64 + header.len;
 
     let mut buf = vec![0; header.len as usize];
     data.read_exact(&mut buf[..])?;
@@ -39,9 +38,8 @@ pub(crate) fn decode_to_app(
 pub(crate) fn decode_to_transport(
     data: &mut impl std::io::Read,
 ) -> Result<(u64, Option<re_protos::log_msg::v1alpha1::log_msg::Msg>), DecodeError> {
-    let mut read_bytes = 0u64;
     let header = MessageHeader::decode(data)?;
-    read_bytes += std::mem::size_of::<MessageHeader>() as u64 + header.len;
+    let read_bytes = std::mem::size_of::<MessageHeader>() as u64 + header.len;
 
     let mut buf = vec![0; header.len as usize];
     data.read_exact(&mut buf[..])?;

crates/store/re_log_encoding/src/decoder/mod.rs

@@ -257,13 +257,7 @@ impl<R: std::io::Read> Decoder<R> {
 
         let (version, options) = options_from_bytes(&data)?;
 
-        Ok(Self {
-            version,
-            options,
-            read: Reader::Raw(read),
-            size_bytes: FileHeader::SIZE as _,
-            app_id_cache: CachingApplicationIdInjector::default(),
-        })
+        Ok(Self::new_with_options(options, version, read))
     }
 
     pub fn new_with_options(options: EncodingOptions, version: CrateVersion, read: R) -> Self {
@@ -394,20 +388,20 @@ impl<R: std::io::Read> Decoder<R> {
             },
         };
 
-        let Some(msg) = msg else {
+        if let Some(msg) = msg {
+            Some(Ok(msg))
+        } else {
             // we might have a concatenated stream, so we peek beyond end of file marker to see
             if self.peek_file_header() {
                 re_log::debug!(
                     "Reached end of stream, but it seems we have a concatenated file, continuing"
                 );
-                return self.next_impl(decoder);
+                self.next_impl(decoder)
+            } else {
+                re_log::trace!("Reached end of stream, iterator complete");
+                None
             }
-
-            re_log::trace!("Reached end of stream, iterator complete");
-            return None;
-        };
-
-        Some(Ok(msg))
+        }
     }
 
     /// Peeks ahead in search of additional `FileHeader`s in the stream.
@@ -424,11 +418,14 @@ impl<R: std::io::Read> Decoder<R> {
                 }
 
                 let mut read = std::io::Cursor::new(read.buffer());
-                if FileHeader::decode(&mut read).is_err() {
-                    return false;
+                match FileHeader::decode(&mut read) {
+                    Ok(_) => true,
+                    Err(DecodeError::Read { .. }) => false,
+                    Err(err) => {
+                        re_log::warn_once!("Error when expecting rrd header: {err}");
+                        false
+                    }
                 }
-
-                true
             }
         }
     }

crates/store/re_log_encoding/src/decoder/stream.rs

@@ -64,6 +64,9 @@ enum State {
     /// Compression is only applied to individual `ArrowMsg`s, instead of
     /// the entire stream.
     Message(crate::codec::file::MessageHeader),
+
+    /// Stop reading
+    Aborted,
 }
 
 impl StreamDecoder {
@@ -109,9 +112,31 @@ impl StreamDecoder {
     fn try_read_impl(&mut self) -> Result<Option<LogMsg>, DecodeError> {
         match self.state {
             State::StreamHeader => {
+                let is_first_header = self.chunks.num_read() == 0;
                 if let Some(header) = self.chunks.try_read(FileHeader::SIZE) {
+                    re_log::trace!(?header, "Decoding StreamHeader");
+
                     // header contains version and compression options
-                    let (version, options) = options_from_bytes(header)?;
+                    let (version, options) = match options_from_bytes(header) {
+                        Ok(ok) => ok,
+                        Err(err) => {
+                            // We expected a header, but didn't find one!
+                            if is_first_header {
+                                return Err(err);
+                            } else {
+                                re_log::error!("Trailing bytes in rrd stream: {header:?} ({err})");
+                                self.state = State::Aborted;
+                                return Ok(None);
+                            }
+                        }
+                    };
+
+                    re_log::trace!(
+                        version = version.to_string(),
+                        ?options,
+                        "Found Stream Header"
+                    );
+
                     self.version = Some(version);
                     self.options = options;
 
@@ -131,6 +156,9 @@ impl StreamDecoder {
                     .try_read(crate::codec::file::MessageHeader::SIZE_BYTES)
                 {
                     let header = crate::codec::file::MessageHeader::from_bytes(bytes)?;
+
+                    re_log::trace!(?header, "MessageHeader");
+
                     self.state = State::Message(header);
                     // we might have data left in the current chunk,
                     // immediately try to read the message content
@@ -139,23 +167,42 @@ impl StreamDecoder {
             }
             State::Message(header) => {
                 if let Some(bytes) = self.chunks.try_read(header.len as usize) {
+                    re_log::trace!(?header, "Read message");
+
                     let message = crate::codec::file::decoder::decode_bytes_to_app(
                         &mut self.app_id_cache,
                         header.kind,
                         bytes,
                     )?;
+
                     if let Some(mut message) = message {
+                        re_log::trace!(
+                            "LogMsg::{}",
+                            match message {
+                                LogMsg::SetStoreInfo { .. } => "SetStoreInfo",
+                                LogMsg::ArrowMsg { .. } => "ArrowMsg",
+                                LogMsg::BlueprintActivationCommand { .. } => {
+                                    "BlueprintActivationCommand"
+                                }
+                            }
+                        );
+
                         propagate_version(&mut message, self.version);
                         self.state = State::MessageHeader;
                         return Ok(Some(message));
                     } else {
+                        re_log::trace!("End of stream - expecting a new Streamheader");
+
                         // `None` means end of stream, but there might be concatenated streams,
                         // so try to read another one.
                         self.state = State::StreamHeader;
                         return self.try_read();
                     }
                 }
             }
+            State::Aborted => {
+                return Ok(None);
+            }
         }
 
         Ok(None)
@@ -182,6 +229,9 @@ struct ChunkBuffer {
 
     /// How many bytes of valid data are currently in `self.buffer`.
     buffer_fill: usize,
+
+    /// How many bytes have been read with [`Self::try_read`] so far?
+    num_read: usize,
 }
 
 impl ChunkBuffer {
@@ -190,6 +240,7 @@ impl ChunkBuffer {
             queue: VecDeque::with_capacity(16),
             buffer: Vec::with_capacity(1024),
             buffer_fill: 0,
+            num_read: 0,
         }
     }
 
@@ -200,6 +251,11 @@ impl ChunkBuffer {
         self.queue.push_back(Chunk::new(chunk));
     }
 
+    /// How many bytes have been read with [`Self::try_read`] so far?
+    fn num_read(&self) -> usize {
+        self.num_read
+    }
+
     /// Attempt to read exactly `n` bytes out of the queued chunks.
     ///
     /// Returns `None` if there is not enough data to return a slice of `n` bytes.
@@ -239,7 +295,7 @@ impl ChunkBuffer {
             // followed by another call to `try_read(N)` with the same `N`
             // won't erroneously return the same bytes
             self.buffer_fill = 0;
-
+            self.num_read += n;
             Some(&self.buffer[..])
         } else {
             None

crates/store/re_log_encoding/src/stream_rrd_from_http.rs

@@ -75,11 +75,14 @@ pub fn stream_rrd_from_http(url: String, on_msg: Arc<HttpMessageCallback>) {
         move |part| match part {
             Ok(part) => match part {
                 ehttp::streaming::Part::Response(ehttp::PartialResponse {
+                    url,
                     ok,
                     status,
                     status_text,
-                    ..
+                    headers,
                 }) => {
+                    re_log::trace!("{url} status: {status} - {status_text}");
+                    re_log::trace!("{url} headers: {headers:#?}");
                     if ok {
                         re_log::debug!("Decoding .rrd file from {url:?}…");
                         ControlFlow::Continue(())