From 2ddbb7eeed87b4fe9e879c357574a4fb83b7e903 Mon Sep 17 00:00:00 2001 From: pedrocarlo Date: Sun, 1 Jun 2025 01:01:35 -0300 Subject: [PATCH] adjust move_to and seek functions to make them truly reentrant + adding return_if_locked_maybe_load in some places so that we read loaded pages --- core/storage/btree.rs | 193 ++++++++++++++++++++++++++++++------ testing/cli_tests/memory.py | 1 - 2 files changed, 160 insertions(+), 34 deletions(-) diff --git a/core/storage/btree.rs b/core/storage/btree.rs index 9b184c6a0..9fb820eda 100644 --- a/core/storage/btree.rs +++ b/core/storage/btree.rs @@ -416,6 +416,9 @@ enum CursorSeekState { /// Max cell_idx max: isize, nearest_matching_cell: Option, + /// Indicates when we have not not found a value in leaf and now will look in the next/prev record. + /// This value is only used for indexbtree + not_found_leaf: bool, }, } @@ -482,9 +485,102 @@ impl CursorSeekState { }; *min } + + /// # Safety + /// method can only be called on CursorSeekState::Seeking + fn set_not_found_leaf(&mut self, not_found: bool) { + let CursorSeekState::Seeking { not_found_leaf, .. } = self else { + unreachable!("method can only be called on CursorSeekState::Seeking") + }; + *not_found_leaf = not_found; + } + + /// # Safety + /// method can only be called on CursorSeekState::Seeking + fn get_not_found_leaf(&mut self) -> bool { + let CursorSeekState::Seeking { not_found_leaf, .. } = self else { + unreachable!("method can only be called on CursorSeekState::Seeking") + }; + *not_found_leaf + } } -type CursorMoveToState = CursorSeekState; +#[derive(Debug)] +/// State used for seeking +enum CursorMoveToState { + Start, + ContinueLoop, + Seeking { + /// Min cell_idx + min: isize, + /// Max cell_idx + max: isize, + nearest_matching_cell: Option, + }, +} + +// These functions below exist to avoid problems with the borrow checker +impl CursorMoveToState { + /// # Safety + /// method can only be called on CursorMoveToState::Seeking + fn set_nearest_matching_cell(&mut self, matching_cell: Option) { + match self { + CursorMoveToState::Seeking { + nearest_matching_cell, + .. + } => *nearest_matching_cell = matching_cell, + _ => unreachable!("method can only be called on CursorMoveToState::Seeking"), + }; + } + + /// # Safety + /// method can only be called on CursorMoveToState::Seeking + fn get_nearest_matching_cell(&mut self) -> Option { + match self { + CursorMoveToState::Seeking { + nearest_matching_cell, + .. + } => nearest_matching_cell.clone(), + _ => unreachable!("method can only be called on CursorMoveToState::Seeking"), + } + } + + /// # Safety + /// method can only be called on CursorMoveToState::Seeking + fn set_max(&mut self, max: isize) { + match self { + CursorMoveToState::Seeking { max: max_state, .. } => *max_state = max, + _ => unreachable!("method can only be called on CursorMoveToState::Seeking"), + } + } + + /// # Safety + /// method can only be called on CursorMoveToState::Seeking + fn get_max(&mut self) -> isize { + match self { + CursorMoveToState::Seeking { max: max_state, .. } => *max_state, + _ => unreachable!("method can only be called on CursorMoveToState::Seeking"), + } + } + + /// # Safety + /// method can only be called on CursorMoveToState::Seeking + fn set_min(&mut self, min: isize) { + match self { + CursorMoveToState::Seeking { min: min_state, .. } => *min_state = min, + _ => unreachable!("method can only be called on CursorMoveToState::Seeking "), + } + } + + /// # Safety + /// method can only be called on CursorMoveToState::Seeking + fn get_min(&mut self) -> isize { + match self { + CursorMoveToState::Seeking { min: min_state, .. } => *min_state, + _ => unreachable!("method can only be called on CursorMoveToState::Seeking"), + } + } +} #[derive(Debug)] struct FindCellState(Option); @@ -1546,7 +1642,7 @@ impl BTreeCursor { let mem_page = self.stack.top(); let page_idx = mem_page.get().get().id; let page = self.read_page(page_idx)?; - return_if_locked!(page.get()); + return_if_locked_maybe_load!(self.pager, page); let page = page.get(); let contents = page.get().contents.as_ref().unwrap(); if contents.is_leaf() { @@ -1577,7 +1673,7 @@ impl BTreeCursor { let iter_dir = seek_op.iteration_direction(); 'outer: loop { let page = self.stack.top(); - return_if_locked!(page.get()); + return_if_locked_maybe_load!(self.pager, page); let page = page.get(); let contents = page.get().contents.as_ref().unwrap(); if contents.is_leaf() { @@ -1585,7 +1681,10 @@ impl BTreeCursor { } let cell_count = contents.cell_count(); - if matches!(self.move_to_state, CursorSeekState::Start) { + if matches!( + self.move_to_state, + CursorMoveToState::Start | CursorMoveToState::ContinueLoop { .. } + ) { let min: isize = 0; let max: isize = cell_count as isize - 1; let leftmost_matching_cell = None; @@ -1620,6 +1719,7 @@ impl BTreeCursor { .set_cell_index(leftmost_matching_cell as i32 + index_change); let mem_page = self.read_page(left_child_page as usize)?; self.stack.push(mem_page); + self.move_to_state = CursorMoveToState::ContinueLoop; continue 'outer; } self.stack.set_cell_index(cell_count as i32 + 1); @@ -1627,6 +1727,7 @@ impl BTreeCursor { Some(right_most_pointer) => { let mem_page = self.read_page(right_most_pointer as usize)?; self.stack.push(mem_page); + self.move_to_state = CursorMoveToState::ContinueLoop; continue 'outer; } None => { @@ -1687,14 +1788,17 @@ impl BTreeCursor { let iter_dir = cmp.iteration_direction(); 'outer: loop { let page = self.stack.top(); - return_if_locked!(page.get()); + return_if_locked_maybe_load!(self.pager, page); let page = page.get(); let contents = page.get().contents.as_ref().unwrap(); if contents.is_leaf() { return Ok(CursorResult::Ok(())); } - if matches!(self.move_to_state, CursorSeekState::Start) { + if matches!( + self.move_to_state, + CursorMoveToState::Start | CursorMoveToState::ContinueLoop + ) { let cell_count = contents.cell_count(); let min: isize = 0; let max: isize = cell_count as isize - 1; @@ -1719,6 +1823,7 @@ impl BTreeCursor { Some(right_most_pointer) => { let mem_page = self.read_page(right_most_pointer as usize)?; self.stack.push(mem_page); + self.move_to_state = CursorMoveToState::ContinueLoop; continue 'outer; } None => { @@ -1759,6 +1864,7 @@ impl BTreeCursor { let mem_page = self.read_page(*left_child_page as usize)?; self.stack.push(mem_page); + self.move_to_state = CursorMoveToState::ContinueLoop; continue 'outer; } @@ -1863,19 +1969,20 @@ impl BTreeCursor { seek_op: SeekOp, ) -> Result> { assert!(self.mv_cursor.is_none()); - self.move_to_root(); - return_if_io!(self.move_to(SeekKey::TableRowId(rowid), seek_op)); - let page = self.stack.top(); - return_if_locked!(page.get()); - let page = page.get(); - let contents = page.get().contents.as_ref().unwrap(); - assert!( - contents.is_leaf(), - "tablebtree_seek() called on non-leaf page" - ); let iter_dir = seek_op.iteration_direction(); if matches!(self.seek_state, CursorSeekState::Start) { + // No need for another move_to_root. Move_to already moves to root + return_if_io!(self.move_to(SeekKey::TableRowId(rowid), seek_op)); + let page = self.stack.top(); + return_if_locked_maybe_load!(self.pager, page); + let page = page.get(); + let contents = page.get().contents.as_ref().unwrap(); + assert!( + contents.is_leaf(), + "tablebtree_seek() called on non-leaf page" + ); + let cell_count = contents.cell_count(); let min: isize = 0; let max: isize = cell_count as isize - 1; @@ -1888,9 +1995,15 @@ impl BTreeCursor { min, max, nearest_matching_cell, + not_found_leaf: false, }; } + let page = self.stack.top(); + return_if_locked_maybe_load!(self.pager, page); + let page = page.get(); + let contents = page.get().contents.as_ref().unwrap(); + loop { let min = self.seek_state.get_min(); let max = self.seek_state.get_max(); @@ -2026,17 +2139,16 @@ impl BTreeCursor { key: &ImmutableRecord, seek_op: SeekOp, ) -> Result> { - self.move_to_root(); - return_if_io!(self.move_to(SeekKey::IndexKey(key), seek_op)); - - let page = self.stack.top(); - return_if_locked!(page.get()); - - let page = page.get(); - let contents = page.get().contents.as_ref().unwrap(); - let cell_count = contents.cell_count(); - if matches!(self.seek_state, CursorSeekState::Start) { + // No need for another move_to_root. Move_to already moves to root + return_if_io!(self.move_to(SeekKey::IndexKey(key), seek_op)); + let page = self.stack.top(); + return_if_locked_maybe_load!(self.pager, page); + + let page = page.get(); + let contents = page.get().contents.as_ref().unwrap(); + let cell_count = contents.cell_count(); + let min: isize = 0; let max: isize = cell_count as isize - 1; @@ -2048,9 +2160,16 @@ impl BTreeCursor { min, max, nearest_matching_cell, + not_found_leaf: false, }; } + let page = self.stack.top(); + return_if_locked_maybe_load!(self.pager, page); + let page = page.get(); + let contents = page.get().contents.as_ref().unwrap(); + let cell_count = contents.cell_count(); + let iter_dir = seek_op.iteration_direction(); loop { @@ -2074,13 +2193,19 @@ impl BTreeCursor { // if we were to return Ok(CursorResult::Ok((None, None))), self.record would be None, which is incorrect, because we already know // that there is a record with a key greater than K (K' = K+2) in the parent interior cell. Hence, we need to move back up the tree // and get the next matching record from there. - match seek_op.iteration_direction() { + match iter_dir { IterationDirection::Forwards => { - self.stack.set_cell_index(cell_count as i32); + if !self.seek_state.get_not_found_leaf() { + self.seek_state.set_not_found_leaf(true); + self.stack.set_cell_index(cell_count as i32); + } return self.get_next_record(Some((SeekKey::IndexKey(key), seek_op))); } IterationDirection::Backwards => { - self.stack.set_cell_index(-1); + if !self.seek_state.get_not_found_leaf() { + self.seek_state.set_not_found_leaf(true); + self.stack.set_cell_index(-1); + } return self.get_prev_record(Some((SeekKey::IndexKey(key), seek_op))); } } @@ -2242,7 +2367,9 @@ impl BTreeCursor { // 5. We scan the leaf cells in the leaf page until we find the cell whose rowid is equal to the rowid we are looking for. // This cell contains the actual data we are looking for. // 6. If we find the cell, we return the record. Otherwise, we return an empty result. - self.move_to_root(); + if matches!(self.move_to_state, CursorMoveToState::Start) { + self.move_to_root(); + } let ret = match key { SeekKey::TableRowId(rowid_key) => self.tablebtree_move_to(rowid_key, cmp), @@ -4047,7 +4174,7 @@ impl BTreeCursor { let mem_page = self.stack.top(); let page_id = mem_page.get().get().id; let page = self.read_page(page_id)?; - return_if_locked!(page.get()); + return_if_locked_maybe_load!(self.pager, page); let page = page.get(); let contents = page.get().contents.as_ref().unwrap(); @@ -4552,7 +4679,7 @@ impl BTreeCursor { let _ = return_if_io!(self.move_to(SeekKey::TableRowId(*int_key), SeekOp::EQ)); let page = self.stack.top(); // TODO(pere): request load - return_if_locked!(page.get()); + return_if_locked_maybe_load!(self.pager, page); let page = page.get(); let contents = page.get().contents.as_ref().unwrap(); @@ -4613,7 +4740,7 @@ impl BTreeCursor { return Err(LimboError::Corrupt("Invalid overflow page number".into())); } let page = self.read_page(next_page as usize)?; - return_if_locked!(page.get()); + return_if_locked_maybe_load!(self.pager, page); let page = page.get(); let contents = page.get().contents.as_ref().unwrap(); diff --git a/testing/cli_tests/memory.py b/testing/cli_tests/memory.py index 5dd50bf37..474fcae7d 100755 --- a/testing/cli_tests/memory.py +++ b/testing/cli_tests/memory.py @@ -17,7 +17,6 @@ def stub_memory_test( vals: int = 100, blobs: bool = True, ): - raise # zero_blob_size = 1024 **2 zero_blob = "0" * blob_size * 2 # vals = 100