From a4efc83f7f29c6ef8ac9b4c57201c9ecc31266b3 Mon Sep 17 00:00:00 2001 From: Jan Tuomi Date: Mon, 7 Oct 2024 12:43:16 +0300 Subject: Add py_bindings lib, move to monorepo structure --- Cargo.lock | 162 ++++++-- Cargo.toml | 27 +- benches/benchmark.rs | 222 ----------- benches/utils.rs | 46 --- log_db/Cargo.toml | 24 ++ log_db/benches/benchmark.rs | 222 +++++++++++ log_db/benches/utils.rs | 46 +++ log_db/src/common.rs | 308 +++++++++++++++ log_db/src/forward_log_reader.rs | 124 ++++++ log_db/src/lib.rs | 794 +++++++++++++++++++++++++++++++++++++++ log_db/src/primary_memtable.rs | 79 ++++ log_db/src/reverse_log_reader.rs | 354 +++++++++++++++++ log_db/src/secondary_memtable.rs | 178 +++++++++ log_db/tests/integration.rs | 551 +++++++++++++++++++++++++++ log_db/tests/resources/test_db1 | Bin 0 -> 31 bytes log_db/tests/resources/test_db2 | Bin 0 -> 46243 bytes log_db/tests/resources/test_db3 | Bin 0 -> 46243 bytes py_bindings/Cargo.toml | 11 + py_bindings/src/lib.rs | 16 + src/common.rs | 308 --------------- src/forward_log_reader.rs | 124 ------ src/lib.rs | 794 --------------------------------------- src/primary_memtable.rs | 79 ---- src/reverse_log_reader.rs | 354 ----------------- src/secondary_memtable.rs | 178 --------- tests/integration.rs | 555 --------------------------- tests/resources/test_db1 | Bin 31 -> 0 bytes tests/resources/test_db2 | Bin 46243 -> 0 bytes tests/resources/test_db3 | Bin 46243 -> 0 bytes 29 files changed, 2846 insertions(+), 2710 deletions(-) delete mode 100644 benches/benchmark.rs delete mode 100644 benches/utils.rs create mode 100644 log_db/Cargo.toml create mode 100644 log_db/benches/benchmark.rs create mode 100644 log_db/benches/utils.rs create mode 100644 log_db/src/common.rs create mode 100644 log_db/src/forward_log_reader.rs create mode 100644 log_db/src/lib.rs create mode 100644 log_db/src/primary_memtable.rs create mode 100644 log_db/src/reverse_log_reader.rs create mode 100644 log_db/src/secondary_memtable.rs create mode 100644 log_db/tests/integration.rs create mode 100644 log_db/tests/resources/test_db1 create mode 100644 log_db/tests/resources/test_db2 create mode 100644 log_db/tests/resources/test_db3 create mode 100644 py_bindings/Cargo.toml create mode 100644 py_bindings/src/lib.rs delete mode 100644 src/common.rs delete mode 100644 src/forward_log_reader.rs delete mode 100644 src/lib.rs delete mode 100644 src/primary_memtable.rs delete mode 100644 src/reverse_log_reader.rs delete mode 100644 src/secondary_memtable.rs delete mode 100644 tests/integration.rs delete mode 100644 tests/resources/test_db1 delete mode 100644 tests/resources/test_db2 delete mode 100644 tests/resources/test_db3 diff --git a/Cargo.lock b/Cargo.lock index fa5b94a..3dad95c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -300,9 +300,9 @@ dependencies = [ [[package]] name = "futures" -version = "0.3.30" +version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "645c6916888f6cb6350d2550b80fb63e734897a8498abe35cfb732b6487804b0" +checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876" dependencies = [ "futures-channel", "futures-core", @@ -315,9 +315,9 @@ dependencies = [ [[package]] name = "futures-channel" -version = "0.3.30" +version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eac8f7d7865dcb88bd4373ab671c8cf4508703796caa2b1985a9ca867b3fcb78" +checksum = "2dff15bf788c671c1934e366d07e30c1814a8ef514e1af724a602e8a2fbe1b10" dependencies = [ "futures-core", "futures-sink", @@ -325,15 +325,15 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.30" +version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dfc6580bb841c5a68e9ef15c77ccc837b40a7504914d52e47b8b0e9bbda25a1d" +checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e" [[package]] name = "futures-executor" -version = "0.3.30" +version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a576fc72ae164fca6b9db127eaa9a9dda0d61316034f33a0a0d4eda41f02b01d" +checksum = "1e28d1d997f585e54aebc3f97d39e72338912123a67330d723fdbb564d646c9f" dependencies = [ "futures-core", "futures-task", @@ -342,27 +342,27 @@ dependencies = [ [[package]] name = "futures-io" -version = "0.3.30" +version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a44623e20b9681a318efdd71c299b6b222ed6f231972bfe2f224ebad6311f0c1" +checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" [[package]] name = "futures-sink" -version = "0.3.30" +version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fb8e00e87438d937621c1c6269e53f536c14d3fbd6a042bb24879e57d474fb5" +checksum = "e575fab7d1e0dcb8d0c7bcf9a63ee213816ab51902e6d244a95819acacf1d4f7" [[package]] name = "futures-task" -version = "0.3.30" +version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38d84fa142264698cdce1a9f9172cf383a0c82de1bddcf3092901442c4097004" +checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" [[package]] name = "futures-util" -version = "0.3.30" +version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d6401deb83407ab3da39eba7e33987a73c3df0c82b4bb5813ee871c19c41d48" +checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81" dependencies = [ "futures-channel", "futures-core", @@ -398,9 +398,15 @@ dependencies = [ [[package]] name = "hashbrown" -version = "0.14.5" +version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" +checksum = "1e087f84d4f86bf4b218b927129862374b72199ae7d8657835f1e89000eea4fb" + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" [[package]] name = "hermit-abi" @@ -416,14 +422,20 @@ checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4" [[package]] name = "indexmap" -version = "2.5.0" +version = "2.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68b900aa2f7301e21c36462b170ee99994de34dff39a4a6a528e80e7376d07e5" +checksum = "707907fe3c25f5424cce2cb7e1cbcafee6bdbe735ca90ef77c29e84591e5b9da" dependencies = [ "equivalent", "hashbrown", ] +[[package]] +name = "indoc" +version = "2.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b248f5224d1d606005e02c97f5aa4e88eeb230488bcc03bc9ca4d7991399f2b5" + [[package]] name = "is-terminal" version = "0.4.13" @@ -514,6 +526,15 @@ version = "2.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -525,9 +546,9 @@ dependencies = [ [[package]] name = "once_cell" -version = "1.19.0" +version = "1.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" +checksum = "1261fe7e33c73b354eab43b1273a57c8f967d0391e80353e51f764ac02cf6775" [[package]] name = "oorandom" @@ -598,6 +619,12 @@ dependencies = [ "plotters-backend", ] +[[package]] +name = "portable-atomic" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc9c68a3f6da06753e9335d63e27f6b9754dd1920d941135b7ea8224f141adb2" + [[package]] name = "ppv-lite86" version = "0.2.20" @@ -627,6 +654,77 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "py_bindings" +version = "0.1.0" +dependencies = [ + "log_db", + "pyo3", +] + +[[package]] +name = "pyo3" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15ee168e30649f7f234c3d49ef5a7a6cbf5134289bc46c29ff3155fa3221c225" +dependencies = [ + "cfg-if", + "indoc", + "libc", + "memoffset", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", + "unindent", +] + +[[package]] +name = "pyo3-build-config" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e61cef80755fe9e46bb8a0b8f20752ca7676dcc07a5277d8b7768c6172e529b3" +dependencies = [ + "once_cell", + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67ce096073ec5405f5ee2b8b31f03a68e02aa10d5d4f565eca04acc41931fa1c" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2440c6d12bc8f3ae39f1e775266fa5122fd0c8891ce7520fa6048e683ad3de28" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1be962f0e06da8f8465729ea2cb71a416d2257dff56cbe40a70d3e62a93ae5d1" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn", +] + [[package]] name = "quote" version = "1.0.37" @@ -688,9 +786,9 @@ dependencies = [ [[package]] name = "redox_syscall" -version = "0.5.6" +version = "0.5.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "355ae415ccd3a04315d3f8246e86d67689ea74d88d915576e1589a351062a13b" +checksum = "9b6dfecf2c74bce2466cabf93f6664d6998a69eb21e39f4207930065b27b771f" dependencies = [ "bitflags", ] @@ -754,9 +852,9 @@ dependencies = [ [[package]] name = "scc" -version = "2.1.17" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c947adb109a8afce5fc9c7bf951f87f146e9147b3a6a58413105628fb1d1e66" +checksum = "836f1e0f4963ef5288b539b643b35e043e76a32d0f4e47e67febf69576527f50" dependencies = [ "sdd", ] @@ -856,6 +954,12 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "target-lexicon" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" + [[package]] name = "tempfile" version = "3.13.0" @@ -885,6 +989,12 @@ version = "1.0.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e91b56cd4cadaeb79bbf1a5645f6b4f8dc5bde8834ad5894a8db35fda9efa1fe" +[[package]] +name = "unindent" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7de7d73e1754487cb58364ee906a499937a0dfabd86bcb980fa99ec8c8fa2ce" + [[package]] name = "utf8parse" version = "0.2.2" diff --git a/Cargo.toml b/Cargo.toml index ca77753..ec5c3ea 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,24 +1,3 @@ -[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 +[workspace] +resolver = "2" +members = ["log_db", "py_bindings"] diff --git a/benches/benchmark.rs b/benches/benchmark.rs deleted file mode 100644 index 384c9b7..0000000 --- a/benches/benchmark.rs +++ /dev/null @@ -1,222 +0,0 @@ -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/benches/utils.rs b/benches/utils.rs deleted file mode 100644 index 63b3d84..0000000 --- a/benches/utils.rs +++ /dev/null @@ -1,46 +0,0 @@ -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 { - 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( - db: &mut DB, - 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/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 { + 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( + db: &mut DB, + 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 { + 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), +} + +impl RecordValue { + pub fn serialize(&self) -> Vec { + 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 { + 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, +} + +impl Record { + pub fn serialize(&self) -> Vec { + 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 { + 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 { + // 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, io::Error> { + // The buffer that stores the bytes read from the file. + let mut read_buf: Vec = Vec::new(); + // The buffer that stores all the bytes of the record read so far in reverse order. + let mut result_buf: Vec = 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 { + let mut special_buf: Vec = 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 { + 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, + segment_size: Option, + memtable_capacity: Option, + fields: Option<&'a Vec<(Field, RecordField)>>, + primary_key: Option, + secondary_keys: Option>, + memtable_evict_policy: Option, + write_durability: Option, +} + +impl<'a, Field: Eq + Clone + Debug> ConfigBuilder<'a, Field> { + pub fn new() -> ConfigBuilder<'a, Field> { + ConfigBuilder:: { + 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) -> &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, io::Error> { + let config = Config:: { + 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 { + 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, + pub memtable_evict_policy: MemtableEvictPolicy, + pub write_durability: WriteDurability, +} + +pub struct DB { + config: Config, + log_path: PathBuf, + log_file: fs::File, + primary_key_index: usize, + primary_memtable: PrimaryMemtable, + secondary_memtables: Vec>, +} + +impl DB { + /// Create a new database configuration builder. + pub fn configure() -> ConfigBuilder<'static, Field> { + ConfigBuilder::new() + } + + fn initialize(config: &Config) -> Result, 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:: { + 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, 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 = 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, 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::>(); + + 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 { + 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 { + 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 { + 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, io::Error> { + // TODO: optimize the vecs out of here + let files = fs::read_dir(&self.config.data_dir)?; + let mut nums: Vec = 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::() + .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, + /// A max heap priority queue of keys. Note: n_operations must + /// be negated upon append to evict oldest values first. + evict_queue: PriorityQueue, + 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, + /// 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 { + 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 { + 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, 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 = 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 { + 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 { + 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) -> Result { + // 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 { + let mut special_buf: Vec = 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 { + 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(&self, state: &mut H) + where + H: Hasher, + { + self.primary_value.hash(state) + } +} + +pub struct SecondaryMemtable { + 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>, + /// A max heap priority queue of keys. Note: n_operations must + /// be negated upon append to evict oldest values first. + evict_queue: PriorityQueue, + evict_policy: MemtableEvictPolicy, +} + +impl SecondaryMemtable { + pub fn new( + field: &Field, + primary_field_index: usize, + capacity: usize, + evict_policy: MemtableEvictPolicy, + ) -> SecondaryMemtable { + 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 new file mode 100644 index 0000000..1f19a86 Binary files /dev/null and b/log_db/tests/resources/test_db1 differ diff --git a/log_db/tests/resources/test_db2 b/log_db/tests/resources/test_db2 new file mode 100644 index 0000000..156659f Binary files /dev/null and b/log_db/tests/resources/test_db2 differ diff --git a/log_db/tests/resources/test_db3 b/log_db/tests/resources/test_db3 new file mode 100644 index 0000000..156659f Binary files /dev/null and b/log_db/tests/resources/test_db3 differ diff --git a/py_bindings/Cargo.toml b/py_bindings/Cargo.toml new file mode 100644 index 0000000..2f7a83f --- /dev/null +++ b/py_bindings/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "py_bindings" +version = "0.1.0" +edition = "2021" + +[lib] +crate-type = ["lib"] + +[dependencies] +log_db = { path = "../log_db" } +pyo3 = "0.22.3" diff --git a/py_bindings/src/lib.rs b/py_bindings/src/lib.rs new file mode 100644 index 0000000..abee9f2 --- /dev/null +++ b/py_bindings/src/lib.rs @@ -0,0 +1,16 @@ +use log_db::*; + +pub fn add(left: u64, right: u64) -> u64 { + left + right +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn it_works() { + let result = add(2, 2); + assert_eq!(result, 4); + } +} diff --git a/src/common.rs b/src/common.rs deleted file mode 100644 index c2af766..0000000 --- a/src/common.rs +++ /dev/null @@ -1,308 +0,0 @@ -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 { - 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), -} - -impl RecordValue { - pub fn serialize(&self) -> Vec { - 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 { - 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, -} - -impl Record { - pub fn serialize(&self) -> Vec { - 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 { - 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 { - // 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/src/forward_log_reader.rs b/src/forward_log_reader.rs deleted file mode 100644 index 10788a7..0000000 --- a/src/forward_log_reader.rs +++ /dev/null @@ -1,124 +0,0 @@ -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, io::Error> { - // The buffer that stores the bytes read from the file. - let mut read_buf: Vec = Vec::new(); - // The buffer that stores all the bytes of the record read so far in reverse order. - let mut result_buf: Vec = 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 { - let mut special_buf: Vec = 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 { - 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/src/lib.rs b/src/lib.rs deleted file mode 100644 index 6787cf1..0000000 --- a/src/lib.rs +++ /dev/null @@ -1,794 +0,0 @@ -#[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, - segment_size: Option, - memtable_capacity: Option, - fields: Option<&'a Vec<(Field, RecordField)>>, - primary_key: Option, - secondary_keys: Option>, - memtable_evict_policy: Option, - write_durability: Option, -} - -impl<'a, Field: Eq + Clone + Debug> ConfigBuilder<'a, Field> { - pub fn new() -> ConfigBuilder<'a, Field> { - ConfigBuilder:: { - 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) -> &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, io::Error> { - let config = Config:: { - 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 { - 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, - pub memtable_evict_policy: MemtableEvictPolicy, - pub write_durability: WriteDurability, -} - -pub struct DB { - config: Config, - log_path: PathBuf, - log_file: fs::File, - primary_key_index: usize, - primary_memtable: PrimaryMemtable, - secondary_memtables: Vec>, -} - -impl DB { - /// Create a new database configuration builder. - pub fn configure() -> ConfigBuilder<'static, Field> { - ConfigBuilder::new() - } - - fn initialize(config: &Config) -> Result, 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:: { - 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, 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 = 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, 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::>(); - - 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 { - 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 { - 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 { - 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, io::Error> { - // TODO: optimize the vecs out of here - let files = fs::read_dir(&self.config.data_dir)?; - let mut nums: Vec = 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::() - .expect("Extension was not a valid number"); - - Some(ext_num) - }) - .collect(); - - nums.sort(); - nums.push(0); - nums.reverse(); - Ok(nums) - } -} diff --git a/src/primary_memtable.rs b/src/primary_memtable.rs deleted file mode 100644 index 0c94ceb..0000000 --- a/src/primary_memtable.rs +++ /dev/null @@ -1,79 +0,0 @@ -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, - /// A max heap priority queue of keys. Note: n_operations must - /// be negated upon append to evict oldest values first. - evict_queue: PriorityQueue, - 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/src/reverse_log_reader.rs b/src/reverse_log_reader.rs deleted file mode 100644 index 5d80dba..0000000 --- a/src/reverse_log_reader.rs +++ /dev/null @@ -1,354 +0,0 @@ -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, - /// 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 { - 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 { - 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, 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 = 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 { - 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 { - 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) -> Result { - // 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 { - let mut special_buf: Vec = 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 { - match self.read_record() { - Ok(Some(record)) => Some(record), - Ok(None) => None, - Err(err) => { - panic!("Error reading record: {:?}", err,) - } - } - } -} diff --git a/src/secondary_memtable.rs b/src/secondary_memtable.rs deleted file mode 100644 index fcf4a7f..0000000 --- a/src/secondary_memtable.rs +++ /dev/null @@ -1,178 +0,0 @@ -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(&self, state: &mut H) - where - H: Hasher, - { - self.primary_value.hash(state) - } -} - -pub struct SecondaryMemtable { - 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>, - /// A max heap priority queue of keys. Note: n_operations must - /// be negated upon append to evict oldest values first. - evict_queue: PriorityQueue, - evict_policy: MemtableEvictPolicy, -} - -impl SecondaryMemtable { - pub fn new( - field: &Field, - primary_field_index: usize, - capacity: usize, - evict_policy: MemtableEvictPolicy, - ) -> SecondaryMemtable { - 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/tests/integration.rs b/tests/integration.rs deleted file mode 100644 index e6be74f..0000000 --- a/tests/integration.rs +++ /dev/null @@ -1,555 +0,0 @@ -extern crate ctor; -extern crate tempfile; - -use ctor::ctor; -use env_logger; -use log::debug; -use log_db::{ - self, ForwardLogReader, RecordField, ACTIVE_LOG_FILENAME, SEQ_LIT_ESCAPE, SEQ_RECORD_SEP, -}; -use log_db::{Record, RecordValue, DB, TEST_RESOURCES_DIR}; -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/tests/resources/test_db1 b/tests/resources/test_db1 deleted file mode 100644 index 1f19a86..0000000 Binary files a/tests/resources/test_db1 and /dev/null differ diff --git a/tests/resources/test_db2 b/tests/resources/test_db2 deleted file mode 100644 index 156659f..0000000 Binary files a/tests/resources/test_db2 and /dev/null differ diff --git a/tests/resources/test_db3 b/tests/resources/test_db3 deleted file mode 100644 index 156659f..0000000 Binary files a/tests/resources/test_db3 and /dev/null differ -- cgit v1.3