xmtp_db/sql_key_store/
transactions.rs

1use super::*;
2use crate::DbConnection;
3use crate::TransactionOutcome;
4use crate::TransactionOutcome::{Continue, Rollback};
5
6/// wrapper around a mutable connection (&mut SqliteConnection)
7/// Requires that all execution/transaction happens in one thread on one connection.
8/// This connection _must only_ be created from starting a transaction
9pub struct MutableTransactionConnection<'a> {
10    // we cannot avoid interior mutability here
11    // because raw_query methods require &self, as do MlsStorage trait methods.
12    // Since we no longer have async transactions, once a transaction is started
13    // we can ensure it occurs all on one thread.
14    pub(crate) conn: parking_lot::Mutex<&'a mut SqliteConnection>,
15}
16
17impl<'a> MutableTransactionConnection<'a> {
18    pub fn new(conn: &'a mut SqliteConnection) -> Self {
19        Self {
20            conn: parking_lot::Mutex::new(conn),
21        }
22    }
23}
24
25impl<'a> ConnectionExt for MutableTransactionConnection<'a> {
26    fn raw_query<T, F>(&self, fun: F) -> Result<T, crate::ConnectionError>
27    where
28        F: FnOnce(&mut SqliteConnection) -> Result<T, diesel::result::Error>,
29        Self: Sized,
30    {
31        let mut conn = self.conn.try_lock().expect("Lock is held somewhere else");
32        fun(&mut conn).map_err(crate::ConnectionError::from)
33    }
34
35    // this should cause a transaction rollback. since reconnect/disconnect is retryable
36    fn disconnect(&self) -> Result<(), crate::ConnectionError> {
37        Err(crate::ConnectionError::DisconnectInTransaction)
38    }
39
40    fn reconnect(&self) -> Result<(), crate::ConnectionError> {
41        Err(crate::ConnectionError::ReconnectInTransaction)
42    }
43}
44
45impl<C: ConnectionExt> XmtpMlsStorageProvider for SqlKeyStore<C> {
46    type Connection = C;
47
48    type TxQuery = SqliteConnection;
49
50    type DbQuery<'a>
51        = DbConnection<&'a C>
52    where
53        Self::Connection: 'a;
54
55    fn db<'a>(&'a self) -> Self::DbQuery<'a> {
56        DbConnection::new(&self.conn)
57    }
58
59    #[xmtp_common::db_span]
60    fn transaction<T, E, F>(&self, f: F) -> Result<TransactionOutcome<T>, E>
61    where
62        F: FnOnce(&mut Self::TxQuery) -> Result<TransactionOutcome<T>, E>,
63        E: From<diesel::result::Error> + From<crate::ConnectionError> + std::error::Error,
64    {
65        let conn = &self.conn;
66
67        // immediate transactions force SQLite to respect BUSY_TIMEOUT
68        // there are a few ways we can get DB Locked Errors:
69        // 1.) A Transaction is already writing
70        //  https://www.sqlite.org/rescode.html#busy
71        // 2.) Promoting a transaction to write:
72        // we start a transaction with BEGIN (read), then later promote the transaction to a write.
73        // another transaction is already writing, so SQLite throws Database Locked.
74        // code: https://www.sqlite.org/rescode.html#busy_snapshot
75        // Solution:
76        // - set BUSY_TIMEOUT. this is effectively a timeout for SQLite to get a lock on the
77        //      write to a table. See [BUSY_TIMEOUT](xmtp_db::configuration::BUSY_TIMEOUT)
78        // - use immediate_transaction to force SQLite to respect busy_timeout as soon as the
79        //      transaction starts. Otherwise, we still run into problem #2, even if BUSY_TIMEOUT is
80        //      set.
81
82        // An intentional `Rollback` is turned into diesel's `RollbackTransaction`
83        // sentinel (the only way to make diesel roll back) and flagged, so below we
84        // can report it as `Ok(Rollback)` without inspecting the opaque `E`; any
85        // other `Err` is a real failure and propagates.
86        let mut rolled_back = false;
87        let inner_result: Result<TransactionOutcome<T>, E> = conn
88            .raw_query(|c| {
89                Ok(c.immediate_transaction(|sqlite_c| match f(sqlite_c) {
90                    Ok(Continue(v)) => Ok(Continue(v)),
91                    Ok(Rollback) => {
92                        rolled_back = true;
93                        Err(E::from(diesel::result::Error::RollbackTransaction))
94                    }
95                    Err(e) => Err(e),
96                }))
97            })
98            .map_err(E::from)?;
99
100        // A failed ROLLBACK after our sentinel is still reported as Ok(Rollback);
101        // that rare rollback-execution error is deliberately swallowed.
102        match inner_result {
103            Ok(outcome) => Ok(outcome),
104            Err(_) if rolled_back => Ok(Rollback),
105            Err(e) => Err(e),
106        }
107    }
108
109    // Same Rollback-sentinel handling as `transaction`; see there for the rationale.
110    fn savepoint<T, E, F>(&self, f: F) -> Result<TransactionOutcome<T>, E>
111    where
112        F: FnOnce(&mut Self::TxQuery) -> Result<TransactionOutcome<T>, E>,
113        E: From<diesel::result::Error> + From<crate::ConnectionError> + std::error::Error,
114    {
115        let mut rolled_back = false;
116        let inner_result: Result<TransactionOutcome<T>, E> = self
117            .conn
118            .raw_query(|c| {
119                Ok(c.transaction(|sqlite_c| match f(sqlite_c) {
120                    Ok(Continue(v)) => Ok(Continue(v)),
121                    Ok(Rollback) => {
122                        rolled_back = true;
123                        Err(E::from(diesel::result::Error::RollbackTransaction))
124                    }
125                    Err(e) => Err(e),
126                }))
127            })
128            .map_err(E::from)?;
129
130        match inner_result {
131            Ok(outcome) => Ok(outcome),
132            Err(_) if rolled_back => Ok(Rollback),
133            Err(e) => Err(e),
134        }
135    }
136
137    fn read<V: Entity<CURRENT_VERSION>>(
138        &self,
139        label: &[u8],
140        key: &[u8],
141    ) -> Result<Option<V>, SqlKeyStoreError> {
142        self.read(label, key)
143    }
144
145    fn read_list<V: Entity<CURRENT_VERSION>>(
146        &self,
147        label: &[u8],
148        key: &[u8],
149    ) -> Result<Vec<V>, <Self as StorageProvider<CURRENT_VERSION>>::Error> {
150        self.read_list(label, key)
151    }
152
153    fn delete(
154        &self,
155        label: &[u8],
156        key: &[u8],
157    ) -> Result<(), <Self as StorageProvider<CURRENT_VERSION>>::Error> {
158        self.delete::<CURRENT_VERSION>(label, key)
159    }
160
161    fn write(
162        &self,
163        label: &[u8],
164        key: &[u8],
165        value: &[u8],
166    ) -> Result<(), <Self as StorageProvider<CURRENT_VERSION>>::Error> {
167        self.write::<CURRENT_VERSION>(label, key, value)
168    }
169
170    #[cfg(feature = "test-utils")]
171    fn hash_all(&self) -> Result<Vec<u8>, SqlKeyStoreError> {
172        self.conn
173            .raw_query(OpenMlsKeyValue::hash_all)
174            .map_err(Into::into)
175    }
176}
177
178#[cfg(test)]
179mod tests {
180
181    #![allow(unused)]
182
183    use crate::{
184        TestDb, TransactionOutcome, XmtpTestDb,
185        group_intent::{IntentKind, IntentState, NewGroupIntent},
186        prelude::QueryGroupIntent,
187    };
188    use xmtp_proto::types::GroupId;
189
190    use super::*;
191
192    // Test to ensure that we can use the transaction() callback without requiring a 'static
193    // lifetimes
194    // This ensures we do not propagate 'static throughout all of our code.
195    // have not figured out a good, ergonomic way to pass SqlKeyStore directly into the
196    // transaction callback
197    struct Foo<C> {
198        key_store: SqlKeyStore<C>,
199    }
200
201    impl<C> Foo<C>
202    where
203        C: ConnectionExt,
204    {
205        async fn long_async_call(&self) {
206            xmtp_common::time::sleep(std::time::Duration::from_millis(10)).await;
207        }
208
209        async fn db_op(&self) {
210            self.long_async_call().await;
211
212            self.key_store
213                .transaction(|conn| {
214                    let storage = conn.key_store();
215                    storage
216                        .db()
217                        .insert_group_intent(NewGroupIntent {
218                            kind: IntentKind::SendMessage,
219                            group_id: GroupId::default(),
220                            data: vec![],
221                            should_push: false,
222                            state: IntentState::ToPublish,
223                        })
224                        .map(Continue)
225                })
226                .unwrap();
227            self.long_async_call().await;
228        }
229    }
230}
xmtp_db/sql_key_store/transactions.rs

xmtp_db/sql_key_store/
transactions.rs
