Fix database importer decoding error (#784)

## Usage and product changes
Stabilize the database import function by eliminating rare decoding
errors that could occur during the import of large datasets. All the
exported files that the import function could not process are still
valid and should be correctly imported after the proposed changes.

## Implementation
* For better control of the manual bytes processing, replace `VecDeque`
with `prost::bytes::*` and isolate the file reader's moving in a single
function (`Iterator::next`).
* Since `read` functions can read fewer bytes than requested even
without meeting EOF, replace single `if` calls with `loop`s <- this was
probably the main issue before.

I was not able to reduce the bug to a single test case, even with
thousands of records (however, we had a report of a failure after
processing 200 entries), so no new behavior tests were introduced. I
decided not to spend more time searching for such cases.

Author: farost

rust/src/database/migration.rs

@@ -19,14 +19,16 @@
 
 use std::{
     cmp::max,
-    collections::VecDeque,
     fs::{File, OpenOptions},
-    io::{BufRead, BufWriter, Read, Write},
+    io::{BufRead, Read, Write},
     marker::PhantomData,
     path::Path,
 };
 
-use prost::Message;
+use prost::{
+    bytes::{Buf, BytesMut},
+    Message,
+};
 use typedb_protocol::migration::Item as MigrationItemProto;
 
 use crate::{error::MigrationError, Error, Result};
@@ -40,35 +42,42 @@ pub(crate) enum DatabaseExportAnswer {
 
 pub struct ProtoMessageIterator<M: Message + Default, R: BufRead> {
     reader: R,
-    buffer: VecDeque<u8>,
+    buffer: BytesMut,
     _phantom_data: PhantomData<M>,
 }
 
 impl<M: Message + Default, R: BufRead> ProtoMessageIterator<M, R> {
-    const BUF_CAPACITY: usize = 1024;
+    const BUF_CAPACITY: usize = 8 * 1024;
     // prost's length delimiters take up to 10 bytes
     const MAX_LENGTH_DELIMITER_LEN: usize = 10;
 
     pub fn new(reader: R) -> Self {
-        Self { reader, buffer: VecDeque::with_capacity(Self::BUF_CAPACITY), _phantom_data: PhantomData }
+        Self { reader, buffer: BytesMut::with_capacity(Self::BUF_CAPACITY), _phantom_data: PhantomData }
     }
 
-    fn try_read_more(&mut self, bytes_to_read: usize) -> std::io::Result<usize> {
-        let mut addition = vec![0; bytes_to_read];
+    fn read_more(&mut self, bytes_to_read: usize) -> std::io::Result<usize> {
+        if self.buffer.capacity() - self.buffer.len() < bytes_to_read {
+            self.buffer.reserve(max(bytes_to_read, Self::BUF_CAPACITY));
+        }
+        let mut addition = vec![0u8; max(bytes_to_read, 1)];
         let bytes_read = self.reader.read(&mut addition)?;
-        self.buffer.extend(&addition[..bytes_read]);
+        self.buffer.extend_from_slice(&addition[..bytes_read]);
         Ok(bytes_read)
     }
 
-    fn try_get_next_message_len(&mut self) -> Result<Option<usize>> {
+    fn decode_next_len(&mut self) -> Result<Option<(usize /*len*/, usize /*consumed*/)>> {
         loop {
-            if let Ok(len) = prost::decode_length_delimiter(&mut self.buffer) {
-                return Ok(Some(len));
-            } else {
-                if self.buffer.len() < Self::MAX_LENGTH_DELIMITER_LEN {
-                    assert!(Self::MAX_LENGTH_DELIMITER_LEN < Self::BUF_CAPACITY);
-                    let to_read = max(Self::MAX_LENGTH_DELIMITER_LEN, Self::BUF_CAPACITY - self.buffer.len());
-                    match self.try_read_more(to_read) {
+            let mut cursor: &[u8] = &self.buffer;
+            match prost::decode_length_delimiter(&mut cursor) {
+                Ok(len) => {
+                    let consumed = self.buffer.len() - cursor.len();
+                    return Ok(Some((len, consumed)));
+                }
+                Err(_) => {
+                    if self.buffer.len() >= Self::MAX_LENGTH_DELIMITER_LEN {
+                        return Err(Error::Migration(MigrationError::CannotDecodeImportedConceptLength));
+                    }
+                    match self.read_more(Self::MAX_LENGTH_DELIMITER_LEN - self.buffer.len()) {
                         Ok(bytes_read) if bytes_read == 0 => {
                             return if self.buffer.is_empty() {
                                 Ok(None)
@@ -79,40 +88,34 @@ impl<M: Message + Default, R: BufRead> ProtoMessageIterator<M, R> {
                         Err(_) => return Err(Error::Migration(MigrationError::CannotDecodeImportedConceptLength)),
                         Ok(_) => continue,
                     }
-                } else {
-                    return Err(Error::Migration(MigrationError::CannotDecodeImportedConceptLength));
                 }
             }
         }
     }
-
-    fn get_message_buf(&mut self, len: usize) -> VecDeque<u8> {
-        let message_buf = self.buffer.split_off(len);
-        std::mem::replace(&mut self.buffer, message_buf)
-    }
 }
 
 impl<M: Message + Default, R: BufRead> Iterator for ProtoMessageIterator<M, R> {
     type Item = Result<M>;
 
     fn next(&mut self) -> Option<Self::Item> {
-        let message_len = match self.try_get_next_message_len() {
-            Ok(Some(len)) => len,
+        let (message_len, consumed) = match self.decode_next_len() {
+            Ok(Some(res)) => res,
             Ok(None) => return None,
             Err(err) => return Some(Err(err)),
         };
 
-        if self.buffer.len() < message_len {
-            let required = message_len - self.buffer.len();
-            let to_read = max(required, Self::BUF_CAPACITY);
-            match self.try_read_more(to_read) {
-                Ok(bytes_read) if bytes_read >= required => {}
-                _ => return Some(Err(Error::Migration(MigrationError::CannotDecodeImportedConcept))),
+        let required = consumed + message_len;
+        while self.buffer.len() < required {
+            let to_read = required - self.buffer.len();
+            match self.read_more(max(to_read, Self::BUF_CAPACITY)) {
+                Ok(0) | Err(_) => return Some(Err(Error::Migration(MigrationError::CannotDecodeImportedConcept))),
+                Ok(_) => {}
             }
         }
 
-        let mut message_buf = self.get_message_buf(message_len);
-        Some(M::decode(&mut message_buf).map_err(|_| Error::Migration(MigrationError::CannotDecodeImportedConcept)))
+        self.buffer.advance(consumed);
+        let message_bytes = self.buffer.split_to(message_len).freeze();
+        Some(M::decode(message_bytes).map_err(|_| Error::Migration(MigrationError::CannotDecodeImportedConcept)))
     }
 }