From 7ea52a3f892df6bb2a5ab550857d11cedec5c3c1 Mon Sep 17 00:00:00 2001
From: PThorpe92 <preston@turso.tech>
Date: Sat, 2 Aug 2025 17:31:01 -0400
Subject: [PATCH] Fix changing page size and initialization for buffer pool

---
 core/lib.rs                 | 20 ++++++++++++--------
 core/storage/buffer_pool.rs |  9 +++++----
 core/storage/pager.rs       |  9 +++++----
 3 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/core/lib.rs b/core/lib.rs
index 4736c6cb4..f79c9358d 100644
--- a/core/lib.rs
+++ b/core/lib.rs
@@ -380,6 +380,11 @@ impl Database {
     }
 
     fn init_pager(&self, page_size: Option<usize>) -> Result<Pager> {
+        let arena_size = if std::env::var("CI").is_ok_and(|v| v.eq_ignore_ascii_case("true")) {
+            BufferPool::TEST_AREA_SIZE
+        } else {
+            BufferPool::DEFAULT_ARENA_SIZE
+        };
         // Open existing WAL file if present
         let mut maybe_shared_wal = self.maybe_shared_wal.write();
         if let Some(shared_wal) = maybe_shared_wal.clone() {
@@ -387,8 +392,8 @@ impl Database {
                 None => unsafe { (*shared_wal.get()).page_size() as usize },
                 Some(size) => size,
             };
-            let buffer_pool = BufferPool::begin_init(&self.io, BufferPool::DEFAULT_ARENA_SIZE)
-                .finalize_page_size(size)?;
+            let buffer_pool =
+                BufferPool::begin_init(&self.io, arena_size).finalize_page_size(size)?;
 
             let db_state = self.db_state.clone();
             let wal = Rc::new(RefCell::new(WalFile::new(
@@ -407,8 +412,7 @@ impl Database {
             )?;
             return Ok(pager);
         }
-
-        let buffer_pool = BufferPool::begin_init(&self.io, BufferPool::DEFAULT_ARENA_SIZE);
+        let buffer_pool = BufferPool::begin_init(&self.io, arena_size);
 
         // No existing WAL; create one.
         let db_state = self.db_state.clone();
@@ -425,16 +429,15 @@ impl Database {
         let size = match page_size {
             Some(size) => size as u32,
             None => {
-                let size = pager
+                pager // if None is passed in, we know that we already initialized so we can safely call `with_header` here
                     .io
                     .block(|| pager.with_header(|header| header.page_size))
                     .unwrap_or_default()
-                    .get();
-                size
+                    .get()
             }
         };
 
-        buffer_pool.finalize_page_size(size as usize);
+        pager.page_size.set(Some(size));
         let wal_path = format!("{}-wal", self.path);
         let file = self.io.open_file(&wal_path, OpenFlags::Create, false)?;
         let real_shared_wal = WalFileShared::new_shared(size, &self.io, file)?;
@@ -1375,6 +1378,7 @@ impl Connection {
         }
 
         *self._db.maybe_shared_wal.write() = None;
+        self.pager.borrow_mut().clear_page_cache();
         let pager = self._db.init_pager(Some(size as usize))?;
         self.pager.replace(Rc::new(pager));
         self.pager.borrow().set_initial_page_size(size);
diff --git a/core/storage/buffer_pool.rs b/core/storage/buffer_pool.rs
index f2de613a0..b0d2b79c0 100644
--- a/core/storage/buffer_pool.rs
+++ b/core/storage/buffer_pool.rs
@@ -137,6 +137,7 @@ impl Default for BufferPool {
 
 impl BufferPool {
     pub const DEFAULT_ARENA_SIZE: usize = 4 * 1024 * 1024; // 4MB arena
+    pub const TEST_AREA_SIZE: usize = 512 * 1024; // 512KB arena for testing
     pub const DEFAULT_PAGE_SIZE: usize = 4096; // 4KB default page size
     const MAX_ARENA_SIZE: usize = 32 * 1024 * 1024; // 32MB max arena
 
@@ -193,9 +194,6 @@ impl BufferPool {
     pub fn finalize_page_size(&self, page_size: usize) -> crate::Result<Arc<BufferPool>> {
         let pool = BUFFER_POOL.get().expect("BufferPool must be initialized");
         let inner = pool.inner_mut();
-        if inner.arena.is_some() {
-            return Ok(pool.clone());
-        }
         tracing::trace!("finalize page size called with size {page_size}");
         if page_size != BufferPool::DEFAULT_PAGE_SIZE {
             // so far we have handed out some temporary buffers, since the page size is not
@@ -204,6 +202,9 @@ impl BufferPool {
                 cache.borrow_mut().reinit_cache(page_size);
             });
         }
+        if inner.arena.is_some() {
+            return Ok(pool.clone());
+        }
         inner.db_page_size.store(page_size, Ordering::Relaxed);
         inner.init_arena()?;
         Ok(pool.clone())
@@ -253,7 +254,7 @@ impl PoolInner {
         ) {
             Ok(arena) => {
                 tracing::trace!(
-                    "Growing buffer pool: added arena {} with size {} MB",
+                    "added arena {} with size {} MB",
                     arena.id,
                     arena_size / (1024 * 1024)
                 );
diff --git a/core/storage/pager.rs b/core/storage/pager.rs
index d685712d1..fc2562ba5 100644
--- a/core/storage/pager.rs
+++ b/core/storage/pager.rs
@@ -1766,6 +1766,8 @@ impl Pager {
                 if let Some(size) = self.page_size.get() {
                     default_header.page_size = PageSize::new(size).expect("page size");
                 }
+                self.buffer_pool
+                    .finalize_page_size(default_header.page_size.get() as usize)?;
                 let page = allocate_new_page(1, &self.buffer_pool, 0);
 
                 let contents = page.get_contents();
@@ -2431,10 +2433,9 @@ mod ptrmap_tests {
         ));
 
         //  Construct interfaces for the pager
-        let buffer_pool = Arc::new(BufferPool::new(Some(page_size as usize)));
-        let page_cache = Arc::new(RwLock::new(DumbLruPageCache::new(
-            (initial_db_pages + 10) as usize,
-        )));
+        let sz = initial_db_pages + 10;
+        let buffer_pool = BufferPool::begin_init(&io, (sz * page_size) as usize);
+        let page_cache = Arc::new(RwLock::new(DumbLruPageCache::new(sz as usize)));
 
         let wal = Rc::new(RefCell::new(WalFile::new(
             io.clone(),