From f123c77ee88237abbfea124cb87b1964f6f3e820 Mon Sep 17 00:00:00 2001
From: meteorgan <meteorite.gan@gmail.com>
Date: Fri, 11 Jul 2025 22:22:16 +0800
Subject: [PATCH] fix set page_size in pager

---
 core/lib.rs              | 17 +++++++++++++++--
 core/storage/pager.rs    | 15 ++++++++++++---
 core/translate/pragma.rs |  7 ++-----
 3 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/core/lib.rs b/core/lib.rs
index 12d758005..540f81761 100644
--- a/core/lib.rs
+++ b/core/lib.rs
@@ -41,6 +41,7 @@ mod numeric;
 #[global_allocator]
 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
 
+use crate::storage::sqlite3_ondisk::is_valid_page_size;
 use crate::storage::{header_accessor, wal::DummyWAL};
 use crate::translate::optimizer::optimize_plan;
 use crate::translate::pragma::TURSO_CDC_DEFAULT_TABLE_NAME;
@@ -510,6 +511,8 @@ pub struct Connection {
     syms: RefCell<SymbolTable>,
     _shared_cache: bool,
     cache_size: Cell<i32>,
+    /// page size used for an uninitialized database or the next vacuum command.
+    /// it's not always equal to the current page size of the database
     page_size: Cell<u32>,
     readonly: Cell<bool>,
     wal_checkpoint_disabled: Cell<bool>,
@@ -833,15 +836,25 @@ impl Connection {
         self.page_size.get()
     }
 
-    /// Reset the page size for the current connection. Can only be called when db is uninitialized.
+    /// Reset the page size for the current connection.
+    ///
+    /// Specifying a new page size does not change the page size immediately.
+    /// Instead, the new page size is remembered and is used to set the page size when the database
+    /// is first created, if it does not already exist when the page_size pragma is issued,
+    /// or at the next VACUUM command that is run on the same database connection while not in WAL mode.
     pub fn reset_page_size(&self, size: u32) -> Result<()> {
-        if self._db.db_state.load(Ordering::SeqCst) != DB_STATE_UNINITIALIZED {
+        if !is_valid_page_size(size) {
             return Ok(());
         }
 
         self.page_size.set(size);
+        if self._db.db_state.load(Ordering::SeqCst) != DB_STATE_UNINITIALIZED {
+            return Ok(());
+        }
+
         let pager = self._db.init_pager(Some(size as usize))?;
         self.pager.replace(Rc::new(pager));
+        self.pager.borrow().set_initial_page_size(size);
 
         Ok(())
     }
diff --git a/core/storage/pager.rs b/core/storage/pager.rs
index b88a92475..2b2d2cd32 100644
--- a/core/storage/pager.rs
+++ b/core/storage/pager.rs
@@ -225,7 +225,7 @@ pub struct Pager {
     /// Cache page_size and reserved_space at Pager init and reuse for subsequent
     /// `usable_space` calls. TODO: Invalidate reserved_space when we add the functionality
     /// to change it.
-    page_size: OnceCell<u32>,
+    page_size: RefCell<Option<u32>>,
     reserved_space: OnceCell<u8>,
 }
 
@@ -291,7 +291,7 @@ impl Pager {
             db_state,
             init_lock,
             allocate_page1_state,
-            page_size: OnceCell::new(),
+            page_size: RefCell::new(None),
             reserved_space: OnceCell::new(),
         })
     }
@@ -586,7 +586,8 @@ impl Pager {
     pub fn usable_space(&self) -> usize {
         let page_size = *self
             .page_size
-            .get_or_init(|| header_accessor::get_page_size(self).unwrap_or_default());
+            .borrow_mut()
+            .get_or_insert_with(|| header_accessor::get_page_size(self).unwrap_or_default());
 
         let reserved_space = *self
             .reserved_space
@@ -595,6 +596,11 @@ impl Pager {
         (page_size as usize) - (reserved_space as usize)
     }
 
+    /// Set the initial page size for the database. Should only be called before the database is initialized
+    pub fn set_initial_page_size(&self, size: u32) {
+        self.page_size.replace(Some(size));
+    }
+
     #[inline(always)]
     #[instrument(skip_all, level = Level::INFO)]
     pub fn begin_read_tx(&self) -> Result<CursorResult<LimboResult>> {
@@ -1057,6 +1063,9 @@ impl Pager {
                 self.db_state.store(DB_STATE_INITIALIZING, Ordering::SeqCst);
                 let mut default_header = DatabaseHeader::default();
                 default_header.database_size += 1;
+                if let Some(size) = *self.page_size.borrow() {
+                    default_header.update_page_size(size);
+                }
                 let page = allocate_page(1, &self.buffer_pool, 0);
 
                 let contents = page.get_contents();
diff --git a/core/translate/pragma.rs b/core/translate/pragma.rs
index d6736ef95..da99412ad 100644
--- a/core/translate/pragma.rs
+++ b/core/translate/pragma.rs
@@ -9,7 +9,7 @@ use turso_sqlite3_parser::ast::{PragmaName, QualifiedName};
 use crate::pragma::pragma_for;
 use crate::schema::Schema;
 use crate::storage::pager::AutoVacuumMode;
-use crate::storage::sqlite3_ondisk::{is_valid_page_size, MIN_PAGE_CACHE_SIZE};
+use crate::storage::sqlite3_ondisk::MIN_PAGE_CACHE_SIZE;
 use crate::storage::wal::CheckpointMode;
 use crate::translate::schema::translate_create_table;
 use crate::util::{normalize_ident, parse_signed_number, parse_string};
@@ -536,11 +536,8 @@ fn turso_cdc_table_columns() -> Vec<ColumnDefinition> {
         },
     ]
 }
-fn update_page_size(connection: Arc<crate::Connection>, page_size: u32) -> crate::Result<()> {
-    if !is_valid_page_size(page_size) {
-        return Ok(());
-    }
 
+fn update_page_size(connection: Arc<crate::Connection>, page_size: u32) -> crate::Result<()> {
     connection.reset_page_size(page_size)?;
     Ok(())
 }