Refactored webext-storage database close function

Author: lougeniaC64

components/webext-storage/src/api.rs

@@ -438,8 +438,9 @@ mod tests {
     #[test]
     fn test_simple() -> Result<()> {
         let ext_id = "x";
-        let mut db = new_mem_db();
-        let tx = db.transaction()?;
+        let db = new_mem_db();
+        let conn = db.get_connection().expect("should retrieve connection");
+        let tx = conn.unchecked_transaction()?;
 
         // an empty store.
         for q in vec![JsonValue::Null, json!("foo"), json!(["foo"])].into_iter() {
@@ -529,8 +530,9 @@ mod tests {
     fn test_check_get_impl() -> Result<()> {
         // This is a port of checkGetImpl in test_ext_storage.js in Desktop.
         let ext_id = "x";
-        let mut db = new_mem_db();
-        let tx = db.transaction()?;
+        let db = new_mem_db();
+        let conn = db.get_connection().expect("should retrieve connection");
+        let tx = conn.unchecked_transaction()?;
 
         let prop = "test-prop";
         let value = "test-value";
@@ -584,8 +586,9 @@ mod tests {
     fn test_bug_1621162() -> Result<()> {
         // apparently Firefox, unlike Chrome, will not optimize the changes.
         // See bug 1621162 for more!
-        let mut db = new_mem_db();
-        let tx = db.transaction()?;
+        let db = new_mem_db();
+        let conn = db.get_connection().expect("should retrieve connection");
+        let tx = conn.unchecked_transaction()?;
         let ext_id = "xyz";
 
         set(&tx, ext_id, json!({"foo": "bar" }))?;
@@ -599,8 +602,9 @@ mod tests {
 
     #[test]
     fn test_quota_maxitems() -> Result<()> {
-        let mut db = new_mem_db();
-        let tx = db.transaction()?;
+        let db = new_mem_db();
+        let conn = db.get_connection().expect("should retrieve connection");
+        let tx = conn.unchecked_transaction()?;
         let ext_id = "xyz";
         for i in 1..SYNC_MAX_ITEMS + 1 {
             set(
@@ -619,8 +623,9 @@ mod tests {
 
     #[test]
     fn test_quota_bytesperitem() -> Result<()> {
-        let mut db = new_mem_db();
-        let tx = db.transaction()?;
+        let db = new_mem_db();
+        let conn = db.get_connection().expect("should retrieve connection");
+        let tx = conn.unchecked_transaction()?;
         let ext_id = "xyz";
         // A string 5 bytes less than the max. This should be counted as being
         // 3 bytes less than the max as the quotes are counted. Plus the length
@@ -645,8 +650,9 @@ mod tests {
 
     #[test]
     fn test_quota_bytes() -> Result<()> {
-        let mut db = new_mem_db();
-        let tx = db.transaction()?;
+        let db = new_mem_db();
+        let conn = db.get_connection().expect("should retrieve connection");
+        let tx = conn.unchecked_transaction()?;
         let ext_id = "xyz";
         let val = "x".repeat(SYNC_QUOTA_BYTES + 1);
 
@@ -682,8 +688,9 @@ mod tests {
 
     #[test]
     fn test_get_bytes_in_use() -> Result<()> {
-        let mut db = new_mem_db();
-        let tx = db.transaction()?;
+        let db = new_mem_db();
+        let conn = db.get_connection().expect("should retrieve connection");
+        let tx = conn.unchecked_transaction()?;
         let ext_id = "xyz";
 
         assert_eq!(get_bytes_in_use(&tx, ext_id, json!(null))?, 0);
@@ -714,8 +721,9 @@ mod tests {
 
     #[test]
     fn test_usage() {
-        let mut db = new_mem_db();
-        let tx = db.transaction().unwrap();
+        let db = new_mem_db();
+        let conn = db.get_connection().expect("should retrieve connection");
+        let tx = conn.unchecked_transaction().unwrap();
         // '{"a":"a","b":"bb","c":"ccc","n":999999}': 39 bytes
         set(&tx, "xyz", json!({ "a": "a" })).unwrap();
         set(&tx, "xyz", json!({ "b": "bb" })).unwrap();
@@ -727,7 +735,7 @@ mod tests {
 
         tx.commit().unwrap();
 
-        let usage = usage(&db).unwrap();
+        let usage = usage(conn).unwrap();
         let expect = [
             UsageInfo {
                 ext_id: "abc".to_string(),

components/webext-storage/src/db.rs

@@ -11,7 +11,7 @@ use rusqlite::Connection;
 use rusqlite::OpenFlags;
 use sql_support::open_database::open_database_with_flags;
 use sql_support::ConnExt;
-use std::ops::{Deref, DerefMut};
+use std::ops::Deref;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use url::Url;
@@ -23,10 +23,16 @@ use url::Url;
 ///
 /// We only support a single writer connection - so that's the only thing we
 /// store. It's still a bit overkill, but there's only so many yaks in a day.
+pub enum WebExtStorageDb {
+    Open(Connection),
+    Closed,
+}
+
 pub struct StorageDb {
-    writer: Connection,
+    pub writer: WebExtStorageDb,
     interrupt_handle: Arc<SqlInterruptHandle>,
 }
+
 impl StorageDb {
     /// Create a new, or fetch an already open, StorageDb backed by a file on disk.
     pub fn new(db_path: impl AsRef<Path>) -> Result<Self> {
@@ -54,7 +60,7 @@ impl StorageDb {
         let conn = open_database_with_flags(db_path, flags, &schema::WebExtMigrationLogin)?;
         Ok(Self {
             interrupt_handle: Arc::new(SqlInterruptHandle::new(&conn)),
-            writer: conn,
+            writer: WebExtStorageDb::Open(conn),
         })
     }
 
@@ -73,29 +79,20 @@ impl StorageDb {
     /// underlying connection so the caller can retry but (a) that's very tricky
     /// in an Arc<Mutex<>> world and (b) we never actually took advantage of
     /// that retry capability.
-    pub fn close(self) -> Result<()> {
-        self.writer.close().map_err(|(writer, err)| {
-            // In rusqlite 0.28.0 and earlier, if we just let `writer` drop,
-            // the close would panic on failure.
-            // Later rusqlite versions will not panic, but this behavior doesn't
-            // hurt there.
-            std::mem::forget(writer);
-            err.into()
-        })
-    }
-}
-
-impl Deref for StorageDb {
-    type Target = Connection;
-
-    fn deref(&self) -> &Self::Target {
-        &self.writer
+    pub fn close(&mut self) -> Result<()> {
+        let conn = match std::mem::replace(&mut self.writer, WebExtStorageDb::Closed) {
+            WebExtStorageDb::Open(conn) => conn,
+            WebExtStorageDb::Closed => return Ok(()),
+        };
+        conn.close().map_err(|(_, y)| Error::SqlError(y))
     }
-}
 
-impl DerefMut for StorageDb {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.writer
+    pub(crate) fn get_connection(&self) -> Result<&Connection> {
+        let db = &self.writer;
+        match db {
+            WebExtStorageDb::Open(y) => Ok(y),
+            WebExtStorageDb::Closed => Err(Error::DatabaseConnectionClosed),
+        }
     }
 }
 
@@ -290,12 +287,13 @@ mod tests {
 
     #[test]
     fn test_meta() -> Result<()> {
-        let writer = new_mem_db();
-        assert_eq!(get_meta::<String>(&writer, "foo")?, None);
-        put_meta(&writer, "foo", &"bar".to_string())?;
-        assert_eq!(get_meta(&writer, "foo")?, Some("bar".to_string()));
-        delete_meta(&writer, "foo")?;
-        assert_eq!(get_meta::<String>(&writer, "foo")?, None);
+        let db = new_mem_db();
+        let conn = &db.get_connection()?;
+        assert_eq!(get_meta::<String>(conn, "foo")?, None);
+        put_meta(conn, "foo", &"bar".to_string())?;
+        assert_eq!(get_meta(conn, "foo")?, Some("bar".to_string()));
+        delete_meta(conn, "foo")?;
+        assert_eq!(get_meta::<String>(conn, "foo")?, None);
         Ok(())
     }
 }

components/webext-storage/src/migration.rs

@@ -346,8 +346,9 @@ mod tests {
         init_source_db(path, f);
 
         // now migrate
-        let mut db = new_mem_db();
-        let tx = db.transaction().expect("tx should work");
+        let db = new_mem_db();
+        let conn = db.get_connection().expect("should retrieve connection");
+        let tx = conn.unchecked_transaction().expect("tx should work");
 
         let mi = migrate(&tx, &tmpdir.path().join("source.db")).expect("migrate should work");
         tx.commit().expect("should work");
@@ -384,17 +385,18 @@ mod tests {
     #[test]
     fn test_happy_paths() {
         // some real data.
-        let conn = do_migrate(HAPPY_PATH_MIGRATION_INFO, |c| {
+        let db = do_migrate(HAPPY_PATH_MIGRATION_INFO, |c| {
             c.execute_batch(HAPPY_PATH_SQL).expect("should populate")
         });
+        let conn = db.get_connection().expect("should retrieve connection");
 
         assert_has(
-            &conn,
+            conn,
             "{e7fefcf3-b39c-4f17-5215-ebfe120a7031}",
             json!({"userWelcomed": 1570659224457u64, "isWho": "4ec8109f"}),
         );
         assert_has(
-            &conn,
+            conn,
             "[email protected]",
             json!({"userRules": [], "ruleActiveStates": {}, "migration_version": 2}),
         );

components/webext-storage/src/schema.rs

@@ -94,22 +94,24 @@ mod tests {
     #[test]
     fn test_create_schema_twice() {
         let db = new_mem_db();
-        db.execute_batch(CREATE_SCHEMA_SQL)
+        let conn = db.get_connection().expect("should retrieve connection");
+        conn.execute_batch(CREATE_SCHEMA_SQL)
             .expect("should allow running twice");
     }
 
     #[test]
     fn test_create_empty_sync_temp_tables_twice() {
         let db = new_mem_db();
-        create_empty_sync_temp_tables(&db).expect("should work first time");
+        let conn = db.get_connection().expect("should retrieve connection");
+        create_empty_sync_temp_tables(conn).expect("should work first time");
         // insert something into our new temp table and check it's there.
-        db.execute_batch(
+        conn.execute_batch(
             "INSERT INTO temp.storage_sync_staging
                             (guid, ext_id) VALUES
                             ('guid', 'ext_id');",
         )
         .expect("should work once");
-        let count = db
+        let count = conn
             .query_row_and_then(
                 "SELECT COUNT(*) FROM temp.storage_sync_staging;",
                 [],
@@ -119,9 +121,9 @@ mod tests {
         assert_eq!(count, 1, "should be one row");
 
         // re-execute
-        create_empty_sync_temp_tables(&db).expect("should second first time");
+        create_empty_sync_temp_tables(conn).expect("should second first time");
         // and it should have deleted existing data.
-        let count = db
+        let count = conn
             .query_row_and_then(
                 "SELECT COUNT(*) FROM temp.storage_sync_staging;",
                 [],

components/webext-storage/src/store.rs

@@ -59,17 +59,19 @@ impl WebExtStorageStore {
     /// Sets one or more JSON key-value pairs for an extension ID. Returns a
     /// list of changes, with existing and new values for each key in `val`.
     pub fn set(&self, ext_id: &str, val: JsonValue) -> Result<StorageChanges> {
-        let db = self.db.lock();
-        let tx = db.unchecked_transaction()?;
+        let db = &self.db.lock();
+        let conn = db.get_connection()?;
+        let tx = conn.unchecked_transaction()?;
         let result = api::set(&tx, ext_id, val)?;
         tx.commit()?;
         Ok(result)
     }
 
     /// Returns information about per-extension usage
     pub fn usage(&self) -> Result<Vec<crate::UsageInfo>> {
-        let db = self.db.lock();
-        api::usage(&db)
+        let db = &self.db.lock();
+        let conn = db.get_connection()?;
+        api::usage(conn)
     }
 
     /// Returns the values for one or more keys `keys` can be:
@@ -90,17 +92,19 @@ impl WebExtStorageStore {
     /// `serde_json::Value::Object`).
     pub fn get(&self, ext_id: &str, keys: JsonValue) -> Result<JsonValue> {
         // Don't care about transactions here.
-        let db = self.db.lock();
-        api::get(&db, ext_id, keys)
+        let db = &self.db.lock();
+        let conn = db.get_connection()?;
+        api::get(conn, ext_id, keys)
     }
 
     /// Deletes the values for one or more keys. As with `get`, `keys` can be
     /// either a single string key, or an array of string keys. Returns a list
     /// of changes, where each change contains the old value for each deleted
     /// key.
     pub fn remove(&self, ext_id: &str, keys: JsonValue) -> Result<StorageChanges> {
-        let db = self.db.lock();
-        let tx = db.unchecked_transaction()?;
+        let db = &self.db.lock();
+        let conn = db.get_connection()?;
+        let tx = conn.unchecked_transaction()?;
         let result = api::remove(&tx, ext_id, keys)?;
         tx.commit()?;
         Ok(result)
@@ -110,8 +114,9 @@ impl WebExtStorageStore {
     /// a list of changes, where each change contains the old value for each
     /// deleted key.
     pub fn clear(&self, ext_id: &str) -> Result<StorageChanges> {
-        let db = self.db.lock();
-        let tx = db.unchecked_transaction()?;
+        let db = &self.db.lock();
+        let conn = db.get_connection()?;
+        let tx = conn.unchecked_transaction()?;
         let result = api::clear(&tx, ext_id)?;
         tx.commit()?;
         Ok(result)
@@ -120,8 +125,9 @@ impl WebExtStorageStore {
     /// Returns the bytes in use for the specified items (which can be null,
     /// a string, or an array)
     pub fn get_bytes_in_use(&self, ext_id: &str, keys: JsonValue) -> Result<usize> {
-        let db = self.db.lock();
-        api::get_bytes_in_use(&db, ext_id, keys)
+        let db = &self.db.lock();
+        let conn = db.get_connection()?;
+        api::get_bytes_in_use(conn, ext_id, keys)
     }
 
     /// Returns a bridged sync engine for Desktop for this store.
@@ -135,8 +141,8 @@ impl WebExtStorageStore {
         // Even though this consumes `self`, the fact we use an Arc<> means
         // we can't guarantee we can actually consume the inner DB - so do
         // the best we can.
-        let shared: ThreadSafeStorageDb = match Arc::try_unwrap(self.db) {
-            Ok(shared) => shared,
+        let shared: ThreadSafeStorageDb = match Arc::into_inner(self.db) {
+            Some(shared) => shared,
             _ => {
                 // The only way this is possible is if the sync engine has an operation
                 // running - but that shouldn't be possible in practice because desktop
@@ -157,7 +163,7 @@ impl WebExtStorageStore {
             }
         };
         // consume the mutex and get back the inner.
-        let db = shared.into_inner();
+        let mut db = shared.into_inner();
         db.close()
     }
 
@@ -177,12 +183,13 @@ impl WebExtStorageStore {
     ///
     /// Note that `filename` isn't normalized or canonicalized.
     pub fn migrate(&self, filename: impl AsRef<Path>) -> Result<()> {
-        let db = self.db.lock();
-        let tx = db.unchecked_transaction()?;
+        let db = &self.db.lock();
+        let conn = db.get_connection()?;
+        let tx = conn.unchecked_transaction()?;
         let result = migrate(&tx, filename.as_ref())?;
         tx.commit()?;
         // Failing to store this information should not cause migration failure.
-        if let Err(e) = result.store(&db) {
+        if let Err(e) = result.store(conn) {
             debug_assert!(false, "Migration error: {:?}", e);
             log::warn!("Failed to record migration telmetry: {}", e);
         }
@@ -192,8 +199,9 @@ impl WebExtStorageStore {
     /// Read-and-delete (e.g. `take` in rust parlance, see Option::take)
     /// operation for any MigrationInfo stored in this database.
     pub fn take_migration_info(&self) -> Result<Option<MigrationInfo>> {
-        let db = self.db.lock();
-        let tx = db.unchecked_transaction()?;
+        let db = &self.db.lock();
+        let conn = db.get_connection()?;
+        let tx = conn.unchecked_transaction()?;
         let result = MigrationInfo::take(&tx)?;
         tx.commit()?;
         Ok(result)