xmtp_mls_core/
worker.rs

1pub mod disappearing_messages;
2pub use crate::hmac_epoch;
3pub mod metrics;
4pub mod tasks;
5
6use crate::context::CoreContext;
7use crate::events::InternalEvent;
8use crate::metrics::SyncMetric;
9pub use crate::reconnect::{LocalNeedsDbReconnect, NeedsDbReconnect};
10use crate::task_channels::TaskWorkerChannels;
11use futures::future::{AbortHandle, Abortable};
12use futures::{StreamExt, stream::FuturesUnordered};
13use metrics::WorkerMetrics;
14use parking_lot::Mutex;
15use std::fmt::Debug;
16use std::pin::Pin;
17use std::{collections::HashMap, hash::Hash, sync::Arc};
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;
23use xmtp_events::{EventFilter, Subscription};
24
25/// Hard cap on how long `WorkerRunner::shutdown` waits for the supervisor
26/// task to drain after cancellation. Anything beyond this gets logged and
27/// the task is allowed to detach — keeps `Client::close` bounded.
28const WORKER_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(2);
29
30/// Worker runner controls for tests. Each assumes process-per-test isolation.
31#[cfg(any(test, feature = "test-hooks"))]
32#[doc(hidden)]
33pub mod test_hooks {
34    use std::sync::Arc;
35    use tokio::sync::Notify;
36
37    type PauseHook = (Vec<u8>, Arc<Notify>, Arc<Notify>);
38    static PAUSE_NEXT_SPAWN: parking_lot::Mutex<Option<PauseHook>> = parking_lot::Mutex::new(None);
39
40    /// Pauses the next worker spawn of `installation_id` once its identity is
41    /// ready and before any worker starts: notifies `entered`, then waits for
42    /// `release`.
43    pub fn pause_next_spawn(installation_id: Vec<u8>, entered: Arc<Notify>, release: Arc<Notify>) {
44        *PAUSE_NEXT_SPAWN.lock() = Some((installation_id, entered, release));
45    }
46
47    pub(super) fn take_spawn_pause(installation_id: &[u8]) -> Option<(Arc<Notify>, Arc<Notify>)> {
48        let mut hook = PAUSE_NEXT_SPAWN.lock();
49        if hook
50            .as_ref()
51            .is_some_and(|(installation, _, _)| installation.as_slice() == installation_id)
52        {
53            hook.take().map(|(_, entered, release)| (entered, release))
54        } else {
55            None
56        }
57    }
58}
59
60#[derive(PartialEq, Eq, Copy, Clone, Hash, Debug)]
61pub enum WorkerKind {
62    AttachmentCleanup,
63    DeviceSync,
64    DisappearingMessages,
65    KeyPackageCleaner,
66    CommitLog,
67    TaskRunner,
68    /// Re-reads what the backend publishes about itself, hourly.
69    ConfigurationRefresh,
70    HmacEpoch,
71}
72
73/// Buffer depth of a worker subscription that may drop events: each wake only
74/// makes the worker read its stored state again.
75pub(crate) const EVENT_BUFFER: usize = 10;
76
77/// Configuration for the cadence and enablement of background workers.
78///
79/// `Default` (all-`None`, empty maps) reproduces the historical behavior:
80/// every worker enabled, each using its own compiled-in interval const, no
81/// jitter. All fields are opt-in so existing callers are unaffected.
82///
83/// Durations are nanoseconds, matching [`crate::config::ForkRecoveryOpts`].
84#[derive(Clone, Debug, Default)]
85pub struct WorkerConfig {
86    /// Global fallback interval (ns) for any worker without a per-kind
87    /// override. `None` => each worker uses its own const default.
88    pub default_interval_ns: Option<u64>,
89    /// Per-worker interval override (ns). Wins over `default_interval_ns`.
90    pub interval_overrides: HashMap<WorkerKind, u64>,
91    /// Per-worker jitter (ns). An absent entry means `0` (deterministic).
92    /// Jitter de-synchronizes a fleet of clients; scoping it per-worker
93    /// avoids blanketing fast workers with a large jitter meant for a slow
94    /// one (e.g. the daily CommitLog worker).
95    pub jitter_overrides: HashMap<WorkerKind, u64>,
96    /// Per-worker enable flag. An absent entry means enabled.
97    pub enabled: HashMap<WorkerKind, bool>,
98}
99
100impl WorkerConfig {
101    /// Resolve `(base, jitter)` for a worker.
102    ///
103    /// Base precedence: per-kind override > global default > `const_default`.
104    /// A resolved base of `0` is clamped to `const_default` to avoid a
105    /// pathological busy-loop. Jitter is the per-kind `jitter_overrides` entry
106    /// (0 if absent).
107    pub fn interval(&self, kind: WorkerKind, const_default: Duration) -> (Duration, Duration) {
108        let base_ns = self
109            .interval_overrides
110            .get(&kind)
111            .copied()
112            .or(self.default_interval_ns);
113        let base = match base_ns {
114            Some(0) | None => const_default,
115            Some(ns) => Duration::from_nanos(ns),
116        };
117        let jitter = self
118            .jitter_overrides
119            .get(&kind)
120            .copied()
121            .map(Duration::from_nanos)
122            .unwrap_or(Duration::ZERO);
123        (base, jitter)
124    }
125
126    /// `true` unless an explicit `false` entry exists for `kind`.
127    pub fn worker_enabled(&self, kind: WorkerKind) -> bool {
128        self.enabled.get(&kind).copied().unwrap_or(true)
129    }
130}
131
132pub struct WorkerRunner {
133    // When this is cloned into the Context this is empty, so the Context and Client have different views
134    factories: Vec<DynFactory>,
135    sync_metrics: SyncMetricsSlot,
136    task_channels: TaskWorkerChannels,
137    handle: Mutex<Option<Box<dyn StreamHandle<StreamOutput = ()>>>>,
138    // Per-worker abort handles. `shutdown` calls `abort()` on each so the
139    // worker future is dropped at its next poll regardless of whether the
140    // outer loop observed the cancellation token. Belt to the token's
141    // suspenders for workers that don't yield to cancellation promptly.
142    abort_handles: Mutex<Vec<AbortHandle>>,
143    subscriptions: HashMap<WorkerKind, Arc<Subscription<InternalEvent>>>,
144}
145
146impl Default for WorkerRunner {
147    fn default() -> Self {
148        Self::new()
149    }
150}
151
152impl WorkerRunner {
153    pub fn new() -> Self {
154        Self {
155            factories: Vec::new(),
156            sync_metrics: SyncMetricsSlot::default(),
157            task_channels: TaskWorkerChannels::default(),
158            handle: Mutex::default(),
159            abort_handles: Mutex::default(),
160            subscriptions: HashMap::new(),
161        }
162    }
163
164    /// The slot the device-sync worker's metrics land in, shared with the context.
165    pub fn sync_metrics_slot(&self) -> &SyncMetricsSlot {
166        &self.sync_metrics
167    }
168
169    pub fn sync_metrics(&self) -> Option<Arc<WorkerMetrics<SyncMetric>>> {
170        self.sync_metrics.lock().clone()
171    }
172
173    pub fn task_channels(&self) -> &TaskWorkerChannels {
174        &self.task_channels
175    }
176
177    /// True while the supervisor handle is held; false after a successful
178    /// `shutdown` or before the first `spawn`. Used in tests and as a
179    /// cheap liveness check.
180    pub fn is_running(&self) -> bool {
181        self.handle.lock().is_some()
182    }
183
184    /// The kinds of every registered worker factory. Used in tests to assert
185    /// which workers a builder configuration registered.
186    #[cfg(any(test, feature = "test-utils"))]
187    pub fn registered_kinds(&self) -> Vec<WorkerKind> {
188        self.factories.iter().map(|f| f.kind()).collect()
189    }
190
191    /// The kinds of every worker that subscribed to events at registration.
192    #[cfg(any(test, feature = "test-utils"))]
193    pub fn subscribed_kinds(&self) -> Vec<WorkerKind> {
194        self.subscriptions.keys().copied().collect()
195    }
196
197    /// The event subscription the worker of `kind` receives, if it has one.
198    #[cfg(any(test, feature = "test-utils"))]
199    pub fn subscription_for_test(
200        &self,
201        kind: WorkerKind,
202    ) -> Option<Arc<Subscription<InternalEvent>>> {
203        self.subscriptions.get(&kind).cloned()
204    }
205}
206
207impl WorkerRunner {
208    pub fn register_new_worker<W: NewWorker<C>, C>(&mut self, ctx: C)
209    where
210        C: CoreContext + 'static,
211    {
212        let events = ctx.events().clone();
213        let factory = W::factory(ctx);
214        let kind = factory.kind();
215        if let Some((filter, depth)) = factory.event_interest() {
216            self.subscriptions
217                .insert(kind, Arc::new(events.subscribe(filter, depth)));
218        }
219        self.factories.push(Arc::new(factory))
220    }
221
222    pub fn spawn<C>(self: &Arc<Self>, ctx: C)
223    where
224        C: CoreContext + 'static,
225    {
226        let mut handle_lock = self.handle.lock();
227        if let Some(handle) = handle_lock.take() {
228            handle.abort_handle().end();
229        }
230        // Force-abort any prior worker futures still alive from a previous spawn.
231        for h in self.abort_handles.lock().drain(..) {
232            h.abort();
233        }
234
235        let this = self.clone();
236        let cancel = ctx.cancellation_token().clone();
237        let handle = xmtp_common::spawn(
238            None,
239            async move {
240                while !ctx.identity().is_ready() {
241                    tokio::select! {
242                        () = xmtp_common::time::sleep(Duration::from_millis(50)) => {},
243                        () = cancel.cancelled() => {
244                            this.close_subscriptions();
245                            return;
246                        }
247                    }
248                }
249
250                #[cfg(any(test, feature = "test-hooks"))]
251                {
252                    let paused = test_hooks::take_spawn_pause(&ctx.installation_id().to_vec());
253                    if let Some((entered, release)) = paused {
254                        entered.notify_one();
255                        release.notified().await;
256                    }
257                }
258
259                let mut futs = FuturesUnordered::new();
260                let mut new_handles = Vec::with_capacity(this.factories.len());
261
262                for factory in &this.factories {
263                    let mut worker = factory.create();
264                    if let Some(subscription) = this.subscriptions.get(&factory.kind()) {
265                        worker.set_subscription(subscription.clone());
266                    }
267
268                    if let Some(metrics) = worker.sync_metrics() {
269                        *this.sync_metrics.lock() = Some(metrics);
270                    }
271                    let (abort_handle, reg) = AbortHandle::new_pair();
272                    new_handles.push(abort_handle);
273                    futs.push(Abortable::new(
274                        xmtp_common::bind_task_hub(worker.spawn(cancel.clone())),
275                        reg,
276                    ));
277                }
278                *this.abort_handles.lock() = new_handles;
279
280                while let Some(outcome) = futs.next().await {
281                    match outcome {
282                        Ok(kind) => tracing::warn!("Worker {kind:?} completed unexpectedly"),
283                        Err(_aborted) => tracing::debug!("worker aborted during shutdown"),
284                    }
285                }
286                if cancel.is_cancelled() {
287                    this.close_subscriptions();
288                }
289            }
290            .instrument(tracing::debug_span!("xmtp_worker_supervisor")),
291        );
292
293        *handle_lock = Some(Box::new(handle));
294    }
295
296    /// Drain the running worker supervisor. Bounded by [`WORKER_SHUTDOWN_TIMEOUT`].
297    /// Cancellation must be signalled separately (via the shared
298    /// `CancellationToken` on the context); this method additionally aborts
299    /// each worker's future to guarantee no further DB writes happen,
300    /// independent of whether the worker's loop respects the token.
301    pub async fn shutdown(&self) {
302        // Hard-kill each worker future at its next poll. Critical for workers
303        // that sit on long sleeps and only check the token between intervals.
304        for h in self.abort_handles.lock().drain(..) {
305            h.abort();
306        }
307        let mut handle = self.handle.lock().take();
308        let Some(handle) = handle.as_mut() else {
309            self.close_subscriptions();
310            return;
311        };
312        match xmtp_common::time::timeout(WORKER_SHUTDOWN_TIMEOUT, handle.end_and_wait()).await {
313            Ok(Ok(())) => {}
314            Ok(Err(e)) => tracing::debug!("worker supervisor ended with: {e:?}"),
315            Err(_) => tracing::warn!(
316                "worker supervisor did not drain within {:?}; abandoning",
317                WORKER_SHUTDOWN_TIMEOUT
318            ),
319        }
320        self.close_subscriptions();
321    }
322
323    fn close_subscriptions(&self) {
324        for subscription in self.subscriptions.values() {
325            subscription.close();
326        }
327    }
328
329    pub async fn wait_for_sync_worker_init(&self) {
330        let handle = self.sync_metrics.lock().clone();
331        if let Some(handle) = handle {
332            let _ = handle.wait_for_init().await;
333        }
334    }
335}
336
337pub type WorkerResult<T> = Result<T, Box<dyn NeedsDbReconnect>>;
338if_native! {
339    type SpawnWorkerFut = dyn Future<Output = WorkerKind> + Send;
340}
341if_wasm! {
342    type SpawnWorkerFut = dyn Future<Output = WorkerKind>;
343}
344
345#[xmtp_common::async_trait]
346pub trait Worker: MaybeSend + MaybeSync + 'static {
347    fn kind(&self) -> WorkerKind;
348
349    fn set_subscription(&mut self, _subscription: Arc<Subscription<InternalEvent>>) {}
350
351    async fn run_tasks(&mut self) -> Result<(), Box<dyn NeedsDbReconnect>>;
352
353    /// The metrics of the device-sync worker; no other worker has any.
354    fn sync_metrics(&self) -> Option<Arc<WorkerMetrics<SyncMetric>>> {
355        None
356    }
357
358    /// Box the worker, erasing its type
359    fn boxed(self) -> Box<dyn Worker>
360    where
361        Self: Sized,
362    {
363        Box::new(self) as Box<_>
364    }
365
366    // Wrap the outer loop (not each `run_tasks` impl) so individual workers
367    // observe cancellation by having their in-flight future dropped at the
368    // next await point — no per-impl plumbing required.
369    fn spawn(
370        mut self: Box<Self>,
371        cancel: CancellationToken,
372    ) -> Instrumented<Pin<Box<SpawnWorkerFut>>> {
373        let kind_str = format!("{:?}", self.kind());
374        let fut = Box::pin(async move {
375            let kind = self.kind();
376            let worker = format!("{:?}", kind);
377            let run = async move {
378                loop {
379                    // No span here: `run_tasks` runs until worker death, so a
380                    // wrapping span exports only on restart with an hours-long
381                    // duration. Per-tick visibility comes from `worker_turn`;
382                    // restarts from the structured logs below.
383                    let outcome = self.run_tasks().await;
384                    if let Err(err) = outcome {
385                        if err.needs_db_reconnect() {
386                            // drop the worker
387                            tracing::debug!("pool disconnected. task will restart on reconnect");
388                            break;
389                        } else {
390                            tracing::error!(worker = %worker, "{:?} worker error: {}", kind, err);
391                            // A restart resends the refused work, so it waits
392                            // out a rate-limit delay first. A delay too large
393                            // for the clock saturates instead of overflowing.
394                            // implements: RATE-041
395                            let delay = err.restart_delay().map_or(WORKER_RESTART_DELAY, |delay| {
396                                delay.max(WORKER_RESTART_DELAY)
397                            });
398                            let now = xmtp_common::time::Instant::now();
399                            let until = xmtp_common::time::saturating_after(now, delay);
400                            xmtp_common::time::sleep(until.saturating_duration_since(now)).await;
401                            tracing::info!(worker = %worker, "Restarting {:?} worker...", kind);
402                        }
403                    }
404                }
405                self.kind()
406            };
407
408            tokio::select! {
409                k = run => k,
410                _ = cancel.cancelled() => {
411                    tracing::debug!("{:?} worker cancelled", kind);
412                    kind
413                }
414            }
415        }) as Pin<Box<SpawnWorkerFut>>;
416        fut.instrument(tracing::debug_span!("libxmtp_worker", kind = kind_str))
417    }
418}
419
420/// A worker that can be built on contexts of type `C`.
421pub trait NewWorker<C> {
422    /// The factory the runner rebuilds this worker with after a restart.
423    fn factory(context: C) -> impl WorkerFactory + 'static;
424}
425
426pub trait WorkerFactory: MaybeSend + MaybeSync {
427    fn kind(&self) -> WorkerKind;
428
429    /// The events that wake the worker: public event kinds plus an internal
430    /// predicate, and the subscription's buffer depth (`None` drops nothing).
431    /// `None` when no event wakes it.
432    fn event_interest(&self) -> Option<(EventFilter<InternalEvent>, Option<usize>)> {
433        None
434    }
435
436    /// Create a new worker
437    fn create(&self) -> BoxedWorker;
438}
439
440pub type BoxedWorker = Box<dyn Worker>;
441pub type DynFactory = Arc<dyn WorkerFactory>;
442
443/// Where the device-sync worker's metrics land once the runner first builds
444/// that worker. Empty while no device-sync worker has started.
445pub type SyncMetricsSlot = Arc<Mutex<Option<Arc<WorkerMetrics<SyncMetric>>>>>;
446
447#[cfg(test)]
448mod worker_config_tests {
449    use super::{WorkerConfig, WorkerKind};
450    use std::time::Duration;
451
452    #[xmtp_common::test]
453    fn default_is_all_enabled_no_overrides() {
454        let cfg = WorkerConfig::default();
455        assert!(cfg.worker_enabled(WorkerKind::DeviceSync));
456        assert!(cfg.worker_enabled(WorkerKind::CommitLog));
457        let (base, jitter) = cfg.interval(WorkerKind::DisappearingMessages, Duration::from_secs(7));
458        assert_eq!(base, Duration::from_secs(7), "falls back to const default");
459        assert_eq!(jitter, Duration::ZERO, "no jitter by default");
460    }
461
462    #[xmtp_common::test]
463    fn per_kind_override_beats_global_default() {
464        let mut cfg = WorkerConfig {
465            default_interval_ns: Some(Duration::from_secs(3).as_nanos() as u64),
466            ..Default::default()
467        };
468        cfg.interval_overrides.insert(
469            WorkerKind::KeyPackageCleaner,
470            Duration::from_secs(9).as_nanos() as u64,
471        );
472        let (base, _) = cfg.interval(WorkerKind::KeyPackageCleaner, Duration::from_secs(99));
473        assert_eq!(base, Duration::from_secs(9), "per-kind override wins");
474        let (other, _) = cfg.interval(WorkerKind::CommitLog, Duration::from_secs(99));
475        assert_eq!(
476            other,
477            Duration::from_secs(3),
478            "global default for un-overridden worker"
479        );
480    }
481
482    #[xmtp_common::test]
483    fn zero_resolved_base_clamps_to_const() {
484        let mut cfg = WorkerConfig::default();
485        cfg.interval_overrides.insert(WorkerKind::CommitLog, 0);
486        let (base, _) = cfg.interval(WorkerKind::CommitLog, Duration::from_secs(60));
487        assert_eq!(
488            base,
489            Duration::from_secs(60),
490            "zero base clamps to const default"
491        );
492    }
493
494    #[xmtp_common::test]
495    fn per_kind_jitter_is_carried() {
496        let mut cfg = WorkerConfig::default();
497        cfg.jitter_overrides.insert(
498            WorkerKind::TaskRunner,
499            Duration::from_secs(2).as_nanos() as u64,
500        );
501        let (_, jitter) = cfg.interval(WorkerKind::TaskRunner, Duration::from_secs(1));
502        assert_eq!(jitter, Duration::from_secs(2));
503    }
504
505    #[xmtp_common::test]
506    fn jitter_is_scoped_per_worker() {
507        let mut cfg = WorkerConfig::default();
508        cfg.jitter_overrides.insert(
509            WorkerKind::CommitLog,
510            Duration::from_secs(6 * 3600).as_nanos() as u64,
511        );
512        // The jittered worker gets its jitter...
513        let (_, commit_log_jitter) = cfg.interval(WorkerKind::CommitLog, Duration::from_secs(300));
514        assert_eq!(commit_log_jitter, Duration::from_secs(6 * 3600));
515        // ...while an un-listed worker stays deterministic.
516        let (_, disappearing_jitter) =
517            cfg.interval(WorkerKind::DisappearingMessages, Duration::from_secs(1));
518        assert_eq!(disappearing_jitter, Duration::ZERO);
519    }
520
521    #[xmtp_common::test]
522    fn disabled_entry_reports_false() {
523        let mut cfg = WorkerConfig::default();
524        cfg.enabled.insert(WorkerKind::DeviceSync, false);
525        assert!(!cfg.worker_enabled(WorkerKind::DeviceSync));
526        assert!(
527            cfg.worker_enabled(WorkerKind::CommitLog),
528            "absent => enabled"
529        );
530    }
531}
xmtp_mls_core/worker.rs

xmtp_mls_core/
worker.rs
