xmtp_mls/worker/
tasks.rs

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

xmtp_mls/worker/
tasks.rs
