Add transaction support

5 files changed, 220 insertions, 108 deletions

diff --git a/log_db/src/common.rs b/log_db/src/common.rs
index 3a5fbfe..d4a95b5 100644
--- a/log_db/src/common.rs
+++ b/log_db/src/common.rs
@@ -41,6 +41,8 @@ pub enum DBError {
     ValidationError(String),
     #[error("consistency check failed: {0}")]
     ConsistencyError(String),
+    #[error("invalid transaction: {0}")]
+    TransactionError(String),
     #[error("unexpected IO error: {0}")]
     IOError(#[from] io::Error),
 }
diff --git a/log_db/src/engine.rs b/log_db/src/engine.rs
index c682704..edf8af6 100644
--- a/log_db/src/engine.rs
+++ b/log_db/src/engine.rs
@@ -8,6 +8,9 @@ pub struct Engine<R: Recordable> {
     primary_key_index: usize,
     refresh_next_logkey: LogKey,
 
+    pub tx_active: bool,
+    pub tx_log: Vec<TxEntry>,
+
     active_metadata_file: fs::File,
     active_data_file: fs::File,
 
@@ -116,6 +119,8 @@ impl<R: Recordable> Engine<R> {
             active_metadata_file,
             active_data_file,
             refresh_next_logkey: LogKey::new(1, 0),
+            tx_active: false,
+            tx_log: vec![],
         };
 
         info!("Rebuilding memtable indexes...");
@@ -221,7 +226,7 @@ impl<R: Recordable> Engine<R> {
         }
     }
 
-    pub fn batch_upsert_records(&mut self, records: impl Iterator<Item = Record>) -> DBResult<()> {
+    pub fn upsert_record(&mut self, record: Record) -> DBResult<()> {
         debug!("Opening file in append mode...");
 
         if !self.ensure_metadata_file_is_active()?
@@ -231,63 +236,14 @@ impl<R: Recordable> Engine<R> {
             )?
         {
             // The log file has been rotated, so we must try again
-            return self.batch_upsert_records(records);
-        }
-
-        let active_symlink_path = self.data_dir_path.join(ACTIVE_SYMLINK_FILENAME);
-        let active_target = fs::read_link(active_symlink_path)?;
-        let segment_num = parse_segment_number(&active_target)?;
-
-        debug!("Appending to log file");
-
-        let mut serialized_data: Vec<u8> = vec![];
-        let mut serialized_metadata: Vec<u8> = vec![];
-        let mut pending_memtable_insertions: Vec<(LogKey, Record)> = vec![];
-        for record in records {
-            // Write the record to the log
-            let serialized = &record.serialize();
-            let record_offset = self.active_data_file.seek(SeekFrom::End(0))?;
-            let record_length = serialized.len() as u64;
-            assert!(record_length > 0);
-
-            serialized_data.extend(serialized);
-
-            let metadata_pos = self.active_metadata_file.seek(SeekFrom::End(0))?;
-            let metadata_index =
-                (metadata_pos - METADATA_FILE_HEADER_SIZE as u64) / METADATA_ROW_LENGTH as u64;
-
-            // Write the record metadata to the metadata file
-            let mut metadata_buf = vec![];
-            metadata_buf.extend(record_offset.to_be_bytes().into_iter());
-            metadata_buf.extend(record_length.to_be_bytes().into_iter());
-
-            assert_eq!(metadata_buf.len(), 16);
-
-            serialized_metadata.extend(metadata_buf);
-
-            let log_key = LogKey::new(segment_num, metadata_index);
-
-            pending_memtable_insertions.push((log_key, record));
+            return self.upsert_record(record);
         }
 
-        self.active_data_file.write_all(&serialized_data)?;
-        self.active_metadata_file.write_all(&serialized_metadata)?;
+        self.tx_log.push(TxEntry::Upsert { record });
 
-        // Flush and sync data and metadata to disk
-        if self.config.write_durability == WriteDurability::Flush {
-            self.active_data_file.flush()?;
-            self.active_metadata_file.flush()?;
-        } else if self.config.write_durability == WriteDurability::FlushSync {
-            self.active_data_file.flush()?;
-            self.active_data_file.sync_all()?;
-            self.active_metadata_file.flush()?;
-            self.active_metadata_file.sync_all()?;
-        }
-
-        debug!("Records appended to log file");
-
-        for (log_key, record) in pending_memtable_insertions {
-            self.insert_record_to_memtables(log_key, record);
+        if !self.tx_active {
+            self.commit_transaction()?;
+            self.tx_log.clear();
         }
 
         Ok(())
@@ -532,44 +488,90 @@ impl<R: Recordable> Engine<R> {
             })
             .collect();
 
+        // TODO: refactor the clone out of here
         for record in &recs {
-            let record_serialized = record.serialize();
+            self.tx_log.push(TxEntry::Delete {
+                record: record.clone(),
+            });
+        }
 
-            let offset = self.active_data_file.seek(SeekFrom::End(0))?;
-            let length = record_serialized.len() as u64;
+        if !self.tx_active {
+            self.commit_transaction()?;
+            self.tx_log.clear();
+        }
 
-            self.active_data_file.write_all(&record_serialized)?;
+        debug!("Records deleted");
 
-            // Flush and sync data to disk
-            if self.config.write_durability == WriteDurability::Flush {
-                self.active_data_file.flush()?;
-            }
-            if self.config.write_durability == WriteDurability::FlushSync {
-                self.active_data_file.flush()?;
-                self.active_data_file.sync_all()?;
-            }
+        Ok(recs)
+    }
 
-            let mut metadata_entry = vec![];
-            metadata_entry.extend(offset.to_be_bytes().into_iter());
-            metadata_entry.extend(length.to_be_bytes().into_iter());
+    pub fn commit_transaction(&mut self) -> DBResult<()> {
+        let active_symlink_path = self.data_dir_path.join(ACTIVE_SYMLINK_FILENAME);
+        let active_target = fs::read_link(active_symlink_path)?;
+        let segment_num = parse_segment_number(&active_target)?;
 
-            self.active_metadata_file.write_all(&metadata_entry)?;
+        let initial_data_offset = self.active_data_file.seek(SeekFrom::End(0))?;
+        let initial_metadata_offset = self.active_metadata_file.seek(SeekFrom::End(0))?;
+        let mut serialized_data: Vec<u8> = vec![];
+        let mut serialized_metadata: Vec<u8> = vec![];
+        let mut pending_memtable_ops: Vec<(LogKey, TxEntry)> = vec![];
 
-            // Flush and sync metadata to disk
-            if self.config.write_durability == WriteDurability::Flush {
-                self.active_metadata_file.flush()?;
-            }
-            if self.config.write_durability == WriteDurability::FlushSync {
-                self.active_metadata_file.flush()?;
-                self.active_metadata_file.sync_all()?;
-            }
+        debug!("Serializing tx_log to byte arrays");
+        for tx_entry in &self.tx_log {
+            let record = match tx_entry {
+                TxEntry::Upsert { record } => record,
+                TxEntry::Delete { record } => record,
+            };
+
+            let serialized = record.serialize();
+            let record_offset = initial_data_offset + serialized_data.len() as u64;
+            let record_length = serialized.len() as u64;
+            assert!(record_length > 0);
+
+            serialized_data.extend(serialized);
+
+            let metadata_pos = initial_metadata_offset + serialized_metadata.len() as u64;
+            let metadata_index =
+                (metadata_pos - METADATA_FILE_HEADER_SIZE as u64) / METADATA_ROW_LENGTH as u64;
+
+            // Write the record metadata to the metadata file
+            let mut metadata_buf = vec![];
+            metadata_buf.extend(record_offset.to_be_bytes().into_iter());
+            metadata_buf.extend(record_length.to_be_bytes().into_iter());
+
+            assert_eq!(metadata_buf.len(), 16);
+
+            serialized_metadata.extend(metadata_buf);
 
-            self.remove_record_from_memtables(&record);
+            let log_key = LogKey::new(segment_num, metadata_index);
+            pending_memtable_ops.push((log_key, tx_entry.clone()));
         }
 
-        debug!("Records deleted");
+        debug!("Writing serialized bytearrays to log files");
+        self.active_data_file.write_all(&serialized_data)?;
+        self.active_metadata_file.write_all(&serialized_metadata)?;
 
-        Ok(recs)
+        // Flush and sync data and metadata to disk
+        if self.config.write_durability == WriteDurability::Flush {
+            self.active_data_file.flush()?;
+            self.active_metadata_file.flush()?;
+        } else if self.config.write_durability == WriteDurability::FlushSync {
+            self.active_data_file.flush()?;
+            self.active_data_file.sync_all()?;
+            self.active_metadata_file.flush()?;
+            self.active_metadata_file.sync_all()?;
+        }
+
+        debug!("Updating memtables");
+        for (log_key, tx_entry) in pending_memtable_ops {
+            match tx_entry {
+                TxEntry::Upsert { record } => self.insert_record_to_memtables(log_key, record),
+                TxEntry::Delete { record } => self.remove_record_from_memtables(&record),
+            }
+        }
+        debug!("Commit done");
+
+        Ok(())
     }
 
     pub fn do_maintenance_tasks(&mut self) -> DBResult<()> {
@@ -719,17 +721,30 @@ impl<R: Recordable> Engine<R> {
         &mut self,
         f: impl FnOnce(&mut Self) -> DBResult<T>,
     ) -> DBResult<T> {
-        self.lock_manager.lock_exclusive()?;
+        // No need to acquire a lock if a transaction is already active
+        // because the lock is already held.
+        if !self.tx_active {
+            self.lock_manager.lock_exclusive()?;
+        }
         let result = f(self);
-        self.lock_manager.unlock()?;
+        if !self.tx_active {
+            self.lock_manager.unlock()?;
+        }
         result
     }
 
     #[inline]
     pub fn with_shared_lock<T>(&mut self, f: impl FnOnce(&mut Self) -> DBResult<T>) -> DBResult<T> {
-        self.lock_manager.lock_shared()?;
+        // No need to acquire a lock if a transaction is already active
+        // because the lock is already held.
+        if !self.tx_active {
+            self.lock_manager.lock_shared()?;
+        }
         let result = f(self);
-        self.lock_manager.unlock()?;
+        if !self.tx_active {
+            self.lock_manager.unlock()?;
+        }
+
         result
     }
 }
diff --git a/log_db/src/lib.rs b/log_db/src/lib.rs
index 698ce64..67a1e0b 100644
--- a/log_db/src/lib.rs
+++ b/log_db/src/lib.rs
@@ -65,29 +65,8 @@ impl<R: Recordable> DB<R> {
         record.validate(&self.engine.config.fields)?;
         debug!("Record is valid");
 
-        self.engine.with_exclusive_lock(move |engine| {
-            engine.batch_upsert_records(std::iter::once(record))
-        })?;
-
-        Ok(())
-    }
-
-    /// Insert a batch of records into the database. If the primary key value for a record already exists,
-    /// the existing record will be replaced by the supplied one. Records are inserted in the order they are given.
-    pub fn batch_upsert(&mut self, recordables: Vec<R>) -> DBResult<()> {
-        let records = recordables
-            .into_iter()
-            .map(|r| Record::from(&r.into_record()))
-            .collect::<Vec<Record>>();
-        debug!("Batch upserting {} records", records.len());
-
-        for record in &records {
-            record.validate(&self.engine.config.fields)?;
-        }
-        debug!("Records are valid");
-
         self.engine
-            .with_exclusive_lock(move |engine| engine.batch_upsert_records(records.into_iter()))?;
+            .with_exclusive_lock(move |engine| engine.upsert_record(record))?;
 
         Ok(())
     }
@@ -209,6 +188,45 @@ impl<R: Recordable> DB<R> {
         self.engine
             .with_exclusive_lock(|engine| engine.refresh_indexes())
     }
+
+    pub fn tx_begin(&mut self) -> DBResult<()> {
+        if self.engine.tx_active {
+            return Err(DBError::TransactionError(
+                "Transaction already active".to_string(),
+            ));
+        }
+
+        self.engine.lock_manager.lock_exclusive()?;
+        self.engine.tx_active = true;
+        Ok(())
+    }
+
+    pub fn tx_commit(&mut self) -> DBResult<()> {
+        if !self.engine.tx_active {
+            return Err(DBError::TransactionError(
+                "No active transaction to commit".to_string(),
+            ));
+        }
+
+        self.engine.commit_transaction()?;
+        self.engine.tx_log.clear();
+        self.engine.tx_active = false;
+        self.engine.lock_manager.unlock()?;
+        Ok(())
+    }
+
+    pub fn tx_rollback(&mut self) -> DBResult<()> {
+        if !self.engine.tx_active {
+            return Err(DBError::TransactionError(
+                "No active transaction to rollback".to_string(),
+            ));
+        }
+
+        self.engine.tx_log.clear();
+        self.engine.tx_active = false;
+        self.engine.lock_manager.unlock()?;
+        Ok(())
+    }
 }
 
 #[cfg(test)]
diff --git a/log_db/src/record.rs b/log_db/src/record.rs
index a464520..6d7cbbf 100644
--- a/log_db/src/record.rs
+++ b/log_db/src/record.rs
@@ -144,3 +144,9 @@ mod tests {
         assert_eq!(record.values, deserialized.values);
     }
 }
