do_allocate_page should return a Result

core/storage/btree.rs

@@ -2314,7 +2314,7 @@ impl BTreeCursor {
                     }
 
                     if !self.stack.has_parent() {
-                        self.balance_root();
+                        self.balance_root()?;
                     }
 
                     let write_info = self.state.mut_write_info().unwrap();
@@ -2890,7 +2890,7 @@ impl BTreeCursor {
                         pages_to_balance_new[i].replace(page.clone());
                     } else {
                         // FIXME: handle page cache is full
-                        let page = self.allocate_page(page_type, 0);
+                        let page = self.allocate_page(page_type, 0)?;
                         pages_to_balance_new[i].replace(page);
                         // Since this page didn't exist before, we can set it to cells length as it
                         // marks them as empty since it is a prefix sum of cells.
@@ -3785,7 +3785,7 @@ impl BTreeCursor {
     /// Balance the root page.
     /// This is done when the root page overflows, and we need to create a new root page.
     /// See e.g. https://en.wikipedia.org/wiki/B-tree
-    fn balance_root(&mut self) {
+    fn balance_root(&mut self) -> Result<()> {
         /* todo: balance deeper, create child and copy contents of root there. Then split root */
         /* if we are in root page then we just need to create a new root and push key there */
 
@@ -3802,7 +3802,7 @@ impl BTreeCursor {
         // FIXME: handle page cache is full
         let child_btree =
             self.pager
-                .do_allocate_page(root_contents.page_type(), 0, BtreePageAllocMode::Any);
+                .do_allocate_page(root_contents.page_type(), 0, BtreePageAllocMode::Any)?;
 
         tracing::debug!(
             "balance_root(root={}, rightmost={}, page_type={:?})",
@@ -3860,6 +3860,7 @@ impl BTreeCursor {
         self.stack.push(root_btree.clone());
         self.stack.set_cell_index(0); // leave parent pointing at the rightmost pointer (in this case 0, as there are no cells), since we will be balancing the rightmost child page.
         self.stack.push(child_btree.clone());
+        Ok(())
     }
 
     fn usable_space(&self) -> usize {
@@ -5153,7 +5154,7 @@ impl BTreeCursor {
         btree_read_page(&self.pager, page_idx)
     }
 
-    pub fn allocate_page(&self, page_type: PageType, offset: usize) -> BTreePage {
+    pub fn allocate_page(&self, page_type: PageType, offset: usize) -> Result<BTreePage> {
         self.pager
             .do_allocate_page(page_type, offset, BtreePageAllocMode::Any)
     }
@@ -7587,11 +7588,11 @@ mod tests {
         let mut cursor = BTreeCursor::new_table(None, pager.clone(), 2);
 
         // Initialize page 2 as a root page (interior)
-        let root_page = cursor.allocate_page(PageType::TableInterior, 0);
+        let root_page = cursor.allocate_page(PageType::TableInterior, 0)?;
 
         // Allocate two leaf pages
-        let page3 = cursor.allocate_page(PageType::TableLeaf, 0);
-        let page4 = cursor.allocate_page(PageType::TableLeaf, 0);
+        let page3 = cursor.allocate_page(PageType::TableLeaf, 0)?;
+        let page4 = cursor.allocate_page(PageType::TableLeaf, 0)?;
 
         // Configure the root page to point to the two leaf pages
         {

core/storage/pager.rs

@@ -480,7 +480,7 @@ impl Pager {
         };
         #[cfg(feature = "omit_autovacuum")]
         {
-            let page = self.do_allocate_page(page_type, 0, BtreePageAllocMode::Any);
+            let page = self.do_allocate_page(page_type, 0, BtreePageAllocMode::Any)?;
             let page_id = page.get().get().id;
             Ok(CursorResult::Ok(page_id as u32))
         }
@@ -491,7 +491,7 @@ impl Pager {
             let auto_vacuum_mode = self.auto_vacuum_mode.borrow();
             match *auto_vacuum_mode {
                 AutoVacuumMode::None => {
-                    let page = self.do_allocate_page(page_type, 0, BtreePageAllocMode::Any);
+                    let page = self.do_allocate_page(page_type, 0, BtreePageAllocMode::Any)?;
                     let page_id = page.get().get().id;
                     Ok(CursorResult::Ok(page_id as u32))
                 }
@@ -515,7 +515,7 @@ impl Pager {
                         page_type,
                         0,
                         BtreePageAllocMode::Exact(root_page_num),
-                    );
+                    )?;
                     let allocated_page_id = page.get().get().id as u32;
                     if allocated_page_id != root_page_num {
                         //  TODO(Zaid): Handle swapping the allocated page with the desired root page
@@ -559,8 +559,8 @@ impl Pager {
         page_type: PageType,
         offset: usize,
         _alloc_mode: BtreePageAllocMode,
-    ) -> BTreePage {
-        let page = self.allocate_page().unwrap();
+    ) -> Result<BTreePage> {
+        let page = self.allocate_page()?;
         let page = Arc::new(BTreePageInner {
             page: RefCell::new(page),
         });
@@ -570,7 +570,7 @@ impl Pager {
             page.get().get().id,
             page.get().get_contents().page_type()
         );
-        page
+        Ok(page)
     }
 
     /// The "usable size" of a database page is the page size specified by the 2-byte integer at offset 16