diff options
| author | Jan Tuomi <jan@jantuomi.fi> | 2024-10-07 12:43:16 +0300 |
|---|---|---|
| committer | Jan Tuomi <jan@jantuomi.fi> | 2024-10-07 12:43:16 +0300 |
| commit | a4efc83f7f29c6ef8ac9b4c57201c9ecc31266b3 (patch) | |
| tree | a18c856a5fc15fe509d3e3277e345176608dee21 /log_db/src | |
| parent | 30295ba5e1bdbfa69ad98fc5908ee4b1aabbd2e5 (diff) | |
Add py_bindings lib, move to monorepo structure
Diffstat (limited to 'log_db/src')
| -rw-r--r-- | log_db/src/common.rs | 308 | ||||
| -rw-r--r-- | log_db/src/forward_log_reader.rs | 124 | ||||
| -rw-r--r-- | log_db/src/lib.rs | 794 | ||||
| -rw-r--r-- | log_db/src/primary_memtable.rs | 79 | ||||
| -rw-r--r-- | log_db/src/reverse_log_reader.rs | 354 | ||||
| -rw-r--r-- | log_db/src/secondary_memtable.rs | 178 |
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 + } +} |
