xmtp_mls/worker/
key_package_maintenance.rs

1//! Key-package maintenance as TaskRunner consumers: payload/seed helpers and the
2//! rotate/sweep work the `KpRotation`/`KpDeletion` dispatch arms call into.
3//! Recurrence + nudging come from the generic layer (TaskOutcome, PullInDeadline).
4
5use crate::context::XmtpSharedContext;
6use crate::identity::IdentityError;
7use crate::state_tx::state_write;
8use crate::worker::NeedsDbReconnect;
9use crate::worker::tasks::enqueue_pull_in;
10use thiserror::Error;
11use tls_codec::Serialize;
12use xmtp_configuration::CREATE_PQ_KEY_PACKAGE_EXTENSION;
13use xmtp_db::StorageError;
14use xmtp_db::TransactionOutcome::Continue;
15use xmtp_db::prelude::*;
16use xmtp_db::sql_key_store::{KEY_PACKAGE_REFERENCES, KEY_PACKAGE_WRAPPER_PRIVATE_KEY};
17use xmtp_db::tasks::{NEVER_EXPIRES, NewTask, TaskDataHash, data_hash_for};
18use xmtp_proto::xmtp::mls::database::{
19    KpDeletion, KpRotation, Task as TaskProto, task::Task as TaskKind,
20};
21
22#[derive(Debug, Error)]
23pub enum KeyPackageMaintenanceError {
24    #[error("generic storage error: {0}")]
25    Storage(#[from] StorageError),
26    #[error("generic identity error: {0}")]
27    Identity(#[from] IdentityError),
28    #[error("metadata error: {0}")]
29    Metadata(StorageError),
30    #[error("failed to fetch expired key packages: {0}")]
31    Fetch(StorageError),
32    #[error("failed to delete key package: {0}")]
33    DeleteKeyPackage(IdentityError),
34    #[error("deletion error: {0}")]
35    Deletion(StorageError),
36    #[error("rotation error: {0}")]
37    Rotation(IdentityError),
38}
39
40impl NeedsDbReconnect for KeyPackageMaintenanceError {
41    fn needs_db_reconnect(&self) -> bool {
42        match self {
43            Self::Storage(s) => s.db_needs_connection(),
44            Self::Identity(s) => s.needs_db_reconnect(),
45            Self::Metadata(s) => s.db_needs_connection(),
46            Self::Fetch(s) => s.db_needs_connection(),
47            Self::DeleteKeyPackage(s) => s.needs_db_reconnect(),
48            Self::Deletion(s) => s.db_needs_connection(),
49            Self::Rotation(s) => s.needs_db_reconnect(),
50        }
51    }
52}
53
54pub(crate) fn kp_rotation_proto() -> TaskProto {
55    TaskProto {
56        task: Some(TaskKind::KpRotation(KpRotation {})),
57    }
58}
59
60pub(crate) fn kp_deletion_proto() -> TaskProto {
61    TaskProto {
62        task: Some(TaskKind::KpDeletion(KpDeletion {})),
63    }
64}
65
66pub(crate) fn kp_rotation_hash() -> TaskDataHash {
67    data_hash_for(&kp_rotation_proto())
68}
69
70pub(crate) fn kp_deletion_hash() -> TaskDataHash {
71    data_hash_for(&kp_deletion_proto())
72}
73
74/// Never-expire recurring seed: the reaper's
75/// `expires_at_ns < now || attempts >= max_attempts` check can never fire.
76pub(crate) fn kp_seed(proto: TaskProto, now: i64) -> Result<NewTask, StorageError> {
77    NewTask::builder()
78        .originating_message_sequence_id(0)
79        .expires_at_ns(NEVER_EXPIRES)
80        .max_attempts(i32::MAX)
81        .next_attempt_at_ns(now)
82        .build(proto)
83}
84
85/// Upload a fresh key package when rotation is due and return whether it occurred.
86/// A confirmed receipt starts retirement of keys published earlier by the backend.
87pub(crate) async fn rotate_if_needed<Context: XmtpSharedContext>(
88    context: &Context,
89) -> Result<bool, KeyPackageMaintenanceError> {
90    if !context
91        .db()
92        .is_identity_needs_rotation()
93        .map_err(KeyPackageMaintenanceError::Metadata)?
94    {
95        return Ok(false);
96    }
97    context
98        .identity()
99        .rotate_and_upload_key_package(
100            context.api(),
101            context.mls_storage(),
102            CREATE_PQ_KEY_PACKAGE_EXTENSION,
103        )
104        .await
105        .map_err(KeyPackageMaintenanceError::Rotation)?;
106    Ok(true)
107}
108
109/// Delete one key package's local material (keystore entry + PQ references).
110#[cfg(test)]
111pub(crate) fn delete_key_package<Context: XmtpSharedContext>(
112    context: &Context,
113    hash_ref: Vec<u8>,
114    pq_pub_key: Option<Vec<u8>>,
115) -> Result<(), IdentityError> {
116    state_write(context.mls_storage(), |tx| {
117        delete_key_package_material(&tx.storage(), &hash_ref, pq_pub_key.as_deref())?;
118        Ok::<_, IdentityError>(Continue(()))
119    })?;
120    Ok(())
121}
122
123/// Delete all parts of one key package under the current database writer.
124fn delete_key_package_material(
125    key_store: &impl XmtpMlsStorageProvider,
126    hash_ref: &[u8],
127    pq_pub_key: Option<&[u8]>,
128) -> Result<(), IdentityError> {
129    let openmls_hash_ref = crate::identity::deserialize_key_package_hash_ref(hash_ref)?;
130    let bundle: Option<openmls::prelude::KeyPackageBundle> =
131        key_store.key_package(&openmls_hash_ref)?;
132    if let Some(bundle) = bundle {
133        let init_key = bundle
134            .key_package()
135            .hpke_init_key()
136            .tls_serialize_detached()?;
137        key_store.delete(KEY_PACKAGE_REFERENCES, &init_key)?;
138    }
139    key_store.delete_key_package(&openmls_hash_ref)?;
140
141    if let Some(pq_pub_key) = pq_pub_key {
142        key_store.delete(
143            KEY_PACKAGE_REFERENCES,
144            crate::identity::pq_key_package_references_key(pq_pub_key)?.as_slice(),
145        )?;
146        key_store.delete(KEY_PACKAGE_WRAPPER_PRIVATE_KEY, hash_ref)?;
147    }
148
149    Ok(())
150}
151
152/// Keep expired keys while any durable Welcome remains pending.
153/// Otherwise, delete each expired key's material and history in one transaction.
154pub(crate) fn sweep_expired<Context: XmtpSharedContext>(
155    context: &Context,
156) -> Result<(), KeyPackageMaintenanceError> {
157    state_write(context.mls_storage(), |tx| {
158        let storage = tx.storage();
159        let conn = storage.db();
160        // The durable Welcome queue owns the key-retention obligation.
161        if conn.has_pending_welcomes()? {
162            return Ok::<_, IdentityError>(Continue(()));
163        }
164        let expired = conn.get_expired_key_packages()?;
165        for kp in &expired {
166            delete_key_package_material(
167                &storage,
168                &kp.key_package_hash_ref,
169                kp.post_quantum_public_key.as_deref(),
170            )?;
171            conn.delete_key_package_entry_with_id(kp.id)?;
172        }
173        Ok(Continue(()))
174    })?;
175    Ok(())
176}
177
178/// Queue rotation within five seconds and persist its task in one transaction.
179/// Create the recurring task before its deadline update, then wake after commit.
180pub(crate) fn queue_key_rotation<Context: XmtpSharedContext>(
181    context: &Context,
182) -> Result<(), StorageError> {
183    state_write(context.mls_storage(), |tx| {
184        queue_key_rotation_in(&tx.storage())?;
185        Ok::<_, StorageError>(Continue(()))
186    })?;
187    // In-memory only; must stay outside the transaction.
188    context.task_channels().wake();
189    Ok(())
190}
191
192/// Queue rotation on the caller's writer so Welcome receipt can commit with it.
193/// The caller wakes TaskRunner only after commit. The task survives a lost wake.
194pub(crate) fn queue_key_rotation_in(
195    storage: &impl XmtpMlsStorageProvider,
196) -> Result<(), StorageError> {
197    let now = xmtp_common::time::now_ns();
198    storage
199        .db()
200        .queue_key_rotation_with_nudge(&kp_rotation_hash(), kp_seed(kp_rotation_proto(), now)?)
201}
202
203/// After anything marks superseded KPs for deletion: ensure the KpDeletion
204/// singleton exists (pull-in against a missing target is a no-op), then pull
205/// it in to the earliest pending delete_at. No-op when nothing is marked, so
206/// it is safe (and idempotent) to call on every dispatch.
207pub(crate) fn nudge_deletion<Context: XmtpSharedContext>(
208    context: &Context,
209) -> Result<(), StorageError> {
210    let db = context.db();
211    let now = xmtp_common::time::now_ns();
212    if let Some(at) = db.min_key_package_delete_at_ns()? {
213        db.create_or_ignore_task(kp_seed(kp_deletion_proto(), now)?)?;
214        enqueue_pull_in(context, kp_deletion_hash(), at, NEVER_EXPIRES)?;
215    }
216    Ok(())
217}
218
219/// Idempotent startup seeding + reconcile: pull-ins only LOWER task deadlines to
220/// the live DB columns, repairing rows stranded by a crash mid-nudge.
221pub(crate) fn seed_and_reconcile_kp_tasks<Context: XmtpSharedContext>(
222    context: &Context,
223) -> Result<(), StorageError> {
224    let db = context.db();
225    let now = xmtp_common::time::now_ns();
226    db.create_or_ignore_task(kp_seed(kp_rotation_proto(), now)?)?;
227    db.create_or_ignore_task(kp_seed(kp_deletion_proto(), now)?)?;
228    // None = pre-registration (no identity row): the seed row already fires at
229    // startup; a pull-in to `now` would be redundant noise.
230    if let Some(rot) = db.next_key_package_rotation_ns()? {
231        enqueue_pull_in(context, kp_rotation_hash(), rot, NEVER_EXPIRES)?;
232    }
233    if let Some(del) = db.min_key_package_delete_at_ns()? {
234        enqueue_pull_in(context, kp_deletion_hash(), del, NEVER_EXPIRES)?;
235    }
236    Ok(())
237}
238
239// These tests use the native connection pool and its error variants.
240#[cfg(all(test, not(target_arch = "wasm32")))]
241mod tests {
242    use super::*;
243    use crate::tester;
244    use crate::worker::tasks::TaskWorker;
245    use crate::worker::{WorkerConfig, WorkerKind};
246    use openmls_traits::storage::StorageProvider;
247    use prost::Message;
248    use xmtp_db::ConnectionExt;
249    use xmtp_proto::xmtp::mls::database::Task as TaskProtoDecode;
250
251    /// A `StorageError` that signals the connection pool was dropped.
252    fn disconnect_storage() -> xmtp_db::StorageError {
253        xmtp_db::StorageError::Platform(xmtp_db::PlatformStorageError::PoolNeedsConnection)
254    }
255
256    /// A storage error that is NOT a disconnect — must never trip the contract.
257    fn benign_storage() -> xmtp_db::StorageError {
258        xmtp_db::StorageError::InvalidHmacLength
259    }
260
261    fn no_runner_cfg() -> WorkerConfig {
262        let mut cfg = WorkerConfig::default();
263        cfg.enabled.insert(WorkerKind::TaskRunner, false);
264        cfg
265    }
266
267    fn row_by_hash(db: &impl QueryTasks, hash: impl AsRef<[u8]>) -> Option<xmtp_db::tasks::Task> {
268        db.get_tasks()
269            .expect("get_tasks should not fail")
270            .into_iter()
271            .find(|t| t.data_hash == hash.as_ref())
272    }
273
274    async fn make_rotation_due(db: &impl QueryIdentity) {
275        db.queue_key_package_rotation()
276            .expect("queue_key_package_rotation should not fail"); // column := now + 5s
277        xmtp_common::time::sleep(std::time::Duration::from_secs(6)).await;
278    }
279
280    #[xmtp_common::test]
281    fn kp_errors_forward_db_reconnect() {
282        use crate::worker::NeedsDbReconnect;
283        use crate::worker::tasks::TaskWorkerError;
284        let e = TaskWorkerError::from(KeyPackageMaintenanceError::Storage(disconnect_storage()));
285        assert!(
286            e.needs_db_reconnect(),
287            "DB outage during KP work must trigger supervisor reconnect, not plain backoff"
288        );
289        let e = TaskWorkerError::from(crate::identity::IdentityError::from(disconnect_storage()));
290        assert!(e.needs_db_reconnect());
291        // Keystore pool loss (rotate/delete paths) must also restart the worker.
292        let e = TaskWorkerError::from(crate::identity::IdentityError::OpenMlsStorageError(
293            xmtp_db::sql_key_store::SqlKeyStoreError::Connection(
294                xmtp_db::ConnectionError::Platform(
295                    xmtp_db::PlatformStorageError::PoolNeedsConnection,
296                ),
297            ),
298        ));
299        assert!(e.needs_db_reconnect());
300        // A non-disconnect storage failure must NOT stop the worker.
301        let e = TaskWorkerError::from(KeyPackageMaintenanceError::Storage(benign_storage()));
302        assert!(
303            !e.needs_db_reconnect(),
304            "benign storage errors must back off, not restart the supervisor"
305        );
306    }
307
308    #[xmtp_common::test(unwrap_try = true)]
309    async fn manual_rotation_nudges_deletion() {
310        tester!(alix, worker_config: no_runner_cfg());
311        let db = alix.context.db();
312        assert!(row_by_hash(&db, kp_deletion_hash()).is_none());
313
314        alix.rotate_and_upload_key_package().await?;
315
316        assert!(
317            row_by_hash(&db, kp_deletion_hash()).is_some(),
318            "manual rotation must self-heal the deletion singleton"
319        );
320        let has_pull_in = db.get_tasks()?.into_iter().any(|t| matches!(
321            TaskProtoDecode::decode(t.data.as_slice()).ok().and_then(|p| p.task),
322            Some(TaskKind::PullInDeadline(p)) if p.target_data_hash == kp_deletion_hash().as_ref()
323        ));
324        assert!(
325            has_pull_in,
326            "manual rotation must enqueue a deletion pull-in"
327        );
328    }
329
330    #[xmtp_common::test(unwrap_try = true)]
331    async fn pending_welcome_preserves_expired_keys_until_completion() {
332        use xmtp_db::diesel::prelude::*;
333        use xmtp_db::incoming_envelope::{
334            IncomingLimits, NetworkEntityKind, NewIncomingEnvelope, PendingBudget, StreamTopic,
335        };
336        use xmtp_db::schema::key_package_history::dsl;
337        use xmtp_proto::types::Cursor;
338
339        tester!(alix, disable_workers);
340        let db = alix.context.db();
341        let history = db
342            .find_key_package_history_entries_before_id(i32::MAX)?
343            .pop()?;
344        let hash =
345            crate::identity::deserialize_key_package_hash_ref(&history.key_package_hash_ref)?;
346        db.raw_query(|conn| {
347            diesel::update(dsl::key_package_history.filter(dsl::id.eq(history.id)))
348                .set(dsl::delete_at_ns.eq(0))
349                .execute(conn)
350        })?;
351        let topic = StreamTopic {
352            entity_id: alix.context.installation_id().to_vec(),
353            kind: NetworkEntityKind::Welcome,
354        };
355        let limit = PendingBudget { rows: 1, bytes: 1 };
356        db.admit_ordered_batch(
357            &topic,
358            Cursor(0),
359            &[NewIncomingEnvelope {
360                sequence_id: Cursor(1),
361                envelope: vec![1],
362            }],
363            IncomingLimits {
364                batch: limit,
365                topic: limit,
366                kind: limit,
367            },
368        )?;
369
370        sweep_expired(&alix.context)?;
371        let retained: Option<openmls::prelude::KeyPackageBundle> =
372            alix.context.mls_storage().key_package(&hash)?;
373        assert!(retained.is_some());
374
375        db.complete_pending_envelope(&topic, Cursor(1))?;
376        sweep_expired(&alix.context)?;
377        let removed: Option<openmls::prelude::KeyPackageBundle> =
378            alix.context.mls_storage().key_package(&hash)?;
379        assert!(removed.is_none());
380        assert!(
381            db.find_key_package_history_entry_by_hash_ref(history.key_package_hash_ref)
382                .is_err()
383        );
384    }
385
386    #[xmtp_common::test(unwrap_try = true)]
387    async fn rotation_task_rotates_and_reschedules() {
388        tester!(alix, worker_config: no_runner_cfg());
389        let db = alix.context.db();
390        let now = xmtp_common::time::now_ns();
391        db.create_or_ignore_task(kp_seed(kp_rotation_proto(), now)?)?;
392        make_rotation_due(&db).await;
393
394        let row = row_by_hash(&db, kp_rotation_hash()).unwrap();
395        TaskWorker::run_and_reschedule_task(row, &alix.context).await?;
396
397        assert!(
398            !db.is_identity_needs_rotation()?,
399            "rotation must have happened"
400        );
401        let after = row_by_hash(&db, kp_rotation_hash()).expect("recurring row survives");
402        let col = db.next_key_package_rotation_ns()?.unwrap();
403        assert_eq!(
404            after.next_attempt_at_ns, col,
405            "reschedule must read the live column"
406        );
407        assert_eq!(after.attempts, 0);
408    }
409
410    #[xmtp_common::test(unwrap_try = true)]
411    async fn rotation_ensures_and_pulls_in_deletion_when_singleton_missing() {
412        tester!(alix, worker_config: no_runner_cfg());
413        let db = alix.context.db();
414        let now = xmtp_common::time::now_ns();
415        db.create_or_ignore_task(kp_seed(kp_rotation_proto(), now)?)?;
416        // Deliberately NO KpDeletion seed: the handler must self-heal it.
417        make_rotation_due(&db).await;
418
419        let row = row_by_hash(&db, kp_rotation_hash()).unwrap();
420        TaskWorker::run_and_reschedule_task(row, &alix.context).await?;
421
422        assert!(
423            row_by_hash(&db, kp_deletion_hash()).is_some(),
424            "rotation must recreate a missing KpDeletion singleton"
425        );
426        let has_pull_in = db.get_tasks()?.iter().any(|t| {
427            matches!(
428                TaskProtoDecode::decode(t.data.as_slice()).ok().and_then(|p| p.task),
429                Some(TaskKind::PullInDeadline(p)) if p.target_data_hash == kp_deletion_hash().as_ref()
430            )
431        });
432        assert!(has_pull_in, "rotation must enqueue a deletion pull-in");
433    }
434
435    #[xmtp_common::test(unwrap_try = true)]
436    async fn deletion_task_sweeps_and_reschedules() {
437        tester!(alix, worker_config: no_runner_cfg());
438        let db = alix.context.db();
439        let now = xmtp_common::time::now_ns();
440        db.create_or_ignore_task(kp_seed(kp_deletion_proto(), now)?)?;
441
442        // A rotation marks the superseded KP delete_at = now + 3s (test cfg).
443        make_rotation_due(&db).await;
444        rotate_if_needed(&alix.context).await?;
445        assert!(db.min_key_package_delete_at_ns()?.is_some());
446        xmtp_common::time::sleep(std::time::Duration::from_secs(4)).await; // pass the grace
447
448        let row = row_by_hash(&db, kp_deletion_hash()).unwrap();
449        TaskWorker::run_and_reschedule_task(row, &alix.context).await?;
450
451        assert!(
452            db.get_expired_key_packages()?.is_empty(),
453            "sweep must delete expired KPs"
454        );
455        let after = row_by_hash(&db, kp_deletion_hash()).expect("recurring row survives");
456        assert!(
457            after.next_attempt_at_ns > xmtp_common::time::now_ns(),
458            "deletion reschedules to next pending deadline or far-future"
459        );
460    }
461
462    #[xmtp_common::test(unwrap_try = true)]
463    async fn kp_tasks_seeded_when_workers_run_absent_when_passive() {
464        tester!(alix); // default: TaskRunner on -> seeds present
465        let db = alix.context.db();
466        assert!(row_by_hash(&db, kp_rotation_hash()).is_some());
467        assert!(row_by_hash(&db, kp_deletion_hash()).is_some());
468
469        tester!(bo, worker_config: no_runner_cfg()); // no TaskRunner -> no seeds
470        let db = bo.context.db();
471        assert!(row_by_hash(&db, kp_rotation_hash()).is_none());
472        assert!(row_by_hash(&db, kp_deletion_hash()).is_none());
473    }
474
475    #[xmtp_common::test(unwrap_try = true)]
476    async fn startup_reconcile_pulls_in_far_scheduled_row() {
477        tester!(alix, worker_config: no_runner_cfg());
478        let db = alix.context.db();
479        let now = xmtp_common::time::now_ns();
480        // Stale persisted row 30d out while the column says due-in-5s
481        // (crash-between-writes scenario).
482        db.create_or_ignore_task(kp_seed(kp_rotation_proto(), now)?)?;
483        let row = row_by_hash(&db, kp_rotation_hash()).unwrap();
484        db.update_task(row.id, 0, now, now + 30 * xmtp_common::NS_IN_DAY)?;
485        db.queue_key_package_rotation()?; // column := now + 5s
486
487        seed_and_reconcile_kp_tasks(&alix.context)?;
488
489        let pull_in = db
490            .get_tasks()?
491            .into_iter()
492            .find(|t| {
493                matches!(
494                    TaskProtoDecode::decode(t.data.as_slice()).ok().and_then(|p| p.task),
495                    Some(TaskKind::PullInDeadline(p)) if p.target_data_hash == kp_rotation_hash().as_ref()
496                )
497            })
498            .expect("reconcile must enqueue a rotation pull-in");
499        TaskWorker::run_and_reschedule_task(pull_in, &alix.context).await?;
500
501        let after = row_by_hash(&db, kp_rotation_hash()).unwrap();
502        let col = db.next_key_package_rotation_ns()?.unwrap();
503        assert_eq!(after.next_attempt_at_ns, col);
504    }
505
506    /// KpRotation firing while NOT due must not rotate or seed deletion — it just
507    /// re-syncs its deadline to the column (spurious-wake safety).
508    #[xmtp_common::test(unwrap_try = true)]
509    async fn rotation_task_not_due_reschedules_without_rotating() {
510        tester!(alix, worker_config: no_runner_cfg());
511        let db = alix.context.db();
512        let now = xmtp_common::time::now_ns();
513        db.create_or_ignore_task(kp_seed(kp_rotation_proto(), now)?)?;
514        // Post-registration column is ~now+30d: not due.
515        let row = row_by_hash(&db, kp_rotation_hash()).unwrap();
516        TaskWorker::run_and_reschedule_task(row, &alix.context).await?;
517
518        assert!(
519            row_by_hash(&db, kp_deletion_hash()).is_none(),
520            "must not seed deletion"
521        );
522        assert!(db.min_key_package_delete_at_ns()?.is_none());
523        let after = row_by_hash(&db, kp_rotation_hash()).unwrap();
524        assert_eq!(
525            after.next_attempt_at_ns,
526            db.next_key_package_rotation_ns()?.unwrap()
527        );
528    }
529
530    /// The welcome nudge must self-heal a missing rotation seed (e.g. startup
531    /// seeding never ran) instead of enqueuing a dropped-on-miss pull-in.
532    #[xmtp_common::test(unwrap_try = true)]
533    async fn welcome_nudge_selfheals_missing_rotation_seed() {
534        tester!(alix, worker_config: no_runner_cfg()); // no TaskRunner -> no seeds
535        let db = alix.context.db();
536        assert!(row_by_hash(&db, kp_rotation_hash()).is_none());
537
538        queue_key_rotation(&alix.context)?;
539
540        assert!(
541            row_by_hash(&db, kp_rotation_hash()).is_some(),
542            "nudge must recreate the missing KpRotation singleton"
543        );
544        let has_pull_in = db.get_tasks()?.iter().any(|t| {
545            matches!(
546                TaskProtoDecode::decode(t.data.as_slice()).ok().and_then(|p| p.task),
547                Some(TaskKind::PullInDeadline(p)) if p.target_data_hash == kp_rotation_hash().as_ref()
548            )
549        });
550        assert!(has_pull_in, "nudge must enqueue a rotation pull-in");
551    }
552
553    /// Regression: welcome nudge must pull the parked rotation task in even when
554    /// the seed dispatched BEFORE the column was lowered (the startup race).
555    #[xmtp_common::test(unwrap_try = true)]
556    async fn welcome_nudge_pulls_in_parked_rotation() {
557        tester!(alix, worker_config: no_runner_cfg());
558        let db = alix.context.db();
559        let now = xmtp_common::time::now_ns();
560        db.create_or_ignore_task(kp_seed(kp_rotation_proto(), now)?)?;
561        // Simulate the seed having already dispatched not-due: park it on the column (~+30d).
562        let parked = row_by_hash(&db, kp_rotation_hash()).unwrap();
563        TaskWorker::run_and_reschedule_task(parked, &alix.context).await?;
564        let parked_at = row_by_hash(&db, kp_rotation_hash())
565            .unwrap()
566            .next_attempt_at_ns;
567        assert!(
568            parked_at > now + xmtp_common::NS_IN_DAY,
569            "precondition: parked far out"
570        );
571
572        queue_key_rotation(&alix.context)?; // welcome: column + pull-in, atomically
573
574        let pull_in = db
575            .get_tasks()?
576            .into_iter()
577            .find(|t| {
578                matches!(
579                    TaskProtoDecode::decode(t.data.as_slice()).ok().and_then(|p| p.task),
580                    Some(TaskKind::PullInDeadline(p)) if p.target_data_hash == kp_rotation_hash().as_ref()
581                )
582            })
583            .expect("nudge must enqueue a durable pull-in");
584        TaskWorker::run_and_reschedule_task(pull_in, &alix.context).await?;
585
586        let after = row_by_hash(&db, kp_rotation_hash()).unwrap();
587        let col = db.next_key_package_rotation_ns()?.unwrap();
588        assert_eq!(
589            after.next_attempt_at_ns, col,
590            "rotation row must be pulled in to the lowered column"
591        );
592        // 5s queue debounce + 2s slack for local ops between `now` and the queue call.
593        assert!(after.next_attempt_at_ns <= now + 7 * xmtp_common::NS_IN_SEC);
594    }
595}
xmtp_mls/worker/key_package_maintenance.rs

xmtp_mls/worker/
key_package_maintenance.rs
