Add exclusive lock request file to prioritize writes

6 files changed, 230 insertions, 95 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 2be0b4d..a96e56c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -47,7 +47,7 @@ version = "1.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6d36fc52c7f6c869915e99412912f22093507da8d9e942ceaf66fe4b7c14422a"
 dependencies = [
- "windows-sys",
+ "windows-sys 0.52.0",
 ]
 
 [[package]]
@@ -57,7 +57,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "5bf74e1b6e971609db8ca7a9ce79fd5768ab6ae46441c572e46cf596f59e57f8"
 dependencies = [
  "anstyle",
- "windows-sys",
+ "windows-sys 0.52.0",
 ]
 
 [[package]]
@@ -85,6 +85,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d3fd119d74b830634cea2a0f58bbd0d54540518a14397557951e79340abc28c0"
 
 [[package]]
+name = "ctor"
+version = "0.2.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "edb49164822f3ee45b17acd4a208cfc1251410cf0cad9a833234c9890774dd9f"
+dependencies = [
+ "quote",
+ "syn",
+]
+
+[[package]]
 name = "env_filter"
 version = "0.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -114,6 +124,22 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5"
 
 [[package]]
+name = "errno"
+version = "0.3.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "534c5cf6194dfab3db3242765c03bbe257cf92f22b38f6bc0c58d59108a820ba"
+dependencies = [
+ "libc",
+ "windows-sys 0.52.0",
+]
+
+[[package]]
+name = "fastrand"
+version = "2.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e8c02a5121d4ea3eb16a80748c74f5549a5665e4c21333c6098f283870fbdea6"
+
+[[package]]
 name = "fs2"
 version = "0.4.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -235,6 +261,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "561d97a539a36e26a9a5fad1ea11a3039a67714694aaa379433e580854bc3dc5"
 
 [[package]]
+name = "linux-raw-sys"
+version = "0.4.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "78b3ae25bc7c8c38cec158d1f2757ee79e9b3740fbc7ccf0e59e4b08d793fa89"
+
+[[package]]
 name = "lock_api"
 version = "0.4.12"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -254,12 +286,14 @@ checksum = "a7a70ba024b9dc04c27ea2f0c0548feb474ec5c54bba33a7f72f873a39d07b24"
 name = "log_db"
 version = "0.1.0"
 dependencies = [
+ "ctor",
  "env_logger",
  "fs2",
  "log",
  "priority-queue",
  "rev_buf_reader",
  "serial_test",
+ "tempfile",
 ]
 
 [[package]]
@@ -386,6 +420,19 @@ dependencies = [
 ]
 
 [[package]]
+name = "rustix"
+version = "0.38.37"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8acb788b847c24f28525660c4d7758620a7210875711f79e7f663cc152726811"
+dependencies = [
+ "bitflags",
+ "errno",
+ "libc",
+ "linux-raw-sys",
+ "windows-sys 0.52.0",
+]
+
+[[package]]
 name = "scc"
 version = "2.1.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -458,6 +505,19 @@ dependencies = [
 ]
 
 [[package]]
+name = "tempfile"
+version = "3.13.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f0f2c9fc62d0beef6951ccffd757e241266a2c833136efbe35af6cd2567dca5b"
+dependencies = [
+ "cfg-if",
+ "fastrand",
+ "once_cell",
+ "rustix",
+ "windows-sys 0.59.0",
+]
+
+[[package]]
 name = "unicode-ident"
 version = "1.0.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -501,6 +561,15 @@ dependencies = [
 ]
 
 [[package]]
+name = "windows-sys"
+version = "0.59.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b"
+dependencies = [
+ "windows-targets",
+]
+
+[[package]]
 name = "windows-targets"
 version = "0.52.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/Cargo.toml b/Cargo.toml
index fd8775b..9c140ae 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -3,15 +3,17 @@ 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"
 rev_buf_reader = "0.3.0"
 
-[lib]
-crate-type = ["lib"]
-
 [dev-dependencies]
+ctor = "0.2.8"
 env_logger = "0.11.5"
 serial_test = "3.1.1"
+tempfile = "3.13.0"
diff --git a/src/common.rs b/src/common.rs
index 4922c73..c4fce39 100644
--- a/src/common.rs
+++ b/src/common.rs
@@ -2,13 +2,16 @@ 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 Unix-like systems
+use std::os::unix::fs::MetadataExt;
 
+// For Windows
 #[cfg(windows)]
-use std::os::windows::fs::MetadataExt; // For 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';
diff --git a/src/lib.rs b/src/lib.rs
index c0cf00c..6d040c4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -8,6 +8,7 @@ mod primary_memtable;
 mod secondary_memtable;
 
 pub use common::*;
+use fs2::lock_contended_error;
 use fs2::FileExt;
 pub use log_reader::ForwardLogReader;
 pub use log_reader::ReverseLogReader;
