LocalStorage: Add table initialisation for in memory databases (#41486)

Fixes #41485

---------

Signed-off-by: Leo Ring <leoring03@gmail.com>

components/storage/shared/mod.rs

@@ -11,6 +11,9 @@ use rusqlite::{Error as RusqliteError, ffi};
 pub const DB_INIT_PRAGMAS: [&str; 2] =
     ["PRAGMA journal_mode = WAL;", "PRAGMA encoding = 'UTF-16';"];
 
+// These pragmas need to be set once for in memory databases
+pub const DB_IN_MEMORY_INIT_PRAGMAS: [&str; 1] = ["PRAGMA encoding = 'UTF-16';"];
+
 // These pragmas need to be run once per connection.
 pub const DB_PRAGMAS: [&str; 4] = [
     "PRAGMA synchronous = NORMAL;",
@@ -19,6 +22,9 @@ pub const DB_PRAGMAS: [&str; 4] = [
     "PRAGMA cache_size = 2000;",
 ];
 
+// These pragmas need to be run once per connection for in memory databases.
+pub const DB_IN_MEMORY_PRAGMAS: [&str; 1] = ["PRAGMA cache_size = 2000;"];
+
 pub(crate) fn is_sqlite_disk_full_error(error: &RusqliteError) -> bool {
     fn has_enospc(mut source: Option<&(dyn StdError + 'static)>) -> bool {
         while let Some(err) = source {

components/storage/tests/webstorage.rs

@@ -26,6 +26,12 @@ fn init_with(dir: &tempfile::TempDir) -> StorageThreads {
     threads.0
 }
 
+fn init_in_memory() -> StorageThreads {
+    let mem_profiler_chan = profile_mem::Profiler::create();
+    let threads = storage::new_storage_threads(mem_profiler_chan, None);
+    threads.0
+}
+
 /// Gracefully shut down the webstorage thread to avoid dangling threads in tests.
 fn shutdown(threads: &StorageThreads) {
     let (sender, receiver) = base_channel::channel().unwrap();
@@ -71,6 +77,41 @@ fn set_and_get_item() {
     shutdown(&threads);
 }
 
+#[test]
+fn set_and_get_item_in_memory() {
+    let threads = init_in_memory();
+    let url = ServoUrl::parse("https://example.com").unwrap();
+
+    // Set a value.
+    let (sender, receiver) = base_channel::channel().unwrap();
+    threads
+        .send(WebStorageThreadMsg::SetItem(
+            sender,
+            WebStorageType::Local,
+            TEST_WEBVIEW_ID,
+            url.clone(),
+            "foo".into(),
+            "bar".into(),
+        ))
+        .unwrap();
+    assert_eq!(receiver.recv().unwrap(), Ok((true, None)));
+
+    // Retrieve the value.
+    let (sender, receiver) = base_channel::channel().unwrap();
+    threads
+        .send(WebStorageThreadMsg::GetItem(
+            sender,
+            WebStorageType::Local,
+            TEST_WEBVIEW_ID,
+            url.clone(),
+            "foo".into(),
+        ))
+        .unwrap();
+    assert_eq!(receiver.recv().unwrap(), Some("bar".into()));
+
+    shutdown(&threads);
+}
+
 #[test]
 fn length_key_and_keys() {
     let (_tmp_dir, threads) = init();

components/storage/webstorage/engines/sqlite.rs

@@ -9,7 +9,7 @@ use base::threadpool::ThreadPool;
 use log::error;
 use rusqlite::Connection;
 
-use crate::shared::{DB_INIT_PRAGMAS, DB_PRAGMAS};
+use crate::shared::{DB_IN_MEMORY_INIT_PRAGMAS, DB_IN_MEMORY_PRAGMAS, DB_INIT_PRAGMAS, DB_PRAGMAS};
 use crate::webstorage::OriginEntry;
 use crate::webstorage::engines::WebStorageEngine;
 
@@ -22,25 +22,39 @@ impl SqliteEngine {
         let connection = match db_dir {
             Some(path) => {
                 let path = path.join("webstorage.sqlite");
-                Self::init_db(&path)?
+                Self::init_db(Some(&path))?
             },
-            None => Connection::open_in_memory()?,
+            None => Self::init_db(None)?,
         };
-        // Initialize the database with necessary pragmas
-        for pragma in DB_PRAGMAS.iter() {
-            let _ = connection.execute(pragma, []);
-        }
         Ok(SqliteEngine { connection })
     }
 
-    pub fn init_db(path: &PathBuf) -> rusqlite::Result<Connection> {
-        if let Some(parent) = path.parent() {
-            let _ = std::fs::create_dir_all(parent);
-        }
-        let connection = Connection::open(path)?;
-        for pragma in DB_INIT_PRAGMAS.iter() {
-            let _ = connection.execute(pragma, []);
-        }
+    pub fn init_db(db_path: Option<&PathBuf>) -> rusqlite::Result<Connection> {
+        let connection = if let Some(path) = db_path {
+            if let Some(parent) = path.parent() {
+                let _ = std::fs::create_dir_all(parent);
+            }
+            let conn = Connection::open(path)?;
+            for pragma in DB_INIT_PRAGMAS.iter() {
+                let _ = conn.execute(pragma, []);
+            }
+            for pragma in DB_PRAGMAS.iter() {
+                let _ = conn.execute(pragma, []);
+            }
+            conn
+        } else {
+            // TODO We probably don't need an in memory implementation at all.
+            // WebStorageEnvironment already keeps all key value pairs in memory via its data field.
+            // A future refactoring could avoid creating a WebStorageEngine entirely when config_dir is None.
+            let conn = Connection::open_in_memory()?;
+            for pragma in DB_IN_MEMORY_INIT_PRAGMAS.iter() {
+                let _ = conn.execute(pragma, []);
+            }
+            for pragma in DB_IN_MEMORY_PRAGMAS.iter() {
+                let _ = conn.execute(pragma, []);
+            }
+            conn
+        };
         connection.execute("CREATE TABLE IF NOT EXISTS data (id INTEGER PRIMARY KEY AUTOINCREMENT, key TEXT, value TEXT);", [])?;
         Ok(connection)
     }