Make pool read-only, with a single write connection.

xmtp_mls/src/client.rs

@@ -1072,7 +1072,7 @@ pub(crate) mod tests {
             .unwrap();
 
         let conn = amal.store().conn().unwrap();
-        conn.raw_query(|conn| diesel::delete(identity_updates::table).execute(conn))
+        conn.raw_query_write(|conn| diesel::delete(identity_updates::table).execute(conn))
             .unwrap();
 
         let members = group.members().await.unwrap();
@@ -1424,6 +1424,7 @@ pub(crate) mod tests {
             .unwrap();
         assert_eq!(amal_group.members().await.unwrap().len(), 1);
         tracing::info!("Syncing bolas welcomes");
+
         // See if Bola can see that they were added to the group
         bola.sync_welcomes(&bola.mls_provider().unwrap())
             .await

xmtp_mls/src/groups/intents.rs

@@ -63,10 +63,12 @@ impl<ScopedClient: ScopedGroupClient> MlsGroup<ScopedClient> {
         intent_kind: IntentKind,
         intent_data: Vec<u8>,
     ) -> Result<StoredGroupIntent, GroupError> {
-        provider.transaction(|provider| {
+        let res = provider.transaction(|provider| {
             let conn = provider.conn_ref();
             self.queue_intent_with_conn(conn, intent_kind, intent_data)
-        })
+        });
+
+        res
     }
 
     fn queue_intent_with_conn(

xmtp_mls/src/groups/mod.rs

@@ -933,6 +933,8 @@ impl<ScopedClient: ScopedGroupClient> MlsGroup<ScopedClient> {
             intent_data.into(),
         )?;
 
+        tracing::warn!("This makes it here?");
+
         self.sync_until_intent_resolved(provider, intent.id).await
     }
 
@@ -2216,7 +2218,7 @@ pub(crate) mod tests {
 
         // The dm shows up
         let alix_groups = alix_conn
-            .raw_query(|conn| groups::table.load::<StoredGroup>(conn))
+            .raw_query_read(|conn| groups::table.load::<StoredGroup>(conn))
             .unwrap();
         assert_eq!(alix_groups.len(), 2);
         // They should have the same ID
@@ -3787,18 +3789,18 @@ pub(crate) mod tests {
         let conn_1: XmtpOpenMlsProvider = bo.store().conn().unwrap().into();
         let conn_2 = bo.store().conn().unwrap();
         conn_2
-            .raw_query(|c| {
+            .raw_query_write(|c| {
                 c.batch_execute("BEGIN EXCLUSIVE").unwrap();
                 Ok::<_, diesel::result::Error>(())
             })
             .unwrap();
 
         let process_result = bo_group.process_messages(bo_messages, &conn_1).await;
         if let Some(GroupError::ReceiveErrors(errors)) = process_result.err() {
-            assert_eq!(errors.len(), 2);
-            assert!(errors
-                .iter()
-                .any(|err| err.to_string().contains("database is locked")));
+            assert_eq!(errors.len(), 1);
+            assert!(errors.iter().any(|err| err
+                .to_string()
+                .contains("cannot start a transaction within a transaction")));
         } else {
             panic!("Expected error")
         }

xmtp_mls/src/storage/encrypted_store/association_state.rs

@@ -109,7 +109,7 @@ impl StoredAssociationState {
             );
 
         let association_states =
-            conn.raw_query(|query_conn| query.load::<StoredAssociationState>(query_conn))?;
+            conn.raw_query_read(|query_conn| query.load::<StoredAssociationState>(query_conn))?;
 
         association_states
             .into_iter()

xmtp_mls/src/storage/encrypted_store/consent_record.rs

@@ -48,7 +48,7 @@ impl DbConnection {
         entity: String,
         entity_type: ConsentType,
     ) -> Result<Option<StoredConsentRecord>, StorageError> {
-        Ok(self.raw_query(|conn| -> diesel::QueryResult<_> {
+        Ok(self.raw_query_read(|conn| -> diesel::QueryResult<_> {
             dsl::consent_records
                 .filter(dsl::entity.eq(entity))
                 .filter(dsl::entity_type.eq(entity_type))
@@ -77,7 +77,7 @@ impl DbConnection {
             );
         }
 
-        let changed = self.raw_query(|conn| -> diesel::QueryResult<_> {
+        let changed = self.raw_query_write(|conn| -> diesel::QueryResult<_> {
             let existing: Vec<StoredConsentRecord> = query.load(conn)?;
             let changed: Vec<_> = records
                 .iter()
@@ -107,7 +107,7 @@ impl DbConnection {
         &self,
         record: &StoredConsentRecord,
     ) -> Result<Option<StoredConsentRecord>, StorageError> {
-        self.raw_query(|conn| {
+        self.raw_query_write(|conn| {
             let maybe_inserted_consent_record: Option<StoredConsentRecord> =
                 diesel::insert_into(dsl::consent_records)
                     .values(record)

xmtp_mls/src/storage/encrypted_store/conversation_list.rs

@@ -139,7 +139,7 @@ impl DbConnection {
                     .select(conversation_list::all_columns())
                     .order(conversation_list_dsl::created_at_ns.asc());
 
-                self.raw_query(|conn| query.load::<ConversationListItem>(conn))?
+                self.raw_query_read(|conn| query.load::<ConversationListItem>(conn))?
             } else {
                 // Only include the specified states
                 let query = query
@@ -153,19 +153,19 @@ impl DbConnection {
                     .select(conversation_list::all_columns())
                     .order(conversation_list_dsl::created_at_ns.asc());
 
-                self.raw_query(|conn| query.load::<ConversationListItem>(conn))?
+                self.raw_query_read(|conn| query.load::<ConversationListItem>(conn))?
             }
         } else {
             // Handle the case where `consent_states` is `None`
-            self.raw_query(|conn| query.load::<ConversationListItem>(conn))?
+            self.raw_query_read(|conn| query.load::<ConversationListItem>(conn))?
         };
 
         // Were sync groups explicitly asked for? Was the include_sync_groups flag set to true?
         // Then query for those separately
         if matches!(conversation_type, Some(ConversationType::Sync)) || *include_sync_groups {
             let query = conversation_list_dsl::conversation_list
                 .filter(conversation_list_dsl::conversation_type.eq(ConversationType::Sync));
-            let mut sync_groups = self.raw_query(|conn| query.load(conn))?;
+            let mut sync_groups = self.raw_query_read(|conn| query.load(conn))?;
             conversations.append(&mut sync_groups);
         }
 

xmtp_mls/src/storage/encrypted_store/db_connection.rs

@@ -1,8 +1,15 @@
+use crate::storage::{xmtp_openmls_provider::XmtpOpenMlsProvider, StorageError};
+use diesel::connection::TransactionManager;
 use parking_lot::Mutex;
-use std::fmt;
-use std::sync::Arc;
+use std::{
+    fmt,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+};
 
-use crate::storage::xmtp_openmls_provider::XmtpOpenMlsProvider;
+use super::XmtpDb;
 
 #[cfg(not(target_arch = "wasm32"))]
 pub type DbConnection = DbConnectionPrivate<super::RawDbConnection>;
@@ -18,38 +25,83 @@ pub type DbConnection = DbConnectionPrivate<sqlite_web::connection::WasmSqliteCo
 // Do not derive clone here.
 // callers should be able to accomplish everything with one conn/reference.
 #[doc(hidden)]
+#[derive(Clone)]
 pub struct DbConnectionPrivate<C> {
-    inner: Arc<Mutex<C>>,
+    read: Arc<Mutex<C>>,
+    write: Option<Arc<Mutex<C>>>,
+    // This field will funnel all reads / writes to the write connection if true.
+    pub(super) in_transaction: Arc<AtomicBool>,
 }
 
 /// Owned DBConnection Methods
 impl<C> DbConnectionPrivate<C> {
     /// Create a new [`DbConnectionPrivate`] from an existing Arc<Mutex<C>>
-    pub(super) fn from_arc_mutex(conn: Arc<Mutex<C>>) -> Self {
-        Self { inner: conn }
+    pub(super) fn from_arc_mutex(read: Arc<Mutex<C>>, write: Option<Arc<Mutex<C>>>) -> Self {
+        Self {
+            read,
+            write,
+            in_transaction: Arc::new(AtomicBool::new(false)),
+        }
     }
 }
 
 impl<C> DbConnectionPrivate<C>
 where
     C: diesel::Connection,
 {
+    pub(crate) fn start_transaction<Db: XmtpDb<Connection = C>>(
+        &self,
+    ) -> Result<TransactionGuard, StorageError> {
+        let mut write = self
+            .write
+            .as_ref()
+            .expect("Tried to open transaction on read-only connection")
+            .lock();
+        <Db as XmtpDb>::TransactionManager::begin_transaction(&mut *write)?;
+
+        if self.in_transaction.swap(true, Ordering::SeqCst) {
+            panic!("Already in transaction.");
+        }
+
+        Ok(TransactionGuard {
+            in_transaction: self.in_transaction.clone(),
+        })
+    }
+
+    fn in_transaction(&self) -> bool {
+        self.in_transaction.load(Ordering::SeqCst)
+    }
+
     /// Do a scoped query with a mutable [`diesel::Connection`]
     /// reference
-    pub(crate) fn raw_query<T, E, F>(&self, fun: F) -> Result<T, E>
+    pub(crate) fn raw_query_read<T, E, F>(&self, fun: F) -> Result<T, E>
     where
         F: FnOnce(&mut C) -> Result<T, E>,
     {
-        let mut lock = self.inner.lock();
+        if self.in_transaction() {
+            if let Some(write) = &self.write {
+                let mut lock = write.lock();
+                return fun(&mut lock);
+            };
+        }
+
+        let mut lock = self.read.lock();
         fun(&mut lock)
     }
 
-    /// Internal-only API to get the underlying `diesel::Connection` reference
-    /// without a scope
-    /// Must be used with care. holding this reference while calling `raw_query`
-    /// will cause a deadlock.
-    pub(super) fn inner_mut_ref(&self) -> parking_lot::MutexGuard<'_, C> {
-        self.inner.lock()
+    /// Do a scoped query with a mutable [`diesel::Connection`]
+    /// reference
+    pub(crate) fn raw_query_write<T, E, F>(&self, fun: F) -> Result<T, E>
+    where
+        F: FnOnce(&mut C) -> Result<T, E>,
+    {
+        if let Some(write_conn) = &self.write {
+            let mut lock = write_conn.lock();
+            return fun(&mut lock);
+        }
+
+        let mut lock = self.read.lock();
+        fun(&mut lock)
     }
 }
 
@@ -71,3 +123,12 @@ impl<C> fmt::Debug for DbConnectionPrivate<C> {
             .finish()
     }
 }
+
+pub struct TransactionGuard {
+    in_transaction: Arc<AtomicBool>,
+}
+impl Drop for TransactionGuard {
+    fn drop(&mut self) {
+        self.in_transaction.store(false, Ordering::SeqCst);
+    }
+}