Merge 'core/mvcc: fix exists to use BTreeCursor as fallback' from Pere Diaz Bou

Closes #4026

Author: jussisaurio

core/mvcc/cursor.rs

@@ -43,10 +43,17 @@ enum PrevState {
         new_position_in_mvcc: CursorPosition,
     },
 }
+
+#[derive(Debug, Clone)]
+enum ExistsState {
+    ExistsBtree,
+}
+
 #[derive(Debug, Clone)]
 enum MvccLazyCursorState {
     Next(NextState),
     Prev(PrevState),
+    Exists(ExistsState),
 }
 
 pub struct MvccLazyCursor<Clock: LogicalClock> {
@@ -648,39 +655,59 @@ impl<Clock: LogicalClock + 'static> CursorTrait for MvccLazyCursor<Clock> {
     }
 
     fn exists(&mut self, key: &Value) -> Result<IOResult<bool>> {
-        self.invalidate_record();
-        let int_key = match key {
-            Value::Integer(i) => i,
-            _ => unreachable!("btree tables are indexed by integers!"),
-        };
-        let rowid = self.db.seek_rowid(
-            Bound::Included(&RowID {
-                table_id: self.table_id,
-                row_id: RowKey::Int(*int_key),
-            }),
-            true,
-            self.tx_id,
-        );
-        tracing::trace!("found {rowid:?}");
-        let exists = if let Some(rowid) = rowid {
-            let RowKey::Int(rowid) = rowid.row_id else {
-                panic!("Rowid is not an integer in mvcc table cursor");
+        if self.state.borrow().is_none() {
+            self.invalidate_record();
+            let int_key = match key {
+                Value::Integer(i) => i,
+                _ => unreachable!("btree tables are indexed by integers!"),
             };
-            rowid == *int_key
-        } else {
-            false
-        };
-        if exists {
-            self.current_pos.replace(CursorPosition::Loaded {
-                row_id: RowID {
+            let rowid = self.db.seek_rowid(
+                Bound::Included(&RowID {
                     table_id: self.table_id,
                     row_id: RowKey::Int(*int_key),
-                },
-                in_btree: false,
-                btree_consumed: false,
-            });
+                }),
+                true,
+                self.tx_id,
+            );
+            tracing::trace!("found {rowid:?}");
+            let exists = if let Some(rowid) = rowid {
+                let RowKey::Int(rowid) = rowid.row_id else {
+                    panic!("Rowid is not an integer in mvcc table cursor");
+                };
+                rowid == *int_key
+            } else {
+                false
+            };
+            if exists {
+                self.current_pos.replace(CursorPosition::Loaded {
+                    row_id: RowID {
+                        table_id: self.table_id,
+                        row_id: RowKey::Int(*int_key),
+                    },
+                    in_btree: false,
+                    btree_consumed: false,
+                });
+                return Ok(IOResult::Done(exists));
+            } else if self.is_btree_allocated() {
+                self.state
+                    .replace(Some(MvccLazyCursorState::Exists(ExistsState::ExistsBtree)));
+            } else {
+                return Ok(IOResult::Done(false));
+            }
         }
-        Ok(IOResult::Done(exists))
+
+        let Some(MvccLazyCursorState::Exists(ExistsState::ExistsBtree)) =
+            self.state.borrow().clone()
+        else {
+            panic!("Invalid state {:?}", self.state.borrow());
+        };
+        assert!(
+            self.is_btree_allocated(),
+            "BTree should be allocated when we are in ExistsBtree state"
+        );
+        self.state.replace(None);
+        let found = return_if_io!(self.btree_cursor.exists(key));
+        Ok(IOResult::Done(found))
     }
 
     fn clear_btree(&mut self) -> Result<IOResult<Option<usize>>> {

core/mvcc/database/tests.rs

@@ -1843,3 +1843,23 @@ pub fn rng_from_time_or_env() -> (ChaCha8Rng, u64) {
     let rng = ChaCha8Rng::seed_from_u64(seed as u64);
     (rng, seed as u64)
 }
+
+#[test]
+fn test_cursor_with_btree_and_mvcc_insert_after_checkpoint_repeated_key() {
+    let mut db = MvccTestDbNoConn::new_with_random_db();
+    // First write some rows and checkpoint so data is flushed to BTree file (.db)
+    {
+        let conn = db.connect();
+        conn.execute("CREATE TABLE t(x integer primary key)")
+            .unwrap();
+        conn.execute("INSERT INTO t VALUES (1)").unwrap();
+        conn.execute("INSERT INTO t VALUES (2)").unwrap();
+        conn.execute("PRAGMA wal_checkpoint(TRUNCATE)").unwrap();
+    }
+    // Now restart so new connection will have to read data from BTree instead of MVCC.
+    db.restart();
+    let conn = db.connect();
+    // Insert a new row so that we have a gap in the BTree.
+    let res = conn.execute("INSERT INTO t VALUES (2)");
+    assert!(res.is_err(), "Expected error because key 2 already exists");
+}