1pub mod device_sync;
2pub mod disappearing_messages;
3pub mod key_package_maintenance;
4pub mod metrics;
5pub(crate) mod notifications;
6pub mod tasks;
7
8use crate::context::XmtpSharedContext;
9use device_sync::worker::SyncMetric;
10use futures::future::{AbortHandle, Abortable};
11use futures::{StreamExt, stream::FuturesUnordered};
12use metrics::WorkerMetrics;
13use parking_lot::Mutex;
14use std::fmt::Debug;
15use std::pin::Pin;
16use std::{any::Any, collections::HashMap, hash::Hash, sync::Arc};
17use tasks::TaskWorkerChannels;
18use tokio_util::sync::CancellationToken;
19use tracing::Instrument;
20use tracing::instrument::Instrumented;
21use xmtp_common::{MaybeSend, MaybeSync, StreamHandle, if_native, if_wasm, time::Duration};
22use xmtp_configuration::WORKER_RESTART_DELAY;
23
24const WORKER_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(2);
28
29#[derive(PartialEq, Eq, Copy, Clone, Hash, Debug)]
30pub enum WorkerKind {
31 DeviceSync,
32 DisappearingMessages,
33 KeyPackageCleaner,
34 CommitLog,
35 TaskRunner,
36 ConfigurationRefresh,
38}
39
40#[derive(Clone, Debug, Default)]
48pub struct WorkerConfig {
49 pub default_interval_ns: Option<u64>,
52 pub interval_overrides: HashMap<WorkerKind, u64>,
54 pub jitter_overrides: HashMap<WorkerKind, u64>,
59 pub enabled: HashMap<WorkerKind, bool>,
61}
62
63impl WorkerConfig {
64 pub fn interval(&self, kind: WorkerKind, const_default: Duration) -> (Duration, Duration) {
71 let base_ns = self
72 .interval_overrides
73 .get(&kind)
74 .copied()
75 .or(self.default_interval_ns);
76 let base = match base_ns {
77 Some(0) | None => const_default,
78 Some(ns) => Duration::from_nanos(ns),
79 };
80 let jitter = self
81 .jitter_overrides
82 .get(&kind)
83 .copied()
84 .map(Duration::from_nanos)
85 .unwrap_or(Duration::ZERO);
86 (base, jitter)
87 }
88
89 pub fn worker_enabled(&self, kind: WorkerKind) -> bool {
91 self.enabled.get(&kind).copied().unwrap_or(true)
92 }
93}
94
95pub struct WorkerRunner {
96 factories: Vec<DynFactory>,
98 metrics: Arc<Mutex<HashMap<WorkerKind, DynMetrics>>>,
99 task_channels: TaskWorkerChannels,
100 handle: Mutex<Option<Box<dyn StreamHandle<StreamOutput = ()>>>>,
101 abort_handles: Mutex<Vec<AbortHandle>>,
106}
107
108impl Default for WorkerRunner {
109 fn default() -> Self {
110 Self::new()
111 }
112}
113
114impl WorkerRunner {
115 pub fn new() -> Self {
116 Self {
117 factories: Vec::new(),
118 metrics: Arc::default(),
119 task_channels: TaskWorkerChannels::default(),
120 handle: Mutex::default(),
121 abort_handles: Mutex::default(),
122 }
123 }
124
125 pub fn metrics(&self) -> &Arc<Mutex<HashMap<WorkerKind, DynMetrics>>> {
126 &self.metrics
127 }
128
129 pub fn sync_metrics(&self) -> Option<Arc<WorkerMetrics<SyncMetric>>> {
130 self.metrics
131 .lock()
132 .get(&WorkerKind::DeviceSync)?
133 .as_sync_metrics()
134 }
135
136 pub fn task_channels(&self) -> &TaskWorkerChannels {
137 &self.task_channels
138 }
139
140 pub fn is_running(&self) -> bool {
144 self.handle.lock().is_some()
145 }
146
147 #[cfg(test)]
150 pub(crate) fn registered_kinds(&self) -> Vec<WorkerKind> {
151 self.factories.iter().map(|f| f.kind()).collect()
152 }
153}
154
155impl WorkerRunner {
156 pub fn register_new_worker<W: Worker, C>(&mut self, ctx: C)
157 where
158 C: XmtpSharedContext + 'static,
159 {
160 let factory = W::factory(ctx);
161 self.factories.push(Arc::new(factory))
162 }
163
164 pub fn spawn<C>(self: &Arc<Self>, ctx: C)
165 where
166 C: XmtpSharedContext + 'static,
167 {
168 let mut handle_lock = self.handle.lock();
169 if let Some(handle) = handle_lock.take() {
170 handle.abort_handle().end();
171 }
172 for h in self.abort_handles.lock().drain(..) {
174 h.abort();
175 }
176
177 let this = self.clone();
178 let cancel = ctx.cancellation_token().clone();
179 let handle = xmtp_common::spawn(
180 None,
181 async move {
182 while !ctx.identity().is_ready() {
183 xmtp_common::time::sleep(Duration::from_millis(50)).await;
184 }
185
186 let mut futs = FuturesUnordered::new();
187 let mut new_handles = Vec::with_capacity(this.factories.len());
188
189 for factory in &this.factories {
190 let metric = this.metrics.lock().get(&factory.kind()).cloned();
191 let (worker, metrics) = factory.create(metric);
192
193 if let Some(metrics) = metrics {
194 this.metrics.lock().insert(factory.kind(), metrics);
195 }
196
197 if let Some(metrics) = worker.metrics() {
198 let mut m = this.metrics.lock();
199 m.insert(worker.kind(), metrics);
200 }
201 let (abort_handle, reg) = AbortHandle::new_pair();
202 new_handles.push(abort_handle);
203 futs.push(Abortable::new(
204 xmtp_common::bind_task_hub(worker.spawn(cancel.clone())),
205 reg,
206 ));
207 }
208 *this.abort_handles.lock() = new_handles;
209
210 while let Some(outcome) = futs.next().await {
211 match outcome {
212 Ok(kind) => tracing::warn!("Worker {kind:?} completed unexpectedly"),
213 Err(_aborted) => tracing::debug!("worker aborted during shutdown"),
214 }
215 }
216 }
217 .instrument(tracing::debug_span!("xmtp_worker_supervisor")),
218 );
219
220 *handle_lock = Some(Box::new(handle));
221 }
222
223 pub async fn shutdown(&self) {
229 for h in self.abort_handles.lock().drain(..) {
232 h.abort();
233 }
234 let mut handle = self.handle.lock().take();
235 let Some(handle) = handle.as_mut() else {
236 return;
237 };
238 match xmtp_common::time::timeout(WORKER_SHUTDOWN_TIMEOUT, handle.end_and_wait()).await {
239 Ok(Ok(())) => {}
240 Ok(Err(e)) => tracing::debug!("worker supervisor ended with: {e:?}"),
241 Err(_) => tracing::warn!(
242 "worker supervisor did not drain within {:?}; abandoning",
243 WORKER_SHUTDOWN_TIMEOUT
244 ),
245 }
246 }
247
248 pub async fn wait_for_sync_worker_init(&self) {
249 let handle = self
250 .metrics
251 .lock()
252 .get(&WorkerKind::DeviceSync)
253 .cloned()
254 .and_then(|h| h.as_sync_metrics());
255 if let Some(handle) = handle {
256 let _ = handle.wait_for_init().await;
257 }
258 }
259}
260
261pub type WorkerResult<T> = Result<T, Box<dyn NeedsDbReconnect>>;
262if_native! {
263 type SpawnWorkerFut = dyn Future<Output = WorkerKind> + Send;
264}
265if_wasm! {
266 type SpawnWorkerFut = dyn Future<Output = WorkerKind>;
267}
268
269#[xmtp_common::async_trait]
270pub trait Worker: MaybeSend + MaybeSync + 'static {
271 fn kind(&self) -> WorkerKind;
272
273 async fn run_tasks(&mut self) -> Result<(), Box<dyn NeedsDbReconnect>>;
274
275 fn metrics(&self) -> Option<DynMetrics> {
276 None
277 }
278
279 fn factory<C>(context: C) -> impl WorkerFactory + 'static
280 where
281 Self: Sized,
282 C: XmtpSharedContext + 'static;
283
284 fn boxed(self) -> Box<dyn Worker>
286 where
287 Self: Sized,
288 {
289 Box::new(self) as Box<_>
290 }
291
292 fn spawn(
296 mut self: Box<Self>,
297 cancel: CancellationToken,
298 ) -> Instrumented<Pin<Box<SpawnWorkerFut>>> {
299 let kind_str = format!("{:?}", self.kind());
300 let fut = Box::pin(async move {
301 let kind = self.kind();
302 let worker = format!("{:?}", kind);
303 let run = async move {
304 loop {
305 let outcome = self.run_tasks().await;
310 if let Err(err) = outcome {
311 if err.needs_db_reconnect() {
312 tracing::debug!("pool disconnected. task will restart on reconnect");
314 break;
315 } else {
316 tracing::error!(worker = %worker, "{:?} worker error: {}", kind, err);
317 xmtp_common::time::sleep(WORKER_RESTART_DELAY).await;
318 tracing::info!(worker = %worker, "Restarting {:?} worker...", kind);
319 }
320 }
321 }
322 self.kind()
323 };
324
325 tokio::select! {
326 k = run => k,
327 _ = cancel.cancelled() => {
328 tracing::debug!("{:?} worker cancelled", kind);
329 kind
330 }
331 }
332 }) as Pin<Box<SpawnWorkerFut>>;
333 fut.instrument(tracing::debug_span!("libxmtp_worker", kind = kind_str))
334 }
335}
336
337#[cfg_attr(not(target_arch = "wasm32"), trait_variant::make(NeedsDbReconnect: Send + Sync))]
338#[cfg_attr(target_arch = "wasm32", trait_variant::make(NeedsDbReconnect: xmtp_common::Wasm))]
339pub trait LocalNeedsDbReconnect: std::error::Error {
340 fn needs_db_reconnect(&self) -> bool;
341}
342
343pub trait WorkerFactory: MaybeSend + MaybeSync {
344 fn kind(&self) -> WorkerKind;
345 fn create(&self, metrics: Option<DynMetrics>) -> (BoxedWorker, Option<DynMetrics>);
347}
348
349pub type BoxedWorker = Box<dyn Worker>;
350pub type DynFactory = Arc<dyn WorkerFactory>;
351
352pub type DynMetrics = Arc<dyn Any + Send + Sync>;
353
354pub trait MetricsCasting {
355 fn as_sync_metrics(&self) -> Option<Arc<WorkerMetrics<SyncMetric>>>;
356}
357
358impl MetricsCasting for DynMetrics {
359 fn as_sync_metrics(&self) -> Option<Arc<WorkerMetrics<SyncMetric>>> {
360 self.clone().downcast().ok()
361 }
362}
363
364#[cfg(all(test, not(target_arch = "wasm32")))]
367mod disconnect_propagation_tests {
368 use super::NeedsDbReconnect;
371 use crate::groups::GroupError;
372 use crate::groups::commit_log::CommitLogError;
373 use crate::mls_store::MlsStoreError;
374 use crate::subscriptions::SubscribeError;
375 use crate::worker::device_sync::DeviceSyncError;
376 use crate::worker::key_package_maintenance::KeyPackageMaintenanceError;
377 use xmtp_db::{ConnectionError, PlatformStorageError, StorageError};
378
379 fn disconnect_storage() -> StorageError {
381 StorageError::Platform(PlatformStorageError::PoolNeedsConnection)
382 }
383
384 fn disconnect_connection() -> ConnectionError {
386 ConnectionError::Platform(PlatformStorageError::PoolNeedsConnection)
387 }
388
389 fn benign_storage() -> StorageError {
391 StorageError::InvalidHmacLength
392 }
393
394 #[xmtp_common::test]
395 fn group_error_forwards_disconnect() {
396 use crate::subscriptions::barrier::{
397 BarrierCause, BarrierError, BarrierFailure, BarrierTopic,
398 };
399 use crate::subscriptions::incoming::IncomingError;
400 use std::sync::Arc;
401 use xmtp_proto::types::{Cursor, Topic};
402
403 assert!(GroupError::Storage(disconnect_storage()).needs_db_reconnect());
404 assert!(GroupError::Db(disconnect_connection()).needs_db_reconnect());
405 assert!(
406 GroupError::MlsStore(MlsStoreError::Connection(disconnect_connection()))
407 .needs_db_reconnect()
408 );
409 assert!(!GroupError::Storage(benign_storage()).needs_db_reconnect());
411 assert!(!GroupError::InvalidGroupMembership.needs_db_reconnect());
412
413 let barrier = |cause| BarrierError::Incomplete {
414 reason: BarrierFailure::Blocked,
415 unfinished: vec![BarrierTopic {
416 topic: Topic::new_group_message([1; 32]),
417 target: Some(Cursor(3)),
418 received: Cursor(2),
419 processed: Cursor(1),
420 unresolved_welcomes: Vec::new(),
421 inactive: false,
422 cause: Some(cause),
423 }],
424 };
425 assert!(
426 GroupError::StreamBarrier(barrier(BarrierCause::Storage(Arc::new(
427 disconnect_storage(),
428 ))))
429 .needs_db_reconnect()
430 );
431 assert!(
432 GroupError::PublishedButUnconfirmed {
433 intent_id: 1,
434 cause: Some(Box::new(barrier(BarrierCause::Receiver(Arc::new(
435 IncomingError::Processing(
436 crate::groups::mls_sync::GroupMessageProcessingError::Storage(
437 disconnect_storage(),
438 ),
439 ),
440 ))))),
441 }
442 .needs_db_reconnect()
443 );
444 assert!(
445 GroupError::Sync(Box::new(crate::groups::summary::SyncSummary::other(
446 GroupError::Storage(disconnect_storage()),
447 )))
448 .needs_db_reconnect()
449 );
450 assert!(
451 !GroupError::StreamBarrier(barrier(BarrierCause::Blocked("unsupported".into())))
452 .needs_db_reconnect()
453 );
454 }
455
456 #[xmtp_common::test]
457 fn mls_store_error_forwards_disconnect() {
458 assert!(MlsStoreError::Storage(disconnect_storage()).needs_db_reconnect());
459 assert!(MlsStoreError::Connection(disconnect_connection()).needs_db_reconnect());
460 assert!(!MlsStoreError::Storage(benign_storage()).needs_db_reconnect());
461 }
462
463 #[xmtp_common::test]
464 fn subscribe_error_forwards_disconnect() {
465 assert!(SubscribeError::Storage(disconnect_storage()).needs_db_reconnect());
466 assert!(SubscribeError::Db(disconnect_connection()).needs_db_reconnect());
467 assert!(
468 SubscribeError::from(GroupError::Storage(disconnect_storage())).needs_db_reconnect()
469 );
470 assert!(!SubscribeError::Storage(benign_storage()).needs_db_reconnect());
471 }
472
473 #[xmtp_common::test]
476 fn task_worker_load_group_forwards_disconnect() {
477 use crate::worker::tasks::TaskWorkerError;
478 assert!(
481 TaskWorkerError::LoadGroup(MlsStoreError::Connection(disconnect_connection()))
482 .needs_db_reconnect()
483 );
484 assert!(
485 TaskWorkerError::Group(GroupError::Storage(disconnect_storage())).needs_db_reconnect()
486 );
487 assert!(!TaskWorkerError::Group(GroupError::InvalidGroupMembership).needs_db_reconnect());
488 }
489
490 #[xmtp_common::test]
491 fn commit_log_error_forwards_disconnect() {
492 assert!(CommitLogError::Connection(disconnect_connection()).needs_db_reconnect());
493 assert!(
494 CommitLogError::GroupError(GroupError::Storage(disconnect_storage()))
495 .needs_db_reconnect()
496 );
497 assert!(
499 !CommitLogError::Connection(ConnectionError::DisconnectInTransaction)
500 .needs_db_reconnect()
501 );
502 }
503
504 #[xmtp_common::test]
505 fn device_sync_error_forwards_disconnect() {
506 assert!(DeviceSyncError::Storage(disconnect_storage()).needs_db_reconnect());
507 assert!(DeviceSyncError::Db(disconnect_connection()).needs_db_reconnect());
508 assert!(
509 DeviceSyncError::Group(GroupError::Storage(disconnect_storage())).needs_db_reconnect()
510 );
511 assert!(
512 DeviceSyncError::MlsStore(MlsStoreError::Connection(disconnect_connection()))
513 .needs_db_reconnect()
514 );
515 assert!(
516 DeviceSyncError::Subscribe(SubscribeError::Db(disconnect_connection()))
517 .needs_db_reconnect()
518 );
519 assert!(!DeviceSyncError::Storage(benign_storage()).needs_db_reconnect());
520 assert!(!DeviceSyncError::InvalidPayload.needs_db_reconnect());
521 }
522
523 #[xmtp_common::test]
524 fn key_package_maintenance_error_forwards_disconnect() {
525 use crate::identity::IdentityError;
526 assert!(KeyPackageMaintenanceError::Storage(disconnect_storage()).needs_db_reconnect());
527 assert!(
530 KeyPackageMaintenanceError::Identity(IdentityError::StorageError(disconnect_storage()))
531 .needs_db_reconnect()
532 );
533 assert!(
534 KeyPackageMaintenanceError::Identity(IdentityError::Db(disconnect_connection()))
535 .needs_db_reconnect()
536 );
537 assert!(!KeyPackageMaintenanceError::Storage(benign_storage()).needs_db_reconnect());
538 }
539}
540
541#[cfg(test)]
542mod worker_config_tests {
543 use super::{WorkerConfig, WorkerKind};
544 use std::time::Duration;
545
546 #[xmtp_common::test]
547 fn default_is_all_enabled_no_overrides() {
548 let cfg = WorkerConfig::default();
549 assert!(cfg.worker_enabled(WorkerKind::DeviceSync));
550 assert!(cfg.worker_enabled(WorkerKind::CommitLog));
551 let (base, jitter) = cfg.interval(WorkerKind::DisappearingMessages, Duration::from_secs(7));
552 assert_eq!(base, Duration::from_secs(7), "falls back to const default");
553 assert_eq!(jitter, Duration::ZERO, "no jitter by default");
554 }
555
556 #[xmtp_common::test]
557 fn per_kind_override_beats_global_default() {
558 let mut cfg = WorkerConfig {
559 default_interval_ns: Some(Duration::from_secs(3).as_nanos() as u64),
560 ..Default::default()
561 };
562 cfg.interval_overrides.insert(
563 WorkerKind::KeyPackageCleaner,
564 Duration::from_secs(9).as_nanos() as u64,
565 );
566 let (base, _) = cfg.interval(WorkerKind::KeyPackageCleaner, Duration::from_secs(99));
567 assert_eq!(base, Duration::from_secs(9), "per-kind override wins");
568 let (other, _) = cfg.interval(WorkerKind::CommitLog, Duration::from_secs(99));
569 assert_eq!(
570 other,
571 Duration::from_secs(3),
572 "global default for un-overridden worker"
573 );
574 }
575
576 #[xmtp_common::test]
577 fn zero_resolved_base_clamps_to_const() {
578 let mut cfg = WorkerConfig::default();
579 cfg.interval_overrides.insert(WorkerKind::CommitLog, 0);
580 let (base, _) = cfg.interval(WorkerKind::CommitLog, Duration::from_secs(60));
581 assert_eq!(
582 base,
583 Duration::from_secs(60),
584 "zero base clamps to const default"
585 );
586 }
587
588 #[xmtp_common::test]
589 fn per_kind_jitter_is_carried() {
590 let mut cfg = WorkerConfig::default();
591 cfg.jitter_overrides.insert(
592 WorkerKind::TaskRunner,
593 Duration::from_secs(2).as_nanos() as u64,
594 );
595 let (_, jitter) = cfg.interval(WorkerKind::TaskRunner, Duration::from_secs(1));
596 assert_eq!(jitter, Duration::from_secs(2));
597 }
598
599 #[xmtp_common::test]
600 fn jitter_is_scoped_per_worker() {
601 let mut cfg = WorkerConfig::default();
602 cfg.jitter_overrides.insert(
603 WorkerKind::CommitLog,
604 Duration::from_secs(6 * 3600).as_nanos() as u64,
605 );
606 let (_, commit_log_jitter) = cfg.interval(WorkerKind::CommitLog, Duration::from_secs(300));
608 assert_eq!(commit_log_jitter, Duration::from_secs(6 * 3600));
609 let (_, disappearing_jitter) =
611 cfg.interval(WorkerKind::DisappearingMessages, Duration::from_secs(1));
612 assert_eq!(disappearing_jitter, Duration::ZERO);
613 }
614
615 #[xmtp_common::test]
616 fn disabled_entry_reports_false() {
617 let mut cfg = WorkerConfig::default();
618 cfg.enabled.insert(WorkerKind::DeviceSync, false);
619 assert!(!cfg.worker_enabled(WorkerKind::DeviceSync));
620 assert!(
621 cfg.worker_enabled(WorkerKind::CommitLog),
622 "absent => enabled"
623 );
624 }
625}