xmtp_mls_core/worker/
tasks.rs

1pub use crate::task_channels::TaskWorkerChannels;
2use crate::{
3    context::CoreContext,
4    events::InternalEvent,
5    worker::{
6        BoxedWorker, EVENT_BUFFER, NeedsDbReconnect, NewWorker, Worker, WorkerFactory, WorkerKind,
7    },
8};
9use std::{marker::PhantomData, sync::Arc};
10use xmtp_common::{MaybeSend, MaybeSync, RetryableError};
11use xmtp_db::tasks::{QueryTasks, Task as DbTask, TaskDataHash};
12use xmtp_db::{StorageError, diesel};
13use xmtp_events::{EventEnvelope, EventFilter, EventKind, Subscription};
14use xmtp_proto::xmtp::mls::database::PullInDeadline;
15
16/// One task-runner turn for handler tests that run outside a running worker.
17#[cfg(any(test, feature = "test-utils"))]
18#[doc(hidden)]
19pub mod test_support {
20    use super::{DbTask, PullInDeadline, StorageError, TaskHandlers, TaskWorker};
21    use crate::context::CoreContext;
22    use xmtp_db::tasks::{NewTask, QueryTasks, TaskDataHash, data_hash_for};
23    use xmtp_proto::xmtp::mls::database::{Task as TaskProto, task::Task as TaskKind};
24
25    /// Run `task` through one turn of `worker`: the expiry and due checks,
26    /// the dispatch, and the delete or reschedule that follows.
27    pub async fn run_and_reschedule_task<Context, Handlers>(
28        worker: &TaskWorker<Context, Handlers>,
29        task: DbTask,
30    ) -> Result<(), Handlers::Error>
31    where
32        Context: CoreContext,
33        Handlers: TaskHandlers<Context>,
34    {
35        worker.run_and_reschedule_task(task).await
36    }
37
38    /// A unique one-shot payload: a PullInDeadline aimed at a random, nonexistent
39    /// target. Applying it is a no-op; its data_hash is unique per call.
40    pub fn unique_proto() -> TaskProto {
41        TaskProto {
42            task: Some(TaskKind::PullInDeadline(PullInDeadline {
43                target_data_hash: xmtp_common::rand_vec::<32>(),
44                not_later_than_ns: 0,
45            })),
46        }
47    }
48
49    /// Insert a task row and return the stored row (found by its data_hash).
50    pub fn seed(
51        db: &impl QueryTasks,
52        proto: TaskProto,
53        next: i64,
54        expires: i64,
55        attempts: i32,
56        max: i32,
57    ) -> DbTask {
58        let hash = data_hash_for(&proto);
59        let task = NewTask::builder()
60            .originating_message_sequence_id(0)
61            .next_attempt_at_ns(next)
62            .expires_at_ns(expires)
63            .attempts(attempts)
64            .max_attempts(max)
65            .build(proto)
66            .expect("a task row");
67        db.create_or_ignore_task(task).expect("a stored task");
68        db.get_tasks()
69            .expect("the stored tasks")
70            .into_iter()
71            .find(|t| t.data_hash == hash.as_ref())
72            .expect("the seeded task")
73    }
74
75    /// Queue a pull-in of the task with `target_data_hash`, as a nudge does.
76    pub fn enqueue_pull_in<Context: CoreContext>(
77        context: &Context,
78        target_data_hash: TaskDataHash,
79        not_later_than_ns: i64,
80        expires_at_ns: i64,
81    ) -> Result<(), StorageError> {
82        super::enqueue_pull_in(context, target_data_hash, not_later_than_ns, expires_at_ns)
83    }
84}
85
86/// `Done` = one-shot, row deleted. `RescheduleAt(ns)` = recurring, row kept and
87/// advanced to that absolute deadline.
88#[derive(Debug, PartialEq, Eq)]
89pub enum TaskOutcome {
90    Done,
91    RescheduleAt(i64),
92}
93
94/// Task-runner controls for tests. The runner reads them; tests set them.
95/// Each assumes process-per-test isolation.
96#[cfg(any(test, feature = "test-hooks"))]
97#[doc(hidden)]
98pub mod test_hooks {
99    use std::sync::Mutex;
100
101    static FAIL_RESCHEDULE: Mutex<bool> = Mutex::new(false);
102    type ScheduledWitness = (Vec<u8>, tokio::sync::oneshot::Sender<()>);
103    static TASK_SCHEDULED_CONSUMED: Mutex<Option<ScheduledWitness>> = Mutex::new(None);
104
105    /// Makes the runner's next retry schedule write fail, once.
106    pub fn fail_next_reschedule() {
107        *FAIL_RESCHEDULE.lock().expect("task test hooks") = true;
108    }
109
110    pub(super) fn take_reschedule_failure() -> bool {
111        std::mem::take(&mut *FAIL_RESCHEDULE.lock().expect("task test hooks"))
112    }
113
114    /// Sends on `consumed` once the task runner of `installation_id` reads a
115    /// `TaskScheduled` event.
116    pub fn on_task_scheduled_consumed(
117        installation_id: Vec<u8>,
118        consumed: tokio::sync::oneshot::Sender<()>,
119    ) {
120        *TASK_SCHEDULED_CONSUMED.lock().expect("task test hooks") =
121            Some((installation_id, consumed));
122    }
123
124    pub(super) fn take_task_scheduled_witness(
125        installation_id: &[u8],
126    ) -> Option<tokio::sync::oneshot::Sender<()>> {
127        let mut hook = TASK_SCHEDULED_CONSUMED.lock().expect("task test hooks");
128        if hook
129            .as_ref()
130            .is_some_and(|(installation, _)| installation.as_slice() == installation_id)
131        {
132            hook.take().map(|(_, consumed)| consumed)
133        } else {
134            None
135        }
136    }
137}
138
139/// The handlers a [`TaskWorker`] runs. The runner owns each task's schedule:
140/// expiry, attempts, backoff and rate-limit holds. The handlers run one row.
141#[xmtp_common::async_trait]
142pub trait TaskHandlers<Context>: MaybeSend + MaybeSync {
143    /// The error of a failed run. The runner reads its rate-limit delay and
144    /// its reconnect signal, and reports its own storage failures as this type.
145    type Error: RetryableError
146        + NeedsDbReconnect
147        + From<StorageError>
148        + std::error::Error
149        + MaybeSend
150        + MaybeSync
151        + 'static;
152
153    /// Run one due task. Gets the whole row: handlers use `task.id`, for
154    /// example to delete it.
155    async fn dispatch(&self, context: &Context, task: &DbTask) -> Result<TaskOutcome, Self::Error>;
156
157    /// Runs when the runner starts, before it reads a task.
158    fn on_start(&self, _context: &Context) -> Result<(), Self::Error> {
159        Ok(())
160    }
161
162    /// Runs once for each batch of events the runner reads that holds an
163    /// event other than its own `TaskScheduled`, and gets the whole batch.
164    /// Public events such as `ConsentChanged` and `ArchiveRestored` arrive here.
165    fn on_events(
166        &self,
167        _context: &Context,
168        _events: &[EventEnvelope<InternalEvent>],
169    ) -> Result<(), Self::Error> {
170        Ok(())
171    }
172}
173
174/// Durably enqueue a `PullInDeadline` for `target_data_hash`, then emit its fact.
175/// The row is committed before the event; duplicates coalesce on data_hash. Lifetime is
176/// bounded by `expires_at_ns` alone (pass `NEVER_EXPIRES` for critical nudges).
177/// Callers must commit the target row FIRST: a pull-in never waits for its target —
178/// a miss is dropped (debug-logged), since a missing target is normally a completed
179/// one-shot and retrying would spin forever on never-expiring nudges.
180pub fn enqueue_pull_in<Context: CoreContext>(
181    context: &Context,
182    target_data_hash: TaskDataHash,
183    not_later_than_ns: i64,
184    expires_at_ns: i64,
185) -> Result<(), xmtp_db::StorageError> {
186    let now = xmtp_common::time::now_ns();
187    let task = xmtp_db::tasks::NewTask::builder()
188        .originating_message_sequence_id(0)
189        .next_attempt_at_ns(now) // the pull-in itself is due immediately
190        .expires_at_ns(expires_at_ns)
191        .max_attempts(i32::MAX) // lifetime bounded by expires_at_ns, not retries
192        .build(xmtp_proto::xmtp::mls::database::Task {
193            task: Some(xmtp_proto::xmtp::mls::database::task::Task::PullInDeadline(
194                PullInDeadline {
195                    target_data_hash: target_data_hash.to_vec(),
196                    not_later_than_ns,
197                },
198            )),
199        })?;
200    context.db().create_or_ignore_task(task)?;
201    use xmtp_events::EventWriter;
202    context
203        .events()
204        .emit(None, Some(InternalEvent::TaskScheduled));
205    Ok(())
206}
207
208/// Run a `PullInDeadline` task: lower its target's deadline to the ceiling.
209pub fn apply_pull_in<Context: CoreContext>(
210    context: &Context,
211    task: &DbTask,
212    pull_in: PullInDeadline,
213) -> Result<(), StorageError> {
214    // Runs on the worker thread — the sole rescheduler of existing
215    // rows' `next_attempt_at_ns` — so no transaction is needed
216    // (inserts happen off-thread; only deadline mutation is guarded).
217    let matched = match TaskDataHash::try_from(pull_in.target_data_hash.as_slice()) {
218        Ok(h) => {
219            // A target in failure backoff keeps its stored schedule,
220            // which holds any rate-limit delay, across restarts too.
221            // implements: RATE-041
222            let backing_off = context
223                .db()
224                .get_tasks()?
225                .iter()
226                .any(|target| target.data_hash == h.as_ref() && target.attempts > 0);
227            backing_off
228                || context
229                    .db()
230                    .pull_in_task_deadline(&h, pull_in.not_later_than_ns)?
231        }
232        Err(_) => false, // malformed length: fall through to the miss log
233    };
234    if !matched {
235        // Completed one-shot target, or a producer broke the
236        // commit-target-first contract. Drop either way (retrying
237        // would spin forever on never-expiring nudges). warn: every
238        // in-tree pull-in targets a never-deleted singleton, so a
239        // miss is always anomalous today.
240        tracing::warn!(
241            target_data_hash = hex::encode(&pull_in.target_data_hash),
242            "pull-in target missing; dropping nudge for task {}",
243            task.id
244        );
245    }
246    Ok(())
247}
248
249/// What wakes the runner: its own `TaskScheduled`, and the changes the
250/// notification task follows.
251fn event_interest() -> (EventFilter<InternalEvent>, Option<usize>) {
252    let filter = EventFilter::new([
253        EventKind::ConsentChanged,
254        EventKind::ConversationJoined,
255        EventKind::ConversationRemoved,
256        EventKind::HmacKeysUpdated,
257        // implements: PUSH-260
258        // An import writes consent without consent.changed.
259        EventKind::ArchiveRestored,
260    ])
261    .with_internal(|event| {
262        matches!(
263            event,
264            InternalEvent::TaskScheduled
265                | InternalEvent::NotificationSettingsChanged
266                | InternalEvent::GroupJoined { .. }
267                | InternalEvent::PreferencesChanged { .. }
268        )
269    });
270    (filter, Some(EVENT_BUFFER))
271}
272
273pub struct Factory<Context, Handlers> {
274    context: Context,
275    handlers: PhantomData<fn() -> Handlers>,
276}
277
278impl<Context, Handlers> WorkerFactory for Factory<Context, Handlers>
279where
280    Context: CoreContext + 'static,
281    Handlers: TaskHandlers<Context> + Default + 'static,
282{
283    fn kind(&self) -> WorkerKind {
284        WorkerKind::TaskRunner
285    }
286
287    fn event_interest(&self) -> Option<(EventFilter<InternalEvent>, Option<usize>)> {
288        Some(event_interest())
289    }
290
291    fn create(&self) -> BoxedWorker {
292        Box::new(TaskWorker::new(self.context.clone(), Handlers::default()))
293    }
294}
295
296pub struct TaskWorker<Context, Handlers> {
297    context: Context,
298    handlers: Handlers,
299    subscription: Option<Arc<Subscription<InternalEvent>>>,
300}
301
302impl<Context, Handlers> NewWorker<Context> for TaskWorker<Context, Handlers>
303where
304    Context: CoreContext + 'static,
305    Handlers: TaskHandlers<Context> + Default + 'static,
306{
307    fn factory(context: Context) -> impl WorkerFactory + 'static {
308        Factory::<Context, Handlers> {
309            context,
310            handlers: PhantomData,
311        }
312    }
313}
314
315#[xmtp_common::async_trait]
316impl<Context, Handlers> Worker for TaskWorker<Context, Handlers>
317where
318    Context: CoreContext + 'static,
319    Handlers: TaskHandlers<Context> + 'static,
320{
321    fn kind(&self) -> WorkerKind {
322        WorkerKind::TaskRunner
323    }
324
325    fn set_subscription(&mut self, subscription: Arc<Subscription<InternalEvent>>) {
326        self.subscription = Some(subscription);
327    }
328
329    async fn run_tasks(&mut self) -> Result<(), Box<dyn NeedsDbReconnect>> {
330        self.run().await.map_err(|e| Box::new(e) as _)
331    }
332}
333
334impl<Context, Handlers> TaskWorker<Context, Handlers>
335where
336    Context: CoreContext,
337    Handlers: TaskHandlers<Context>,
338{
339    pub fn new(context: Context, handlers: Handlers) -> Self {
340        Self {
341            context,
342            handlers,
343            subscription: None,
344        }
345    }
346
347    #[cfg(any(test, feature = "test-hooks"))]
348    fn record_task_fact(&self, event: &xmtp_events::EventEnvelope<InternalEvent>) {
349        if !matches!(&event.internal, Some(InternalEvent::TaskScheduled)) {
350            return;
351        }
352        let witness =
353            test_hooks::take_task_scheduled_witness(&self.context.installation_id().to_vec());
354        if let Some(received) = witness {
355            let _ = received.send(());
356        }
357    }
358
359    pub async fn run(&mut self) -> Result<(), Handlers::Error> {
360        if self.subscription.is_none() {
361            let (filter, depth) = event_interest();
362            self.subscription = Some(Arc::new(self.context.events().subscribe(filter, depth)));
363        }
364        let subscription = self
365            .subscription
366            .as_ref()
367            .expect("subscription is initialized")
368            .clone();
369        self.handlers.on_start(&self.context)?;
370        loop {
371            let drained = subscription.drain();
372            #[cfg(any(test, feature = "test-hooks"))]
373            for event in &drained {
374                self.record_task_fact(event);
375            }
376            self.forward_events(&drained)?;
377            let next_task = self.context.db().get_next_task()?;
378            let next_wakeup = Self::next_wakeup(
379                next_task.as_ref().map(|t| t.next_attempt_at_ns),
380                xmtp_common::time::now_ns(),
381            );
382            tokio::select! {
383                event = subscription.next() => {
384                    let Some(event) = event else { break Ok(()); };
385                    #[cfg(any(test, feature = "test-hooks"))]
386                    self.record_task_fact(&event);
387                    let mut events = vec![event];
388                    events.extend(subscription.drain());
389                    self.forward_events(&events)?;
390                }
391                () = xmtp_common::time::sleep(next_wakeup) => {
392                    if let Some(task) = next_task {
393                        self.run_and_reschedule_task(task).await?;
394                    }
395                }
396            }
397        }
398    }
399
400    /// Hand a batch of events to the handlers, unless it holds only the
401    /// runner's own `TaskScheduled` facts.
402    fn forward_events(
403        &self,
404        events: &[EventEnvelope<InternalEvent>],
405    ) -> Result<(), Handlers::Error> {
406        if events
407            .iter()
408            .any(|event| !matches!(event.internal, Some(InternalEvent::TaskScheduled)))
409        {
410            self.handlers.on_events(&self.context, events)?;
411        }
412        Ok(())
413    }
414
415    #[tracing::instrument(skip_all, fields(worker = "TaskRunner", operation = "worker_turn"))]
416    pub(crate) async fn run_and_reschedule_task(
417        &self,
418        task: DbTask,
419    ) -> Result<(), Handlers::Error> {
420        let context = &self.context;
421        let now = xmtp_common::time::now_ns();
422        if task.expires_at_ns < now || task.attempts >= task.max_attempts {
423            context.db().delete_task(task.id)?;
424            return Ok(());
425        }
426        if task.next_attempt_at_ns > now {
427            // This will get called again — expected scheduler behavior, not a
428            // warning (it was ~3k warns/day in prod).
429            tracing::debug!(
430                task_id = task.id,
431                "Task {} called before next attempt at {}. Now: {now}",
432                task.id,
433                task.next_attempt_at_ns
434            );
435            return Ok(());
436        }
437        match self.handlers.dispatch(context, &task).await {
438            Ok(TaskOutcome::Done) => {
439                context.db().delete_task(task.id)?;
440            }
441            Ok(TaskOutcome::RescheduleAt(t)) => {
442                // Plain advance + attempts=0. A MIN floor here would pin a just-run
443                // (past-due) row and hot-loop it.
444                match context.db().update_task(task.id, 0, now, t) {
445                    Ok(_) => {}
446                    Err(StorageError::DieselResult(diesel::result::Error::NotFound)) => {
447                        tracing::debug!("Task {} vanished before reschedule; skipping", task.id);
448                    }
449                    Err(e) => return Err(e.into()),
450                }
451            }
452            Err(error) => {
453                // Every kind backs off from when its run finished, so a long
454                // run does not use up its own backoff.
455                let finished = xmtp_common::time::now_ns();
456                let mut attempts = task.attempts + 1;
457                let attempt_scaling_factor = (task.backoff_scaling_factor as f64).powi(attempts);
458                let next_attempt_duration = (((task.initial_backoff_duration_ns as f64)
459                    * attempt_scaling_factor) as i64)
460                    .min(task.max_backoff_duration_ns);
461                let backoff_at_ns = finished.saturating_add(next_attempt_duration);
462                let mut next_attempt_at_ns = backoff_at_ns;
463                let mut expires_at_ns = task.expires_at_ns;
464                // A rate-limited request waits at least the delay, counted from
465                // when the rejection arrived, not from when the turn started.
466                // Neither the wait it adds past the backoff nor the refused turn
467                // counts against the task's lifetime or its attempts, so a rate
468                // limit never expires or deletes the refused work. At least one
469                // attempt is recorded, which marks the task as backing off, so
470                // a pull-in keeps it on its schedule.
471                // implements: RATE-041
472                if let Some(delay) = error.rate_limit_delay() {
473                    attempts = task.attempts.max(1);
474                    let until = finished
475                        .saturating_add(i64::try_from(delay.as_nanos()).unwrap_or(i64::MAX));
476                    next_attempt_at_ns = next_attempt_at_ns.max(until);
477                    expires_at_ns = expires_at_ns
478                        .saturating_add(next_attempt_at_ns.saturating_sub(backoff_at_ns));
479                }
480                let next_attempt_duration = next_attempt_at_ns.saturating_sub(finished);
481                tracing::warn!(%error, "Task {} retry failed. Retrying in {next_attempt_duration}ns", task.id);
482                #[cfg(any(test, feature = "test-hooks"))]
483                let fail_reschedule = test_hooks::take_reschedule_failure();
484                #[cfg(not(any(test, feature = "test-hooks")))]
485                let fail_reschedule = false;
486                let stored = if fail_reschedule {
487                    Err(StorageError::DieselResult(
488                        diesel::result::Error::BrokenTransactionManager,
489                    ))
490                } else {
491                    context.db().update_task_and_expiry(
492                        task.id,
493                        attempts,
494                        now,
495                        next_attempt_at_ns,
496                        expires_at_ns,
497                    )
498                };
499                match stored {
500                    Ok(_) => {}
501                    // The row was concurrently deleted (e.g. a dead-row cleanup in
502                    // upsert_pending_self_remove_task crossed this retry). Nothing
503                    // left to reschedule — don't abort the worker loop.
504                    Err(StorageError::DieselResult(diesel::result::Error::NotFound)) => {
505                        tracing::debug!("Task {} vanished before reschedule; skipping", task.id);
506                    }
507                    // The task stays due, so a restarted worker would resend
508                    // it at once. Return the rejection instead: the worker then
509                    // restarts only after the delay.
510                    // implements: RATE-041
511                    Err(e) if error.rate_limit_delay().is_some() => {
512                        tracing::warn!(%e, "Task {} rate-limit hold not stored", task.id);
513                        return Err(error);
514                    }
515                    Err(e) => return Err(e.into()),
516                }
517            }
518        }
519        Ok(())
520    }
521
522    fn next_wakeup(
523        next_attempt_at_ns: Option<i64>,
524        // these are passed in for testing
525        now: i64,
526    ) -> xmtp_common::time::Duration {
527        use xmtp_common::r#const::NS_IN_DAY;
528        let now_plus_one_day = now.saturating_add(NS_IN_DAY);
529        let next_task_wakeup = next_attempt_at_ns.unwrap_or(i64::MAX).min(now_plus_one_day);
530        if now > next_task_wakeup {
531            xmtp_common::time::Duration::from_nanos(0)
532        } else {
533            std::time::Duration::from_nanos((next_task_wakeup - now) as u64)
534        }
535    }
536}
537
538#[cfg(test)]
539mod tests {
540    use super::test_support::{seed, unique_proto};
541    use super::*;
542    use crate::identity::IdentityError;
543    use crate::test_context::real_context;
544    use prost::Message;
545    use xmtp_db::tasks::data_hash_for;
546    use xmtp_proto::xmtp::mls::database::{Task as TaskProto, task::Task as TaskKind};
547
548    /// Handlers that apply each task's pull-in, as the client's dispatch
549    /// does, and finish the task.
550    struct PullInHandler;
551
552    #[xmtp_common::async_trait]
553    impl<Context: CoreContext> TaskHandlers<Context> for PullInHandler {
554        type Error = IdentityError;
555
556        async fn dispatch(
557            &self,
558            context: &Context,
559            task: &DbTask,
560        ) -> Result<TaskOutcome, IdentityError> {
561            let task_proto = TaskProto::decode(task.data.as_slice()).expect("a stored task");
562            let Some(TaskKind::PullInDeadline(pull_in)) = task_proto.task else {
563                panic!("only pull-in tasks are seeded");
564            };
565            apply_pull_in(context, task, pull_in)?;
566            Ok(TaskOutcome::Done)
567        }
568    }
569
570    /// Handlers that reschedule the task with `data_hash` to `target`.
571    struct RescheduleAtHandler {
572        data_hash: Vec<u8>,
573        target: i64,
574    }
575
576    #[xmtp_common::async_trait]
577    impl<Context: CoreContext> TaskHandlers<Context> for RescheduleAtHandler {
578        type Error = IdentityError;
579
580        async fn dispatch(
581            &self,
582            _context: &Context,
583            task: &DbTask,
584        ) -> Result<TaskOutcome, IdentityError> {
585            assert_eq!(task.data_hash, self.data_hash, "only one task is seeded");
586            Ok(TaskOutcome::RescheduleAt(self.target))
587        }
588    }
589
590    /// The failure a [`FailAfter`] run ends with: not retryable, no
591    /// rate-limit delay and no reconnect, as an invalid task hash is.
592    #[derive(Debug, thiserror::Error)]
593    enum TaskFixtureError {
594        #[error(transparent)]
595        Storage(#[from] StorageError),
596        #[error("invalid hash for {id}, expected: {expected}, got: {got}")]
597        InvalidHash {
598            id: i64,
599            expected: String,
600            got: String,
601        },
602    }
603
604    impl RetryableError for TaskFixtureError {
605        fn is_retryable(&self) -> bool {
606            false
607        }
608    }
609
610    impl NeedsDbReconnect for TaskFixtureError {
611        fn needs_db_reconnect(&self) -> bool {
612            false
613        }
614    }
615
616    /// Handlers whose every run works for the given time, then fails.
617    struct FailAfter(xmtp_common::time::Duration);
618
619    #[xmtp_common::async_trait]
620    impl<Context: CoreContext> TaskHandlers<Context> for FailAfter {
621        type Error = TaskFixtureError;
622
623        async fn dispatch(
624            &self,
625            _context: &Context,
626            task: &DbTask,
627        ) -> Result<TaskOutcome, TaskFixtureError> {
628            xmtp_common::time::sleep(self.0).await;
629            Err(TaskFixtureError::InvalidHash {
630                id: task.id.into(),
631                expected: String::new(),
632                got: String::new(),
633            })
634        }
635    }
636
637    #[xmtp_common::test(unwrap_try = true)]
638    async fn done_deletes() {
639        let context = real_context().await;
640        let db = context.db();
641        let now = xmtp_common::time::now_ns();
642        let row = seed(&db, unique_proto(), now - 1, i64::MAX, 0, i32::MAX);
643        TaskWorker::new(context.clone(), PullInHandler)
644            .run_and_reschedule_task(row)
645            .await?;
646        assert!(db.get_tasks()?.is_empty(), "Done task must be deleted");
647    }
648
649    /// A failed task backs off from when its run finished, so a long run
650    /// does not use up its own backoff.
651    #[xmtp_common::test(unwrap_try = true)]
652    async fn failed_task_backs_off_from_run_finish() {
653        use xmtp_common::time::Duration;
654        let context = real_context().await;
655        let db = context.db();
656        let proto = unique_proto();
657        let hash = data_hash_for(&proto);
658        let row = seed(
659            &db,
660            proto,
661            xmtp_common::time::now_ns() - 1,
662            i64::MAX,
663            0,
664            i32::MAX,
665        );
666        let backoff = (row.initial_backoff_duration_ns as f64
667            * (row.backoff_scaling_factor as f64).powi(1)) as i64;
668        let work = Duration::from_secs(2);
669        let started = xmtp_common::time::now_ns();
670        TaskWorker::new(context.clone(), FailAfter(work))
671            .run_and_reschedule_task(row)
672            .await?;
673        let failed = db
674            .get_tasks()?
675            .into_iter()
676            .find(|t| t.data_hash == hash.as_ref())
677            .expect("the failed task is kept");
678        assert_eq!(failed.attempts, 1);
679        assert!(failed.next_attempt_at_ns >= started + work.as_nanos() as i64 + backoff);
680    }
681
682    #[xmtp_common::test(unwrap_try = true)]
683    async fn recurring_task_advances_and_does_not_hot_loop() {
684        let context = real_context().await;
685        let db = context.db();
686        let now = xmtp_common::time::now_ns();
687        let proto = unique_proto();
688        // High attempts + past-due: only the i64::MAX/i32::MAX seed keeps the
689        // reaper (expires < now || attempts >= max) from deleting it.
690        let row = seed(&db, proto.clone(), now - 1, i64::MAX, 1_000_000, i32::MAX);
691        let target = now + xmtp_common::NS_IN_DAY;
692        let handler = RescheduleAtHandler {
693            data_hash: data_hash_for(&proto).to_vec(),
694            target,
695        };
696
697        TaskWorker::new(context.clone(), handler)
698            .run_and_reschedule_task(row)
699            .await?;
700
701        let after = db.get_tasks()?.pop().expect("recurring row must survive");
702        assert_eq!(
703            after.next_attempt_at_ns, target,
704            "deadline must ADVANCE, not stay past-due"
705        );
706        assert_eq!(after.attempts, 0, "success resets the backoff counter");
707    }
708
709    #[xmtp_common::test(unwrap_try = true)]
710    async fn not_yet_due_task_is_not_run_early() {
711        let context = real_context().await;
712        let db = context.db();
713        let now = xmtp_common::time::now_ns();
714        let proto = unique_proto();
715        let far = now + 30 * xmtp_common::NS_IN_DAY;
716        let row = seed(&db, proto.clone(), far, i64::MAX, 0, i32::MAX);
717        // If the guard failed and the task ran, the handler would rewrite next_attempt.
718        let handler = RescheduleAtHandler {
719            data_hash: data_hash_for(&proto).to_vec(),
720            target: now,
721        };
722
723        TaskWorker::new(context.clone(), handler)
724            .run_and_reschedule_task(row)
725            .await?;
726
727        let after = db.get_tasks()?.pop().unwrap();
728        assert_eq!(
729            after.next_attempt_at_ns, far,
730            "not-yet-due task must not run on a 1-day-cap wake"
731        );
732        assert_eq!(
733            after.attempts, 0,
734            "not-yet-due task must not have its attempts incremented"
735        );
736    }
737}
xmtp_mls_core/worker/tasks.rs

xmtp_mls_core/worker/
tasks.rs
