Implement forward log reader, rebuild memtable at init

6 files changed, 300 insertions, 78 deletions

diff --git a/src/common.rs b/src/common.rs
index f7aa755..4922c73 100644
--- a/src/common.rs
+++ b/src/common.rs
@@ -12,9 +12,28 @@ pub const ACTIVE_LOG_FILENAME: &str = "db";
 pub const DEFAULT_READ_BUF_SIZE: usize = 1024 * 1024; // 1 MB
 pub const FIELD_SEPARATOR: u8 = b'\x1C';
 pub const ESCAPE_CHARACTER: u8 = b'\x1D';
-pub const SEQ_RECORD_SEP: &[u8] = &[FIELD_SEPARATOR, FIELD_SEPARATOR, ESCAPE_CHARACTER];
-pub const SEQ_LIT_ESCAPE: &[u8] = &[ESCAPE_CHARACTER, ESCAPE_CHARACTER, ESCAPE_CHARACTER];
-pub const SEQ_LIT_FIELD_SEP: &[u8] = &[ESCAPE_CHARACTER, FIELD_SEPARATOR, ESCAPE_CHARACTER];
+
+// Special sequences. Note: these must have the same length!
+// Since the log is read both forwards and backwards, we must have a signal
+// character (ESCAPE_CHARACTER) on both sides of the special sequence.
+pub const SEQ_RECORD_SEP: &[u8] = &[
+    ESCAPE_CHARACTER,
+    FIELD_SEPARATOR,
+    FIELD_SEPARATOR,
+    ESCAPE_CHARACTER,
+];
+pub const SEQ_LIT_ESCAPE: &[u8] = &[
+    ESCAPE_CHARACTER,
+    ESCAPE_CHARACTER,
+    ESCAPE_CHARACTER,
+    ESCAPE_CHARACTER,
+];
+pub const SEQ_LIT_FIELD_SEP: &[u8] = &[
+    ESCAPE_CHARACTER,
+    ESCAPE_CHARACTER,
+    FIELD_SEPARATOR,
+    ESCAPE_CHARACTER,
+];
 
 #[derive(Debug, Eq, PartialEq)]
 pub enum SpecialSequence {
diff --git a/src/lib.rs b/src/lib.rs
index 3b249d8..c0cf00c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -9,7 +9,8 @@ mod secondary_memtable;
 
 pub use common::*;
 use fs2::FileExt;
-pub use log_reader::LogReader;
+pub use log_reader::ForwardLogReader;
+pub use log_reader::ReverseLogReader;
 use primary_memtable::PrimaryMemtable;
 use secondary_memtable::SecondaryMemtable;
 use std::fmt::Debug;
@@ -146,7 +147,7 @@ impl<Field: Eq + Clone + Debug> DB<Field> {
     }
 
     fn initialize(config: &Config<Field>) -> Result<DB<Field>, io::Error> {
-        info!("Initializing DB");
+        info!("Initializing DB...");
         // If data_dir does not exist, create it
         if !fs::exists(&config.data_dir)? {
             fs::create_dir_all(&config.data_dir)?;
@@ -214,13 +215,25 @@ impl<Field: Eq + Clone + Debug> DB<Field> {
             })
             .collect();
 
-        let db = DB::<Field> {
+        let mut db = DB::<Field> {
             config: config.clone(),
-            log_path,
+            log_path: log_path.clone(),
             primary_key_index,
             primary_memtable,
             secondary_memtables,
         };
+
+        info!("Rebuilding memtable indexes...");
+        let mut file = fs::OpenOptions::new().read(true).open(&log_path)?;
+
+        let forward_log_reader = ForwardLogReader::new(&mut file);
+        for record in forward_log_reader {
+            db.update_primary_index(&record);
+            db.update_secondary_indexes(&record);
+        }
+
+        info!("Database ready.");
+
         Ok(db)
     }
 
@@ -295,34 +308,12 @@ impl<Field: Eq + Clone + Debug> DB<Field> {
         file.unlock()?;
 
         debug!("Record appended to log file, lock released");
-        debug!("Updating primary memtable");
 
-        let primary_value =
-            &record.values[self.primary_key_index]
-                .as_indexable()
-                .ok_or(io::Error::new(
-                    io::ErrorKind::InvalidInput,
-                    "Primary key must be an IndexableValue",
-                ))?;
-
-        self.primary_memtable.set(primary_value, record);
+        debug!("Updating primary memtable");
+        self.update_primary_index(record);
 
-        self.secondary_memtables
-            .iter_mut()
-            .for_each(|secondary_memtable| {
-                debug!(
-                    "Updating memtable for index on {:?}",
-                    &secondary_memtable.field
-                );
-                for (index, (schema_field, _)) in self.config.fields.iter().enumerate() {
-                    if schema_field == &secondary_memtable.field {
-                        let key = record.values[index]
-                            .as_indexable()
-                            .expect("Secondary index key was not indexable");
-                        secondary_memtable.set(&key, record);
-                    }
-                }
-            });
+        debug!("Updating secondary memtables");
+        self.update_secondary_indexes(record);
 
         Ok(())
     }
@@ -372,8 +363,8 @@ impl<Field: Eq + Clone + Debug> DB<Field> {
 
         debug!("Lock acquired, searching log file for record");
 
-        let mut log_reader = LogReader::new(&mut file)?;
-        let result = log_reader.find(|record| {
+        let mut reverse_log_reader = ReverseLogReader::new(&mut file)?;
+        let result = reverse_log_reader.find(|record| {
             let record_key = record.values[self.primary_key_index]
                 .as_indexable()
                 .expect("A non-indexable value was stored at key index");
@@ -391,8 +382,13 @@ impl<Field: Eq + Clone + Debug> DB<Field> {
             }
         };
 
-        debug!("Found matching record in log file. Storing result in primary memtable.");
-        self.primary_memtable.set(&query_key, &result_value);
+        debug!("Found matching record in log file.");
+
+        debug!("Updating primary memtable");
+        self.update_primary_index(&result_value);
+
+        debug!("Updating secondary memtables");
+        self.update_secondary_indexes(&result_value);
 
         Ok(Some(result_value))
     }
@@ -472,7 +468,7 @@ impl<Field: Eq + Clone + Debug> DB<Field> {
 
         debug!("Lock acquired, searching log file for record");
 
-        let result = LogReader::new(&mut file)?
+        let result = ReverseLogReader::new(&mut file)?
             .filter(|record| {
                 let record_key = record.values[key_index]
                     .as_indexable()
@@ -496,4 +492,30 @@ impl<Field: Eq + Clone + Debug> DB<Field> {
 
         Ok(result)
     }
+
+    fn update_primary_index(&mut self, record: &Record) {
+        let key = record.values[self.primary_key_index]
+            .as_indexable()
+            .expect("A non-indexable value was stored at key index");
+        self.primary_memtable.set(&key, record);
+    }
+
+    fn update_secondary_indexes(&mut self, record: &Record) {
+        self.secondary_memtables
+            .iter_mut()
+            .for_each(|secondary_memtable| {
+                debug!(
+                    "Updating memtable for index on {:?}",
+                    &secondary_memtable.field
+                );
+                for (index, (schema_field, _)) in self.config.fields.iter().enumerate() {
+                    if schema_field == &secondary_memtable.field {
+                        let key = record.values[index]
+                            .as_indexable()
+                            .expect("Secondary index key was not indexable");
+                        secondary_memtable.set(&key, record);
+                    }
+                }
+            });
+    }
 }
diff --git a/src/log_reader.rs b/src/log_reader.rs
index 0ee0a81..e32ef53 100644
--- a/src/log_reader.rs
+++ b/src/log_reader.rs
@@ -1,16 +1,30 @@
 use super::common::*;
 use rev_buf_reader::RevBufReader;
 use std::fs::{self};
-use std::io::{self, BufRead, Read, Seek};
+use std::io::{self, BufRead, Read, Seek, SeekFrom};
 
-pub struct LogReader<'a> {
+/// There are three special characters that need to be handled:
+/// Here: SC = escape char, FS = field separator.
+/// - FS FS SC  -> actual record separator
+/// - SC FS SC  -> literal FS
+/// - SC SC SC  -> literal SC
+fn validate_special(buf: &[u8]) -> Option<SpecialSequence> {
+    match buf {
+        SEQ_RECORD_SEP => Some(SpecialSequence::RecordSeparator),
+        SEQ_LIT_FIELD_SEP => Some(SpecialSequence::LiteralFieldSeparator),
+        SEQ_LIT_ESCAPE => Some(SpecialSequence::LiteralEscape),
+        _ => None,
+    }
+}
+
+pub struct ReverseLogReader<'a> {
     rev_reader: RevBufReader<&'a mut fs::File>,
 }
 
-impl<'a> LogReader<'a> {
-    pub fn new(file: &mut fs::File) -> Result<LogReader, io::Error> {
+impl<'a> ReverseLogReader<'a> {
+    pub fn new(file: &mut fs::File) -> Result<ReverseLogReader, io::Error> {
         let rev_reader = RevBufReader::new(file);
-        Ok(LogReader { rev_reader })
+        Ok(ReverseLogReader { rev_reader })
     }
 
     fn read_record(&mut self) -> Result<Option<Record>, io::Error> {
@@ -48,7 +62,7 @@ impl<'a> LogReader<'a> {
                 SpecialSequence::RecordSeparator => {
                     // The record is complete, so we can break out of the loop.
                     // Move the cursor back to the beginning of the special sequence.
-                    self.rev_reader.seek_relative(3)?;
+                    self.rev_reader.seek_relative(SEQ_RECORD_SEP.len() as i64)?;
                     break;
                 }
                 SpecialSequence::LiteralFieldSeparator => {
@@ -68,7 +82,7 @@ impl<'a> LogReader<'a> {
     }
 
     fn read_special_sequence(&mut self) -> Result<SpecialSequence, io::Error> {
-        let mut special_buf: Vec<u8> = vec![0; 3];
+        let mut special_buf: Vec<u8> = vec![0; SEQ_RECORD_SEP.len()];
         self.rev_reader.read_exact(&mut special_buf)?;
 
         match validate_special(&special_buf.as_slice()) {
@@ -81,7 +95,7 @@ impl<'a> LogReader<'a> {
     }
 }
 
-impl Iterator for LogReader<'_> {
+impl Iterator for ReverseLogReader<'_> {
     type Item = Record;
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -93,16 +107,87 @@ impl Iterator for LogReader<'_> {
     }
 }
 
-/// There are three special characters that need to be handled:
-/// Here: SC = escape char, FS = field separator.
-/// - FS FS SC  -> actual record separator
-/// - SC FS SC  -> literal FS
-/// - SC SC SC  -> literal SC
-fn validate_special(buf: &[u8]) -> Option<SpecialSequence> {
-    match buf {
-        SEQ_RECORD_SEP => Some(SpecialSequence::RecordSeparator),
-        SEQ_LIT_FIELD_SEP => Some(SpecialSequence::LiteralFieldSeparator),
-        SEQ_LIT_ESCAPE => Some(SpecialSequence::LiteralEscape),
-        _ => None,
+pub struct ForwardLogReader<'a> {
+    reader: io::BufReader<&'a mut fs::File>,
+}
+
+impl<'a> ForwardLogReader<'a> {
+    pub fn new(file: &mut fs::File) -> ForwardLogReader {
+        let reader = io::BufReader::new(file);
+        ForwardLogReader { reader }
+    }
+
+    fn read_record(&mut self) -> Result<Option<Record>, io::Error> {
+        // The buffer that stores the bytes read from the file.
+        let mut read_buf: Vec<u8> = Vec::new();
+        // The buffer that stores all the bytes of the record read so far in reverse order.
+        let mut result_buf: Vec<u8> = Vec::new();
+
+        // Try reading a byte from the file.
+        // If we've reached the end of the file, return None.
+        let mut peek_buf = vec![0, 1];
+        match self.reader.read_exact(&mut peek_buf) {
+            Ok(_) => {
+                // Go back one byte (not sure why you need to seek by -2 here?)
+                self.reader.seek_relative(-2)?;
+            }
+            Err(ref e) if e.kind() == io::ErrorKind::UnexpectedEof => {
+                return Ok(None);
+            }
+            Err(e) => {
+                return Err(e);
+            }
+        }
+
+        loop {
+            read_buf.clear();
+            self.reader.read_until(ESCAPE_CHARACTER, &mut read_buf)?;
+            self.reader.seek_relative(-1)?;
+            result_buf.extend(&read_buf[..read_buf.len() - 1]);
+
+            // Otherwise, we must have encountered an escape character.
+            match self.read_special_sequence()? {
+                SpecialSequence::RecordSeparator => {
+                    // The record is complete, so we can break out of the loop.
+                    break;
+                }
+                SpecialSequence::LiteralFieldSeparator => {
+                    // The field separator is escaped, so we need to add it to the result buffer.
+                    result_buf.push(FIELD_SEPARATOR);
+                }
+                SpecialSequence::LiteralEscape => {
+                    // The escape character is escaped, so we need to add it to the result buffer.
+                    result_buf.push(ESCAPE_CHARACTER);
+                }
+            }
+        }
+
+        let record = Record::deserialize(&result_buf);
+        Ok(Some(record))
+    }
+
+    fn read_special_sequence(&mut self) -> Result<SpecialSequence, io::Error> {
+        let mut special_buf: Vec<u8> = vec![0; SEQ_RECORD_SEP.len()];
+        self.reader.read_exact(&mut special_buf)?;
+
+        match validate_special(&special_buf.as_slice()) {
+            Some(special) => Ok(special),
+            None => Err(io::Error::new(
+                io::ErrorKind::InvalidData,
+                "Not a special sequence",
+            )),
+        }
+    }
+}
+
+impl Iterator for ForwardLogReader<'_> {
+    type Item = Record;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.read_record() {
+            Ok(Some(record)) => Some(record),
+            Ok(None) => None,
+            Err(err) => panic!("Error reading record: {:?}", err),
+        }
     }
 }
diff --git a/tests/integration.rs b/tests/integration.rs
index 3e132bc..2aada21 100644
--- a/tests/integration.rs
+++ b/tests/integration.rs
@@ -1,6 +1,6 @@
 use env_logger;
 use log_db;
-use log_db::{LogReader, Record, RecordFieldType, RecordValue, DB};
+use log_db::{ForwardLogReader, Record, RecordFieldType, RecordValue, ReverseLogReader, DB};
 use serial_test::serial;
 use std::fs;
 use std::path::Path;
@@ -18,6 +18,8 @@ enum Field {
 #[test]
 #[serial]
 fn test_initialize() {
+    init_test();
+
     let _db = DB::configure()
         .data_dir(TEST_DATA_DIR)
         .fields(&vec![
@@ -29,13 +31,14 @@ fn test_initialize() {
         .initialize()
         .expect("Failed to initialize DB instance");
 
-    // Clean up
-    std::fs::remove_dir_all(TEST_DATA_DIR.to_string()).unwrap();
+    cleanup_test();
 }
 
 #[test]
 #[serial]
-fn test_upsert_and_get_with_memtable() {
+fn test_upsert_and_get_with_primary_memtable() {
+    init_test();
+
     let mut db = DB::configure()
         .data_dir(TEST_DATA_DIR)
         .fields(&vec![
@@ -64,14 +67,14 @@ fn test_upsert_and_get_with_memtable() {
         _ => false,
     });
 
-    // Clean up
-    std::fs::remove_dir_all(TEST_DATA_DIR.to_string()).unwrap();
+    cleanup_test();
 }
 
 #[test]
 #[serial]
 fn test_upsert_and_get_without_memtable() {
-    let _ = env_logger::builder().is_test(true).try_init();
+    init_test();
+
     let mut db = DB::configure()
         .data_dir(TEST_DATA_DIR)
         .memtable_capacity(0)
@@ -147,13 +150,14 @@ fn test_upsert_and_get_without_memtable() {
         _ => false,
     });
 
-    // Clean up
-    std::fs::remove_dir_all(TEST_DATA_DIR.to_string()).unwrap();
+    cleanup_test();
 }
 
 #[test]
 #[serial]
 fn test_upsert_fails_on_invalid_number_of_values() {
+    init_test();
+
     let mut db = DB::configure()
         .data_dir(TEST_DATA_DIR)
         .fields(&vec![
@@ -174,13 +178,14 @@ fn test_upsert_fails_on_invalid_number_of_values() {
     };
     assert!(db.upsert(&record).is_err());
 
-    // Clean up
-    std::fs::remove_dir_all(TEST_DATA_DIR.to_string()).unwrap();
+    cleanup_test();
 }
 
 #[test]
 #[serial]
 fn test_upsert_fails_on_invalid_value_type() {
+    init_test();
+
     let mut db = DB::configure()
         .data_dir(TEST_DATA_DIR)
         .fields(&vec![
@@ -201,41 +206,85 @@ fn test_upsert_fails_on_invalid_value_type() {
     };
     assert!(db.upsert(&record).is_err());
 
-    // Clean up
-    std::fs::remove_dir_all(TEST_DATA_DIR.to_string()).unwrap();
+    cleanup_test();
 }
 
 #[test]
 #[serial]
-fn test_log_reader_fixture_db1() {
-    let _ = env_logger::builder().is_test(true).try_init();
+fn test_reverse_log_reader_fixture_db1() {
+    init_test();
+
     let db_path = Path::new(TEST_RESOURCES_DIR).join("test_db1");
     let mut file = fs::OpenOptions::new()
         .read(true)
         .open(&db_path)
         .expect("Failed to open file");
-    let mut log_reader = LogReader::new(&mut file).unwrap();
+    let mut reverse_log_reader = ReverseLogReader::new(&mut file).unwrap();
 
     // There are two records in the log with "schema": Int, Null
 
-    let last_record = log_reader.next().expect("Failed to read the last record");
+    let last_record = reverse_log_reader
+        .next()
+        .expect("Failed to read the last record");
     assert!(match last_record.values.as_slice() {
         [RecordValue::Int(10), RecordValue::Null] => true,
         _ => false,
     });
 
-    let first_record = log_reader.next().expect("Failed to read the first record");
+    let first_record = reverse_log_reader
+        .next()
+        .expect("Failed to read the first record");
     assert!(match first_record.values.as_slice() {
         // Note: the int value is equal to the escape byte
         [RecordValue::Int(0x1D), RecordValue::Null] => true,
         _ => false,
     });
+
+    assert!(reverse_log_reader.next().is_none());
+
+    cleanup_test();
+}
+
+#[test]
+#[serial]
+fn test_forward_log_reader_fixture_db1() {
+    init_test();
+
+    let db_path = Path::new(TEST_RESOURCES_DIR).join("test_db1");
+    let mut file = fs::OpenOptions::new()
+        .read(true)
+        .open(&db_path)
+        .expect("Failed to open file");
+    let mut forward_log_reader = ForwardLogReader::new(&mut file);
+
+    // There are two records in the log with "schema": Int, Null
+
+    let first_record = forward_log_reader
+        .next()
+        .expect("Failed to read the first record");
+    assert!(match first_record.values.as_slice() {
+        [RecordValue::Int(0x1D), RecordValue::Null] => true,
+        _ => false,
+    });
+
+    let last_record = forward_log_reader
+        .next()
+        .expect("Failed to read the last record");
+    assert!(match last_record.values.as_slice() {
+        [RecordValue::Int(10), RecordValue::Null] => true,
+        _ => false,
+    });
+
+    assert!(forward_log_reader.next().is_none());
+
+    cleanup_test();
 }
 
 #[test]
 #[serial]
 fn test_upsert_and_get_from_secondary_memtable() {
-    let _ = env_logger::builder().is_test(true).try_init();
+    init_test();
+
     let mut db = DB::configure()
         .data_dir(TEST_DATA_DIR)
         .fields(&vec![
@@ -286,6 +335,53 @@ fn test_upsert_and_get_from_secondary_memtable() {
 
     assert_eq!(johns.len(), 2);
 
-    // Clean up
-    std::fs::remove_dir_all(TEST_DATA_DIR.to_string()).unwrap();
+    cleanup_test();
+}
+
+#[test]
+#[serial]
+fn test_initialize_and_read_from_primary_memtable_fixture_db2() {
+    init_test();
+
+    // Copy the fixture DB to the test data directory
+    fs::create_dir_all(TEST_DATA_DIR).expect("Failed to create the test data directory");
+    fs::copy(
+        &Path::new(TEST_RESOURCES_DIR).join("test_db2"),
+        &Path::new(TEST_DATA_DIR).join("db"),
+    )
+    .expect("Failed to copy the fixture DB");
+
+    let mut db = DB::configure()
+        .data_dir(TEST_DATA_DIR)
+        .fields(&vec![
+            (Field::Id, RecordFieldType::Int),
+            (Field::Name, RecordFieldType::String),
+            (Field::Data, RecordFieldType::Bytes),
+        ])
+        .primary_key(Field::Id)
+        .initialize()
+        .expect("Failed to initialize DB instance");
+
+    // Delete the DB so that any results must come from a memtable
+    fs::remove_file(Path::new(TEST_DATA_DIR).join("db")).expect("Failed to delete the DB log file");
+
+    let result = db.get(&RecordValue::Int(1)).unwrap().unwrap();
+
+    // Check that the IDs match
+    let expected = RecordValue::Int(1);
+    assert!(match (&result.values[0], &expected) {
+        (RecordValue::Int(a), RecordValue::Int(b)) => a == b,
+        _ => false,
+    });
+
+    cleanup_test();
+}
+
+fn init_test() {
+    let _ = env_logger::builder().is_test(true).try_init();
+}
+
+fn cleanup_test() {
+    std::fs::remove_dir_all(TEST_DATA_DIR.to_string())
+        .unwrap_or_else(|e| eprintln!("Failed to delete the test data directory: {:?}", e));
 }
diff --git a/tests/resources/test_db1 b/tests/resources/test_db1
index e75edf5..1f19a86 100644
--- a/tests/resources/test_db1
+++ b/tests/resources/test_db1
diff --git a/tests/resources/test_db2 b/tests/resources/test_db2
new file mode 100644
index 0000000..2448640
--- /dev/null
+++ b/tests/resources/test_db2