xmtp_mls/worker/
tasks.rs

1use crate::{
2    context::XmtpSharedContext,
3    subscriptions::internal::InternalEvent,
4    worker::{
5        NeedsDbReconnect, Worker, WorkerFactory, WorkerKind, device_sync::DeviceSyncError,
6        key_package_maintenance as kp,
7    },
8};
9use prost::Message;
10use std::sync::Arc;
11use xmtp_configuration::KEY_PACKAGE_ROTATION_INTERVAL_NS;
12use xmtp_db::prelude::{QueryIdentity, QueryKeyPackageHistory};
13use xmtp_db::tasks::{QueryTasks, Task as DbTask, TaskDataHash};
14use xmtp_db::{StorageError, diesel};
15use xmtp_events::Subscription;
16use xmtp_proto::xmtp::mls::database::Task as TaskProto;
17
18/// How far out to push a task whose kind this build does not understand. Long
19/// enough that an old client sharing the database does not spin on it, short
20/// enough that a newer client picks it up soon after it starts.
21const UNKNOWN_TASK_DEFER_NS: i64 = 60 * 60 * 1_000_000_000;
22
23/// `Done` = one-shot, row deleted. `RescheduleAt(ns)` = recurring, row kept and
24/// advanced to that absolute deadline.
25#[derive(Debug, PartialEq, Eq)]
26pub(crate) enum TaskOutcome {
27    Done,
28    RescheduleAt(i64),
29}
30
31#[cfg(test)]
32pub(crate) mod test_hooks {
33    use std::sync::Mutex;
34    /// `(target_data_hash, deadline)` → `run_task` returns `RescheduleAt(deadline)`
35    /// for the matching task. Reset at test end; assumes process-per-test isolation.
36    pub(crate) static RESCHEDULE_OVERRIDE: Mutex<Option<(Vec<u8>, i64)>> = Mutex::new(None);
37    pub(crate) static TASK_SCHEDULED_CONSUMED: Mutex<
38        Option<(Vec<u8>, tokio::sync::oneshot::Sender<()>)>,
39    > = Mutex::new(None);
40}
41
42#[derive(thiserror::Error, Debug)]
43pub enum TaskWorkerError {
44    #[error("generic storage error: {0}")]
45    Storage(#[from] xmtp_db::StorageError),
46    #[error("group error: {0}")]
47    Group(#[from] crate::groups::GroupError),
48    #[error("device sync error: {0}")]
49    DeviceSync(#[from] DeviceSyncError),
50    #[error("failed to load MLS group from store: {0}")]
51    LoadGroup(#[from] crate::mls_store::MlsStoreError),
52    #[error("invalid task data for {id}: {error}")]
53    InvalidTaskData { id: i64, error: prost::DecodeError },
54    #[error("invalid hash for {id}, expected: {expected}, got: {got}")]
55    InvalidHash {
56        id: i64,
57        expected: String,
58        got: String,
59    },
60    #[error(transparent)]
61    Conversion(#[from] xmtp_proto::ConversionError),
62    #[error("identity error: {0}")]
63    Identity(#[from] crate::identity::IdentityError),
64    #[error("key package maintenance error: {0}")]
65    KeyPackageMaintenance(
66        #[from] crate::worker::key_package_maintenance::KeyPackageMaintenanceError,
67    ),
68    #[error("notification task failed")]
69    Notification(#[from] crate::client::notifications::NotificationError),
70}
71
72impl NeedsDbReconnect for TaskWorkerError {
73    fn needs_db_reconnect(&self) -> bool {
74        match self {
75            TaskWorkerError::Storage(s)
76            | TaskWorkerError::DeviceSync(DeviceSyncError::Storage(s)) => s.db_needs_connection(),
77            TaskWorkerError::LoadGroup(e) => e.needs_db_reconnect(),
78            // Forward through GroupError's own classifier so a dropped pool hiding
79            // in a `Db`/`MlsStore` (not just `Storage`) variant still restarts the
80            // worker instead of being retried on a dead connection.
81            TaskWorkerError::Group(e) => e.needs_db_reconnect(),
82            TaskWorkerError::DeviceSync(_) => false,
83            TaskWorkerError::InvalidTaskData { .. } => false,
84            TaskWorkerError::InvalidHash { .. } => false,
85            TaskWorkerError::Conversion(_) => false,
86            TaskWorkerError::Identity(e) => e.needs_db_reconnect(),
87            TaskWorkerError::KeyPackageMaintenance(e) => e.needs_db_reconnect(),
88            TaskWorkerError::Notification(
89                crate::client::notifications::NotificationError::Storage(e),
90            ) => e.db_needs_connection(),
91            TaskWorkerError::Notification(
92                crate::client::notifications::NotificationError::Group(e),
93            ) => e.needs_db_reconnect(),
94            TaskWorkerError::Notification(_) => false,
95        }
96    }
97}
98
99#[derive(Clone)]
100pub struct TaskWorkerChannels {
101    /// Serializes notification requests across inline calls and task turns.
102    pub(crate) notification_request: Arc<tokio::sync::Mutex<()>>,
103    /// Confirmed backend topics retained while local notifications are disabled.
104    /// This memory state does not survive a client restart.
105    /// Lock before the database writer. Never hold across an await.
106    pub(crate) notification_pending_topics: Arc<
107        parking_lot::Mutex<
108            std::collections::BTreeMap<Vec<u8>, xmtp_db::notifications::UploadedTopic>,
109        >,
110    >,
111    notification_revision: Arc<std::sync::atomic::AtomicUsize>,
112}
113
114impl Default for TaskWorkerChannels {
115    fn default() -> Self {
116        Self::new()
117    }
118}
119
120impl TaskWorkerChannels {
121    pub fn new() -> Self {
122        Self {
123            notification_request: Default::default(),
124            notification_pending_topics: Default::default(),
125            notification_revision: Default::default(),
126        }
127    }
128    /// Invalidate a prepared notification batch after a local change.
129    pub fn mark_notification_changed(&self) {
130        use std::sync::atomic::Ordering;
131        self.notification_revision.fetch_add(1, Ordering::AcqRel);
132    }
133
134    pub(crate) fn notification_revision(&self) -> usize {
135        self.notification_revision
136            .load(std::sync::atomic::Ordering::Acquire)
137    }
138}
139
140/// Durably enqueue a `PullInDeadline` for `target_data_hash`, then emit its fact.
141/// The row is committed before the event; duplicates coalesce on data_hash. Lifetime is
142/// bounded by `expires_at_ns` alone (pass `NEVER_EXPIRES` for critical nudges).
143/// Callers must commit the target row FIRST: a pull-in never waits for its target —
144/// a miss is dropped (debug-logged), since a missing target is normally a completed
145/// one-shot and retrying would spin forever on never-expiring nudges.
146pub(crate) fn enqueue_pull_in<Context: XmtpSharedContext>(
147    context: &Context,
148    target_data_hash: TaskDataHash,
149    not_later_than_ns: i64,
150    expires_at_ns: i64,
151) -> Result<(), xmtp_db::StorageError> {
152    let now = xmtp_common::time::now_ns();
153    let task = xmtp_db::tasks::NewTask::builder()
154        .originating_message_sequence_id(0)
155        .next_attempt_at_ns(now) // the pull-in itself is due immediately
156        .expires_at_ns(expires_at_ns)
157        .max_attempts(i32::MAX) // lifetime bounded by expires_at_ns, not retries
158        .build(xmtp_proto::xmtp::mls::database::Task {
159            task: Some(xmtp_proto::xmtp::mls::database::task::Task::PullInDeadline(
160                xmtp_proto::xmtp::mls::database::PullInDeadline {
161                    target_data_hash: target_data_hash.to_vec(),
162                    not_later_than_ns,
163                },
164            )),
165        })?;
166    context.db().create_or_ignore_task(task)?;
167    use xmtp_events::EventWriter;
168    context
169        .events()
170        .emit(None, Some(InternalEvent::TaskScheduled));
171    Ok(())
172}
173
174pub struct Factory<Context> {
175    context: Context,
176}
177
178impl<Context> WorkerFactory for Factory<Context>
179where
180    Context: XmtpSharedContext + 'static,
181{
182    fn kind(&self) -> WorkerKind {
183        WorkerKind::TaskRunner
184    }
185
186    fn create(
187        &self,
188        _metrics: Option<crate::worker::DynMetrics>,
189    ) -> (
190        crate::worker::BoxedWorker,
191        Option<crate::worker::DynMetrics>,
192    ) {
193        let worker = TaskWorker::new(self.context.clone());
194        (Box::new(worker) as Box<_>, None)
195    }
196}
197
198pub struct TaskWorker<Context> {
199    context: Context,
200    subscription: Option<Arc<Subscription<InternalEvent>>>,
201}
202
203#[xmtp_common::async_trait]
204impl<Context> Worker for TaskWorker<Context>
205where
206    Context: XmtpSharedContext + 'static,
207{
208    fn kind(&self) -> WorkerKind {
209        WorkerKind::TaskRunner
210    }
211
212    fn set_subscription(&mut self, subscription: Arc<Subscription<InternalEvent>>) {
213        self.subscription = Some(subscription);
214    }
215
216    async fn run_tasks(&mut self) -> Result<(), Box<dyn NeedsDbReconnect>> {
217        self.run().await.map_err(|e| Box::new(e) as _)
218    }
219
220    fn factory<C>(context: C) -> impl WorkerFactory + 'static
221    where
222        Self: Sized,
223        C: XmtpSharedContext + 'static,
224    {
225        Factory { context }
226    }
227}
228
229impl<Context> TaskWorker<Context>
230where
231    Context: XmtpSharedContext + 'static,
232{
233    pub fn new(context: Context) -> Self {
234        Self {
235            context,
236            subscription: None,
237        }
238    }
239
240    #[cfg(test)]
241    fn record_task_fact(&self, event: &xmtp_events::EventEnvelope<InternalEvent>) {
242        if !matches!(&event.internal, Some(InternalEvent::TaskScheduled)) {
243            return;
244        }
245        let witness = {
246            let mut hook = test_hooks::TASK_SCHEDULED_CONSUMED.lock().unwrap();
247            if hook.as_ref().is_some_and(|(installation, _)| {
248                installation.as_slice() == self.context.installation_id().to_vec().as_slice()
249            }) {
250                hook.take()
251            } else {
252                None
253            }
254        };
255        if let Some((_, received)) = witness {
256            let _ = received.send(());
257        }
258    }
259
260    pub async fn run(&mut self) -> Result<(), TaskWorkerError> {
261        if self.subscription.is_none() {
262            let (filter, depth) = crate::worker::worker_event_filter(WorkerKind::TaskRunner)
263                .expect("task runner has an event filter");
264            self.subscription = Some(Arc::new(self.context.events().subscribe(filter, depth)));
265        }
266        let subscription = self
267            .subscription
268            .as_ref()
269            .expect("subscription is initialized")
270            .clone();
271        crate::worker::notifications::wake(&self.context)?;
272        loop {
273            let drained = subscription.drain();
274            #[cfg(test)]
275            for event in &drained {
276                self.record_task_fact(event);
277            }
278            if drained
279                .into_iter()
280                .any(|event| !matches!(event.internal, Some(InternalEvent::TaskScheduled)))
281            {
282                crate::worker::notifications::wake(&self.context)?;
283            }
284            let next_task = self.context.db().get_next_task()?;
285            let next_wakeup = Self::next_wakeup(
286                next_task.as_ref().map(|t| t.next_attempt_at_ns),
287                xmtp_common::time::now_ns(),
288            );
289            tokio::select! {
290                event = subscription.next() => {
291                    let Some(event) = event else { break Ok(()); };
292                    #[cfg(test)]
293                    self.record_task_fact(&event);
294                    let more_notification_changes = subscription.drain().into_iter().any(|event| {
295                            !matches!(event.internal, Some(InternalEvent::TaskScheduled))
296                        });
297                    let notification_changed = !matches!(event.internal, Some(InternalEvent::TaskScheduled))
298                        || more_notification_changes;
299                    if notification_changed {
300                        crate::worker::notifications::wake(&self.context)?;
301                    }
302                }
303                () = xmtp_common::time::sleep(next_wakeup) => {
304                    if let Some(task) = next_task {
305                        Self::run_and_reschedule_task(task, &self.context).await?;
306                    }
307                }
308            }
309        }
310    }
311    #[tracing::instrument(skip_all, fields(worker = "TaskRunner", operation = "worker_turn"))]
312    pub(crate) async fn run_and_reschedule_task(
313        task: DbTask,
314        context: &Context,
315    ) -> Result<(), TaskWorkerError> {
316        let now = xmtp_common::time::now_ns();
317        if task.expires_at_ns < now || task.attempts >= task.max_attempts {
318            context.db().delete_task(task.id)?;
319            return Ok(());
320        }
321        if task.next_attempt_at_ns > now {
322            // This will get called again — expected scheduler behavior, not a
323            // warning (it was ~3k warns/day in prod).
324            tracing::debug!(
325                task_id = task.id,
326                "Task {} called before next attempt at {}. Now: {now}",
327                task.id,
328                task.next_attempt_at_ns
329            );
330            return Ok(());
331        }
332        match Self::run_task(&task, context).await {
333            Ok(TaskOutcome::Done) => {
334                context.db().delete_task(task.id)?;
335            }
336            Ok(TaskOutcome::RescheduleAt(t)) => {
337                // Plain advance + attempts=0. A MIN floor here would pin a just-run
338                // (past-due) row and hot-loop it.
339                match context.db().update_task(task.id, 0, now, t) {
340                    Ok(_) => {}
341                    Err(StorageError::DieselResult(diesel::result::Error::NotFound)) => {
342                        tracing::debug!("Task {} vanished before reschedule; skipping", task.id);
343                    }
344                    Err(e) => return Err(e.into()),
345                }
346            }
347            Err(error) => {
348                let attempts = task.attempts + 1;
349                let attempt_scaling_factor = (task.backoff_scaling_factor as f64).powi(attempts);
350                let next_attempt_duration = (((task.initial_backoff_duration_ns as f64)
351                    * attempt_scaling_factor) as i64)
352                    .min(task.max_backoff_duration_ns);
353                let retry_from =
354                    if task.data_hash == crate::worker::notifications::task_hash().as_ref() {
355                        xmtp_common::time::now_ns()
356                    } else {
357                        now
358                    };
359                let next_attempt_at_ns = retry_from.saturating_add(next_attempt_duration);
360                tracing::warn!(%error, "Task {} retry failed. Retrying in {next_attempt_duration}ns", task.id);
361                match context
362                    .db()
363                    .update_task(task.id, attempts, now, next_attempt_at_ns)
364                {
365                    Ok(_) => {}
366                    // The row was concurrently deleted (e.g. a dead-row cleanup in
367                    // upsert_pending_self_remove_task crossed this retry). Nothing
368                    // left to reschedule — don't abort the worker loop.
369                    Err(StorageError::DieselResult(diesel::result::Error::NotFound)) => {
370                        tracing::debug!("Task {} vanished before reschedule; skipping", task.id);
371                    }
372                    Err(e) => return Err(e.into()),
373                }
374            }
375        }
376        Ok(())
377    }
378    async fn run_task(task: &DbTask, context: &Context) -> Result<TaskOutcome, TaskWorkerError> {
379        #[cfg(test)]
380        if let Some((hash, t)) = test_hooks::RESCHEDULE_OVERRIDE.lock().unwrap().clone()
381            && task.data_hash == hash
382        {
383            return Ok(TaskOutcome::RescheduleAt(t));
384        }
385        let data_hash = xmtp_common::sha256_bytes(&task.data);
386        if task.data_hash != data_hash {
387            let expected = hex::encode(&data_hash);
388            let got = hex::encode(&task.data_hash);
389            tracing::warn!(
390                "Task {} data hash mismatch. Expected {expected}, got {got}",
391                task.id,
392            );
393        }
394        let task_proto = match TaskProto::decode(task.data.as_slice()) {
395            Ok(task_proto) => task_proto,
396            Err(e) => {
397                context.db().delete_task(task.id)?;
398                tracing::warn!("Task {} data decode error: {}", task.id, e);
399                return Ok(TaskOutcome::Done);
400            }
401        };
402        match task_proto.task {
403            Some(xmtp_proto::xmtp::mls::database::task::Task::NotificationSync(_)) => {
404                return Ok(crate::worker::notifications::run(context).await?);
405            }
406            Some(xmtp_proto::xmtp::mls::database::task::Task::ProcessPendingSelfRemove(
407                pending,
408            )) => {
409                Self::process_pending_self_remove(task, pending, context).await?;
410            }
411            Some(xmtp_proto::xmtp::mls::database::task::Task::PullInDeadline(p)) => {
412                // Runs on the worker thread — the sole rescheduler of existing
413                // rows' `next_attempt_at_ns` — so no transaction is needed
414                // (inserts happen off-thread; only deadline mutation is guarded).
415                let matched = match TaskDataHash::try_from(p.target_data_hash.as_slice()) {
416                    Ok(h) => context
417                        .db()
418                        .pull_in_task_deadline(&h, p.not_later_than_ns)?,
419                    Err(_) => false, // malformed length: fall through to the miss log
420                };
421                if !matched {
422                    // Completed one-shot target, or a producer broke the
423                    // commit-target-first contract. Drop either way (retrying
424                    // would spin forever on never-expiring nudges). warn: every
425                    // in-tree pull-in targets a never-deleted singleton, so a
426                    // miss is always anomalous today.
427                    tracing::warn!(
428                        target_data_hash = hex::encode(&p.target_data_hash),
429                        "pull-in target missing; dropping nudge for task {}",
430                        task.id
431                    );
432                }
433            }
434            Some(xmtp_proto::xmtp::mls::database::task::Task::KpRotation(_)) => {
435                let now = xmtp_common::time::now_ns();
436                kp::rotate_if_needed(context).await?;
437                // Unconditional (no-op when nothing is marked): a backoff retry
438                // after rotate-succeeded/nudge-failed must still re-nudge.
439                kp::nudge_deletion(context)?;
440                // Advance to the LIVE rotation column (rotate reset it to +30d; a
441                // welcome nudge may have re-lowered it). Do NOT hardcode +30d.
442                let next = context
443                    .db()
444                    .next_key_package_rotation_ns()?
445                    .unwrap_or(now + KEY_PACKAGE_ROTATION_INTERVAL_NS);
446                return Ok(TaskOutcome::RescheduleAt(next));
447            }
448            Some(xmtp_proto::xmtp::mls::database::task::Task::KpDeletion(_)) => {
449                kp::sweep_expired(context)?;
450                let now = xmtp_common::time::now_ns();
451                // Nothing pending: park one rotation interval out (deletions only
452                // arise from rotations; same constant as the KpRotation arm — the
453                // exact value isn't load-bearing, nudges pull the task in sooner).
454                let next = context
455                    .db()
456                    .min_key_package_delete_at_ns()?
457                    .unwrap_or(now + KEY_PACKAGE_ROTATION_INTERVAL_NS);
458                return Ok(TaskOutcome::RescheduleAt(next));
459            }
460            Some(xmtp_proto::xmtp::mls::database::task::Task::AddMissingInstallations(add)) => {
461                Self::run_add_missing_installations(task, add, context).await?;
462            }
463            Some(xmtp_proto::xmtp::mls::database::task::Task::KpLiveness(_)) => {
464                // The variant exists in the regenerated protos but nothing in
465                // this crate schedules it yet, so a row can only appear from a
466                // newer client sharing this database. Leave it for that client
467                // rather than deleting it: falling through to `Done` would drop
468                // work this build simply cannot see, and the owning feature will
469                // add real handling. `expires_at_ns` still bounds how long an
470                // unclaimed row can sit here, so deferring cannot leak rows.
471                tracing::warn!(
472                    "Task {} is a KpLiveness task, which this version does not handle. Deferring.",
473                    task.id
474                );
475                return Ok(TaskOutcome::RescheduleAt(
476                    xmtp_common::time::now_ns() + UNKNOWN_TASK_DEFER_NS,
477                ));
478            }
479            None => {
480                tracing::error!("Task {} has no data. Deleting.", task.id);
481                context.db().delete_task(task.id)?;
482            }
483        }
484        Ok(TaskOutcome::Done)
485    }
486    fn next_wakeup(
487        next_attempt_at_ns: Option<i64>,
488        // these are passed in for testing
489        now: i64,
490    ) -> xmtp_common::time::Duration {
491        use xmtp_common::r#const::NS_IN_DAY;
492        let now_plus_one_day = now.saturating_add(NS_IN_DAY);
493        let next_task_wakeup = next_attempt_at_ns.unwrap_or(i64::MAX).min(now_plus_one_day);
494        if now > next_task_wakeup {
495            xmtp_common::time::Duration::from_nanos(0)
496        } else {
497            std::time::Duration::from_nanos((next_task_wakeup - now) as u64)
498        }
499    }
500
501    /// Run a `ProcessPendingSelfRemove` task: load the group and remove members
502    /// who requested to leave (a no-op unless this client is super-admin).
503    async fn process_pending_self_remove(
504        task: &DbTask,
505        pending: xmtp_proto::xmtp::mls::database::ProcessPendingSelfRemove,
506        context: &Context,
507    ) -> Result<(), TaskWorkerError> {
508        // A malformed group_id can never succeed — drop the task, don't retry.
509        let Ok(group_id) = xmtp_proto::types::GroupId::try_from(pending.group_id.as_slice()) else {
510            tracing::warn!(
511                "Task {} has a malformed group_id for ProcessPendingSelfRemove. Deleting.",
512                task.id
513            );
514            context.db().delete_task(task.id)?;
515            return Ok(());
516        };
517        match crate::mls_store::MlsStore::new(context.clone()).group(&group_id) {
518            Ok(group) => {
519                // No-op unless super-admin; idempotent, so retries are safe.
520                group.process_pending_self_removals().await?;
521                Ok(())
522            }
523            Err(crate::mls_store::MlsStoreError::NotFound(_)) => {
524                tracing::debug!(
525                    "Task {} targets a group that no longer exists. Deleting.",
526                    task.id
527                );
528                context.db().delete_task(task.id)?;
529                Ok(())
530            }
531            // A DB/connection error is transient — let it retry.
532            Err(e) => Err(e.into()),
533        }
534    }
535
536    /// Run an `AddMissingInstallations` task: load the group and reconcile its
537    /// membership with the inbox's latest identity state. Idempotent — the
538    /// underlying group method no-ops when membership is already current, so
539    /// retries and duplicate rows are safe.
540    async fn run_add_missing_installations(
541        task: &DbTask,
542        add: xmtp_proto::xmtp::mls::database::AddMissingInstallations,
543        context: &Context,
544    ) -> Result<(), TaskWorkerError> {
545        // A malformed group_id can never succeed — drop the task, don't retry.
546        let Ok(group_id) = xmtp_proto::types::GroupId::try_from(add.group_id.as_slice()) else {
547            tracing::warn!(
548                "Task {} has a malformed group_id for AddMissingInstallations. Deleting.",
549                task.id
550            );
551            context.db().delete_task(task.id)?;
552            return Ok(());
553        };
554        match crate::mls_store::MlsStore::new(context.clone()).group(&group_id) {
555            Ok(group) => {
556                group.add_missing_installations().await?;
557                Ok(())
558            }
559            Err(crate::mls_store::MlsStoreError::NotFound(_)) => {
560                tracing::debug!(
561                    "Task {} targets a group that no longer exists. Deleting.",
562                    task.id
563                );
564                context.db().delete_task(task.id)?;
565                Ok(())
566            }
567            // A DB/connection error is transient — let it retry.
568            Err(e) => Err(e.into()),
569        }
570    }
571}
572
573#[cfg(test)]
574mod tests {
575    use super::*;
576    use crate::tester;
577    use crate::worker::{WorkerConfig, WorkerKind};
578    use xmtp_db::tasks::{NewTask, data_hash_for};
579    use xmtp_proto::xmtp::mls::database::{Task as TaskProto, task::Task as TaskKind};
580
581    /// A unique one-shot payload: a PullInDeadline aimed at a random, nonexistent
582    /// target. Applying it is a no-op; its data_hash is unique per call.
583    fn unique_proto() -> TaskProto {
584        TaskProto {
585            task: Some(TaskKind::PullInDeadline(
586                xmtp_proto::xmtp::mls::database::PullInDeadline {
587                    target_data_hash: xmtp_common::rand_vec::<32>(),
588                    not_later_than_ns: 0,
589                },
590            )),
591        }
592    }
593
594    fn no_runner_cfg() -> WorkerConfig {
595        let mut cfg = WorkerConfig::default();
596        cfg.enabled.insert(WorkerKind::TaskRunner, false);
597        cfg
598    }
599
600    /// Insert a task row and return the stored row (found by its data_hash).
601    fn seed(
602        db: &impl QueryTasks,
603        proto: TaskProto,
604        next: i64,
605        expires: i64,
606        attempts: i32,
607        max: i32,
608    ) -> xmtp_db::tasks::Task {
609        let hash = data_hash_for(&proto);
610        let task = NewTask::builder()
611            .originating_message_sequence_id(0)
612            .next_attempt_at_ns(next)
613            .expires_at_ns(expires)
614            .attempts(attempts)
615            .max_attempts(max)
616            .build(proto)
617            .unwrap();
618        db.create_or_ignore_task(task).unwrap();
619        db.get_tasks()
620            .unwrap()
621            .into_iter()
622            .find(|t| t.data_hash == hash.as_ref())
623            .unwrap()
624    }
625
626    #[xmtp_common::test(unwrap_try = true)]
627    async fn done_deletes() {
628        tester!(alix, worker_config: no_runner_cfg());
629        let db = alix.context.db();
630        let now = xmtp_common::time::now_ns();
631        let row = seed(&db, unique_proto(), now - 1, i64::MAX, 0, i32::MAX);
632        TaskWorker::run_and_reschedule_task(row, &alix.context).await?;
633        assert!(db.get_tasks()?.is_empty(), "Done task must be deleted");
634    }
635
636    #[xmtp_common::test(unwrap_try = true)]
637    async fn add_missing_installations_missing_group_deletes_task() {
638        tester!(alix, worker_config: no_runner_cfg());
639        let db = alix.context.db();
640        let now = xmtp_common::time::now_ns();
641        let proto = TaskProto {
642            task: Some(TaskKind::AddMissingInstallations(
643                xmtp_proto::xmtp::mls::database::AddMissingInstallations {
644                    group_id: xmtp_common::rand_vec::<16>(),
645                },
646            )),
647        };
648        let row = seed(&db, proto, now - 1, i64::MAX, 0, 3);
649        TaskWorker::run_and_reschedule_task(row, &alix.context).await?;
650        assert!(
651            db.get_tasks()?.is_empty(),
652            "task for a nonexistent group must be deleted, not retried"
653        );
654    }
655
656    #[xmtp_common::test(unwrap_try = true)]
657    async fn recurring_task_advances_and_does_not_hot_loop() {
658        tester!(alix, worker_config: no_runner_cfg());
659        let db = alix.context.db();
660        let now = xmtp_common::time::now_ns();
661        let proto = unique_proto();
662        let row = seed(&db, proto.clone(), now - 1, i64::MAX, 5, i32::MAX);
663        let target = now + xmtp_common::NS_IN_DAY;
664        *test_hooks::RESCHEDULE_OVERRIDE.lock().unwrap() =
665            Some((data_hash_for(&proto).to_vec(), target));
666
667        TaskWorker::run_and_reschedule_task(row, &alix.context).await?;
668
669        *test_hooks::RESCHEDULE_OVERRIDE.lock().unwrap() = None;
670        let after = db.get_tasks()?.pop().expect("recurring row must survive");
671        assert_eq!(
672            after.next_attempt_at_ns, target,
673            "deadline must ADVANCE, not stay past-due"
674        );
675        assert_eq!(after.attempts, 0, "success resets the backoff counter");
676    }
677
678    #[xmtp_common::test(unwrap_try = true)]
679    async fn never_expire_seed_survives_reaper() {
680        tester!(alix, worker_config: no_runner_cfg());
681        let db = alix.context.db();
682        let now = xmtp_common::time::now_ns();
683        let proto = unique_proto();
684        // High attempts + past-due: only the i64::MAX/i32::MAX seed keeps the
685        // reaper (expires < now || attempts >= max) from deleting it.
686        let row = seed(&db, proto.clone(), now - 1, i64::MAX, 1_000_000, i32::MAX);
687        *test_hooks::RESCHEDULE_OVERRIDE.lock().unwrap() =
688            Some((data_hash_for(&proto).to_vec(), now + xmtp_common::NS_IN_DAY));
689
690        TaskWorker::run_and_reschedule_task(row, &alix.context).await?;
691
692        *test_hooks::RESCHEDULE_OVERRIDE.lock().unwrap() = None;
693        let after = db.get_tasks()?;
694        assert_eq!(after.len(), 1, "never-expire seed must not be reaped");
695        assert_eq!(
696            after[0].next_attempt_at_ns,
697            now + xmtp_common::NS_IN_DAY,
698            "seed must have been rescheduled to the target deadline"
699        );
700    }
701
702    #[xmtp_common::test(unwrap_try = true)]
703    async fn not_yet_due_task_is_not_run_early() {
704        tester!(alix, worker_config: no_runner_cfg());
705        let db = alix.context.db();
706        let now = xmtp_common::time::now_ns();
707        let proto = unique_proto();
708        let far = now + 30 * xmtp_common::NS_IN_DAY;
709        let row = seed(&db, proto.clone(), far, i64::MAX, 0, i32::MAX);
710        // If the guard failed and the task ran, the hook would rewrite next_attempt.
711        *test_hooks::RESCHEDULE_OVERRIDE.lock().unwrap() =
712            Some((data_hash_for(&proto).to_vec(), now));
713
714        TaskWorker::run_and_reschedule_task(row, &alix.context).await?;
715
716        *test_hooks::RESCHEDULE_OVERRIDE.lock().unwrap() = None;
717        let after = db.get_tasks()?.pop().unwrap();
718        assert_eq!(
719            after.next_attempt_at_ns, far,
720            "not-yet-due task must not run on a 1-day-cap wake"
721        );
722        assert_eq!(
723            after.attempts, 0,
724            "not-yet-due task must not have its attempts incremented"
725        );
726    }
727
728    #[xmtp_common::test(unwrap_try = true)]
729    async fn pull_in_arm_lowers_existing_target() {
730        tester!(alix, worker_config: no_runner_cfg());
731        let db = alix.context.db();
732        let now = xmtp_common::time::now_ns();
733        // Far-future recurring target.
734        let target_proto = unique_proto();
735        let target_hash = data_hash_for(&target_proto);
736        let far = now + 30 * xmtp_common::NS_IN_DAY;
737        seed(&db, target_proto, far, i64::MAX, 0, i32::MAX);
738        // Due pull-in aimed at it.
739        enqueue_pull_in(&alix.context, target_hash, now + 1_000, i64::MAX)?;
740        let pull_in_row = db
741            .get_tasks()?
742            .into_iter()
743            .find(|t| t.data_hash != target_hash.as_ref())
744            .expect("pull-in row exists");
745
746        TaskWorker::run_and_reschedule_task(pull_in_row, &alix.context).await?;
747
748        let rows = db.get_tasks()?;
749        let target = rows
750            .iter()
751            .find(|t| t.data_hash == target_hash.as_ref())
752            .expect("target survives");
753        assert_eq!(
754            target.next_attempt_at_ns,
755            now + 1_000,
756            "arm must lower the target to the ceiling"
757        );
758        assert_eq!(rows.len(), 1, "applied pull-in must self-delete");
759    }
760
761    #[xmtp_common::test(unwrap_try = true)]
762    async fn pull_in_task_runs_and_pulls_in() {
763        tester!(alix); // live TaskRunner
764        let db = alix.context.db();
765        let now = xmtp_common::time::now_ns();
766        let proto = unique_proto();
767        // Ceiling 1 day out: the exact-equality assert below proves lowering
768        // regardless of magnitude, and a generous ceiling means the target can't
769        // become due (and get dispatched/deleted) even on a pathologically slow
770        // CI runner — only constraints are "> test wall time" and "!= far".
771        let ceiling = now + xmtp_common::NS_IN_DAY;
772        let far = now + 30 * xmtp_common::NS_IN_DAY;
773        seed(&db, proto.clone(), far, i64::MAX, 0, i32::MAX);
774        let hash = data_hash_for(&proto);
775
776        enqueue_pull_in(&alix.context, hash, ceiling, i64::MAX)?;
777
778        // Poll up to ~10s (wasm-safe): the worker dispatches the due pull-in,
779        // which lowers the target and self-deletes.
780        let mut pulled = false;
781        for _ in 0..50u32 {
782            xmtp_common::time::sleep(std::time::Duration::from_millis(200)).await;
783            let rows = db.get_tasks()?;
784            let target_ok = rows
785                .iter()
786                .any(|t| t.data_hash == hash.as_ref() && t.next_attempt_at_ns == ceiling);
787            let pull_in_gone = !rows.iter().any(|t| {
788                t.data_hash != hash.as_ref() && t.next_attempt_at_ns <= xmtp_common::time::now_ns()
789            });
790            if target_ok && pull_in_gone {
791                pulled = true;
792                break;
793            }
794        }
795        assert!(
796            pulled,
797            "worker must apply the pull-in and lower the target deadline"
798        );
799    }
800}
xmtp_mls/worker/tasks.rs

xmtp_mls/worker/
tasks.rs
