xmtp_mls_core/
state_tx.rs

1//! Synchronous access to mutable MLS state under the database writer.
2
3use openmls::group::MlsGroup as OpenMlsGroup;
4use xmtp_db::{
5    NotFound, StorageError, TransactionOutcome, TransactionalKeyStore, XmtpMlsStorageProvider,
6    sql_key_store::SqlKeyStoreError,
7};
8use xmtp_events::{EventBuffer, EventBus};
9use xmtp_proto::types::GroupId;
10
11use crate::events::InternalEvent;
12
13/// A capability created only after an immediate write transaction starts.
14///
15/// Keep network requests and callbacks outside this capability. A group is
16/// loaded for one synchronous operation and is dropped before the transaction
17/// ends. A retry must enter a new transaction and load the group again.
18pub struct StateTx<'a, T> {
19    query: &'a mut T,
20}
21
22impl<T: TransactionalKeyStore> StateTx<'_, T> {
23    /// Load current state after the database writer has been acquired.
24    ///
25    /// The operation borrows the group. It cannot return the mutable group or
26    /// a future that borrows it. All storage access uses this transaction.
27    pub fn with_group<R, E>(
28        &mut self,
29        group_id: GroupId,
30        operation: impl FnOnce(&mut OpenMlsGroup, &T::Store<'_>) -> Result<R, E>,
31    ) -> Result<R, E>
32    where
33        E: From<StorageError> + From<SqlKeyStoreError>,
34    {
35        let storage = self.query.key_store();
36        let mut group = OpenMlsGroup::load(&storage, &group_id.to_openmls())?
37            .ok_or_else(|| StorageError::from(NotFound::MlsGroup(group_id)))?;
38        operation(&mut group, &storage)
39    }
40
41    /// Access keys and records for a state operation that does not load a group.
42    pub fn storage(&mut self) -> T::Store<'_> {
43        self.query.key_store()
44    }
45
46    /// Roll back a failed attempt while retaining the outer writer lock.
47    ///
48    /// The failed attempt's group is dropped before the caller can load state
49    /// to record a rejection. This also discards OpenMLS ratchet caches.
50    pub fn savepoint<R, E>(
51        &mut self,
52        operation: impl FnOnce(
53            &mut StateTx<'_, <T::Store<'_> as XmtpMlsStorageProvider>::TxQuery>,
54        ) -> Result<TransactionOutcome<R>, E>,
55    ) -> Result<TransactionOutcome<R>, E>
56    where
57        E: From<xmtp_db::diesel::result::Error>
58            + From<xmtp_db::ConnectionError>
59            + std::error::Error,
60    {
61        self.query
62            .key_store()
63            .savepoint(|query| operation(&mut StateTx { query }))
64    }
65}
66
67/// Acquire the cross-process writer before exposing state or keys.
68///
69/// An error rolls back all writes. An explicit rollback is useful for trials;
70/// the caller must keep only immutable requirements from such a trial.
71pub fn state_write<S, R, E>(
72    storage: &S,
73    operation: impl FnOnce(&mut StateTx<'_, S::TxQuery>) -> Result<TransactionOutcome<R>, E>,
74) -> Result<TransactionOutcome<R>, E>
75where
76    S: XmtpMlsStorageProvider,
77    E: From<xmtp_db::diesel::result::Error> + From<xmtp_db::ConnectionError> + std::error::Error,
78{
79    storage.transaction(|query| {
80        let outcome = operation(&mut StateTx { query })?;
81        #[cfg(all(any(test, feature = "test-hooks"), not(target_arch = "wasm32")))]
82        if matches!(outcome, TransactionOutcome::Continue(_)) && precommit_test_hook::is_set() {
83            use xmtp_db::ConnectionExt;
84            query.key_store().db().raw_query(|conn| {
85                precommit_test_hook::run(conn);
86                Ok(())
87            })?;
88        }
89        Ok(outcome)
90    })
91}
92
93/// Sync the write before returning queued work or publishing it.
94pub fn state_write_durable<S, R, E>(
95    storage: &S,
96    operation: impl FnOnce(&mut StateTx<'_, S::TxQuery>) -> Result<TransactionOutcome<R>, E>,
97) -> Result<TransactionOutcome<R>, E>
98where
99    S: XmtpMlsStorageProvider,
100    E: From<xmtp_db::diesel::result::Error> + From<xmtp_db::ConnectionError> + std::error::Error,
101{
102    storage.durable_transaction(|query| {
103        let outcome = operation(&mut StateTx { query })?;
104        #[cfg(all(any(test, feature = "test-hooks"), not(target_arch = "wasm32")))]
105        if matches!(outcome, TransactionOutcome::Continue(_)) && precommit_test_hook::is_set() {
106            use xmtp_db::ConnectionExt;
107            query.key_store().db().raw_query(|conn| {
108                precommit_test_hook::run(conn);
109                Ok(())
110            })?;
111        }
112        Ok(outcome)
113    })
114}
115
116/// Flush this write's events only after its database transaction commits.
117pub fn state_write_with_events<S, R, E>(
118    storage: &S,
119    bus: &EventBus<InternalEvent>,
120    operation: impl FnOnce(
121        &mut StateTx<'_, S::TxQuery>,
122        &EventBuffer<'_, InternalEvent>,
123    ) -> Result<TransactionOutcome<R>, E>,
124) -> Result<TransactionOutcome<R>, E>
125where
126    S: XmtpMlsStorageProvider,
127    E: From<xmtp_db::diesel::result::Error> + From<xmtp_db::ConnectionError> + std::error::Error,
128{
129    let mut result = None;
130    let _: Result<(), ()> = bus.with_buffer(|events| {
131        result = Some(state_write(storage, |tx| operation(tx, events)));
132        if matches!(&result, Some(Ok(TransactionOutcome::Continue(_)))) {
133            Ok(())
134        } else {
135            Err(())
136        }
137    });
138    result.expect("state write produced a result")
139}
140
141/// [`state_write_with_events`] with a durable transaction.
142pub fn state_write_with_events_durable<S, R, E>(
143    storage: &S,
144    bus: &EventBus<InternalEvent>,
145    operation: impl FnOnce(
146        &mut StateTx<'_, S::TxQuery>,
147        &EventBuffer<'_, InternalEvent>,
148    ) -> Result<TransactionOutcome<R>, E>,
149) -> Result<TransactionOutcome<R>, E>
150where
151    S: XmtpMlsStorageProvider,
152    E: From<xmtp_db::diesel::result::Error> + From<xmtp_db::ConnectionError> + std::error::Error,
153{
154    bus.with_buffer(
155        |events| match state_write_durable(storage, |tx| operation(tx, events)) {
156            result @ Ok(TransactionOutcome::Continue(_)) => Ok(result),
157            result => Err(result),
158        },
159    )
160    .unwrap_or_else(|result| result)
161}
162
163/// A test hook that runs on the writer's connection after a state write's
164/// operation succeeds and before its transaction commits.
165#[cfg(all(any(test, feature = "test-hooks"), not(target_arch = "wasm32")))]
166#[doc(hidden)]
167pub mod precommit_test_hook {
168    use std::{cell::RefCell, marker::PhantomData, rc::Rc};
169    use xmtp_db::diesel::SqliteConnection;
170
171    type Hook = Box<dyn FnMut(&mut SqliteConnection)>;
172
173    thread_local! {
174        static HOOK: RefCell<Option<Hook>> = const { RefCell::new(None) };
175    }
176
177    /// Keeps the hook installed on this thread until dropped, then restores
178    /// the previous one. The hook stays on the thread that owns the
179    /// synchronous state operation.
180    pub struct Guard {
181        previous: Option<Hook>,
182        _same_thread: PhantomData<Rc<()>>,
183    }
184
185    /// Runs `hook` before each state write on this thread commits, until the
186    /// returned guard drops.
187    pub fn install(hook: impl FnMut(&mut SqliteConnection) + 'static) -> Guard {
188        Guard {
189            previous: HOOK.with(|slot| slot.replace(Some(Box::new(hook)))),
190            _same_thread: PhantomData,
191        }
192    }
193
194    pub(super) fn is_set() -> bool {
195        HOOK.with(|slot| slot.borrow().is_some())
196    }
197
198    pub(super) fn run(conn: &mut SqliteConnection) {
199        HOOK.with(|slot| {
200            if let Some(hook) = slot.borrow_mut().as_mut() {
201                hook(conn);
202            }
203        });
204    }
205
206    impl Drop for Guard {
207        fn drop(&mut self) {
208            HOOK.with(|slot| {
209                slot.replace(self.previous.take());
210            });
211        }
212    }
213}
214
215#[cfg(test)]
216mod event_tests {
217    use super::*;
218    use crate::{context::CoreContext, events::emit_consent_changes, test_context::real_context};
219    use xmtp_db::{
220        consent_record::{ConsentKey, ConsentState, ConsentWrite, Decision, Owed, write_consent},
221        prelude::*,
222    };
223    use xmtp_events::{ClientEvent, EventFilter, EventKind};
224
225    // verifies: EVENT-010, EVENT-012
226    #[xmtp_common::test(unwrap_try = true)]
227    async fn event_write_flushes_after_commit_and_clears_on_rollback() {
228        let context = real_context().await;
229        let bus = context.events();
230        let subscription = bus.subscribe(EventFilter::new([EventKind::ConsentChanged]), Some(10));
231        #[cfg(not(target_arch = "wasm32"))]
232        let before_commit = bus.subscribe(EventFilter::new([EventKind::ConsentChanged]), Some(10));
233        let committed = ConsentKey::inbox("committed");
234        #[cfg(not(target_arch = "wasm32"))]
235        let fired = std::rc::Rc::new(std::cell::Cell::new(0));
236        #[cfg(not(target_arch = "wasm32"))]
237        let hook = precommit_test_hook::install({
238            let fired = fired.clone();
239            move |_| {
240                fired.set(fired.get() + 1);
241                assert!(before_commit.drain().is_empty());
242            }
243        });
244        state_write_with_events(context.mls_storage(), bus, |tx, events| {
245            let storage = tx.storage();
246            let db = storage.db();
247            let decide = ConsentWrite::Decide(Decision::Set(ConsentState::Allowed));
248            let changes: Vec<_> = write_consent(&db, &committed, decide, Owed::Mark)?
249                .into_iter()
250                .collect();
251            emit_consent_changes(&context, events, &db, &changes)?;
252            Ok::<_, StorageError>(TransactionOutcome::Continue(()))
253        })?;
254        #[cfg(not(target_arch = "wasm32"))]
255        drop(hook);
256        #[cfg(not(target_arch = "wasm32"))]
257        assert!(fired.get() > 0);
258        assert_eq!(
259            context
260                .db()
261                .consent_row(&committed)?
262                .map(|record| record.state),
263            Some(ConsentState::Allowed)
264        );
265        assert!(
266            matches!(subscription.drain().pop().and_then(|item| item.client), Some(ClientEvent::ConsentChanged(change)) if change.entity == "committed")
267        );
268
269        let rolled_back = ConsentKey::inbox("rolled_back");
270        state_write_with_events::<_, (), StorageError>(
271            context.mls_storage(),
272            bus,
273            |tx, events| {
274                let storage = tx.storage();
275                let db = storage.db();
276                let decide = ConsentWrite::Decide(Decision::Set(ConsentState::Denied));
277                let changes: Vec<_> = write_consent(&db, &rolled_back, decide, Owed::Mark)?
278                    .into_iter()
279                    .collect();
280                emit_consent_changes(&context, events, &db, &changes)?;
281                Ok(TransactionOutcome::Rollback)
282            },
283        )?;
284        assert!(context.db().consent_row(&rolled_back)?.is_none());
285        assert!(subscription.drain().is_empty());
286    }
287}
xmtp_mls_core/state_tx.rs

xmtp_mls_core/
state_tx.rs
