fixed transaction scan

07CalC · 07CalC · commit 87a52cc3bef6 · 2025-12-18T21:14:32.000+05:30
diff --git a/kv/src/core/db.rs b/kv/src/core/db.rs
@@ -367,6 +367,17 @@ impl Db {
 
         DbIterator::new(memtable, immutables, sstables, start_bound, end_bound)
     }
+
+    pub fn scan_seq(&self, start: Option<&[u8]>, end: Option<&[u8]>, seq: u64) -> DbIterator {
+        let memtable = Arc::clone(&self.memtable.load());
+        let immutables = (**self.immutable_memtables.load()).clone();
+        let sstables = (**self.sstables.load()).clone();
+
+        let start_bound = start.map(|s| s.to_vec());
+        let end_bound = end.map(|e| e.to_vec());
+
+        DbIterator::new_with_seq(memtable, immutables, sstables, start_bound, end_bound, Some(seq))
+    }
 }
 
 // custom memory drop implementation to join all the threads (i.e. the compaction and flush queue
diff --git a/kv/src/core/iterator.rs b/kv/src/core/iterator.rs
@@ -76,6 +76,8 @@ pub struct DbIterator {
     last_key: Option<Vec<u8>>,
     start_bound: Option<Vec<u8>>,
     end_bound: Option<Vec<u8>>,
+    max_seq: Option<u64>, // for snapshot isolation, we should only see the values with seq less
+                          // than this
 }
 
 impl DbIterator {
@@ -85,6 +87,24 @@ impl DbIterator {
         sstables: Vec<SSTReader>,
         start_bound: Option<Vec<u8>>,
         end_bound: Option<Vec<u8>>,
+    ) -> Self {
+        Self::new_with_seq(
+            memtable,
+            immutable_memtable,
+            sstables,
+            start_bound,
+            end_bound,
+            None,
+        )
+    }
+
+    pub fn new_with_seq(
+        memtable: Arc<Memtable>,
+        immutable_memtable: Vec<Arc<Memtable>>,
+        sstables: Vec<SSTReader>,
+        start_bound: Option<Vec<u8>>,
+        end_bound: Option<Vec<u8>>,
+        max_seq: Option<u64>,
     ) -> Self {
         let mut sources = Vec::new();
         let mut heap = BinaryHeap::new();
@@ -131,7 +151,9 @@ impl DbIterator {
 
         // preload one entry from each source
         for i in 0..sources.len() {
-            if let Some(entry) = Self::advance_source(&mut sources, i, &start_bound, &end_bound) {
+            if let Some(entry) =
+                Self::advance_source(&mut sources, i, &start_bound, &end_bound, &max_seq)
+            {
                 heap.push(entry);
             }
         }
@@ -142,6 +164,7 @@ impl DbIterator {
             last_key: None,
             start_bound,
             end_bound,
+            max_seq,
         }
     }
 
@@ -150,6 +173,7 @@ impl DbIterator {
         source_idx: usize,
         start_bound: &Option<Vec<u8>>,
         end_bound: &Option<Vec<u8>>,
+        max_seq: &Option<u64>,
     ) -> Option<IterEntry> {
         loop {
             let (key, value, seq, priority) = match &mut sources[source_idx] {
@@ -172,6 +196,14 @@ impl DbIterator {
                 },
             };
 
+            // kkip entries with sequence numbers >= max_seq, for snapshot isolation
+            // we use >= to match the strict inequality in get_seq (seq < snapshot_seq)
+            if let Some(max) = max_seq {
+                if seq >= *max {
+                    continue;
+                }
+            }
+
             if let Some(start) = start_bound {
                 if &key < start {
                     continue;
@@ -215,6 +247,7 @@ impl Iterator for DbIterator {
                         idx,
                         &self.start_bound,
                         &self.end_bound,
+                        &self.max_seq,
                     ) {
                         self.heap.push(next_entry);
                     }
diff --git a/kv/src/transaction/mod.rs b/kv/src/transaction/mod.rs
@@ -60,8 +60,8 @@ impl<'a> Transaction<'a> {
     }
 
     pub fn scan(&self, start: Option<&[u8]>, end: Option<&[u8]>) -> TransactionIterator {
-        // Get underlying DB iterator
-        let db_iter = self.db.scan(start, end);
+        // Get underlying DB iterator with snapshot isolation at transaction's sequence
+        let db_iter = self.db.scan_seq(start, end, self.seq);
 
         // Collect transaction buffer entries within the range
         let mut txn_entries = Vec::new();
@@ -88,6 +88,7 @@ impl<'a> Transaction<'a> {
             txn_entries,
             txn_pos: 0,
             last_key: None,
+            peeked_db_entry: None,
         }
     }
 }
@@ -97,58 +98,64 @@ pub struct TransactionIterator {
     txn_entries: Vec<(Vec<u8>, Vec<u8>)>,
     txn_pos: usize,
     last_key: Option<Vec<u8>>,
+    peeked_db_entry: Option<(Vec<u8>, Vec<u8>)>, // Store peeked DB entry
 }
 
 impl Iterator for TransactionIterator {
     type Item = (Vec<u8>, Vec<u8>);
 
     fn next(&mut self) -> Option<Self::Item> {
         loop {
-            // Peek at both sources
+            // peek at both sources first
             let txn_entry = if self.txn_pos < self.txn_entries.len() {
                 Some(&self.txn_entries[self.txn_pos])
             } else {
                 None
             };
 
-            let db_entry = self.db_iter.next();
+            // get or peek DB entry
+            if self.peeked_db_entry.is_none() {
+                self.peeked_db_entry = self.db_iter.next();
+            }
 
-            // Determine which entry to return
-            let (key, val) = match (txn_entry, db_entry) {
-                (Some((tk, tv)), Some((dk, dv))) => {
+            // determine which entry to be returned
+            let (key, val) = match (txn_entry, &self.peeked_db_entry) {
+                (Some((tk, tv)), Some((dk, _))) => {
                     // Both have entries, pick the smaller key
                     // Transaction entries take precedence on equal keys
-                    if tk <= &dk {
+                    if tk <= dk {
                         self.txn_pos += 1;
                         (tk.clone(), tv.clone())
                     } else {
-                        (dk, dv)
+                        let entry = self.peeked_db_entry.take().unwrap();
+                        entry
                     }
                 }
                 (Some((tk, tv)), None) => {
-                    // Only transaction has entry
+                    // only transaction has entry
                     self.txn_pos += 1;
                     (tk.clone(), tv.clone())
                 }
-                (None, Some((dk, dv))) => {
-                    // Only DB has entry
-                    (dk, dv)
+                (None, Some(_)) => {
+                    // only DB has entry
+                    let entry = self.peeked_db_entry.take().unwrap();
+                    entry
                 }
                 (None, None) => {
-                    // Both exhausted
+                    // none has a entry
                     return None;
                 }
             };
 
-            // Skip duplicates
+            // skip duplicates
             if let Some(ref last) = self.last_key {
                 if &key == last {
                     continue;
                 }
             }
             self.last_key = Some(key.clone());
 
-            // Skip tombstones (empty values)
+            // skip tombstones, i.e. empty values or deleted values
             if val.is_empty() {
                 continue;
             }
diff --git a/kv/tests/transaction_test.rs b/kv/tests/transaction_test.rs
@@ -1,5 +1,5 @@
 use keylite_kv::core::Db;
-use std::fs;
+use std::{fs, string};
 
 #[test]
 fn test_basic_transaction_commit() {
@@ -98,29 +98,30 @@ fn test_transaction_scan_all() {
     let _ = fs::remove_dir_all(test_dir);
 
     let db = Db::open(test_dir).unwrap();
-    
-    // Put some data in DB
+
     db.put(b"a1", b"db_a1").unwrap();
     db.put(b"b2", b"db_b2").unwrap();
     db.put(b"c3", b"db_c3").unwrap();
-    
-    // Create transaction with additional data
+
     let mut txn = db.begin();
     txn.put(b"a2", b"txn_a2");
     txn.put(b"b1", b"txn_b1");
     txn.put(b"d1", b"txn_d1");
-    
-    // Scan all
+
     let results: Vec<_> = txn.scan(None, None).collect();
-    
+    for res in results.iter() {
+        println!("{:?}", String::from_utf8(res.1.clone()));
+    }
+    println!("{:?}", results);
+
     assert_eq!(results.len(), 6);
     assert_eq!(results[0], (b"a1".to_vec(), b"db_a1".to_vec()));
     assert_eq!(results[1], (b"a2".to_vec(), b"txn_a2".to_vec()));
     assert_eq!(results[2], (b"b1".to_vec(), b"txn_b1".to_vec()));
     assert_eq!(results[3], (b"b2".to_vec(), b"db_b2".to_vec()));
     assert_eq!(results[4], (b"c3".to_vec(), b"db_c3".to_vec()));
     assert_eq!(results[5], (b"d1".to_vec(), b"txn_d1".to_vec()));
-    
+
     let _ = fs::remove_dir_all(test_dir);
 }
 
@@ -130,27 +131,25 @@ fn test_transaction_scan_with_range() {
     let _ = fs::remove_dir_all(test_dir);
 
     let db = Db::open(test_dir).unwrap();
-    
-    // Put some data in DB
+
     db.put(b"a1", b"db_a1").unwrap();
     db.put(b"b2", b"db_b2").unwrap();
     db.put(b"c3", b"db_c3").unwrap();
     db.put(b"d4", b"db_d4").unwrap();
-    
-    // Create transaction with additional data
+
     let mut txn = db.begin();
     txn.put(b"b1", b"txn_b1");
     txn.put(b"c1", b"txn_c1");
-    
-    // Scan with range [b, d)
+
     let results: Vec<_> = txn.scan(Some(b"b"), Some(b"d")).collect();
-    
+    println!("{:?}", results);
+
     assert_eq!(results.len(), 4);
     assert_eq!(results[0], (b"b1".to_vec(), b"txn_b1".to_vec()));
     assert_eq!(results[1], (b"b2".to_vec(), b"db_b2".to_vec()));
     assert_eq!(results[2], (b"c1".to_vec(), b"txn_c1".to_vec()));
     assert_eq!(results[3], (b"c3".to_vec(), b"db_c3".to_vec()));
-    
+
     let _ = fs::remove_dir_all(test_dir);
 }
 
@@ -160,24 +159,24 @@ fn test_transaction_scan_with_overwrites() {
     let _ = fs::remove_dir_all(test_dir);
 
     let db = Db::open(test_dir).unwrap();
-    
+
     // Put some data in DB
     db.put(b"key1", b"db_value1").unwrap();
     db.put(b"key2", b"db_value2").unwrap();
     db.put(b"key3", b"db_value3").unwrap();
-    
+
     // Create transaction that overwrites some keys
     let mut txn = db.begin();
     txn.put(b"key2", b"txn_value2");
-    
+
     // Scan all - should see transaction value for key2
     let results: Vec<_> = txn.scan(None, None).collect();
-    
+
     assert_eq!(results.len(), 3);
     assert_eq!(results[0], (b"key1".to_vec(), b"db_value1".to_vec()));
     assert_eq!(results[1], (b"key2".to_vec(), b"txn_value2".to_vec()));
     assert_eq!(results[2], (b"key3".to_vec(), b"db_value3".to_vec()));
-    
+
     let _ = fs::remove_dir_all(test_dir);
 }
 
@@ -187,23 +186,23 @@ fn test_transaction_scan_with_deletes() {
     let _ = fs::remove_dir_all(test_dir);
 
     let db = Db::open(test_dir).unwrap();
-    
+
     // Put some data in DB
     db.put(b"key1", b"db_value1").unwrap();
     db.put(b"key2", b"db_value2").unwrap();
     db.put(b"key3", b"db_value3").unwrap();
-    
+
     // Create transaction that deletes a key
     let mut txn = db.begin();
     txn.del(b"key2");
-    
+
     // Scan all - should not see key2
     let results: Vec<_> = txn.scan(None, None).collect();
-    
+
     assert_eq!(results.len(), 2);
     assert_eq!(results[0], (b"key1".to_vec(), b"db_value1".to_vec()));
     assert_eq!(results[1], (b"key3".to_vec(), b"db_value3".to_vec()));
-    
+
     let _ = fs::remove_dir_all(test_dir);
 }
 
@@ -213,21 +212,21 @@ fn test_transaction_scan_start_only() {
     let _ = fs::remove_dir_all(test_dir);
 
     let db = Db::open(test_dir).unwrap();
-    
+
     db.put(b"a", b"val_a").unwrap();
     db.put(b"b", b"val_b").unwrap();
     db.put(b"c", b"val_c").unwrap();
     db.put(b"d", b"val_d").unwrap();
-    
+
     let txn = db.begin();
-    
+
     // Scan from 'c' onwards
     let results: Vec<_> = txn.scan(Some(b"c"), None).collect();
-    
+
     assert_eq!(results.len(), 2);
     assert_eq!(results[0], (b"c".to_vec(), b"val_c".to_vec()));
     assert_eq!(results[1], (b"d".to_vec(), b"val_d".to_vec()));
-    
+
     let _ = fs::remove_dir_all(test_dir);
 }
 
@@ -237,21 +236,20 @@ fn test_transaction_scan_end_only() {
     let _ = fs::remove_dir_all(test_dir);
 
     let db = Db::open(test_dir).unwrap();
-    
+
     db.put(b"a", b"val_a").unwrap();
     db.put(b"b", b"val_b").unwrap();
     db.put(b"c", b"val_c").unwrap();
     db.put(b"d", b"val_d").unwrap();
-    
+
     let txn = db.begin();
-    
+
     // Scan up to 'c' (exclusive)
     let results: Vec<_> = txn.scan(None, Some(b"c")).collect();
-    
+
     assert_eq!(results.len(), 2);
     assert_eq!(results[0], (b"a".to_vec(), b"val_a".to_vec()));
     assert_eq!(results[1], (b"b".to_vec(), b"val_b".to_vec()));
-    
+
     let _ = fs::remove_dir_all(test_dir);
 }
-
