aboutsummaryrefslogtreecommitdiffstats
path: root/log_db/src
diff options
context:
space:
mode:
authorJan Tuomi <jan@jantuomi.fi>2024-10-07 12:43:16 +0300
committerJan Tuomi <jan@jantuomi.fi>2024-10-07 12:43:16 +0300
commita4efc83f7f29c6ef8ac9b4c57201c9ecc31266b3 (patch)
treea18c856a5fc15fe509d3e3277e345176608dee21 /log_db/src
parent30295ba5e1bdbfa69ad98fc5908ee4b1aabbd2e5 (diff)
Add py_bindings lib, move to monorepo structure
Diffstat (limited to 'log_db/src')
-rw-r--r--log_db/src/common.rs308
-rw-r--r--log_db/src/forward_log_reader.rs124
-rw-r--r--log_db/src/lib.rs794
-rw-r--r--log_db/src/primary_memtable.rs79
-rw-r--r--log_db/src/reverse_log_reader.rs354
-rw-r--r--log_db/src/secondary_memtable.rs178
6 files changed, 1837 insertions, 0 deletions
diff --git a/log_db/src/common.rs b/log_db/src/common.rs
new file mode 100644
index 0000000..c2af766
--- /dev/null
+++ b/log_db/src/common.rs
@@ -0,0 +1,308 @@
+use std::fmt::Display;
+use std::fs::{metadata, File};
+use std::io::{self};
+use std::path::PathBuf;
+
+// For Unix-like systems
+#[cfg(unix)]
+use std::os::unix::fs::MetadataExt;
+
+// For Windows
+#[cfg(windows)]
+use std::os::windows::fs::MetadataExt;
+
+pub const ACTIVE_LOG_FILENAME: &str = "db";
+pub const EXCL_LOCK_REQUEST_FILENAME: &str = "excl_lock_req";
+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 TEST_RESOURCES_DIR: &str = "tests/resources";
+
+// 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,
+];
+
+/// There are three special sequences that need to be handled:
+/// Here: SC = escape char, FS = field separator.
+/// - SC FS FS SC -> actual record separator
+/// - SC SC FS SC -> literal FS
+/// - SC SC SC SC -> literal SC
+///
+/// Returns SpecialSequence or None if not valid.
+pub 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,
+ }
+}
+
+#[derive(Debug, Eq, PartialEq)]
+pub enum SpecialSequence {
+ RecordSeparator,
+ LiteralFieldSeparator,
+ LiteralEscape,
+}
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub enum MemtableEvictPolicy {
+ LeastWritten,
+ LeastRead,
+ LeastReadOrWritten,
+}
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub enum WriteDurability {
+ /// Changes are written to an application-level write buffer without flushing to the OS write buffer or syncing to disk.
+ /// The buffered writer will batch writes to the OS buffer for maximum performance.
+ /// Offers the lowest durability guarantees but is very fast.
+ Async,
+ /// Changes are written to the OS write buffer but not immediately synced to disk.
+ /// Offers better durability guarantees than Async but is slower.
+ /// This is generally recommended. Most OSes will sync the write buffer to disk within a few seconds.
+ Flush,
+ /// Changes are written to the OS write buffer and synced to disk immediately.
+ /// Offers the best durability guarantees but is the slowest.
+ FlushSync,
+}
+
+impl Display for WriteDurability {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
+ write!(f, "{:?}", self)?;
+ Ok(())
+ }
+}
+
+#[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq, Hash)]
+pub enum IndexableValue {
+ Int(i64),
+ String(String),
+}
+
+#[derive(Debug, Clone)]
+pub enum RecordFieldType {
+ Int,
+ Float,
+ String,
+ Bytes,
+}
+
+#[derive(Debug, Clone)]
+pub struct RecordField {
+ pub field_type: RecordFieldType,
+ pub nullable: bool,
+}
+
+impl RecordField {
+ pub fn int() -> Self {
+ RecordField {
+ field_type: RecordFieldType::Int,
+ nullable: false,
+ }
+ }
+
+ pub fn float() -> Self {
+ RecordField {
+ field_type: RecordFieldType::Float,
+ nullable: false,
+ }
+ }
+
+ pub fn string() -> Self {
+ RecordField {
+ field_type: RecordFieldType::String,
+ nullable: false,
+ }
+ }
+
+ pub fn bytes() -> Self {
+ RecordField {
+ field_type: RecordFieldType::Bytes,
+ nullable: false,
+ }
+ }
+
+ pub fn nullable(&mut self) -> Self {
+ let mut new = self.clone();
+ new.nullable = true;
+ new
+ }
+}
+
+#[derive(Debug, Clone)]
+pub enum RecordValue {
+ Null,
+ Int(i64),
+ Float(f64),
+ String(String),
+ Bytes(Vec<u8>),
+}
+
+impl RecordValue {
+ pub fn serialize(&self) -> Vec<u8> {
+ match self {
+ RecordValue::Null => {
+ vec![0] // Tag for Null
+ }
+ RecordValue::Int(i) => {
+ let mut bytes = vec![1]; // Tag for Int
+ let data_bytes = escape_bytes(&i.to_be_bytes());
+ bytes.extend(&data_bytes);
+ bytes
+ }
+ RecordValue::Float(f) => {
+ let mut bytes = vec![2]; // Tag for Float
+ let data_bytes = escape_bytes(&f.to_be_bytes());
+ bytes.extend(&data_bytes);
+ bytes
+ }
+ RecordValue::String(s) => {
+ let mut bytes = vec![3]; // Tag for String
+ let length = s.len() as u64;
+ let length_bytes = escape_bytes(&length.to_be_bytes());
+ bytes.extend(&length_bytes);
+ let data_bytes = escape_bytes(s.as_bytes());
+ bytes.extend(&data_bytes);
+ bytes
+ }
+ RecordValue::Bytes(b) => {
+ let mut bytes = vec![4]; // Tag for Bytes
+ let length = b.len() as u64;
+ let length_bytes = escape_bytes(&length.to_be_bytes());
+ bytes.extend(&length_bytes);
+ let data_bytes = escape_bytes(b);
+ bytes.extend(&data_bytes);
+ bytes
+ }
+ }
+ }
+
+ /// Deserialize a RecordValue from a byte slice.
+ /// Returns the deserialized RecordValue and the number of bytes consumed.
+ pub fn deserialize(bytes: &[u8]) -> (RecordValue, usize) {
+ match bytes[0] {
+ 0 => (RecordValue::Null, 1),
+ 1 => {
+ let mut int_bytes = [0; 8];
+ int_bytes.copy_from_slice(&bytes[1..1 + 8]);
+ (RecordValue::Int(i64::from_be_bytes(int_bytes)), 1 + 8)
+ }
+ 2 => {
+ let mut float_bytes = [0; 8];
+ float_bytes.copy_from_slice(&bytes[1..1 + 8]);
+ (RecordValue::Float(f64::from_be_bytes(float_bytes)), 1 + 8)
+ }
+ 3 => {
+ let length_bytes = &bytes[1..1 + 8];
+ let length = u64::from_be_bytes(length_bytes.try_into().unwrap()) as usize;
+ (
+ RecordValue::String(
+ String::from_utf8(bytes[1 + 8..1 + 8 + length].to_vec()).unwrap(),
+ ),
+ 1 + 8 + length,
+ )
+ }
+ 4 => {
+ let length_bytes = &bytes[1..1 + 8];
+ let length = u64::from_be_bytes(length_bytes.try_into().unwrap()) as usize;
+ (
+ RecordValue::Bytes(bytes[1 + 8..1 + 8 + length].to_vec()),
+ 1 + 8 + length,
+ )
+ }
+ _ => panic!("Invalid tag: {}", bytes[0]),
+ }
+ }
+
+ pub fn as_indexable(&self) -> Option<IndexableValue> {
+ match self {
+ RecordValue::Int(i) => Some(IndexableValue::Int(*i)),
+ RecordValue::String(s) => Some(IndexableValue::String(s.clone())),
+ _ => None,
+ }
+ }
+}
+
+#[derive(Debug, Clone)]
+pub struct Record {
+ pub values: Vec<RecordValue>,
+}
+
+impl Record {
+ pub fn serialize(&self) -> Vec<u8> {
+ let mut bytes = Vec::new();
+ for value in &self.values {
+ bytes.extend(value.serialize());
+ }
+ bytes
+ }
+
+ pub fn deserialize(bytes: &[u8]) -> Record {
+ let mut values = Vec::new();
+ let mut start = 0;
+ while start < bytes.len() {
+ let (rv, consumed) = RecordValue::deserialize(&bytes[start..]);
+ values.push(rv);
+ start += consumed;
+ }
+ Record { values }
+ }
+}
+
+pub fn escape_bytes(buf: &[u8]) -> Vec<u8> {
+ let mut result = Vec::new();
+ for byte in buf {
+ match byte {
+ &FIELD_SEPARATOR => {
+ result.extend(SEQ_LIT_FIELD_SEP);
+ }
+ &ESCAPE_CHARACTER => {
+ result.extend(SEQ_LIT_ESCAPE);
+ }
+ _ => result.push(*byte),
+ }
+ }
+ result
+}
+
+pub fn is_file_same_as_path(file: &File, path: &PathBuf) -> io::Result<bool> {
+ // Get the metadata for the open file handle
+ let file_metadata = file.metadata()?;
+
+ // Get the metadata for the file at the specified path
+ let path_metadata = metadata(path)?;
+
+ // Platform-specific comparison
+ #[cfg(unix)]
+ {
+ Ok(
+ file_metadata.dev() == path_metadata.dev()
+ && file_metadata.ino() == path_metadata.ino(),
+ )
+ }
+
+ #[cfg(windows)]
+ {
+ Ok(file_metadata.file_index() == path_metadata.file_index()
+ && file_metadata.volume_serial_number() == path_metadata.volume_serial_number())
+ }
+}
diff --git a/log_db/src/forward_log_reader.rs b/log_db/src/forward_log_reader.rs
new file mode 100644
index 0000000..10788a7
--- /dev/null
+++ b/log_db/src/forward_log_reader.rs
@@ -0,0 +1,124 @@
+use super::common::*;
+use std::fs::{self};
+use std::io::{self, BufRead, Read};
+
+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];
+ match self.reader.read_exact(&mut peek_buf) {
+ Ok(_) => {
+ // Go back one byte
+ self.reader.seek_relative(-1)?;
+ }
+ 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),
+ }
+ }
+}
+
+#[cfg(test)]
+mod reverse_reader_tests {
+ use super::*;
+ use std::path::Path;
+
+ #[test]
+ fn test_forward_log_reader_fixture_db1() {
+ 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());
+ }
+}
diff --git a/log_db/src/lib.rs b/log_db/src/lib.rs
new file mode 100644
index 0000000..6787cf1
--- /dev/null
+++ b/log_db/src/lib.rs
@@ -0,0 +1,794 @@
+#[macro_use]
+extern crate log;
+
+mod common;
+mod forward_log_reader;
+mod primary_memtable;
+mod reverse_log_reader;
+mod secondary_memtable;
+
+pub use common::*;
+pub use forward_log_reader::ForwardLogReader;
+use fs2::lock_contended_error;
+use fs2::FileExt;
+use primary_memtable::PrimaryMemtable;
+pub use reverse_log_reader::ReverseLogReader;
+use secondary_memtable::SecondaryMemtable;
+use std::fmt::Debug;
+use std::fs::{self};
+use std::io::{self, Write};
+use std::os::unix::fs::MetadataExt;
+use std::path::{Path, PathBuf};
+use std::thread;
+
+pub struct ConfigBuilder<'a, Field: Eq + Clone + Debug> {
+ data_dir: Option<String>,
+ segment_size: Option<usize>,
+ memtable_capacity: Option<usize>,
+ fields: Option<&'a Vec<(Field, RecordField)>>,
+ primary_key: Option<Field>,
+ secondary_keys: Option<Vec<Field>>,
+ memtable_evict_policy: Option<MemtableEvictPolicy>,
+ write_durability: Option<WriteDurability>,
+}
+
+impl<'a, Field: Eq + Clone + Debug> ConfigBuilder<'a, Field> {
+ pub fn new() -> ConfigBuilder<'a, Field> {
+ ConfigBuilder::<Field> {
+ data_dir: None,
+ segment_size: None,
+ memtable_capacity: None,
+ fields: None,
+ primary_key: None,
+ secondary_keys: None,
+ memtable_evict_policy: None,
+ write_durability: None,
+ }
+ }
+
+ /// Directory where the database will store its data.
+ pub fn data_dir(&mut self, data_dir: &str) -> &mut Self {
+ self.data_dir = Some(data_dir.to_string());
+ self
+ }
+
+ /// The maximum size of a segment file in bytes.
+ /// Once a segment file reaches this size, it is closed and a new one is created.
+ /// Closed segment files can be compacted.
+ pub fn segment_size(&mut self, segment_size: usize) -> &mut Self {
+ self.segment_size = Some(segment_size);
+ self
+ }
+
+ /// The maximum size of a single memtable in terms of records.
+ /// Note that each secondary index will have its own memtable.
+ pub fn memtable_capacity(&mut self, memtable_capacity: usize) -> &mut Self {
+ self.memtable_capacity = Some(memtable_capacity);
+ self
+ }
+
+ /// The field schema of the database.
+ pub fn fields(&mut self, fields: &'a Vec<(Field, RecordField)>) -> &mut Self {
+ self.fields = Some(fields);
+ self
+ }
+
+ /// The primary key of the database, used to construct
+ /// the primary memtable index. This should be the field
+ /// that is most frequently queried.
+ pub fn primary_key(&mut self, primary_key: Field) -> &mut Self {
+ self.primary_key = Some(primary_key);
+ self
+ }
+
+ /// The secondary keys of the database, used to construct
+ /// the secondary memtable indexes.
+ pub fn secondary_keys(&mut self, secondary_keys: Vec<Field>) -> &mut Self {
+ self.secondary_keys = Some(secondary_keys);
+ self
+ }
+
+ /// The eviction policy for the memtables. Determines which
+ /// record will be dropped from a memtable when it reaches
+ /// capacity.
+ pub fn memtable_evict_policy(
+ &mut self,
+ memtable_evict_policy: MemtableEvictPolicy,
+ ) -> &mut Self {
+ self.memtable_evict_policy = Some(memtable_evict_policy);
+ self
+ }
+
+ /// The write durability policy for the database.
+ /// This determines how writes are persisted to disk.
+ /// The default is WriteDurability::Flush.
+ pub fn write_durability(&mut self, write_durability: WriteDurability) -> &mut Self {
+ self.write_durability = Some(write_durability);
+ self
+ }
+
+ pub fn initialize(&self) -> Result<DB<Field>, io::Error> {
+ let config = Config::<Field> {
+ data_dir: self.data_dir.clone().unwrap_or("db_data".to_string()),
+ segment_size: self.segment_size.unwrap_or(4 * 1024 * 1024), // 4MB
+ memtable_capacity: self.memtable_capacity.unwrap_or(1_000_000),
+ fields: self
+ .fields
+ .ok_or(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "Required config value \"fields\" is not set",
+ ))?
+ .clone(),
+ primary_key: self.primary_key.clone().ok_or(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "Required config value \"primary_key\" is not set",
+ ))?,
+ secondary_keys: self.secondary_keys.clone().unwrap_or(Vec::new()),
+ memtable_evict_policy: self
+ .memtable_evict_policy
+ .clone()
+ .unwrap_or(MemtableEvictPolicy::LeastReadOrWritten),
+ write_durability: self
+ .write_durability
+ .clone()
+ .unwrap_or(WriteDurability::Flush),
+ };
+
+ DB::initialize(&config)
+ }
+}
+
+#[derive(Clone)]
+struct Config<Field: Eq + Clone> {
+ pub data_dir: String,
+ pub segment_size: usize,
+ pub memtable_capacity: usize,
+ pub fields: Vec<(Field, RecordField)>,
+ pub primary_key: Field,
+ pub secondary_keys: Vec<Field>,
+ pub memtable_evict_policy: MemtableEvictPolicy,
+ pub write_durability: WriteDurability,
+}
+
+pub struct DB<Field: Eq + Clone + Debug> {
+ config: Config<Field>,
+ log_path: PathBuf,
+ log_file: fs::File,
+ primary_key_index: usize,
+ primary_memtable: PrimaryMemtable,
+ secondary_memtables: Vec<SecondaryMemtable<Field>>,
+}
+
+impl<Field: Eq + Clone + Debug> DB<Field> {
+ /// Create a new database configuration builder.
+ pub fn configure() -> ConfigBuilder<'static, Field> {
+ ConfigBuilder::new()
+ }
+
+ fn initialize(config: &Config<Field>) -> Result<DB<Field>, io::Error> {
+ info!("Initializing DB...");
+ // If data_dir does not exist, create it
+ if !fs::exists(&config.data_dir)? {
+ fs::create_dir_all(&config.data_dir)?;
+ }
+
+ let log_path = Path::new(&config.data_dir).join(ACTIVE_LOG_FILENAME);
+
+ // Create the log file if it does not exist
+ let log_file_file = fs::OpenOptions::new()
+ .create(true)
+ .read(true)
+ .append(true)
+ .open(&log_path)?;
+
+ // Create the exclusive lock request file if it does not exist
+ fs::OpenOptions::new()
+ .create(true)
+ .write(true)
+ .open(&Path::new(&config.data_dir).join(EXCL_LOCK_REQUEST_FILENAME))?;
+
+ // Calculate the index of the primary value in a record
+ let primary_key_index = config
+ .fields
+ .iter()
+ .position(|(field, _)| field == &config.primary_key)
+ .ok_or(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "Primary key not found in schema after initialize",
+ ))?;
+
+ // Join primary key and secondary keys vec into a single vec
+ let mut all_keys = vec![&config.primary_key];
+ all_keys.extend(&config.secondary_keys);
+
+ // If any of the keys is not in the schema or
+ // is not an IndexableValue, return an error
+ for &key in &all_keys {
+ let (_, RecordField { field_type, .. }) = config
+ .fields
+ .iter()
+ .find(|(field, _)| field == key)
+ .ok_or(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "Secondary key must be present in the field schema",
+ ))?;
+
+ match field_type {
+ RecordFieldType::Int | RecordFieldType::String => {}
+ _ => {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "Secondary key must be an IndexableValue",
+ ))
+ }
+ }
+ }
+ let primary_memtable = PrimaryMemtable::new(
+ config.memtable_capacity,
+ config.memtable_evict_policy.clone(),
+ );
+ let secondary_memtables = config
+ .secondary_keys
+ .iter()
+ .map(|key| {
+ SecondaryMemtable::new(
+ key,
+ primary_key_index,
+ config.memtable_capacity,
+ config.memtable_evict_policy.clone(),
+ )
+ })
+ .collect();
+
+ let mut db = DB::<Field> {
+ config: config.clone(),
+ log_path: log_path.clone(),
+ log_file: log_file_file,
+ 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)
+ }
+
+ /// Insert a record into the database. If the primary key value already exists,
+ /// the existing record will be replaced by the supplied one.
+ pub fn upsert(&mut self, record: &Record) -> Result<(), io::Error> {
+ debug!("Upserting record: {:?}", record);
+ // Validate the record length
+ if record.values.len() != self.config.fields.len() {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ format!(
+ "Record has an incorrect number of fields: {}, expected {}",
+ record.values.len(),
+ self.config.fields.len()
+ ),
+ ));
+ }
+
+ // Validate that record fields match schema types
+ for (i, (_, field)) in self.config.fields.iter().enumerate() {
+ match (&record.values[i], field) {
+ (
+ RecordValue::Null,
+ RecordField {
+ nullable: true,
+ field_type: _,
+ },
+ ) => {}
+ (
+ RecordValue::Int(_),
+ RecordField {
+ field_type: RecordFieldType::Int,
+ ..
+ },
+ ) => {}
+ (
+ RecordValue::String(_),
+ RecordField {
+ field_type: RecordFieldType::String,
+ ..
+ },
+ ) => {}
+ (
+ RecordValue::Bytes(_),
+ RecordField {
+ field_type: RecordFieldType::Bytes,
+ ..
+ },
+ ) => {}
+ _ => {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ format!(
+ "Record field {} has incorrect type: {:?}, expected {:?}",
+ &i, &record.values[i], &field.field_type
+ ),
+ ))
+ }
+ }
+ }
+
+ debug!("Record is valid");
+ debug!("Opening file in append mode and acquiring exclusive lock...");
+
+ // Acquire an exclusive lock for writing
+ self.request_exclusive_lock()?;
+
+ if self.ensure_correct_file_is_open()? {
+ // The log file has been rotated, so we must try again
+ return self.upsert(record);
+ }
+
+ debug!("Lock acquired, appending to log file");
+
+ // Write the record to the log
+ // Each serialized row is suffixed with the field separator character sequence
+ let mut serialized_record = record.serialize();
+ serialized_record.extend(SEQ_RECORD_SEP);
+ self.log_file.write_all(&serialized_record)?;
+
+ // Flush and sync to disk
+ if self.config.write_durability == WriteDurability::Flush {
+ self.log_file.flush()?;
+ }
+ if self.config.write_durability == WriteDurability::FlushSync {
+ self.log_file.flush()?;
+ self.log_file.sync_all()?;
+ }
+
+ self.log_file.unlock()?;
+
+ debug!("Record appended to log file, lock released");
+
+ debug!("Updating primary memtable");
+ self.update_primary_index(record);
+
+ debug!("Updating secondary memtables");
+ self.update_secondary_indexes(record);
+
+ Ok(())
+ }
+
+ /// Get a record by its primary index value.
+ /// E.g. `db.get(RecordValue::Int(10))`.
+ pub fn get(&mut self, query_key: &RecordValue) -> Result<Option<Record>, io::Error> {
+ let query_key_original = query_key;
+ debug!(
+ "Getting record with field {:?} = {:?}",
+ &self.config.primary_key, query_key
+ );
+ let query_key = query_key_original.as_indexable().ok_or(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "Queried value must be indexable",
+ ))?;
+
+ debug!("Looking up key {:?} in primary memtable", query_key);
+ let found = self.primary_memtable.get(&query_key);
+ if let Some(record) = found {
+ debug!("Found record in primary memtable: {:?}", record);
+ return Ok(Some(record.clone()));
+ }
+
+ debug!(
+ "No memtable entry found, looking up key {:?} in log file",
+ query_key
+ );
+
+ debug!(
+ "Matching records based on value at primary key index ({})",
+ &self.primary_key_index
+ );
+ debug!("Opening file in read mode and acquiring shared lock...");
+
+ // Open the file and acquire a shared lock for reading
+ let mut file = fs::OpenOptions::new().read(true).open(&self.log_path)?;
+
+ self.request_shared_lock(&mut file)?;
+
+ if !is_file_same_as_path(&file, &self.log_path)? {
+ // The log file has been rotated, so we must try again
+ debug!("Lock acquired, but the log file has been rotated. Retrying get...");
+ file.unlock()?;
+ drop(file);
+ return self.get(query_key_original);
+ }
+
+ debug!("Lock acquired, searching log files for record");
+
+ let segment_numbers = self.segment_numbers()?;
+ let mut result: Option<Record> = None;
+ for &n in &segment_numbers {
+ if n == 0 {
+ debug!("Searching the active log file...");
+ result = ReverseLogReader::new(&mut file)?.find(|record| {
+ let record_key = record.values[self.primary_key_index]
+ .as_indexable()
+ .expect("A non-indexable value was stored at key index");
+ record_key == query_key
+ });
+
+ debug!("Active log file searched, releasing shared lock...");
+ file.unlock()?;
+ } else {
+ debug!("Locking and searching rotated log segment file {}...", n);
+ let path = Path::new(&self.config.data_dir)
+ .join(ACTIVE_LOG_FILENAME)
+ .with_extension(n.to_string());
+
+ let mut segm_file = fs::OpenOptions::new().read(true).open(&path)?;
+ self.request_shared_lock(&mut segm_file)?;
+ result = ReverseLogReader::new(&mut segm_file)?.find(|record| {
+ let record_key = record.values[self.primary_key_index]
+ .as_indexable()
+ .expect("A non-indexable value was stored at key index");
+ record_key == query_key
+ });
+
+ debug!("Segment file searched, releasing shared lock...");
+ segm_file.unlock()?;
+ };
+
+ if result.is_some() {
+ break;
+ }
+ }
+
+ debug!("Record search complete");
+
+ let result_value = match &result {
+ Some(record) => record,
+ None => {
+ debug!("No record found for key {:?}", query_key);
+ return Ok(None);
+ }
+ };
+
+ 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(result)
+ }
+
+ /// Get a collection of records based on a field value.
+ /// Indexes will be used if they contain the requested key.
+ pub fn find_all(
+ &mut self,
+ field: &Field,
+ query_key: &RecordValue,
+ ) -> Result<Vec<Record>, io::Error> {
+ // If querying by primary key, return the result of `get` wrapped in a vec.
+ if field == &self.config.primary_key {
+ return match self.get(query_key)? {
+ Some(record) => Ok(vec![record.clone()]),
+ None => Ok(vec![]),
+ };
+ }
+
+ // Otherwise, continue with querying secondary indexes.
+ let query_key_original = query_key;
+ debug!(
+ "Finding all records with field {:?} = {:?}",
+ field, query_key
+ );
+ let query_key = query_key_original.as_indexable().ok_or(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "Queried value must be indexable",
+ ))?;
+
+ // Try to find a memtable with the queried key
+ let found_memtable_index = self
+ .secondary_memtables
+ .iter_mut()
+ .position(|mt| &mt.field == field);
+
+ if let Some(memtable_index) = found_memtable_index {
+ debug!(
+ "Found suitable secondary index. Looking up key {:?} in the memtable",
+ query_key
+ );
+ let records = self.secondary_memtables[memtable_index].find_all(&query_key);
+ debug!("Found matching key");
+ return Ok(records.iter().map(|&record| record.clone()).collect());
+ }
+
+ debug!(
+ "No memtable entry found, looking up key {:?} in log file",
+ query_key
+ );
+
+ // Get the index of the requested field
+ let key_index = self
+ .config
+ .fields
+ .iter()
+ .position(|(schema_field, _)| schema_field == field)
+ .ok_or(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "Key not found in schema after initialize",
+ ))?;
+
+ debug!("Matching key index {}", key_index);
+ debug!("Acquiring shared lock...");
+
+ // Acquire a shared lock for reading
+ self.log_file.lock_shared()?;
+
+ if self.ensure_correct_file_is_open()? {
+ // The log file has been rotated, so we must try again
+ return self.find_all(field, query_key_original);
+ }
+
+ debug!("Lock acquired, searching log file for record");
+
+ let result = ReverseLogReader::new(&mut self.log_file)?
+ .filter(|record| {
+ let record_key = record.values[key_index]
+ .as_indexable()
+ .expect("A non-indexable value was stored at key index");
+ record_key == query_key
+ })
+ .collect::<Vec<Record>>();
+
+ self.log_file.unlock()?;
+ debug!("Record search complete, lock released");
+
+ debug!(
+ "Number of matching records found in log file: {}",
+ result.len()
+ );
+
+ if let Some(memtable_index) = found_memtable_index {
+ debug!("Inserting result set into secondary index");
+ self.secondary_memtables[memtable_index].set_all(&query_key, &result);
+ }
+
+ Ok(result)
+ }
+
+ /// Ensures that the `self.log_file` handle is still pointing to the correct file.
+ /// If the file has been rotated, the handle will be closed and reopened.
+ /// Returns `true` if the file has been rotated and the handle has been reopened.
+ fn ensure_correct_file_is_open(&mut self) -> Result<bool, io::Error> {
+ if !is_file_same_as_path(&self.log_file, &self.log_path)? {
+ // The log file has been rotated, so we must try again
+ debug!(
+ "Lock acquired, but the log file has been rotated. Reopening file and retrying..."
+ );
+ self.log_file.unlock()?;
+
+ self.log_file = fs::OpenOptions::new()
+ .create(true)
+ .read(true)
+ .append(true)
+ .open(&self.log_path)?;
+
+ Ok(true)
+ } else {
+ Ok(false)
+ }
+ }
+
+ 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);
+ }
+ }
+ });
+ }
+
+ fn request_exclusive_lock(&mut self) -> Result<(), io::Error> {
+ // Create a lock on the exclusive lock request file to signal to readers that they should wait
+ let lock_request_path = Path::new(&self.config.data_dir).join(EXCL_LOCK_REQUEST_FILENAME);
+ let lock_request_file = fs::OpenOptions::new()
+ .create(true)
+ .write(true) // When requesting a lock, we need to have either read or write permissions
+ .open(&lock_request_path)?;
+
+ // Attempt to acquire an exclusive lock on the lock request file
+ // This will block until the lock is acquired
+ lock_request_file.lock_exclusive()?;
+
+ // Check that the exclusive lock request file is still the same as the one we opened
+ // NOTE: this isn't strictly necessary, but it's a good sanity check. Disabled for now.
+ // if !is_file_same_as_path(&lock_request_file, &lock_request_path)? {
+ // // The lock request file has been removed
+ // return Err(io::Error::new(
+ // io::ErrorKind::Other,
+ // "Lock request file was removed unexpectedly",
+ // ));
+ // }
+
+ // Acquire an exclusive lock on the log file
+ self.log_file.lock_exclusive()?;
+
+ // Unlock the request file
+ lock_request_file.unlock()?;
+
+ Ok(())
+ }
+
+ fn is_exclusive_lock_requested(&self, data_dir: &str) -> Result<bool, io::Error> {
+ let lock_request_path = Path::new(data_dir).join(EXCL_LOCK_REQUEST_FILENAME);
+ let lock_request_file = fs::OpenOptions::new()
+ .create(true)
+ .write(true) // When requesting a lock, we need to have either read or write permissions
+ .open(&lock_request_path)?;
+
+ // Attempt to acquire a shared lock on the lock request file
+ // If the file is already locked, return false
+ match lock_request_file.try_lock_shared() {
+ Err(e) => {
+ if e.kind() == lock_contended_error().kind() {
+ return Ok(true);
+ }
+ return Err(e);
+ }
+ Ok(_) => {
+ // Check that the exclusive lock request file is still the same as the one we opened
+ if !is_file_same_as_path(&lock_request_file, &lock_request_path)? {
+ // The lock request file has been removed
+ return Err(io::Error::new(
+ io::ErrorKind::Other,
+ "Lock request file was removed unexpectedly",
+ ));
+ }
+
+ lock_request_file.unlock()?;
+ return Ok(false);
+ }
+ }
+ }
+
+ fn request_shared_lock(&self, file: &mut fs::File) -> Result<(), io::Error> {
+ const SHARED_LOCK_WAIT_MAX_MS: u64 = 100;
+ let mut timeout = 5;
+ loop {
+ if self.is_exclusive_lock_requested(&self.config.data_dir)? {
+ debug!("Exclusive lock requested, waiting for {}ms before requesting a shared lock again", timeout);
+ thread::sleep(std::time::Duration::from_millis(timeout));
+ timeout = std::cmp::min(timeout * 2, SHARED_LOCK_WAIT_MAX_MS);
+ } else {
+ file.lock_shared()?;
+ return Ok(());
+ }
+ }
+ }
+
+ /// Check if there are any pending tasks and do them. Tasks include:
+ /// - Rotating the active log file if it has reached capacity and compacting it.
+ ///
+ /// This function should be called periodically to ensure that the database remains in an optimal state.
+ /// Note that this function is synchronous and may block for a relatively long time.
+ /// You may call this function in a separate thread or process to avoid blocking the main thread.
+ /// However, the database will be exclusively locked, so all writes will be blocked during the tasks.
+ pub fn do_maintenance_tasks(&mut self) -> Result<(), io::Error> {
+ let active_log_path = Path::new(&self.config.data_dir).join(ACTIVE_LOG_FILENAME);
+ let active_log_md = fs::metadata(&active_log_path)?;
+
+ if active_log_md.size() >= self.config.segment_size as u64 {
+ // Rotate the active log file
+
+ debug!("Starting rotation, requesting exclusive lock...");
+ self.request_exclusive_lock()?;
+
+ debug!("Exclusive lock acquired, rotating active log file...");
+ let next_segment_number = self.next_segment_number()?;
+ let next_segment_path =
+ &active_log_path.with_extension(next_segment_number.to_string());
+
+ debug!("Renaming active log file to {:?}", &next_segment_path);
+ fs::rename(&active_log_path, &next_segment_path)?;
+
+ // Create a new active log file
+ self.log_file = fs::OpenOptions::new()
+ .create(true)
+ .write(true)
+ .append(true)
+ .open(&active_log_path)?;
+
+ // The new active log file is not locked by this client so it cannot be touched.
+ // Compact the rotated segment.
+
+ debug!("Active log file rotated");
+
+ // TODO compaction of rotated segments
+ }
+
+ Ok(())
+ }
+
+ fn next_segment_number(&self) -> Result<u64, io::Error> {
+ match self.segment_numbers()?.iter().max() {
+ Some(greatest) => Ok(greatest + 1),
+ None => Ok(1),
+ }
+ }
+
+ /// Query the filesystem to get the numbers of existing segments
+ /// in the intended reading order: first the active log (signaled with 0),
+ /// then the segments from the greatest ordinal (newest) to the least (oldest).
+ /// E.g. `vec![0, 4, 3, 2, 1]`.
+ fn segment_numbers(&self) -> Result<Vec<u64>, io::Error> {
+ // TODO: optimize the vecs out of here
+ let files = fs::read_dir(&self.config.data_dir)?;
+ let mut nums: Vec<u64> = files
+ .filter_map(|f| {
+ let f_path = match f {
+ Ok(f) => f.path(),
+ Err(_) => return None,
+ };
+
+ if !f_path.is_file() {
+ return None;
+ }
+
+ let name = &f_path
+ .with_extension("")
+ .file_name()
+ .expect("File did not have a name?")
+ .to_str()
+ .expect("Failed to convert file name to string")
+ .to_string();
+
+ if name != ACTIVE_LOG_FILENAME {
+ return None;
+ }
+
+ let ext = match f_path.extension() {
+ Some(ext) => ext,
+ None => return None,
+ };
+
+ let ext_num = ext
+ .to_str()
+ .expect("Extension was not a valid UTF-8 string")
+ .parse::<u64>()
+ .expect("Extension was not a valid number");
+
+ Some(ext_num)
+ })
+ .collect();
+
+ nums.sort();
+ nums.push(0);
+ nums.reverse();
+ Ok(nums)
+ }
+}
diff --git a/log_db/src/primary_memtable.rs b/log_db/src/primary_memtable.rs
new file mode 100644
index 0000000..0c94ceb
--- /dev/null
+++ b/log_db/src/primary_memtable.rs
@@ -0,0 +1,79 @@
+use super::common::*;
+use priority_queue::PriorityQueue;
+use std::collections::BTreeMap;
+
+pub struct PrimaryMemtable {
+ capacity: usize,
+ /// Running counter of memtable operations, used as priority
+ /// in evict_queue.
+ n_operations: u64,
+ records: BTreeMap<IndexableValue, Record>,
+ /// A max heap priority queue of keys. Note: n_operations must
+ /// be negated upon append to evict oldest values first.
+ evict_queue: PriorityQueue<IndexableValue, i64>,
+ evict_policy: MemtableEvictPolicy,
+}
+
+impl PrimaryMemtable {
+ pub fn new(capacity: usize, evict_policy: MemtableEvictPolicy) -> PrimaryMemtable {
+ PrimaryMemtable {
+ capacity,
+ n_operations: 0,
+ records: BTreeMap::new(),
+ evict_queue: PriorityQueue::new(),
+ evict_policy,
+ }
+ }
+
+ pub fn set(&mut self, key: &IndexableValue, value: &Record) {
+ if self.capacity == 0 {
+ return;
+ }
+
+ debug!(
+ "Inserting/updating record in primary memtable with key {:?} = {:?}",
+ &key, &value,
+ );
+
+ if self.records.len() >= self.capacity {
+ let (evict_key, _prio) = self.evict_queue.pop().expect("Evict queue was empty");
+ self.records.remove(&evict_key);
+ }
+
+ self.records.insert(key.clone(), value.clone());
+
+ if self.evict_policy == MemtableEvictPolicy::LeastWritten
+ || self.evict_policy == MemtableEvictPolicy::LeastReadOrWritten
+ {
+ self.set_priority(&key);
+ }
+ }
+
+ pub fn get(&mut self, key: &IndexableValue) -> Option<&Record> {
+ if self.evict_policy == MemtableEvictPolicy::LeastRead
+ || self.evict_policy == MemtableEvictPolicy::LeastReadOrWritten
+ {
+ self.set_priority(&key);
+ }
+
+ self.records.get(key)
+ }
+
+ fn set_priority(&mut self, key: &IndexableValue) {
+ let priority = self.get_and_increment_current_priority();
+ match self.evict_queue.get(key) {
+ Some(_) => {
+ self.evict_queue.change_priority(key, priority);
+ }
+ None => {
+ self.evict_queue.push(key.clone(), priority);
+ }
+ }
+ }
+
+ fn get_and_increment_current_priority(&mut self) -> i64 {
+ let ret = -(self.n_operations as i64);
+ self.n_operations += 1;
+ ret
+ }
+}
diff --git a/log_db/src/reverse_log_reader.rs b/log_db/src/reverse_log_reader.rs
new file mode 100644
index 0000000..5d80dba
--- /dev/null
+++ b/log_db/src/reverse_log_reader.rs
@@ -0,0 +1,354 @@
+use super::common::*;
+use std::fs::{self};
+use std::io::{self, Read, Seek, SeekFrom};
+
+pub struct ReverseLogReader<'a> {
+ /// The file to read from end to beginning.
+ file: &'a mut fs::File,
+ /// The internal buffer used to read from the file.
+ /// It is populated with the last INTERNAL_BUF_SIZE bytes read from the file
+ /// and is used to read records in reverse order
+ internal_buf: Vec<u8>,
+ /// The current position in the internal buffer. It is decremented as bytes are read
+ /// from the buffer. When a read is requested and the internal position is 0, the buffer
+ /// is populated with the next (= closer to the start of the file) INTERNAL_BUF_SIZE bytes from the file.
+ /// Note: This is the index of the next byte to be read from the internal buffer + 1
+ internal_pos: usize,
+ /// A flag indicating whether the record separator at the cursor position has been consumed.
+ /// Useful to avoid consuming the separator once when reading until an escape character, and
+ /// a second time when reading a new record and validating it ends in a separator.
+ consumed_record_sep: bool,
+}
+
+// This value is based on the reverse_read_file_with_various_buffer_sizes benchmark.
+// Greater values yield little to no performance improvement.
+const DEFAULT_INTERNAL_BUF_SIZE: usize = 32768;
+impl<'a> ReverseLogReader<'a> {
+ pub fn new(file: &mut fs::File) -> Result<ReverseLogReader, io::Error> {
+ file.seek(SeekFrom::End(0))?;
+ Ok(ReverseLogReader {
+ file,
+ internal_buf: vec![0; DEFAULT_INTERNAL_BUF_SIZE],
+ internal_pos: 0,
+ consumed_record_sep: false,
+ })
+ }
+
+ pub fn new_with_size(
+ file: &mut fs::File,
+ internal_buf_size: usize,
+ ) -> Result<ReverseLogReader, io::Error> {
+ file.seek(SeekFrom::End(0))?;
+ Ok(ReverseLogReader {
+ file,
+ internal_buf: vec![0; internal_buf_size],
+ internal_pos: 0,
+ consumed_record_sep: false,
+ })
+ }
+
+ pub fn read_record(&mut self) -> Result<Option<Record>, io::Error> {
+ if self.file.stream_position()? == 0 && self.internal_pos == 0 {
+ return Ok(None);
+ }
+
+ if !self.consumed_record_sep {
+ // Check that record ends with a record separator
+ match self.read_special_sequence()? {
+ SpecialSequence::RecordSeparator => {}
+ _ => {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidData,
+ "Record does not end in record separator",
+ ));
+ }
+ }
+ }
+ self.consumed_record_sep = false;
+
+ let mut result_buf: Vec<u8> = vec![];
+ let mut read_buf = Vec::with_capacity(self.internal_buf.len());
+ loop {
+ read_buf.clear();
+ let read = self.read_until(ESCAPE_CHARACTER, &mut read_buf)?;
+
+ result_buf.extend(&read_buf[..read]);
+
+ if self.file.stream_position()? == 0 && self.internal_pos == 0 {
+ // We read until the start of the file, we are done
+ break;
+ }
+
+ match self.read_special_sequence()? {
+ SpecialSequence::LiteralEscape => {
+ result_buf.push(ESCAPE_CHARACTER);
+ }
+ SpecialSequence::LiteralFieldSeparator => {
+ result_buf.push(FIELD_SEPARATOR);
+ }
+ SpecialSequence::RecordSeparator => {
+ self.consumed_record_sep = true;
+ break;
+ }
+ }
+ }
+
+ result_buf.reverse();
+ Ok(Some(Record::deserialize(&result_buf)))
+ }
+
+ /// Read exactly `buf.len()` bytes from the file, return an error if the file is exhausted.
+ /// The bytes are returned in start -> end order.
+ /// If an error is returned, the contents of `buf` are in an undefined state.
+ fn read_exact(&mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
+ let mut read = 0;
+ while read < buf.len() {
+ if self.internal_pos == 0 {
+ let populated_n = self.populate_internal_buf()?;
+ if populated_n == 0 {
+ return Err(io::Error::new(
+ io::ErrorKind::UnexpectedEof,
+ "Unexpected end of file",
+ ));
+ }
+ }
+
+ let end = buf.len() - read;
+ let n = std::cmp::min(self.internal_pos, end);
+ buf[end - n..end]
+ .copy_from_slice(&self.internal_buf[self.internal_pos - n..self.internal_pos]);
+ read += n;
+ self.internal_pos -= n;
+ }
+
+ Ok(read)
+ }
+
+ fn populate_internal_buf(&mut self) -> Result<usize, io::Error> {
+ let current_seek_pos = self.file.stream_position()? as usize;
+
+ // Seek back by the size of the internal buffer or to the beginning of the file
+ let seek_length = if current_seek_pos > self.internal_buf.len() {
+ self.internal_pos = self.internal_buf.len();
+ self.internal_buf.len()
+ } else {
+ self.internal_buf = vec![0; current_seek_pos as usize];
+ self.internal_pos = current_seek_pos as usize;
+ current_seek_pos
+ };
+
+ self.file.seek_relative(-(seek_length as i64))?;
+ self.file.read_exact(&mut self.internal_buf)?;
+ self.file.seek_relative(-(seek_length as i64))?;
+
+ Ok(seek_length as usize)
+ }
+
+ /// Iterate over the internal buffer with internal_pos as the index.
+ /// If the byte is found or the file has been exhausted, we return the number of bytes read.
+ /// The `buf` parameter is used to store the bytes read from the internal buffer, excluding the found byte,
+ /// in reverse order.
+ fn read_until(&mut self, byte: u8, buf: &mut Vec<u8>) -> Result<usize, io::Error> {
+ // TODO: optimize this
+ let mut read = 0;
+ loop {
+ // If we reach internal_pos == 0, we need to populate the internal buffer.
+ if self.internal_pos == 0 {
+ let populated_n = self.populate_internal_buf()?;
+ if populated_n == 0 {
+ return Ok(read);
+ }
+ }
+
+ while self.internal_pos > 0 {
+ let index = self.internal_pos - 1;
+ if self.internal_buf[index] == byte {
+ return Ok(read);
+ }
+ buf.push(self.internal_buf[index]);
+ read += 1;
+ self.internal_pos -= 1;
+ }
+ }
+ }
+
+ fn read_special_sequence(&mut self) -> Result<SpecialSequence, io::Error> {
+ let mut special_buf: Vec<u8> = vec![0; SEQ_RECORD_SEP.len()];
+ self.read_exact(&mut special_buf)?;
+
+ match validate_special(&special_buf.as_slice()) {
+ Some(special) => Ok(special),
+ None => {
+ let pos = self.file.stream_position().unwrap() + self.internal_pos as u64;
+
+ Err(io::Error::new(
+ io::ErrorKind::InvalidData,
+ format!(
+ "Not a special sequence: {:?} at pos: {:x}",
+ special_buf, pos,
+ ),
+ ))
+ }
+ }
+ }
+}
+
+#[cfg(test)]
+mod reverse_reader_tests {
+ use super::*;
+ use std::io::Write;
+ use std::path::Path;
+
+ #[test]
+ fn test_read_until_found() {
+ let mut file = tempfile::tempfile().unwrap();
+ file.write_all(b"hello,world").unwrap();
+ let mut reader = ReverseLogReader::new(&mut file).unwrap();
+ let mut buf = vec![];
+ assert_eq!(reader.read_until(b',', &mut buf).unwrap(), 5);
+ assert_eq!(buf, b"dlrow");
+ }
+
+ #[test]
+ fn test_read_until_not_found() {
+ let mut file = tempfile::tempfile().unwrap();
+ file.write_all(b"hello,world").unwrap();
+ let mut reader = ReverseLogReader::new(&mut file).unwrap();
+ let mut buf = vec![];
+ let read = reader.read_until(b'!', &mut buf).unwrap();
+ assert_eq!(buf, b"dlrow,olleh");
+ assert_eq!(read, 11);
+ }
+
+ #[test]
+ fn test_read_special_sequence() {
+ let mut file = tempfile::tempfile().unwrap();
+ let mut buf = vec![];
+ buf.extend(SEQ_RECORD_SEP);
+ buf.extend(SEQ_LIT_ESCAPE);
+ buf.extend(SEQ_LIT_FIELD_SEP);
+ // Note: written in start -> end order, read end -> start
+ file.write_all(&buf).unwrap();
+
+ let mut reader = ReverseLogReader::new(&mut file).unwrap();
+ assert_eq!(
+ reader.read_special_sequence().unwrap(),
+ SpecialSequence::LiteralFieldSeparator
+ );
+ assert_eq!(
+ reader.read_special_sequence().unwrap(),
+ SpecialSequence::LiteralEscape
+ );
+ assert_eq!(
+ reader.read_special_sequence().unwrap(),
+ SpecialSequence::RecordSeparator
+ );
+ }
+
+ #[test]
+ fn test_populate_internal_buf() {
+ let mut file = tempfile::tempfile().unwrap();
+ file.write_all(b"hello,world").unwrap();
+ let mut reader = ReverseLogReader::new_with_size(&mut file, 3).unwrap();
+
+ reader.populate_internal_buf().unwrap();
+ assert_eq!(
+ String::from_utf8(reader.internal_buf.clone()).unwrap(),
+ "rld".to_string()
+ );
+
+ reader.populate_internal_buf().unwrap();
+ assert_eq!(
+ String::from_utf8(reader.internal_buf.clone()).unwrap(),
+ ",wo".to_string()
+ );
+
+ reader.populate_internal_buf().unwrap();
+ assert_eq!(
+ String::from_utf8(reader.internal_buf.clone()).unwrap(),
+ "llo".to_string()
+ );
+
+ reader.populate_internal_buf().unwrap();
+ assert_eq!(
+ String::from_utf8(reader.internal_buf.clone()).unwrap(),
+ "he".to_string()
+ );
+ }
+
+ #[test]
+ fn test_reverse_log_reader_fixture_db1() {
+ 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 reverse_log_reader = ReverseLogReader::new(&mut file).unwrap();
+
+ // There are two records in the log with "schema": Int, Null
+
+ 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 = 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());
+ }
+
+ #[test]
+ fn test_read_exact() {
+ let mut file = tempfile::tempfile().unwrap();
+ file.write_all(b"hello,world").unwrap();
+ let mut reader = ReverseLogReader::new(&mut file).unwrap();
+ let mut buf = vec![0; 3];
+
+ let read = reader.read_exact(&mut buf).unwrap();
+ assert_eq!(buf, b"rld");
+ assert_eq!(read, 3);
+
+ let read = reader.read_exact(&mut buf).unwrap();
+ assert_eq!(buf, b",wo");
+ assert_eq!(read, 3);
+
+ let read = reader.read_exact(&mut buf).unwrap();
+ assert_eq!(buf, b"llo");
+ assert_eq!(read, 3);
+
+ assert!(reader.read_exact(&mut buf).unwrap_err().kind() == io::ErrorKind::UnexpectedEof);
+ }
+
+ #[test]
+ fn test_read_exact_insufficient_bytes() {
+ let mut file = tempfile::tempfile().unwrap();
+ file.write_all(b"hello").unwrap();
+ let mut reader = ReverseLogReader::new(&mut file).unwrap();
+ let mut buf = vec![0; 10];
+ assert!(reader.read_exact(&mut buf).unwrap_err().kind() == io::ErrorKind::UnexpectedEof);
+ }
+}
+
+impl Iterator for ReverseLogReader<'_> {
+ 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/log_db/src/secondary_memtable.rs b/log_db/src/secondary_memtable.rs
new file mode 100644
index 0000000..fcf4a7f
--- /dev/null
+++ b/log_db/src/secondary_memtable.rs
@@ -0,0 +1,178 @@
+use super::common::*;
+use priority_queue::PriorityQueue;
+use std::collections::BTreeMap;
+use std::collections::HashSet;
+use std::fmt::Debug;
+use std::hash::{Hash, Hasher};
+
+struct UniqueRecord {
+ /// The value of the record's primary key field
+ primary_value: IndexableValue,
+ /// The record itself.
+ record: Record,
+}
+
+impl PartialEq for UniqueRecord {
+ fn eq(&self, other: &UniqueRecord) -> bool {
+ self.primary_value == other.primary_value
+ }
+}
+
+impl Eq for UniqueRecord {}
+
+impl Hash for UniqueRecord {
+ fn hash<H>(&self, state: &mut H)
+ where
+ H: Hasher,
+ {
+ self.primary_value.hash(state)
+ }
+}
+
+pub struct SecondaryMemtable<Field: Eq + Clone + Debug> {
+ pub field: Field,
+ pub primary_field_index: usize,
+ capacity: usize,
+ /// Running counter of memtable operations, used as priority
+ /// in evict_queue.
+ n_operations: u64,
+ records: BTreeMap<IndexableValue, HashSet<UniqueRecord>>,
+ /// A max heap priority queue of keys. Note: n_operations must
+ /// be negated upon append to evict oldest values first.
+ evict_queue: PriorityQueue<IndexableValue, i64>,
+ evict_policy: MemtableEvictPolicy,
+}
+
+impl<Field: Eq + Clone + Debug> SecondaryMemtable<Field> {
+ pub fn new(
+ field: &Field,
+ primary_field_index: usize,
+ capacity: usize,
+ evict_policy: MemtableEvictPolicy,
+ ) -> SecondaryMemtable<Field> {
+ SecondaryMemtable {
+ field: field.clone(),
+ primary_field_index,
+ capacity,
+ n_operations: 0,
+ records: BTreeMap::new(),
+ evict_queue: PriorityQueue::new(),
+ evict_policy,
+ }
+ }
+
+ pub fn set(&mut self, key: &IndexableValue, value: &Record) {
+ if self.capacity == 0 {
+ return;
+ }
+
+ debug!(
+ "Inserting/updating record in secondary memtable with key {:?} = {:?}",
+ &key, &value,
+ );
+
+ if self.records.len() >= self.capacity {
+ let (evict_key, _prio) = self.evict_queue.pop().expect("Evict queue was empty");
+ self.records.remove(&evict_key);
+ }
+
+ let unique_record = UniqueRecord {
+ primary_value: value.values[self.primary_field_index]
+ .as_indexable()
+ .expect("Value at primary field index was not indexable"),
+ record: value.clone(),
+ };
+
+ match self.records.get_mut(key) {
+ Some(existing) => {
+ debug!(
+ "Existing entry found with {} records in the set",
+ &existing.len()
+ );
+ existing.insert(unique_record);
+ }
+ None => {
+ debug!("No existing entry found, creating one.");
+ let mut set = HashSet::with_capacity(1);
+ set.insert(unique_record);
+ self.records.insert(key.clone(), set);
+ }
+ }
+
+ if self.evict_policy == MemtableEvictPolicy::LeastWritten
+ || self.evict_policy == MemtableEvictPolicy::LeastReadOrWritten
+ {
+ self.set_priority(&key);
+ }
+ }
+
+ pub fn set_all(&mut self, key: &IndexableValue, values: &[Record]) {
+ if self.capacity == 0 {
+ return;
+ }
+
+ debug!(
+ "Replacing set of records in secondary memtable with key {:?} ({} values)",
+ &key,
+ &values.len(),
+ );
+
+ if self.records.len() >= self.capacity {
+ let (evict_key, _prio) = self.evict_queue.pop().expect("Evict queue was empty");
+ self.records.remove(&evict_key);
+ }
+
+ let mut set = HashSet::with_capacity(values.len());
+ values.iter().for_each(|value| {
+ let unique_record = UniqueRecord {
+ primary_value: value.values[self.primary_field_index]
+ .as_indexable()
+ .expect("Value at primary field index was not indexable"),
+ record: value.clone(),
+ };
+ set.insert(unique_record);
+ });
+
+ self.records.insert(key.clone(), set);
+
+ if self.evict_policy == MemtableEvictPolicy::LeastWritten
+ || self.evict_policy == MemtableEvictPolicy::LeastReadOrWritten
+ {
+ self.set_priority(&key);
+ }
+ }
+
+ pub fn find_all(&mut self, key: &IndexableValue) -> Vec<&Record> {
+ if self.evict_policy == MemtableEvictPolicy::LeastRead
+ || self.evict_policy == MemtableEvictPolicy::LeastReadOrWritten
+ {
+ self.set_priority(&key);
+ }
+
+ match self.records.get(key) {
+ None => vec![],
+ Some(set) => set
+ .iter()
+ .map(|unique_record| &unique_record.record)
+ .collect(),
+ }
+ }
+
+ fn set_priority(&mut self, key: &IndexableValue) {
+ let priority = self.get_and_increment_current_priority();
+ match self.evict_queue.get(key) {
+ Some(_) => {
+ self.evict_queue.change_priority(key, priority);
+ }
+ None => {
+ self.evict_queue.push(key.clone(), priority);
+ }
+ }
+ }
+
+ fn get_and_increment_current_priority(&mut self) -> i64 {
+ let ret = -(self.n_operations as i64);
+ self.n_operations += 1;
+ ret
+ }
+}