@@ -17,6 +18,7 @@ use std::fmt::Debug;
 use std::fs::{self};
 use std::io::{self, Write};
 use std::path::{Path, PathBuf};
+use std::thread;
 
 pub struct ConfigBuilder<'a, Field: Eq + Clone + Debug> {
     data_dir: Option<String>,
@@ -156,11 +158,17 @@ impl<Field: Eq + Clone + Debug> DB<Field> {
         let log_path = Path::new(&config.data_dir).join(ACTIVE_LOG_FILENAME);
 
         // Create the log file if it does not exist
-        let _file = fs::OpenOptions::new()
+        fs::OpenOptions::new()
             .create(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
@@ -283,7 +291,7 @@ impl<Field: Eq + Clone + Debug> DB<Field> {
             .open(&self.log_path)?;
 
         // Acquire an exclusive lock for writing
-        file.lock_exclusive()?;
+        self.request_exclusive_lock(&self.config.data_dir, &mut file)?;
 
         if !is_file_same_as_path(&file, &self.log_path)? {
             // The log file has been rotated, so we must try again
@@ -351,7 +359,8 @@ impl<Field: Eq + Clone + Debug> DB<Field> {
 
         // Open the file and acquire a shared lock for reading
         let mut file = fs::OpenOptions::new().read(true).open(&self.log_path)?;
-        file.lock_shared()?;
+
+        self.request_shared_lock(&self.config.data_dir, &mut file)?;
 
         if !is_file_same_as_path(&file, &self.log_path)? {
             // The log file has been rotated, so we must try again
@@ -518,4 +527,81 @@ impl<Field: Eq + Clone + Debug> DB<Field> {
                 }
             });
     }
+
+    fn request_exclusive_lock(&self, data_dir: &str, file: &mut fs::File) -> 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(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
+        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
+        file.lock_exclusive()?;
+
+        // Unlock the request file
+        lock_request_file.unlock()?;
+
+        Ok(())
+    }
+
+    fn is_exclusive_lock_requested(&self, data_dir: &str) -> Result<bool, io::Error> {
+        let lock_request_path = Path::new(data_dir).join(EXCL_LOCK_REQUEST_FILENAME);
+        let lock_request_file = fs::OpenOptions::new()
+            .create(true)
+            .write(true) // When requesting a lock, we need to have either read or write permissions
+            .open(&lock_request_path)?;
+
+        // Attempt to acquire a shared lock on the lock request file
+        // If the file is already locked, return false
+        match lock_request_file.try_lock_shared() {
+            Err(e) => {
+                if e.kind() == lock_contended_error().kind() {
+                    return Ok(true);
+                }
+                return Err(e);
+            }
+            Ok(_) => {
+                // Check that the exclusive lock request file is still the same as the one we opened
+                if !is_file_same_as_path(&lock_request_file, &lock_request_path)? {
+                    // The lock request file has been removed
+                    return Err(io::Error::new(
+                        io::ErrorKind::Other,
+                        "Lock request file was removed unexpectedly",
+                    ));
+                }
+
+                lock_request_file.unlock()?;
+                return Ok(false);
+            }
+        }
+    }
+
+    fn request_shared_lock(&self, data_dir: &str, 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(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(());
+            }
+        }
+    }
 }
diff --git a/src/log_reader.rs b/src/log_reader.rs
index e32ef53..d2b364b 100644
--- a/src/log_reader.rs
+++ b/src/log_reader.rs
@@ -3,11 +3,11 @@ use rev_buf_reader::RevBufReader;
 use std::fs::{self};
 use std::io::{self, BufRead, Read, Seek, SeekFrom};
 
-/// There are three special characters that need to be handled:
+/// There are three special sequences that need to be handled:
 /// Here: SC = escape char, FS = field separator.
-/// - FS FS SC  -> actual record separator
-/// - SC FS SC  -> literal FS
-/// - SC SC SC  -> literal SC
+/// - SC FS FS SC  -> actual record separator
+/// - SC SC FS SC  -> literal FS
+/// - SC SC SC SC  -> literal SC
 fn validate_special(buf: &[u8]) -> Option<SpecialSequence> {
     match buf {
         SEQ_RECORD_SEP => Some(SpecialSequence::RecordSeparator),
diff --git a/tests/integration.rs b/tests/integration.rs
index f5cdd19..fdca59f 100644
--- a/tests/integration.rs
+++ b/tests/integration.rs
@@ -1,3 +1,7 @@
+extern crate ctor;
+extern crate tempfile;
+
+use ctor::ctor;
 use env_logger;
 use log_db;
 use log_db::{ForwardLogReader, Record, RecordFieldType, RecordValue, ReverseLogReader, DB};
@@ -5,17 +9,25 @@ use serial_test::serial;
 use std::fs;
 use std::path::Path;
 use std::thread;
+use std::time::Duration;
+use tempfile::tempdir;
 
-const TEST_DATA_DIR: &str = "test_db_data";
 const TEST_RESOURCES_DIR: &str = "tests/resources";
 
-fn init_test() {
+#[ctor]
+fn init_logger() {
     let _ = env_logger::builder().is_test(true).try_init();
 }
 
-fn cleanup_test() {
-    std::fs::remove_dir_all(TEST_DATA_DIR.to_string())
-        .unwrap_or_else(|e| eprintln!("Failed to delete the test data directory: {:?}", e));
+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)]
@@ -26,12 +38,10 @@ enum Field {
 }
 
 #[test]
-#[serial]
 fn test_initialize() {
-    init_test();
-
+    let data_dir = tmp_dir();
     let _db = DB::configure()
-        .data_dir(TEST_DATA_DIR)
+        .data_dir(&data_dir)
         .fields(&vec![
             (Field::Id, RecordFieldType::Int),
             (Field::Name, RecordFieldType::String),
@@ -40,17 +50,13 @@ fn test_initialize() {
         .primary_key(Field::Id)
         .initialize()
         .expect("Failed to initialize DB instance");
-
-    cleanup_test();
 }
 
 #[test]
-#[serial]
 fn test_upsert_and_get_with_primary_memtable() {
-    init_test();
-
+    let data_dir = tmp_dir();
     let mut db = DB::configure()
-        .data_dir(TEST_DATA_DIR)
+        .data_dir(&data_dir)
         .fields(&vec![
             (Field::Id, RecordFieldType::Int),
             (Field::Name, RecordFieldType::String),
@@ -76,17 +82,13 @@ fn test_upsert_and_get_with_primary_memtable() {
         (RecordValue::Int(a), RecordValue::Int(b)) => a == b,
         _ => false,
     });
-
-    cleanup_test();
 }
 
 #[test]
-#[serial]
 fn test_upsert_and_get_without_memtable() {
-    init_test();
-
+    let data_dir = tmp_dir();
     let mut db = DB::configure()
-        .data_dir(TEST_DATA_DIR)
+        .data_dir(&data_dir)
         .memtable_capacity(0)
         .fields(&vec![
             (Field::Id, RecordFieldType::Int),
@@ -159,17 +161,13 @@ fn test_upsert_and_get_without_memtable() {
         (RecordValue::String(a), RecordValue::String(b)) => a == b,
         _ => false,
     });
-
-    cleanup_test();
 }
 
 #[test]
-#[serial]
 fn test_upsert_fails_on_invalid_number_of_values() {
-    init_test();
-
+    let data_dir = tmp_dir();
     let mut db = DB::configure()
-        .data_dir(TEST_DATA_DIR)
+        .data_dir(&data_dir)
         .fields(&vec![
             (Field::Id, RecordFieldType::Int),
             (Field::Name, RecordFieldType::String),
@@ -187,17 +185,13 @@ fn test_upsert_fails_on_invalid_number_of_values() {
         ],
     };
     assert!(db.upsert(&record).is_err());
-
-    cleanup_test();
 }
 
 #[test]
-#[serial]
 fn test_upsert_fails_on_invalid_value_type() {
-    init_test();
-
+    let data_dir = tmp_dir();
     let mut db = DB::configure()
-        .data_dir(TEST_DATA_DIR)
+        .data_dir(&data_dir)
         .fields(&vec![
             (Field::Id, RecordFieldType::Int),
             (Field::Name, RecordFieldType::String),
@@ -215,15 +209,10 @@ fn test_upsert_fails_on_invalid_value_type() {
         ],
     };
     assert!(db.upsert(&record).is_err());
-
-    cleanup_test();
 }
 
 #[test]
-#[serial]
 fn test_reverse_log_reader_fixture_db1() {
-    init_test();
-
     let db_path = Path::new(TEST_RESOURCES_DIR).join("test_db1");
     let mut file = fs::OpenOptions::new()
         .read(true)
@@ -251,15 +240,10 @@ fn test_reverse_log_reader_fixture_db1() {
     });
 
     assert!(reverse_log_reader.next().is_none());
-
-    cleanup_test();
 }
 
 #[test]
-#[serial]
 fn test_forward_log_reader_fixture_db1() {
-    init_test();
-
     let db_path = Path::new(TEST_RESOURCES_DIR).join("test_db1");
     let mut file = fs::OpenOptions::new()
         .read(true)
@@ -286,17 +270,13 @@ fn test_forward_log_reader_fixture_db1() {
     });
 
     assert!(forward_log_reader.next().is_none());
-
-    cleanup_test();
 }
 
 #[test]
-#[serial]
 fn test_upsert_and_get_from_secondary_memtable() {
-    init_test();
-
+    let data_dir = tmp_dir();
     let mut db = DB::configure()
-        .data_dir(TEST_DATA_DIR)
+        .data_dir(&data_dir)
         .fields(&vec![
             (Field::Id, RecordFieldType::Int),
             (Field::Name, RecordFieldType::String),
@@ -336,7 +316,7 @@ fn test_upsert_and_get_from_secondary_memtable() {
     db.upsert(&record2).unwrap();
 
     // Delete the DB so that any results must come from a memtable
-    fs::remove_file(Path::new(TEST_DATA_DIR).join("db")).expect("Failed to delete the DB log file");
+    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
@@ -344,25 +324,21 @@ fn test_upsert_and_get_from_secondary_memtable() {
         .expect("Failed to find all Johns");
 
     assert_eq!(johns.len(), 2);
-
-    cleanup_test();
 }
 
 #[test]
-#[serial]
 fn test_initialize_and_read_from_primary_memtable_fixture_db2() {
-    init_test();
-
+    let data_dir = tmp_dir();
     // Copy the fixture DB to the test data directory
-    fs::create_dir_all(TEST_DATA_DIR).expect("Failed to create the test data directory");
+    fs::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(TEST_DATA_DIR).join("db"),
+        &Path::new(&data_dir).join("db"),
     )
     .expect("Failed to copy the fixture DB");
 
     let mut db = DB::configure()
-        .data_dir(TEST_DATA_DIR)
+        .data_dir(&data_dir)
         .fields(&vec![
             (Field::Id, RecordFieldType::Int),
             (Field::Name, RecordFieldType::String),
@@ -373,7 +349,7 @@ fn test_initialize_and_read_from_primary_memtable_fixture_db2() {
         .expect("Failed to initialize DB instance");
 
     // Delete the DB so that any results must come from a memtable
-    fs::remove_file(Path::new(TEST_DATA_DIR).join("db")).expect("Failed to delete the DB log file");
+    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();
 
@@ -383,20 +359,19 @@ fn test_initialize_and_read_from_primary_memtable_fixture_db2() {
         (RecordValue::Int(a), RecordValue::Int(b)) => a == b,
         _ => false,
     });
-
-    cleanup_test();
 }
 
 #[test]
-#[serial]
 fn test_multiple_writing_threads() {
-    init_test();
-
+    let data_dir = tmp_dir();
     let mut threads = vec![];
-    for i in 0..10 {
+    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(TEST_DATA_DIR)
+                .data_dir(&data_dir)
                 .fields(&vec![(Field::Id, RecordFieldType::Int)])
                 .primary_key(Field::Id)
                 .initialize()
@@ -415,13 +390,13 @@ fn test_multiple_writing_threads() {
 
     // Read the records
     let mut db = DB::configure()
-        .data_dir(TEST_DATA_DIR)
+        .data_dir(&data_dir)
         .fields(&vec![(Field::Id, RecordFieldType::Int)])
         .primary_key(Field::Id)
         .initialize()
         .expect("Failed to initialize DB instance");
 
-    for i in 0..10 {
+    for i in 0..threads_n {
         let result = db
             .get(&RecordValue::Int(i))
             .expect("Failed to get record")
@@ -432,32 +407,34 @@ fn test_multiple_writing_threads() {
             _ => false,
         });
     }
-
-    cleanup_test();
 }
 
 #[test]
-#[serial]
 fn test_one_writer_and_multiple_reading_threads() {
-    init_test();
-
+    let data_dir = tmp_dir();
     let mut threads = vec![];
+    let threads_n = 20;
 
     // Add readers that poll for the records
-    for i in 0..10 {
+    for i in 0..threads_n {
+        let data_dir = data_dir.clone();
         threads.push(thread::spawn(move || {
             let mut db = DB::configure()
-                .data_dir(TEST_DATA_DIR)
+                .data_dir(&data_dir)
                 .fields(&vec![(Field::Id, RecordFieldType::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 => continue,
+                    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) {
@@ -474,13 +451,13 @@ fn test_one_writer_and_multiple_reading_threads() {
     // Add a writer that inserts the records
     threads.push(thread::spawn(move || {
         let mut db = DB::configure()
-            .data_dir(TEST_DATA_DIR)
+            .data_dir(&data_dir)
             .fields(&vec![(Field::Id, RecordFieldType::Int)])
             .primary_key(Field::Id)
             .initialize()
             .expect("Failed to initialize DB instance");
 
-        for i in 0..10 {
+        for i in 0..threads_n {
             let record = Record {
                 values: vec![RecordValue::Int(i)],
             };
@@ -491,6 +468,4 @@ fn test_one_writer_and_multiple_reading_threads() {
     for thread in threads {
         thread.join().expect("Failed to join thread");
     }
-
-    cleanup_test();
 }