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 | |
| parent | 30295ba5e1bdbfa69ad98fc5908ee4b1aabbd2e5 (diff) | |
Add py_bindings lib, move to monorepo structure
Diffstat (limited to 'log_db')
| -rw-r--r-- | log_db/Cargo.toml | 24 | ||||
| -rw-r--r-- | log_db/benches/benchmark.rs | 222 | ||||
| -rw-r--r-- | log_db/benches/utils.rs | 46 | ||||
| -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 | ||||
| -rw-r--r-- | log_db/tests/integration.rs | 551 | ||||
| -rw-r--r-- | log_db/tests/resources/test_db1 | bin | 0 -> 31 bytes | |||
| -rw-r--r-- | log_db/tests/resources/test_db2 | bin | 0 -> 46243 bytes | |||
| -rw-r--r-- | log_db/tests/resources/test_db3 | bin | 0 -> 46243 bytes |
13 files changed, 2680 insertions, 0 deletions
diff --git a/log_db/Cargo.toml b/log_db/Cargo.toml new file mode 100644 index 0000000..ca77753 --- /dev/null +++ b/log_db/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "log_db" +version = "0.1.0" +edition = "2021" + +[lib] +crate-type = ["lib"] + +[dependencies] +fs2 = "0.4.3" +log = "0.4.22" +priority-queue = "2.1.1" + +[dev-dependencies] +ctor = "0.2.8" +env_logger = "0.11.5" +serial_test = "3.1.1" +criterion = { version = "0.5", features = ["html_reports"] } +rand = "0.8.5" +tempfile = "3.13.0" + +[[bench]] +name = "benchmark" +harness = false diff --git a/log_db/benches/benchmark.rs b/log_db/benches/benchmark.rs new file mode 100644 index 0000000..384c9b7 --- /dev/null +++ b/log_db/benches/benchmark.rs @@ -0,0 +1,222 @@ +mod utils; + +use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion}; +use log_db::*; +use std::fs::OpenOptions; +use std::path::Path; +use tempfile; +use utils::*; + +#[derive(Eq, PartialEq, Clone, Debug)] +enum Field { + Id, + Name, + Data, +} + +pub fn upsert_various_initial_sizes(c: &mut Criterion) { + let mut group = c.benchmark_group("upsert_various_initial_sizes"); + + for size in [100, 1000, 10000, 100_000, 1_000_000, 10_000_000] { + let data_dir_obj = tempfile::tempdir().expect("Failed to get tmpdir"); + let data_dir = &data_dir_obj + .path() + .to_str() + .expect("Failed to convert tmpdir path to str"); + let mut db = DB::configure() + .data_dir(&data_dir) + .fields(&vec![ + (Field::Id, RecordField::int()), + (Field::Name, RecordField::string()), + (Field::Data, RecordField::bytes()), + ]) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB"); + prefill_db(&mut db, size).expect("Failed to prefill DB"); + + group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &_size| { + b.iter(|| { + let record = random_record(0, size as i64 + 1); + let _ = db.upsert(black_box(&record)); + }); + }); + } +} + +pub fn upsert_write_durability(c: &mut Criterion) { + let mut group = c.benchmark_group("upsert_write_durability"); + + for mode in [ + WriteDurability::Async, + WriteDurability::Flush, + WriteDurability::FlushSync, + ] { + group.bench_with_input(BenchmarkId::from_parameter(&mode), &mode, |b, _mode| { + let data_dir_obj = tempfile::tempdir().expect("Failed to get tmpdir"); + let data_dir = &data_dir_obj + .path() + .to_str() + .expect("Failed to convert tmpdir path to str"); + let mut db = DB::configure() + .data_dir(&data_dir) + .fields(&vec![ + (Field::Id, RecordField::int()), + (Field::Name, RecordField::string()), + (Field::Data, RecordField::bytes()), + ]) + .write_durability(mode.clone()) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB"); + + b.iter(|| { + let record = random_record(0, 1000); + let _ = db.upsert(black_box(&record)); + }); + }); + } +} + +pub fn get_from_disk_various_initial_sizes(c: &mut Criterion) { + let mut group = c.benchmark_group("get_from_disk_various_initial_sizes"); + + for size in [0, 10, 100, 1000, 3300, 6700, 10000, 50000, 100_000] { + let data_dir_obj = tempfile::tempdir().expect("Failed to get tmpdir"); + let data_dir = &data_dir_obj + .path() + .to_str() + .expect("Failed to convert tmpdir path to str"); + let mut db = DB::configure() + .data_dir(&data_dir) + .memtable_capacity(0) + .fields(&vec![ + (Field::Id, RecordField::int()), + (Field::Name, RecordField::string()), + (Field::Data, RecordField::bytes()), + ]) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB"); + prefill_db(&mut db, size).expect("Failed to prefill DB"); + + group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &_size| { + b.iter(|| { + let id = random_int(0, size as i64 + 1); + let _ = db.get(black_box(&RecordValue::Int(id))); + }); + }); + } +} + +pub fn get_various_memtable_capacities(c: &mut Criterion) { + let mut group = c.benchmark_group("get_various_memtable_capacities"); + + const PREFILL_N: usize = 10000; + let data_dir_obj = tempfile::tempdir().expect("Failed to get tmpdir"); + let data_dir = &data_dir_obj + .path() + .to_str() + .expect("Failed to convert tmpdir path to str"); + + // Create a db instance for prefilling + let mut db = DB::configure() + .data_dir(&data_dir) + .fields(&vec![ + (Field::Id, RecordField::int()), + (Field::Name, RecordField::string()), + (Field::Data, RecordField::bytes()), + ]) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB"); + + prefill_db(&mut db, PREFILL_N).expect("Failed to prefill DB"); + drop(db); + + // prefill_db generates IDs between 0..1000, so having memtable_capacity = 1000 + // effectively indexes the whole DB + for size in (0..).map(|x| x * 100).take_while(|&x| x <= 1000) { + group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &_size| { + let mut db = DB::configure() + .data_dir(&data_dir) + .fields(&vec![ + (Field::Id, RecordField::int()), + (Field::Name, RecordField::string()), + (Field::Data, RecordField::bytes()), + ]) + .memtable_capacity(size) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB"); + + b.iter(|| { + let id = random_int(0, 1000 + 1); + let _ = db.get(black_box(&RecordValue::Int(id))); + }); + }); + } +} + +fn reverse_read_file_with_various_buffer_sizes(c: &mut Criterion) { + let mut group = c.benchmark_group("reverse_read_file_with_various_buffer_sizes"); + group.sample_size(50); + + // odd powers of 2 + let buffer_sizes = [128, 512, 2048, 8192, 32768, 131_072, 524_288]; + const PREFILL_N: usize = 100_000; + + let data_dir_obj = tempfile::tempdir().expect("Failed to get tmpdir"); + let data_dir = &data_dir_obj + .path() + .to_str() + .expect("Failed to convert tmpdir path to str"); + + // Create a db instance for prefilling + let mut db = DB::configure() + .data_dir(&data_dir) + .fields(&vec![ + (Field::Id, RecordField::int()), + (Field::Name, RecordField::string()), + (Field::Data, RecordField::bytes()), + ]) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB"); + + prefill_db(&mut db, PREFILL_N).expect("Failed to prefill DB"); + drop(db); + + for size in buffer_sizes { + group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &_size| { + let mut file = OpenOptions::new() + .read(true) + .open(Path::new(data_dir).join("db")) + .expect("Failed to open log file"); + + b.iter(|| { + let mut rev_reader = ReverseLogReader::new_with_size(&mut file, size) + .expect("Failed to create ReverseLogReader"); + + // This is to avoid optimizing out the loop + let mut i = 0; + for _ in &mut rev_reader { + i += 1; + } + + i + }); + }); + } +} + +// Register the benchmark group +criterion_group!( + benches, + upsert_various_initial_sizes, + upsert_write_durability, + get_from_disk_various_initial_sizes, + get_various_memtable_capacities, + reverse_read_file_with_various_buffer_sizes, +); +criterion_main!(benches); diff --git a/log_db/benches/utils.rs b/log_db/benches/utils.rs new file mode 100644 index 0000000..63b3d84 --- /dev/null +++ b/log_db/benches/utils.rs @@ -0,0 +1,46 @@ +use log_db::*; +use rand::distributions::Alphanumeric; +use rand::Rng; +use std::fmt::Debug; +use std::io; + +// Function to generate a random integer +pub fn random_int(from: i64, to: i64) -> i64 { + let mut rng = rand::thread_rng(); + rng.gen_range(from..to) +} + +// Function to generate a random string +pub fn random_string(len: usize) -> String { + let mut rng = rand::thread_rng(); + (0..len).map(|_| rng.sample(Alphanumeric) as char).collect() +} + +// Function to generate random bytes +pub fn random_bytes(len: usize) -> Vec<u8> { + let mut rng = rand::thread_rng(); + (0..len).map(|_| rng.gen()).collect() +} + +// Function to generate a random record +pub fn random_record(from_id: i64, to_id: i64) -> Record { + Record { + values: vec![ + RecordValue::Int(random_int(from_id, to_id)), // Random int value between 0..1000 + RecordValue::String(random_string(5)), // Random string of length 5 + RecordValue::Bytes(random_bytes(10)), // Random bytes of length 10 + ], + } +} + +pub fn prefill_db<T: Eq + Clone + Debug>( + db: &mut DB<T>, + n_records: usize, +) -> Result<(), io::Error> { + for _ in 0..n_records { + let record = random_record(0, n_records as i64); + db.upsert(&record)?; + } + + Ok(()) +} 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 + } +} diff --git a/log_db/tests/integration.rs b/log_db/tests/integration.rs new file mode 100644 index 0000000..a2453b2 --- /dev/null +++ b/log_db/tests/integration.rs @@ -0,0 +1,551 @@ +extern crate ctor; +extern crate tempfile; + +use ctor::ctor; +use env_logger; +use log_db::*; +use std::fs::{self, OpenOptions}; +use std::path::Path; +use std::thread; +use std::time::Duration; +use tempfile::tempdir; + +#[ctor] +fn init_logger() { + let _ = env_logger::builder().is_test(true).try_init(); +} + +fn tmp_dir() -> String { + let dir = tempdir() + .expect("Failed to create temporary directory") + .path() + .to_str() + .expect("Failed to convert temporary directory path to string") + .to_string(); + fs::create_dir_all(&dir).expect("Failed to create temporary directory"); + dir +} + +#[derive(Eq, PartialEq, Clone, Debug)] +enum Field { + Id, + Name, + Data, +} + +#[test] +fn test_initialize() { + let data_dir = tmp_dir(); + let _db = DB::configure() + .data_dir(&data_dir) + .fields(&vec![ + (Field::Id, RecordField::int()), + (Field::Name, RecordField::string()), + (Field::Data, RecordField::bytes()), + ]) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB instance"); +} + +#[test] +fn test_upsert_and_get_with_primary_memtable() { + let data_dir = tmp_dir(); + let mut db = DB::configure() + .data_dir(&data_dir) + .fields(&vec![ + (Field::Id, RecordField::int()), + (Field::Name, RecordField::string()), + (Field::Data, RecordField::bytes()), + ]) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB instance"); + + let record = Record { + values: vec![ + RecordValue::Int(1), + RecordValue::String("Alice".to_string()), + RecordValue::Bytes(vec![0, 1, 2]), + ], + }; + db.upsert(&record).unwrap(); + + let result = db.get(&RecordValue::Int(1)).unwrap().unwrap(); + + // Check that the IDs match + assert!(match (&result.values[0], &record.values[0]) { + (RecordValue::Int(a), RecordValue::Int(b)) => a == b, + _ => false, + }); +} + +#[test] +fn test_upsert_and_get_without_memtable() { + let data_dir = tmp_dir(); + let mut db = DB::configure() + .data_dir(&data_dir) + .memtable_capacity(0) + .fields(&vec![ + (Field::Id, RecordField::int()), + (Field::Name, RecordField::string().nullable()), + (Field::Data, RecordField::bytes()), + ]) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB instance"); + + // Insert some records + let record0 = Record { + values: vec![ + RecordValue::Int(0), + RecordValue::Null, + RecordValue::Bytes(vec![3, 4, 5]), + ], + }; + db.upsert(&record0).unwrap(); + + let record1 = Record { + values: vec![ + RecordValue::Int(1), + RecordValue::String("Alice".to_string()), + RecordValue::Bytes(vec![0, 1, 2]), + ], + }; + db.upsert(&record1).unwrap(); + + let record2 = Record { + values: vec![ + RecordValue::Int(1), + RecordValue::String("Bob".to_string()), + RecordValue::Bytes(vec![0, 1, 2]), + ], + }; + db.upsert(&record2).unwrap(); + + let record3 = Record { + values: vec![ + RecordValue::Int(2), + RecordValue::String("George".to_string()), + RecordValue::Bytes(vec![]), + ], + }; + db.upsert(&record3).unwrap(); + + // Get with ID = 0 + let result = db.get(&RecordValue::Int(0)).unwrap().unwrap(); + + // Should match record0 + assert!(match (&result.values[0], &record0.values[0]) { + (RecordValue::Int(a), RecordValue::Int(b)) => a == b, + _ => false, + }); + assert!(match (&result.values[1], &record0.values[1]) { + (RecordValue::Null, RecordValue::Null) => true, + _ => false, + }); + + // Get with ID = 1 + let result = db.get(&RecordValue::Int(1)).unwrap().unwrap(); + + // Should match record2 + assert!(match (&result.values[0], &record2.values[0]) { + (RecordValue::Int(a), RecordValue::Int(b)) => a == b, + _ => false, + }); + assert!(match (&result.values[1], &record2.values[1]) { + (RecordValue::String(a), RecordValue::String(b)) => a == b, + _ => false, + }); +} + +#[test] +fn test_upsert_fails_on_null_in_non_nullable_field() { + let data_dir = tmp_dir(); + let mut db = DB::configure() + .data_dir(&data_dir) + .fields(&vec![(Field::Id, RecordField::int())]) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB instance"); + + let record = Record { + // Null value + values: vec![RecordValue::Null], + }; + assert!(db.upsert(&record).is_err()); +} + +#[test] +fn test_upsert_fails_on_invalid_number_of_values() { + let data_dir = tmp_dir(); + let mut db = DB::configure() + .data_dir(&data_dir) + .fields(&vec![ + (Field::Id, RecordField::int()), + (Field::Name, RecordField::string()), + (Field::Data, RecordField::bytes()), + ]) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB instance"); + + let record = Record { + // Missing primary key + values: vec![ + RecordValue::String("Alice".to_string()), + RecordValue::Bytes(vec![0, 1, 2]), + ], + }; + assert!(db.upsert(&record).is_err()); +} + +#[test] +fn test_upsert_fails_on_invalid_value_type() { + let data_dir = tmp_dir(); + let mut db = DB::configure() + .data_dir(&data_dir) + .fields(&vec![ + (Field::Id, RecordField::int()), + (Field::Name, RecordField::string()), + (Field::Data, RecordField::bytes()), + ]) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB instance"); + + let record = Record { + values: vec![ + RecordValue::String("foo".to_string()), + RecordValue::String("bar".to_string()), + RecordValue::String("baz".to_string()), + ], + }; + assert!(db.upsert(&record).is_err()); +} + +#[test] +fn test_upsert_and_get_from_secondary_memtable() { + let data_dir = tmp_dir(); + let mut db = DB::configure() + .data_dir(&data_dir) + .fields(&vec![ + (Field::Id, RecordField::int()), + (Field::Name, RecordField::string()), + (Field::Data, RecordField::bytes()), + ]) + .primary_key(Field::Id) + .secondary_keys(vec![Field::Name]) + .initialize() + .expect("Failed to initialize DB instance"); + + // Insert some records + let record0 = Record { + values: vec![ + RecordValue::Int(0), + RecordValue::String("John".to_string()), + RecordValue::Bytes(vec![3, 4, 5]), + ], + }; + db.upsert(&record0).unwrap(); + + let record1 = Record { + values: vec![ + RecordValue::Int(1), + RecordValue::String("John".to_string()), + RecordValue::Bytes(vec![1, 2, 3]), + ], + }; + db.upsert(&record1).unwrap(); + + let record2 = Record { + values: vec![ + RecordValue::Int(2), + RecordValue::String("George".to_string()), + RecordValue::Bytes(vec![1, 2, 3]), + ], + }; + db.upsert(&record2).unwrap(); + + // Delete the DB so that any results must come from a memtable + fs::remove_file(Path::new(&data_dir).join("db")).expect("Failed to delete the DB log file"); + + // There should be 2 Johns + let johns = db + .find_all(&Field::Name, &RecordValue::String("John".to_string())) + .expect("Failed to find all Johns"); + + assert_eq!(johns.len(), 2); +} + +#[test] +fn test_initialize_and_read_from_primary_memtable_fixture_db2() { + let data_dir = tmp_dir(); + // Copy the fixture DB to the test data directory + fs::create_dir_all(&data_dir).expect("Failed to create the test data directory"); + fs::copy( + &Path::new(TEST_RESOURCES_DIR).join("test_db2"), + &Path::new(&data_dir).join("db"), + ) + .expect("Failed to copy the fixture DB"); + + let mut db = DB::configure() + .data_dir(&data_dir) + .fields(&vec![ + (Field::Id, RecordField::int()), + (Field::Name, RecordField::string()), + (Field::Data, RecordField::bytes()), + ]) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB instance"); + + // Delete the DB so that any results must come from a memtable + fs::remove_file(Path::new(&data_dir).join("db")).expect("Failed to delete the DB log file"); + + let result = db.get(&RecordValue::Int(1)).unwrap().unwrap(); + + // Check that the IDs match + let expected = RecordValue::Int(1); + assert!(match (&result.values[0], &expected) { + (RecordValue::Int(a), RecordValue::Int(b)) => a == b, + _ => false, + }); +} + +#[test] +fn test_initialize_without_memtables_fixture_db3() { + let data_dir = tmp_dir(); + // Copy the fixture DB to the test data directory + fs::create_dir_all(&data_dir).expect("Failed to create the test data directory"); + fs::copy( + &Path::new(TEST_RESOURCES_DIR).join("test_db3"), + &Path::new(&data_dir).join("db"), + ) + .expect("Failed to copy the fixture DB"); + + let mut db = DB::configure() + .data_dir(&data_dir) + .fields(&vec![ + (Field::Id, RecordField::int()), + (Field::Name, RecordField::string()), + (Field::Data, RecordField::bytes()), + ]) + .memtable_capacity(0) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB instance"); + + let result = db.get(&RecordValue::Int(1)).unwrap().unwrap(); + + // Check that the IDs match + let expected = RecordValue::Int(1); + assert!(match (&result.values[0], &expected) { + (RecordValue::Int(a), RecordValue::Int(b)) => a == b, + _ => false, + }); +} + +#[test] +fn test_multiple_writing_threads() { + let data_dir = tmp_dir(); + let mut threads = vec![]; + let threads_n = 100; + + for i in 0..threads_n { + let data_dir = data_dir.clone(); + threads.push(thread::spawn(move || { + let mut db = DB::configure() + .data_dir(&data_dir) + .fields(&vec![(Field::Id, RecordField::int())]) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB instance"); + + let record = Record { + values: vec![RecordValue::Int(i)], + }; + db.upsert(&record).expect("Failed to upsert record"); + })); + } + + for thread in threads { + thread.join().expect("Failed to join thread"); + } + + // Read the records + let mut db = DB::configure() + .data_dir(&data_dir) + .fields(&vec![(Field::Id, RecordField::int())]) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB instance"); + + for i in 0..threads_n { + let result = db + .get(&RecordValue::Int(i)) + .expect("Failed to get record") + .expect("Record not found"); + let expected = RecordValue::Int(i); + assert!(match (&result.values[0], &expected) { + (RecordValue::Int(a), RecordValue::Int(b)) => a == b, + _ => false, + }); + } +} + +#[test] +fn test_one_writer_and_multiple_reading_threads() { + let data_dir = tmp_dir(); + let mut threads = vec![]; + let threads_n = 100; + + // Add readers that poll for the records + for i in 0..threads_n { + let data_dir = data_dir.clone(); + threads.push(thread::spawn(move || { + let mut db = DB::configure() + .data_dir(&data_dir) + .fields(&vec![(Field::Id, RecordField::int())]) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB instance"); + + let mut timeout = 5; + loop { + let result = db.get(&RecordValue::Int(i)).expect("Failed to get record"); + match result { + None => { + thread::sleep(Duration::from_millis(timeout)); + timeout = std::cmp::min(timeout * 2, 100); + continue; + } + Some(result) => { + let expected = RecordValue::Int(i); + assert!(match (&result.values[0], &expected) { + (RecordValue::Int(a), RecordValue::Int(b)) => a == b, + _ => false, + }); + break; + } + }; + } + })); + } + + // Add a writer that inserts the records + threads.push(thread::spawn(move || { + let mut db = DB::configure() + .data_dir(&data_dir) + .fields(&vec![(Field::Id, RecordField::int())]) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB instance"); + + for i in 0..threads_n { + let record = Record { + values: vec![RecordValue::Int(i)], + }; + db.upsert(&record).expect("Failed to upsert record"); + } + })); + + for thread in threads { + thread.join().expect("Failed to join thread"); + } +} + +#[test] +fn test_literal_escape_is_escaped() { + let data_dir = tmp_dir(); + + let mut db = DB::configure() + .data_dir(&data_dir) + .memtable_capacity(0) // disable memtables + .fields(&vec![ + (Field::Id, RecordField::int()), + (Field::Data, RecordField::bytes()), + ]) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB instance"); + + let record = Record { + values: vec![ + RecordValue::Int(1), + RecordValue::Bytes(vec![0x1A, 0x1B, 0x1C, 0x1D]), + ], + }; + + db.upsert(&record).expect("Failed to upsert record"); + + let found = db + .get(&RecordValue::Int(1)) + .expect("Failed to get record") + .expect("Record not found"); + + let received = match &found.values[1] { + RecordValue::Bytes(bytes) => bytes, + _ => panic!("Unexpected record value type"), + }; + + assert_eq!(received, &vec![0x1A, 0x1B, 0x1C, 0x1D]); +} + +#[test] +fn test_log_is_rotated_when_capacity_reached() { + let data_dir = tmp_dir(); + + let record = Record { + values: vec![RecordValue::Int(1), RecordValue::Bytes(vec![1, 2, 3, 4])], + }; + let record_len = &record.serialize().len() + SEQ_RECORD_SEP.len(); + + let mut db = DB::configure() + .data_dir(&data_dir) + .memtable_capacity(0) // disable memtables + .segment_size(10 * record_len) // small log segment size + .fields(&vec![ + (Field::Id, RecordField::int()), + (Field::Data, RecordField::bytes()), + ]) + .primary_key(Field::Id) + .initialize() + .expect("Failed to initialize DB instance"); + + // Insert more records than fits the capacity + for _ in 0..25 { + db.upsert(&record).expect("Failed to upsert record"); + + db.do_maintenance_tasks() + .expect("Failed to do maintenance tasks"); + } + + // Check that the rotated segments exist + assert!(Path::new(&data_dir) + .join(ACTIVE_LOG_FILENAME) + .with_extension("1") + .exists()); + + assert!(Path::new(&data_dir) + .join(ACTIVE_LOG_FILENAME) + .with_extension("2") + .exists()); + + assert!(!Path::new(&data_dir) + .join(ACTIVE_LOG_FILENAME) + .with_extension("3") + .exists()); + + // Check that the active file only contains two rows + let mut file = OpenOptions::new() + .read(true) + .open(Path::new(&data_dir).join(ACTIVE_LOG_FILENAME)) + .expect("File could not be opened"); + let records_in_active_log = ForwardLogReader::new(&mut file).count(); + assert_eq!(records_in_active_log, 5); + + // Look for nonexistant record to scan all segment files + let found = db.get(&RecordValue::Int(2)).expect("Failed to get record"); + assert!(found.is_none()); +} diff --git a/log_db/tests/resources/test_db1 b/log_db/tests/resources/test_db1 Binary files differnew file mode 100644 index 0000000..1f19a86 --- /dev/null +++ b/log_db/tests/resources/test_db1 diff --git a/log_db/tests/resources/test_db2 b/log_db/tests/resources/test_db2 Binary files differnew file mode 100644 index 0000000..156659f --- /dev/null +++ b/log_db/tests/resources/test_db2 diff --git a/log_db/tests/resources/test_db3 b/log_db/tests/resources/test_db3 Binary files differnew file mode 100644 index 0000000..156659f --- /dev/null +++ b/log_db/tests/resources/test_db3 |
