xmtp_mls/
context.rs

1#[cfg(test)]
2use crate::GroupCommitLock;
3use crate::builder::{DeviceSyncMode, ForkRecoveryOpts};
4use crate::client::DeviceSync;
5use crate::groups::change_callbacks::UnstableChangeCallbacks;
6use crate::server_configuration::ServerConfigurationHandle;
7use crate::subscriptions::{LocalEvents, SyncWorkerEvent};
8use crate::utils::VersionInfo;
9use crate::worker::device_sync::worker::SyncMetric;
10use crate::worker::disappearing_messages::DisappearingChannels;
11use crate::worker::metrics::WorkerMetrics;
12use crate::worker::tasks::TaskWorkerChannels;
13use crate::worker::{DynMetrics, MetricsCasting, WorkerConfig, WorkerKind};
14use crate::{
15    identity::{Identity, IdentityError},
16    mutex_registry::MutexRegistry,
17};
18use parking_lot::Mutex;
19use std::collections::HashMap;
20use std::sync::Arc;
21use std::sync::atomic::{AtomicBool, Ordering};
22use tokio::sync::broadcast;
23use tokio_util::sync::CancellationToken;
24use xmtp_api::{ApiClientWrapper, XmtpApi};
25use xmtp_common::{MaybeSend, MaybeSync};
26use xmtp_db::XmtpDb;
27use xmtp_db::XmtpMlsStorageProvider;
28use xmtp_db::xmtp_openmls_provider::XmtpOpenMlsProviderRef;
29use xmtp_id::scw_verifier::SmartContractSignatureVerifier;
30use xmtp_id::{InboxIdRef, associations::builder::SignatureRequest};
31use xmtp_proto::types::InstallationId;
32
33#[cfg(any(test, feature = "test-utils"))]
34use crate::worker::device_sync::DeviceSyncClient;
35
36/// The local context a XMTP MLS needs to function:
37/// - Sqlite Database
38/// - Identity for the User
39pub struct XmtpMlsLocalContext<ApiClient, Db, S> {
40    /// XMTP Identity
41    pub(crate) identity: Identity,
42    /// The XMTP Api Client
43    pub(crate) api_client: ApiClientWrapper<ApiClient>,
44    /// XMTP Local Storage
45    pub(crate) store: Db,
46    pub(crate) mls_storage: S,
47    pub(crate) mutexes: MutexRegistry,
48    #[cfg(test)]
49    pub(crate) mls_commit_lock: Arc<GroupCommitLock>,
50    pub(crate) version_info: VersionInfo,
51    /// What this deployment published about itself, resolved once at build
52    /// (spec 006 CFG-030), plus the latch a refresh may set.
53    pub(crate) server_configuration: ServerConfigurationHandle,
54    pub(crate) local_events: broadcast::Sender<LocalEvents>,
55    pub(crate) delivery_owner: Arc<Mutex<Option<xmtp_db::delivery::DeliveryOwner>>>,
56    pub(crate) worker_events: broadcast::Sender<SyncWorkerEvent>,
57    pub(crate) scw_verifier: Arc<Box<dyn SmartContractSignatureVerifier>>,
58    pub(crate) device_sync: DeviceSync,
59    pub(crate) fork_recovery_opts: ForkRecoveryOpts,
60    /// Unstable: SDK-registered notifications for group-state changes. Empty
61    /// unless the host opted in at build time.
62    pub(crate) change_callbacks: UnstableChangeCallbacks,
63    pub(crate) incoming_runtime: Arc<crate::subscriptions::incoming::IncomingRuntime>,
64    pub(crate) identity_resolutions: Arc<crate::identity_updates::IdentityResolutionRegistry>,
65    pub(crate) worker_config: WorkerConfig,
66    // pub(crate) workers: Arc<WorkerRunner>,
67    pub(crate) worker_metrics: Arc<Mutex<HashMap<WorkerKind, DynMetrics>>>,
68    pub(crate) task_channels: TaskWorkerChannels,
69    pub(crate) disappearing_channels: DisappearingChannels,
70    pub(crate) cancellation_token: CancellationToken,
71    // Set only after a successful `Client::close` (workers stopped + DB
72    // disconnected). The cancellation token tracks "shutdown initiated";
73    // this tracks "shutdown completed cleanly" — distinct semantics so
74    // a `disconnect()` failure mid-`close` doesn't silently short-circuit
75    // future retries.
76    pub(crate) shutdown_complete: Arc<AtomicBool>,
77}
78
79impl<ApiClient, Db, S> XmtpMlsLocalContext<ApiClient, Db, S>
80where
81    Db: XmtpDb + 'static,
82    ApiClient: XmtpApi + 'static,
83    S: XmtpMlsStorageProvider + 'static,
84{
85    /// get a reference to the monolithic Database object where
86    /// higher-level queries are defined
87    pub fn db(&self) -> Db::DbQuery {
88        self.store.db()
89    }
90
91    /// Creates a new MLS Provider
92    pub fn mls_provider(&'_ self) -> XmtpOpenMlsProviderRef<'_, S> {
93        XmtpOpenMlsProviderRef::new(&self.mls_storage)
94    }
95
96    pub fn store(&self) -> &Db {
97        &self.store
98    }
99
100    pub fn scw_verifier(&self) -> &Arc<Box<dyn SmartContractSignatureVerifier>> {
101        &self.scw_verifier
102    }
103
104    pub fn device_sync_worker_enabled(&self) -> bool {
105        !matches!(self.device_sync.mode, DeviceSyncMode::Disabled)
106    }
107
108    /// Reconstructs the DeviceSyncClient from the context
109    /// used in tests
110    #[cfg(any(test, feature = "test-utils"))]
111    pub fn device_sync_client(
112        self: &Arc<XmtpMlsLocalContext<ApiClient, Db, S>>,
113    ) -> DeviceSyncClient<Arc<Self>> {
114        let metrics = self.sync_metrics();
115        DeviceSyncClient::new(
116            Arc::clone(self),
117            metrics.unwrap_or(Arc::new(WorkerMetrics::new(self.installation_id()))),
118        )
119    }
120}
121
122impl<ApiClient, Db, S> XmtpMlsLocalContext<ApiClient, Db, S> {
123    pub fn replace_mls_store<S2>(self, mls_store: S2) -> XmtpMlsLocalContext<ApiClient, Db, S2> {
124        XmtpMlsLocalContext::<ApiClient, Db, S2> {
125            identity: self.identity,
126            api_client: self.api_client,
127            store: self.store,
128            mls_storage: mls_store,
129            mutexes: self.mutexes,
130            #[cfg(test)]
131            mls_commit_lock: self.mls_commit_lock,
132            version_info: self.version_info,
133            server_configuration: self.server_configuration,
134            local_events: self.local_events,
135            delivery_owner: self.delivery_owner,
136            worker_events: self.worker_events,
137            scw_verifier: self.scw_verifier,
138            device_sync: self.device_sync,
139            fork_recovery_opts: self.fork_recovery_opts,
140            change_callbacks: self.change_callbacks,
141            incoming_runtime: self.incoming_runtime,
142            identity_resolutions: self.identity_resolutions,
143            worker_config: self.worker_config,
144            worker_metrics: self.worker_metrics,
145            task_channels: self.task_channels,
146            disappearing_channels: self.disappearing_channels,
147            cancellation_token: self.cancellation_token,
148            shutdown_complete: self.shutdown_complete,
149        }
150    }
151}
152
153impl<ApiClient, Db, S> XmtpMlsLocalContext<ApiClient, Db, S> {
154    /// The installation public key is the primary identifier for an installation
155    pub fn installation_public_key(&self) -> InstallationId {
156        (*self.identity.installation_keys.public_bytes()).into()
157    }
158
159    /// The installation public key is the primary identifier for an installation
160    pub fn installation_id(&self) -> InstallationId {
161        self.identity.installation_id()
162    }
163
164    /// Get the account address of the blockchain account associated with this client
165    pub fn inbox_id(&self) -> InboxIdRef<'_> {
166        self.identity.inbox_id()
167    }
168
169    /// Integrators should always check the `signature_request` return value of this function before calling `register_identity`.
170    /// If `signature_request` returns `None`, then the wallet signature is not required and `register_identity` can be called with None as an argument.
171    pub fn signature_request(&self) -> Option<SignatureRequest> {
172        self.identity.signature_request()
173    }
174
175    pub fn sign_with_public_context(
176        &self,
177        text: impl AsRef<str>,
178    ) -> Result<Vec<u8>, IdentityError> {
179        self.identity.sign_with_public_context(text)
180    }
181
182    #[cfg(test)]
183    pub fn mls_commit_lock(&self) -> &Arc<GroupCommitLock> {
184        &self.mls_commit_lock
185    }
186
187    pub fn sync_metrics(&self) -> Option<Arc<WorkerMetrics<SyncMetric>>> {
188        self.worker_metrics
189            .lock()
190            .get(&WorkerKind::DeviceSync)?
191            .as_sync_metrics()
192    }
193
194    pub fn cancellation_token(&self) -> &CancellationToken {
195        &self.cancellation_token
196    }
197
198    pub fn shutdown_complete(&self) -> bool {
199        self.shutdown_complete.load(Ordering::Acquire)
200    }
201
202    pub fn mark_shutdown_complete(&self) {
203        self.shutdown_complete.store(true, Ordering::Release);
204    }
205}
206
207pub trait XmtpSharedContext
208where
209    Self: MaybeSend + MaybeSync + Sized + Clone,
210{
211    type Db: XmtpDb;
212    type ApiClient: XmtpApi;
213    type MlsStorage: XmtpMlsStorageProvider;
214    /// Owned, cloneable handle for background work; it cannot borrow short-lived resources.
215    type ContextReference: XmtpSharedContext<Db = Self::Db, ApiClient = Self::ApiClient, MlsStorage = Self::MlsStorage>
216        + 'static;
217
218    /// Return the handle to clone when work can outlive this context borrow.
219    fn context_ref(&self) -> &Self::ContextReference;
220    fn db(&self) -> <Self::Db as XmtpDb>::DbQuery;
221    fn api(&self) -> &ApiClientWrapper<Self::ApiClient>;
222    fn scw_verifier(&self) -> Arc<Box<dyn SmartContractSignatureVerifier>>;
223
224    fn device_sync(&self) -> &DeviceSync;
225
226    fn device_sync_worker_enabled(&self) -> bool {
227        !matches!(self.device_sync().mode, DeviceSyncMode::Disabled)
228    }
229
230    fn fork_recovery_opts(&self) -> &ForkRecoveryOpts;
231
232    fn worker_config(&self) -> &WorkerConfig;
233
234    /// Resolve `(base, jitter)` for a worker from the configured
235    /// [`WorkerConfig`], falling back to the worker's compiled-in const.
236    fn worker_interval(
237        &self,
238        kind: WorkerKind,
239        const_default: std::time::Duration,
240    ) -> (std::time::Duration, std::time::Duration) {
241        self.worker_config().interval(kind, const_default)
242    }
243
244    /// Creates a new MLS Provider
245    fn mls_provider(&'_ self) -> XmtpOpenMlsProviderRef<'_, Self::MlsStorage> {
246        XmtpOpenMlsProviderRef::new(self.mls_storage())
247    }
248
249    fn mls_storage(&self) -> &Self::MlsStorage;
250    fn identity(&self) -> &Identity;
251
252    fn signature_request(&self) -> Option<SignatureRequest> {
253        self.identity().signature_request()
254    }
255
256    fn inbox_id(&self) -> InboxIdRef<'_> {
257        self.identity().inbox_id()
258    }
259
260    fn installation_id(&self) -> InstallationId {
261        (*self.identity().installation_keys.public_bytes()).into()
262    }
263
264    fn version_info(&self) -> &VersionInfo;
265    /// The configuration snapshot every consumer in spec 006 section 6.4 reads.
266    fn server_configuration(&self) -> &ServerConfigurationHandle;
267    fn worker_events(&self) -> &broadcast::Sender<SyncWorkerEvent>;
268    fn local_events(&self) -> &broadcast::Sender<LocalEvents>;
269    /// This context's default-consumer token; the database is the ownership authority.
270    fn delivery_owner(&self) -> &Mutex<Option<xmtp_db::delivery::DeliveryOwner>>;
271
272    /// Release this context's message consumer before disconnecting its database.
273    fn close_message_delivery(&self) -> Result<(), xmtp_db::StorageError> {
274        use xmtp_db::delivery::QueryDelivery;
275        let mut registered = self.delivery_owner().lock();
276        if let Some(owner) = *registered {
277            self.db().release_delivery_owner(owner)?;
278            *registered = None;
279        }
280        Ok(())
281    }
282    fn task_channels(&self) -> &TaskWorkerChannels;
283    fn disappearing_channels(&self) -> &DisappearingChannels;
284    /// Unstable: the host's registered group-change callbacks.
285    fn change_callbacks(&self) -> &UnstableChangeCallbacks;
286    /// Shared incoming runtime. Its limits and transport are internal client policy.
287    fn incoming_runtime(&self) -> &crate::subscriptions::incoming::IncomingRuntime;
288    /// Coalesces exact identity lookups without treating a newer snapshot as the requested one.
289    fn identity_resolution_registry(&self) -> &crate::identity_updates::IdentityResolutionRegistry;
290    fn sync_metrics(&self) -> Option<Arc<WorkerMetrics<SyncMetric>>>;
291    #[cfg(test)]
292    fn mls_commit_lock(&self) -> &Arc<GroupCommitLock>;
293    fn mutexes(&self) -> &MutexRegistry;
294    fn cancellation_token(&self) -> &CancellationToken;
295
296    /// Returns `true` once `Client::close` has been called on this context.
297    fn is_closed(&self) -> bool {
298        self.cancellation_token().is_cancelled()
299    }
300
301    /// Returns `true` only after `Client::close` has fully torn the client
302    /// down (workers drained + DB disconnected). Distinct from [`Self::is_closed`]
303    /// which fires the moment shutdown begins — this guards `close`'s
304    /// idempotency check so a mid-shutdown failure stays retryable.
305    fn shutdown_complete(&self) -> bool;
306
307    fn mark_shutdown_complete(&self);
308}
309
310impl<XApiClient, XDb, XMls> XmtpSharedContext for Arc<XmtpMlsLocalContext<XApiClient, XDb, XMls>>
311where
312    XApiClient: XmtpApi + 'static,
313    XDb: XmtpDb + 'static,
314    XMls: XmtpMlsStorageProvider + 'static,
315{
316    type Db = XDb;
317    type ApiClient = XApiClient;
318    type MlsStorage = XMls;
319    type ContextReference = Self;
320
321    fn context_ref(&self) -> &Self::ContextReference {
322        self
323    }
324
325    fn db(&self) -> <Self::Db as XmtpDb>::DbQuery {
326        self.store.db()
327    }
328
329    fn api(&self) -> &ApiClientWrapper<Self::ApiClient> {
330        &self.api_client
331    }
332
333    fn scw_verifier(&self) -> Arc<Box<dyn SmartContractSignatureVerifier>> {
334        self.scw_verifier.clone()
335    }
336
337    fn device_sync(&self) -> &DeviceSync {
338        &self.device_sync
339    }
340
341    fn fork_recovery_opts(&self) -> &ForkRecoveryOpts {
342        &self.fork_recovery_opts
343    }
344
345    fn worker_config(&self) -> &WorkerConfig {
346        &self.worker_config
347    }
348
349    /// a reference to the MLS Storage Type
350    /// This can be related to 'db()' but may also be separate
351    fn mls_storage(&self) -> &Self::MlsStorage {
352        &self.mls_storage
353    }
354
355    fn identity(&self) -> &Identity {
356        &self.identity
357    }
358
359    fn version_info(&self) -> &VersionInfo {
360        &self.version_info
361    }
362
363    fn server_configuration(&self) -> &ServerConfigurationHandle {
364        &self.server_configuration
365    }
366
367    fn worker_events(&self) -> &broadcast::Sender<SyncWorkerEvent> {
368        &self.worker_events
369    }
370
371    fn local_events(&self) -> &broadcast::Sender<LocalEvents> {
372        &self.local_events
373    }
374
375    fn delivery_owner(&self) -> &Mutex<Option<xmtp_db::delivery::DeliveryOwner>> {
376        &self.delivery_owner
377    }
378
379    #[cfg(test)]
380    fn mls_commit_lock(&self) -> &Arc<GroupCommitLock> {
381        &self.mls_commit_lock
382    }
383
384    fn task_channels(&self) -> &TaskWorkerChannels {
385        &self.task_channels
386    }
387
388    fn disappearing_channels(&self) -> &DisappearingChannels {
389        &self.disappearing_channels
390    }
391
392    fn change_callbacks(&self) -> &UnstableChangeCallbacks {
393        &self.change_callbacks
394    }
395
396    fn incoming_runtime(&self) -> &crate::subscriptions::incoming::IncomingRuntime {
397        &self.incoming_runtime
398    }
399
400    fn identity_resolution_registry(&self) -> &crate::identity_updates::IdentityResolutionRegistry {
401        &self.identity_resolutions
402    }
403
404    fn sync_metrics(&self) -> Option<Arc<WorkerMetrics<SyncMetric>>> {
405        self.worker_metrics
406            .lock()
407            .get(&WorkerKind::DeviceSync)?
408            .as_sync_metrics()
409    }
410
411    fn mutexes(&self) -> &MutexRegistry {
412        &self.mutexes
413    }
414
415    fn cancellation_token(&self) -> &CancellationToken {
416        &self.cancellation_token
417    }
418
419    fn shutdown_complete(&self) -> bool {
420        self.shutdown_complete.load(Ordering::Acquire)
421    }
422
423    fn mark_shutdown_complete(&self) {
424        self.shutdown_complete.store(true, Ordering::Release);
425    }
426}
427
428impl<T> XmtpSharedContext for &T
429where
430    T: XmtpSharedContext,
431{
432    type Db = <T as XmtpSharedContext>::Db;
433    type ApiClient = <T as XmtpSharedContext>::ApiClient;
434    type MlsStorage = <T as XmtpSharedContext>::MlsStorage;
435    type ContextReference = <T as XmtpSharedContext>::ContextReference;
436
437    fn context_ref(&self) -> &Self::ContextReference {
438        <T as XmtpSharedContext>::context_ref(self)
439    }
440
441    fn db(&self) -> <Self::Db as XmtpDb>::DbQuery {
442        <T as XmtpSharedContext>::db(self)
443    }
444
445    fn api(&self) -> &ApiClientWrapper<Self::ApiClient> {
446        <T as XmtpSharedContext>::api(self)
447    }
448
449    fn scw_verifier(&self) -> Arc<Box<dyn SmartContractSignatureVerifier>> {
450        <T as XmtpSharedContext>::scw_verifier(self)
451    }
452
453    fn device_sync(&self) -> &DeviceSync {
454        <T as XmtpSharedContext>::device_sync(self)
455    }
456
457    fn device_sync_worker_enabled(&self) -> bool {
458        <T as XmtpSharedContext>::device_sync_worker_enabled(self)
459    }
460
461    fn fork_recovery_opts(&self) -> &ForkRecoveryOpts {
462        <T as XmtpSharedContext>::fork_recovery_opts(self)
463    }
464
465    fn worker_config(&self) -> &WorkerConfig {
466        <T as XmtpSharedContext>::worker_config(self)
467    }
468
469    fn mls_storage(&self) -> &Self::MlsStorage {
470        <T as XmtpSharedContext>::mls_storage(self)
471    }
472
473    fn identity(&self) -> &Identity {
474        <T as XmtpSharedContext>::identity(self)
475    }
476
477    fn version_info(&self) -> &VersionInfo {
478        <T as XmtpSharedContext>::version_info(self)
479    }
480
481    fn server_configuration(&self) -> &ServerConfigurationHandle {
482        <T as XmtpSharedContext>::server_configuration(self)
483    }
484
485    fn worker_events(&self) -> &broadcast::Sender<SyncWorkerEvent> {
486        <T as XmtpSharedContext>::worker_events(self)
487    }
488
489    fn local_events(&self) -> &broadcast::Sender<LocalEvents> {
490        <T as XmtpSharedContext>::local_events(self)
491    }
492
493    fn delivery_owner(&self) -> &Mutex<Option<xmtp_db::delivery::DeliveryOwner>> {
494        <T as XmtpSharedContext>::delivery_owner(self)
495    }
496
497    #[cfg(test)]
498    fn mls_commit_lock(&self) -> &Arc<GroupCommitLock> {
499        <T as XmtpSharedContext>::mls_commit_lock(self)
500    }
501
502    fn task_channels(&self) -> &TaskWorkerChannels {
503        <T as XmtpSharedContext>::task_channels(self)
504    }
505
506    fn disappearing_channels(&self) -> &DisappearingChannels {
507        <T as XmtpSharedContext>::disappearing_channels(self)
508    }
509
510    fn change_callbacks(&self) -> &UnstableChangeCallbacks {
511        <T as XmtpSharedContext>::change_callbacks(self)
512    }
513
514    fn incoming_runtime(&self) -> &crate::subscriptions::incoming::IncomingRuntime {
515        <T as XmtpSharedContext>::incoming_runtime(self)
516    }
517
518    fn identity_resolution_registry(&self) -> &crate::identity_updates::IdentityResolutionRegistry {
519        <T as XmtpSharedContext>::identity_resolution_registry(self)
520    }
521
522    fn sync_metrics(&self) -> Option<Arc<WorkerMetrics<SyncMetric>>> {
523        <T as XmtpSharedContext>::sync_metrics(self)
524    }
525
526    fn mutexes(&self) -> &MutexRegistry {
527        <T as XmtpSharedContext>::mutexes(self)
528    }
529
530    fn cancellation_token(&self) -> &CancellationToken {
531        <T as XmtpSharedContext>::cancellation_token(self)
532    }
533
534    fn shutdown_complete(&self) -> bool {
535        <T as XmtpSharedContext>::shutdown_complete(self)
536    }
537
538    fn mark_shutdown_complete(&self) {
539        <T as XmtpSharedContext>::mark_shutdown_complete(self)
540    }
541}
xmtp_mls/context.rs

xmtp_mls/
context.rs