+
+#[derive(Clone, Debug)]
+pub enum TxEntry {
+    Upsert { record: Record },
+    Delete { record: Record },
+}
diff --git a/log_db/tests/integration.rs b/log_db/tests/integration.rs
index 5535adf..5d56fbf 100644
--- a/log_db/tests/integration.rs
+++ b/log_db/tests/integration.rs
@@ -27,6 +27,10 @@ pub fn tmp_dir() -> String {
 #[ctor]
 fn init_logger() {
     let _ = env_logger::builder().is_test(true).try_init();
+
+    // todo add panic hook stuff
+    // - https://stackoverflow.com/questions/54917373/retrieving-backtrace-from-a-panic-in-hook-in-rust
+    // - https://github.com/sndels/yuki/blob/e86b379165ec657197b1c14b78164bd09a8aa1dc/yuki/src/main.rs#L74
 }
 
 #[derive(Eq, PartialEq, Clone, Debug)]
@@ -36,6 +40,7 @@ enum Field {
     Data,
 }
 
+#[derive(Debug)]
 struct Inst {
     pub id: i64,
     pub name: Option<String>,
@@ -647,3 +652,69 @@ fn test_batch_find_by() {
         vec![2, 3, 4]
     );
 }
+
+#[test]
+fn test_commit_transaction() {
+    let data_dir = tmp_dir();
+    let mut db = DB::<Inst>::configure()
+        .data_dir(&data_dir)
+        .initialize()
+        .expect("Failed to initialize DB instance");
+
+    db.tx_begin().expect("Failed to begin transaction");
+
+    db.upsert(Inst {
+        id: 0,
+        name: Some("John".to_string()),
+        data: vec![3, 4, 5],
+    })
+    .unwrap();
+
+    db.upsert(Inst {
+        id: 1,
+        name: Some("John".to_string()),
+        data: vec![1, 2, 3],
+    })
+    .unwrap();
+
+    db.tx_commit().expect("Failed to commit transaction");
+
+    let johns = db
+        .find_by(&Field::Name, &Value::String("John".to_string()))
+        .unwrap();
+
+    assert_eq!(johns.len(), 2);
+}
+
+#[test]
+fn test_rollback_transaction() {
+    let data_dir = tmp_dir();
+    let mut db = DB::<Inst>::configure()
+        .data_dir(&data_dir)
+        .initialize()
+        .expect("Failed to initialize DB instance");
+
+    db.tx_begin().expect("Failed to begin transaction");
+
+    db.upsert(Inst {
+        id: 0,
+        name: Some("John".to_string()),
+        data: vec![3, 4, 5],
+    })
+    .unwrap();
+
+    db.upsert(Inst {
+        id: 1,
+        name: Some("John".to_string()),
+        data: vec![1, 2, 3],
+    })
+    .unwrap();
+
+    db.tx_rollback().expect("Failed to rollback transaction");
+
+    let johns = db
+        .find_by(&Field::Name, &Value::String("John".to_string()))
+        .unwrap();
+
+    assert_eq!(johns.len(), 0);
+}