xmtp_db/encrypted_store/
tasks.rs

1use super::{ConnectionExt, db_connection::DbConnection, schema::tasks};
2use crate::StorageError;
3use derive_builder::Builder;
4use diesel::prelude::*;
5use prost::Message;
6use xmtp_common::{NS_IN_DAY, NS_IN_SEC, time::now_ns};
7use xmtp_proto::types::GroupId;
8use xmtp_proto::xmtp::mls::database::{Task as TaskProto, task::Task as TaskKind};
9
10#[derive(Queryable, Identifiable, Debug, Clone)]
11#[diesel(table_name = tasks)]
12#[diesel(primary_key(id))]
13pub struct Task {
14    pub id: i32,
15    pub originating_message_sequence_id: i64,
16
17    pub created_at_ns: i64,
18    pub expires_at_ns: i64,
19    pub attempts: i32,
20    pub max_attempts: i32,
21    pub last_attempted_at_ns: i64,
22    pub backoff_scaling_factor: f32,
23    pub max_backoff_duration_ns: i64,
24    pub initial_backoff_duration_ns: i64,
25    pub next_attempt_at_ns: i64,
26    pub data_hash: Vec<u8>,
27    pub data: Vec<u8>,
28}
29
30#[derive(Insertable, Debug, PartialEq, Clone, Builder)]
31#[diesel(table_name = tasks)]
32#[builder(build_fn(skip))]
33pub struct NewTask {
34    pub originating_message_sequence_id: i64,
35
36    pub created_at_ns: i64,
37    pub expires_at_ns: i64,
38    pub attempts: i32,
39    pub max_attempts: i32,
40    pub last_attempted_at_ns: i64,
41    pub backoff_scaling_factor: f32,
42    pub max_backoff_duration_ns: i64,
43    pub initial_backoff_duration_ns: i64,
44    pub next_attempt_at_ns: i64,
45    #[builder(setter(skip))]
46    pub data_hash: Vec<u8>,
47    #[builder(setter(skip))]
48    pub data: Vec<u8>,
49}
50
51impl NewTask {
52    pub fn builder() -> NewTaskBuilder {
53        NewTaskBuilder::default()
54    }
55}
56
57impl NewTaskBuilder {
58    pub fn build(&mut self, task: TaskProto) -> Result<NewTask, StorageError> {
59        use derive_builder::UninitializedFieldError;
60        let err = |s: &'static str| UninitializedFieldError::new(s);
61        let data = task.encode_to_vec();
62        let data_hash = xmtp_common::sha256_array(&data).to_vec();
63        let new_task = NewTask {
64            originating_message_sequence_id: self
65                .originating_message_sequence_id
66                .ok_or_else(|| err("originating_message_sequence_id"))?,
67            created_at_ns: self.created_at_ns.unwrap_or_else(now_ns),
68            expires_at_ns: self
69                .expires_at_ns
70                .unwrap_or_else(|| now_ns() + NS_IN_DAY * 3),
71            attempts: self.attempts.unwrap_or(0),
72            max_attempts: self.max_attempts.unwrap_or(20),
73            last_attempted_at_ns: self.last_attempted_at_ns.unwrap_or_else(now_ns),
74            backoff_scaling_factor: self.backoff_scaling_factor.unwrap_or(1.5),
75            max_backoff_duration_ns: self.max_backoff_duration_ns.unwrap_or(60 * NS_IN_SEC),
76            initial_backoff_duration_ns: self.initial_backoff_duration_ns.unwrap_or(2 * NS_IN_SEC),
77            next_attempt_at_ns: self.next_attempt_at_ns.unwrap_or_else(now_ns),
78            data_hash,
79            data,
80        };
81        Ok(new_task)
82    }
83}
84
85// impl_store_or_ignore!(Task, tasks);
86
87/// A task row's identity: sha256 over the prost-encoded payload. Payload
88/// encodings must stay canonical — never add protobuf map fields to task
89/// messages (map entry order is nondeterministic); see the pinned-encoding test.
90#[derive(Clone, Copy, PartialEq, Eq)]
91pub struct TaskDataHash([u8; 32]);
92
93impl TaskDataHash {
94    pub fn to_vec(&self) -> Vec<u8> {
95        self.0.to_vec()
96    }
97}
98
99impl AsRef<[u8]> for TaskDataHash {
100    fn as_ref(&self) -> &[u8] {
101        &self.0
102    }
103}
104
105impl std::fmt::Debug for TaskDataHash {
106    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
107        write!(f, "TaskDataHash({})", hex::encode(self.0))
108    }
109}
110
111impl TryFrom<&[u8]> for TaskDataHash {
112    type Error = std::array::TryFromSliceError;
113    fn try_from(v: &[u8]) -> Result<Self, Self::Error> {
114        Ok(Self(v.try_into()?))
115    }
116}
117
118/// Compute a task payload's `data_hash` exactly as `NewTaskBuilder::build` does.
119pub fn data_hash_for(task: &TaskProto) -> TaskDataHash {
120    let bytes = task.encode_to_vec();
121    // Hash-as-identity requires deterministic encoding; a map field in any task
122    // payload would break this (HashMap iteration order varies per process).
123    debug_assert_eq!(
124        bytes,
125        task.encode_to_vec(),
126        "task payload encoding is nondeterministic — hashes cannot identify rows"
127    );
128    TaskDataHash(xmtp_common::sha256_array(&bytes))
129}
130
131/// Never reaped by expiry: the row's lifetime is bounded by other means (a
132/// recurring row lives forever; an applied pull-in self-deletes).
133pub const NEVER_EXPIRES: i64 = i64::MAX;
134
135pub trait QueryTasks {
136    fn create_task(&self, task: NewTask) -> Result<Task, StorageError>;
137
138    /// Idempotent enqueue: a payload-identical duplicate is a no-op (the existing
139    /// row wins; OR IGNORE swallows any constraint hit, not just data_hash UNIQUE).
140    fn create_or_ignore_task(&self, task: NewTask) -> Result<(), StorageError>;
141
142    /// Lower a task's `next_attempt_at_ns` to `MIN(current, at_ns)` — never raises.
143    /// Returns whether a row matched; a missing target is a no-op (`false`).
144    /// TaskWorker dispatch thread only (sole rescheduler).
145    fn pull_in_task_deadline(
146        &self,
147        target_data_hash: &TaskDataHash,
148        at_ns: i64,
149    ) -> Result<bool, StorageError>;
150
151    fn get_tasks(&self) -> Result<Vec<Task>, StorageError>;
152
153    fn get_next_task(&self) -> Result<Option<Task>, StorageError>;
154
155    /// Ensure exactly one live `ProcessPendingSelfRemove` task exists for
156    /// `group_id`. Clears only dead rows (expired / attempts-exhausted) then
157    /// insert-or-ignores, so a live retrying task keeps its backoff and is never
158    /// deleted out from under the TaskRunner, while a stale dead row can't block
159    /// a fresh retry via the `data_hash` unique constraint.
160    fn upsert_pending_self_remove_task(
161        &self,
162        group_id: &GroupId,
163        task: NewTask,
164    ) -> Result<(), StorageError>;
165
166    fn update_task(
167        &self,
168        id: i32,
169        attempts: i32,
170        last_attempted_at_ns: i64,
171        next_attempt_at_ns: i64,
172    ) -> Result<Task, StorageError>;
173
174    fn delete_task(&self, id: i32) -> Result<bool, StorageError>;
175}
176
177impl<T: QueryTasks> QueryTasks for &'_ T {
178    fn create_task(&self, task: NewTask) -> Result<Task, StorageError> {
179        (**self).create_task(task)
180    }
181
182    fn create_or_ignore_task(&self, task: NewTask) -> Result<(), StorageError> {
183        (**self).create_or_ignore_task(task)
184    }
185
186    fn pull_in_task_deadline(
187        &self,
188        target_data_hash: &TaskDataHash,
189        at_ns: i64,
190    ) -> Result<bool, StorageError> {
191        (**self).pull_in_task_deadline(target_data_hash, at_ns)
192    }
193
194    fn get_tasks(&self) -> Result<Vec<Task>, StorageError> {
195        (**self).get_tasks()
196    }
197
198    fn get_next_task(&self) -> Result<Option<Task>, StorageError> {
199        (**self).get_next_task()
200    }
201
202    fn upsert_pending_self_remove_task(
203        &self,
204        group_id: &GroupId,
205        task: NewTask,
206    ) -> Result<(), StorageError> {
207        (**self).upsert_pending_self_remove_task(group_id, task)
208    }
209
210    fn update_task(
211        &self,
212        id: i32,
213        attempts: i32,
214        last_attempted_at_ns: i64,
215        next_attempt_at_ns: i64,
216    ) -> Result<Task, StorageError> {
217        (**self).update_task(id, attempts, last_attempted_at_ns, next_attempt_at_ns)
218    }
219
220    fn delete_task(&self, id: i32) -> Result<bool, StorageError> {
221        (**self).delete_task(id)
222    }
223}
224
225impl<C: ConnectionExt> QueryTasks for DbConnection<C> {
226    fn create_task(&self, task: NewTask) -> Result<Task, StorageError> {
227        self.raw_query(|conn| {
228            diesel::insert_into(tasks::table)
229                .values(task)
230                .get_result::<Task>(conn)
231        })
232        .map_err(Into::into)
233    }
234
235    fn create_or_ignore_task(&self, task: NewTask) -> Result<(), StorageError> {
236        // A single INSERT OR IGNORE is atomic; no explicit transaction needed.
237        self.raw_query(|conn| {
238            diesel::insert_or_ignore_into(tasks::table)
239                .values(task)
240                .execute(conn)
241        })?;
242        Ok(())
243    }
244
245    fn pull_in_task_deadline(
246        &self,
247        target_data_hash: &TaskDataHash,
248        at_ns: i64,
249    ) -> Result<bool, StorageError> {
250        use diesel::dsl::sql;
251        use diesel::sql_types::BigInt;
252        let matched = self.raw_query(|conn| {
253            diesel::update(tasks::table.filter(tasks::data_hash.eq(target_data_hash.as_ref())))
254                .set(
255                    tasks::next_attempt_at_ns.eq(sql::<BigInt>("MIN(next_attempt_at_ns, ")
256                        .bind::<BigInt, _>(at_ns)
257                        .sql(")")),
258                )
259                .execute(conn)
260        })?;
261        Ok(matched > 0)
262    }
263
264    fn get_tasks(&self) -> Result<Vec<Task>, StorageError> {
265        self.raw_query(|conn| tasks::table.load::<Task>(conn))
266            .map_err(Into::into)
267    }
268
269    fn get_next_task(&self) -> Result<Option<Task>, StorageError> {
270        self.raw_query(|conn| {
271            tasks::table
272                .order(tasks::next_attempt_at_ns)
273                .first::<Task>(conn)
274                .optional()
275        })
276        .map_err(Into::into)
277    }
278
279    fn upsert_pending_self_remove_task(
280        &self,
281        group_id: &GroupId,
282        task: NewTask,
283    ) -> Result<(), StorageError> {
284        let now = now_ns();
285        self.raw_query(|conn| {
286            conn.transaction(|conn| {
287                // Clear only DEAD rows for this group (expired or attempts
288                // exhausted), then insert-or-ignore. We deliberately leave a LIVE
289                // row untouched: deleting it would reset the TaskRunner's backoff
290                // (resurrecting an intentionally-delayed task) and could race the
291                // worker into calling update_task on a now-deleted id. The new
292                // task carries the same data (group_id only), so the unique
293                // data_hash constraint dedups it against any live row; clearing
294                // dead rows first frees that hash so a fresh retry can take over.
295                let rows: Vec<(i32, i32, i32, i64, Vec<u8>)> = tasks::table
296                    .select((
297                        tasks::id,
298                        tasks::attempts,
299                        tasks::max_attempts,
300                        tasks::expires_at_ns,
301                        tasks::data,
302                    ))
303                    .load(conn)?;
304                for (id, attempts, max_attempts, expires_at_ns, data) in rows {
305                    let is_self_remove = matches!(
306                        TaskProto::decode(data.as_slice()).ok().and_then(|t| t.task),
307                        Some(TaskKind::ProcessPendingSelfRemove(p)) if p.group_id == group_id.as_slice()
308                    );
309                    let is_dead = expires_at_ns < now || attempts >= max_attempts;
310                    if is_self_remove && is_dead {
311                        diesel::delete(tasks::table.filter(tasks::id.eq(id))).execute(conn)?;
312                    }
313                }
314                diesel::insert_or_ignore_into(tasks::table)
315                    .values(task)
316                    .execute(conn)?;
317                Ok(())
318            })
319        })
320        .map_err(Into::into)
321    }
322
323    fn update_task(
324        &self,
325        id: i32,
326        attempts: i32,
327        last_attempted_at_ns: i64,
328        next_attempt_at_ns: i64,
329    ) -> Result<Task, StorageError> {
330        self.raw_query(|conn| {
331            diesel::update(tasks::table.filter(tasks::id.eq(id)))
332                .set((
333                    tasks::attempts.eq(attempts),
334                    tasks::last_attempted_at_ns.eq(last_attempted_at_ns),
335                    tasks::next_attempt_at_ns.eq(next_attempt_at_ns),
336                ))
337                .get_result::<Task>(conn)
338        })
339        .map_err(Into::into)
340    }
341
342    fn delete_task(&self, id: i32) -> Result<bool, StorageError> {
343        let num_deleted = self.raw_query(|conn| {
344            diesel::delete(tasks::table.filter(tasks::id.eq(id))).execute(conn)
345        })?;
346        Ok(num_deleted == 1)
347    }
348}
349
350#[cfg(test)]
351pub(crate) mod tests {
352    use super::*;
353    use crate::test_utils::with_connection;
354
355    #[xmtp_common::test]
356    fn get_tasks_returns_empty_list_initially() {
357        with_connection(|conn| {
358            let tasks = conn.get_tasks().unwrap();
359            assert!(tasks.is_empty());
360        })
361    }
362
363    #[xmtp_common::test]
364    fn update_task_returns_error_when_not_found() {
365        with_connection(|conn| {
366            // Try to update a task that doesn't exist
367            let result = conn.update_task(999, 5, 1000, 2000);
368            // The update should fail when the task doesn't exist
369            assert!(result.is_err());
370        })
371    }
372
373    #[xmtp_common::test]
374    fn delete_task_returns_false_when_not_found() {
375        with_connection(|conn| {
376            let deleted = conn.delete_task(999).unwrap();
377            assert!(!deleted);
378        })
379    }
380
381    // Generate a random task data for testing to ensure that the hashes are unique
382    fn gen_task_data() -> TaskProto {
383        TaskProto {
384            task: Some(xmtp_proto::xmtp::mls::database::task::Task::PullInDeadline(
385                xmtp_proto::xmtp::mls::database::PullInDeadline {
386                    target_data_hash: xmtp_common::rand_vec::<32>(),
387                    not_later_than_ns: 1000,
388                },
389            )),
390        }
391    }
392
393    #[xmtp_common::test]
394    fn all_task_operations_work_together() {
395        with_connection(|conn| {
396            let now = xmtp_common::time::now_ns();
397
398            // 1. Create first task (should be next to run)
399            let task1 = NewTaskBuilder::default()
400                .originating_message_sequence_id(1)
401                .created_at_ns(now)
402                .expires_at_ns(now + 3_600_000_000_000)
403                .attempts(0)
404                .max_attempts(5)
405                .last_attempted_at_ns(0)
406                .backoff_scaling_factor(1.5)
407                .max_backoff_duration_ns(600_000_000_000)
408                .initial_backoff_duration_ns(2_000_000_000)
409                .next_attempt_at_ns(now + 1000) // Later attempt time
410                .build(gen_task_data())
411                .unwrap();
412
413            // 2. Create second task (should be first to run)
414            let task2 = NewTaskBuilder::default()
415                .originating_message_sequence_id(2)
416                .created_at_ns(now)
417                .expires_at_ns(now + 7_200_000_000_000) // 2 hours from now
418                .attempts(0)
419                .max_attempts(3)
420                .last_attempted_at_ns(0)
421                .backoff_scaling_factor(2.0)
422                .max_backoff_duration_ns(300_000_000_000)
423                .initial_backoff_duration_ns(1_000_000_000)
424                .next_attempt_at_ns(now + 500) // Earlier attempt time - should be next
425                .build(gen_task_data())
426                .unwrap();
427
428            // 3. Verify no tasks initially
429            assert!(conn.get_next_task().unwrap().is_none());
430            assert!(conn.get_tasks().unwrap().is_empty());
431
432            // 4. Create both tasks
433            let created_task1 = conn.create_task(task1).unwrap();
434            let created_task2 = conn.create_task(task2).unwrap();
435
436            let task1_id = created_task1.id;
437            let task2_id = created_task2.id;
438            assert!(task1_id >= 0, "task1_id: {task1_id}");
439            assert!(task2_id >= 0, "task2_id: {task2_id}");
440            assert_ne!(task1_id, task2_id);
441
442            // 5. Verify both tasks appear in get_tasks
443            let all_tasks = conn.get_tasks().unwrap();
444            assert_eq!(all_tasks.len(), 2);
445
446            // 6. Verify get_next_task returns the task with earlier next_attempt_at_ns (task2)
447            let next_task = conn.get_next_task().unwrap();
448            assert!(next_task.is_some());
449            let next_task = next_task.unwrap();
450            assert_eq!(next_task.id, task2_id);
451            assert_eq!(next_task.next_attempt_at_ns, now + 500);
452
453            // 7. Update task1 to have an even earlier next_attempt_at_ns
454            let updated_task1 = conn
455                .update_task(
456                    task1_id,
457                    1,          // attempts
458                    now + 2000, // last_attempted_at_ns
459                    now + 200,  // next_attempt_at_ns - now earliest
460                )
461                .unwrap();
462
463            // Verify the update
464            assert_eq!(updated_task1.id, task1_id);
465            assert_eq!(updated_task1.attempts, 1);
466            assert_eq!(updated_task1.next_attempt_at_ns, now + 200);
467
468            // 8. Verify get_next_task now returns task1 (earliest next_attempt_at_ns)
469            let next_task = conn.get_next_task().unwrap();
470            assert!(next_task.is_some());
471            let next_task = next_task.unwrap();
472            assert_eq!(next_task.id, task1_id);
473            assert_eq!(next_task.next_attempt_at_ns, now + 200);
474
475            // 9. Verify both tasks appear in get_tasks with correct data
476            let all_tasks_after_update = conn.get_tasks().unwrap();
477            assert_eq!(all_tasks_after_update.len(), 2);
478
479            // Find each task by ID
480            let updated_task1_in_list = all_tasks_after_update
481                .iter()
482                .find(|t| t.id == task1_id)
483                .unwrap();
484            let task2_in_list = all_tasks_after_update
485                .iter()
486                .find(|t| t.id == task2_id)
487                .unwrap();
488
489            assert_eq!(updated_task1_in_list.attempts, 1);
490            assert_eq!(updated_task1_in_list.next_attempt_at_ns, now + 200);
491            assert_eq!(task2_in_list.attempts, 0);
492            assert_eq!(task2_in_list.next_attempt_at_ns, now + 500);
493
494            // 10. Delete task1
495            let deleted = conn.delete_task(task1_id).unwrap();
496            assert!(deleted);
497
498            // 11. Verify get_next_task now returns task2
499            let next_task = conn.get_next_task().unwrap();
500            assert!(next_task.is_some());
501            let next_task = next_task.unwrap();
502            assert_eq!(next_task.id, task2_id);
503
504            // 12. Verify only task2 remains in get_tasks
505            let remaining_tasks = conn.get_tasks().unwrap();
506            assert_eq!(remaining_tasks.len(), 1);
507            assert_eq!(remaining_tasks[0].id, task2_id);
508
509            // 13. Delete task2
510            let deleted = conn.delete_task(task2_id).unwrap();
511            assert!(deleted);
512
513            // 14. Verify no tasks remain
514            let all_tasks_after_delete = conn.get_tasks().unwrap();
515            assert!(all_tasks_after_delete.is_empty());
516            assert!(conn.get_next_task().unwrap().is_none());
517
518            // 15. Verify delete returns false for non-existent task
519            let deleted_again = conn.delete_task(task1_id).unwrap();
520            assert!(!deleted_again);
521        })
522    }
523
524    #[xmtp_common::test]
525    fn data_hash_for_matches_builder() {
526        let proto = gen_task_data();
527        let task = NewTask::builder()
528            .originating_message_sequence_id(0)
529            .build(proto.clone())
530            .unwrap();
531        assert_eq!(task.data_hash, data_hash_for(&proto).as_ref());
532    }
533
534    /// data_hash values live in persisted rows and must match across app
535    /// upgrades. If this test fails, prost's encoding of these payloads drifted:
536    /// that ORPHANS every existing recurring/pull-in row. Do NOT update the
537    /// constants without a row-migration story.
538    #[xmtp_common::test]
539    fn data_hash_encoding_is_pinned() {
540        use xmtp_proto::xmtp::mls::database::{
541            AddMissingInstallations, KpDeletion, KpRotation, PullInDeadline,
542        };
543        let rotation = TaskProto {
544            task: Some(TaskKind::KpRotation(KpRotation {})),
545        };
546        let deletion = TaskProto {
547            task: Some(TaskKind::KpDeletion(KpDeletion {})),
548        };
549        // Empty singleton payloads: one tag byte + zero length.
550        assert_eq!(rotation.encode_to_vec(), [0x2a, 0x00]);
551        assert_eq!(deletion.encode_to_vec(), [0x32, 0x00]);
552        assert_eq!(
553            hex::encode(data_hash_for(&rotation)),
554            "17d5f5a33ab5f6aed0395d2bc0a4e5df61d92441ea8d77b0952c01bc8aa8bde0"
555        );
556        assert_eq!(
557            hex::encode(data_hash_for(&deletion)),
558            "913da1f8df6f8fd47593840d533ba0458cc9873996bf310460abb495b34c232a"
559        );
560        // Non-empty payload sample (bytes + i64 fields).
561        let pull_in = TaskProto {
562            task: Some(TaskKind::PullInDeadline(PullInDeadline {
563                target_data_hash: vec![0x11; 32],
564                not_later_than_ns: 1_234_567_890,
565            })),
566        };
567        assert_eq!(
568            hex::encode(data_hash_for(&pull_in)),
569            "16b424873a34096e5157ab9f0a31e80dff1d23ebd8b1aab2b948a4732abfc849"
570        );
571        // AddMissingInstallations (bytes group_id): outer tag for oneof field 7
572        // is (7<<3)|2 = 0x3a, LEN 0x22 (34 = 2-byte inner header + 32 bytes);
573        // inner field 1 tag 0x0a, LEN 0x20.
574        let add = TaskProto {
575            task: Some(TaskKind::AddMissingInstallations(AddMissingInstallations {
576                group_id: vec![0x22; 32],
577            })),
578        };
579        assert_eq!(add.encode_to_vec()[..4], [0x3a, 0x22, 0x0a, 0x20]);
580        assert_eq!(
581            hex::encode(data_hash_for(&add)),
582            "2180ffa08703d0dcc600a070da9d162da6821e29176af24e7a9955af2cf43764"
583        );
584        // Determinism under repetition and a decode round-trip.
585        let bytes = pull_in.encode_to_vec();
586        for _ in 0..100 {
587            assert_eq!(pull_in.encode_to_vec(), bytes);
588        }
589        let decoded = TaskProto::decode(bytes.as_slice()).unwrap();
590        assert_eq!(decoded.encode_to_vec(), bytes);
591    }
592
593    #[xmtp_common::test]
594    fn create_or_ignore_task_is_idempotent() {
595        with_connection(|conn| {
596            let proto = gen_task_data();
597            let mk = || {
598                NewTask::builder()
599                    .originating_message_sequence_id(0)
600                    .build(proto.clone())
601                    .unwrap()
602            };
603            conn.create_or_ignore_task(mk()).unwrap();
604            // Second byte-identical insert must be a silent no-op, NOT a
605            // unique-constraint error (plain create_task would error here).
606            conn.create_or_ignore_task(mk()).unwrap();
607            assert_eq!(conn.get_tasks().unwrap().len(), 1);
608        })
609    }
610
611    #[xmtp_common::test]
612    fn pull_in_lowers_deadline() {
613        with_connection(|conn| {
614            let proto = gen_task_data();
615            let now = now_ns();
616            let task = NewTask::builder()
617                .originating_message_sequence_id(0)
618                .next_attempt_at_ns(now + NS_IN_DAY)
619                .build(proto.clone())
620                .unwrap();
621            conn.create_or_ignore_task(task).unwrap();
622            let hash = data_hash_for(&proto);
623
624            // Lowers a far-out deadline.
625            assert!(conn.pull_in_task_deadline(&hash, now + 5).unwrap());
626            assert_eq!(
627                conn.get_next_task().unwrap().unwrap().next_attempt_at_ns,
628                now + 5
629            );
630
631            // Never raises (MIN): a later ceiling keeps the row but not the value.
632            assert!(conn.pull_in_task_deadline(&hash, now + NS_IN_DAY).unwrap());
633            assert_eq!(
634                conn.get_next_task().unwrap().unwrap().next_attempt_at_ns,
635                now + 5
636            );
637
638            // Missing target: no-op reported as false, no error.
639            let absent = TaskDataHash::try_from([0xAAu8; 32].as_slice()).unwrap();
640            assert!(!conn.pull_in_task_deadline(&absent, now).unwrap());
641            assert_eq!(
642                conn.get_next_task().unwrap().unwrap().next_attempt_at_ns,
643                now + 5
644            );
645        })
646    }
647
648    #[xmtp_common::test(unwrap_try = true)]
649    fn upsert_pending_self_remove_dedups_per_group() {
650        use xmtp_proto::xmtp::mls::database::ProcessPendingSelfRemove;
651        let build = |gid: &GroupId| {
652            let proto = TaskProto {
653                task: Some(TaskKind::ProcessPendingSelfRemove(
654                    ProcessPendingSelfRemove {
655                        group_id: gid.to_vec(),
656                    },
657                )),
658            };
659            NewTask::builder()
660                .originating_message_sequence_id(0)
661                .build(proto)
662                .unwrap()
663        };
664        with_connection(|conn| {
665            // First upsert inserts; a second for the same group dedups, not piles up.
666            conn.upsert_pending_self_remove_task(&GroupId::ONE, build(&GroupId::ONE))?;
667            conn.upsert_pending_self_remove_task(&GroupId::ONE, build(&GroupId::ONE))?;
668            assert_eq!(conn.get_tasks()?.len(), 1);
669
670            // A different group gets its own task.
671            conn.upsert_pending_self_remove_task(&GroupId::TWO, build(&GroupId::TWO))?;
672            assert_eq!(conn.get_tasks()?.len(), 2);
673        })
674    }
675
676    #[xmtp_common::test(unwrap_try = true)]
677    fn upsert_preserves_live_task_but_replaces_dead_one() {
678        use xmtp_proto::xmtp::mls::database::ProcessPendingSelfRemove;
679        let proto = |gid: &GroupId| TaskProto {
680            task: Some(TaskKind::ProcessPendingSelfRemove(
681                ProcessPendingSelfRemove {
682                    group_id: gid.to_vec(),
683                },
684            )),
685        };
686        with_connection(|conn| {
687            // A live task that has already retried twice and backed off.
688            let now = now_ns();
689            let live = NewTask::builder()
690                .originating_message_sequence_id(0)
691                .attempts(2)
692                .next_attempt_at_ns(now + NS_IN_DAY)
693                .build(proto(&GroupId::ONE))?;
694            conn.create_task(live)?;
695
696            // Re-upsert must NOT reset its backoff: the live row is left in place.
697            conn.upsert_pending_self_remove_task(&GroupId::ONE, {
698                NewTask::builder()
699                    .originating_message_sequence_id(0)
700                    .next_attempt_at_ns(now)
701                    .build(proto(&GroupId::ONE))?
702            })?;
703            let tasks = conn.get_tasks()?;
704            assert_eq!(tasks.len(), 1);
705            assert_eq!(tasks[0].attempts, 2);
706            assert_eq!(tasks[0].next_attempt_at_ns, now + NS_IN_DAY);
707
708            // A dead task (attempts exhausted) IS replaced with a fresh retry.
709            let dead = NewTask::builder()
710                .originating_message_sequence_id(0)
711                .attempts(20)
712                .max_attempts(20)
713                .build(proto(&GroupId::TWO))?;
714            conn.create_task(dead)?;
715            conn.upsert_pending_self_remove_task(&GroupId::TWO, {
716                NewTask::builder()
717                    .originating_message_sequence_id(0)
718                    .attempts(0)
719                    .build(proto(&GroupId::TWO))?
720            })?;
721            let two: Vec<_> = conn
722                .get_tasks()?
723                .into_iter()
724                .filter(|t| {
725                    matches!(
726                        TaskProto::decode(t.data.as_slice()).ok().and_then(|p| p.task),
727                        Some(TaskKind::ProcessPendingSelfRemove(p)) if p.group_id == GroupId::TWO.as_slice()
728                    )
729                })
730                .collect();
731            assert_eq!(two.len(), 1);
732            assert_eq!(two[0].attempts, 0);
733        })
734    }
735}
xmtp_db/encrypted_store/tasks.rs

xmtp_db/encrypted_store/
tasks.rs
