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 emit after commit.
180pub(crate) fn queue_key_rotation<Context: XmtpSharedContext>(
181    context: &Context,
182) -> Result<(), StorageError> {
183    crate::state_tx::state_write_with_events(
184        context.mls_storage(),
185        context.events(),
186        |tx, events| {
187            queue_key_rotation_in(&tx.storage())?;
188            xmtp_events::EventWriter::emit(
189                events,
190                None,
191                Some(crate::subscriptions::internal::InternalEvent::TaskScheduled),
192            );
193            Ok::<_, StorageError>(Continue(()))
194        },
195    )?;
196    Ok(())
197}
198
199/// Queue rotation on the caller's writer so Welcome receipt can commit with it.
200/// The caller emits the task fact in its event buffer.
201pub(crate) fn queue_key_rotation_in(
202    storage: &impl XmtpMlsStorageProvider,
203) -> Result<(), StorageError> {
204    let now = xmtp_common::time::now_ns();
205    storage
206        .db()
207        .queue_key_rotation_with_nudge(&kp_rotation_hash(), kp_seed(kp_rotation_proto(), now)?)
208}
209
210/// After anything marks superseded KPs for deletion: ensure the KpDeletion
211/// singleton exists (pull-in against a missing target is a no-op), then pull
212/// it in to the earliest pending delete_at. No-op when nothing is marked, so
213/// it is safe (and idempotent) to call on every dispatch.
214pub(crate) fn nudge_deletion<Context: XmtpSharedContext>(
215    context: &Context,
216) -> Result<(), StorageError> {
217    let db = context.db();
218    let now = xmtp_common::time::now_ns();
219    if let Some(at) = db.min_key_package_delete_at_ns()? {
220        db.create_or_ignore_task(kp_seed(kp_deletion_proto(), now)?)?;
221        enqueue_pull_in(context, kp_deletion_hash(), at, NEVER_EXPIRES)?;
222    }
223    Ok(())
224}
225
226/// Idempotent startup seeding + reconcile: pull-ins only LOWER task deadlines to
227/// the live DB columns, repairing rows stranded by a crash mid-nudge.
228pub(crate) fn seed_and_reconcile_kp_tasks<Context: XmtpSharedContext>(
229    context: &Context,
230) -> Result<(), StorageError> {
231    let db = context.db();
232    let now = xmtp_common::time::now_ns();
233    db.create_or_ignore_task(kp_seed(kp_rotation_proto(), now)?)?;
234    db.create_or_ignore_task(kp_seed(kp_deletion_proto(), now)?)?;
235    // None = pre-registration (no identity row): the seed row already fires at
236    // startup; a pull-in to `now` would be redundant noise.
237    if let Some(rot) = db.next_key_package_rotation_ns()? {
238        enqueue_pull_in(context, kp_rotation_hash(), rot, NEVER_EXPIRES)?;
239    }
240    if let Some(del) = db.min_key_package_delete_at_ns()? {
241        enqueue_pull_in(context, kp_deletion_hash(), del, NEVER_EXPIRES)?;
242    }
243    Ok(())
244}
245
246// These tests use the native connection pool and its error variants.
247#[cfg(all(test, not(target_arch = "wasm32")))]
248mod tests {
249    use super::*;
250    use crate::tester;
251    use crate::worker::tasks::TaskWorker;
252    use crate::worker::{WorkerConfig, WorkerKind};
253    use openmls_traits::storage::StorageProvider;
254    use prost::Message;
255    use xmtp_db::ConnectionExt;
256    use xmtp_proto::xmtp::mls::database::Task as TaskProtoDecode;
257
258    /// A `StorageError` that signals the connection pool was dropped.
259    fn disconnect_storage() -> xmtp_db::StorageError {
260        xmtp_db::StorageError::Platform(xmtp_db::PlatformStorageError::PoolNeedsConnection)
261    }
262
263    /// A storage error that is NOT a disconnect — must never trip the contract.
264    fn benign_storage() -> xmtp_db::StorageError {
265        xmtp_db::StorageError::InvalidHmacLength
266    }
267
268    fn no_runner_cfg() -> WorkerConfig {
269        let mut cfg = WorkerConfig::default();
270        cfg.enabled.insert(WorkerKind::TaskRunner, false);
271        cfg
272    }
273
274    fn row_by_hash(db: &impl QueryTasks, hash: impl AsRef<[u8]>) -> Option<xmtp_db::tasks::Task> {
275        db.get_tasks()
276            .expect("get_tasks should not fail")
277            .into_iter()
278            .find(|t| t.data_hash == hash.as_ref())
279    }
280
281    async fn make_rotation_due(db: &impl QueryIdentity) {
282        db.queue_key_package_rotation()
283            .expect("queue_key_package_rotation should not fail"); // column := now + 5s
284        xmtp_common::time::sleep(std::time::Duration::from_secs(6)).await;
285    }
286
287    #[xmtp_common::test]
288    fn kp_errors_forward_db_reconnect() {
289        use crate::worker::NeedsDbReconnect;
290        use crate::worker::tasks::TaskWorkerError;
291        let e = TaskWorkerError::from(KeyPackageMaintenanceError::Storage(disconnect_storage()));
292        assert!(
293            e.needs_db_reconnect(),
294            "DB outage during KP work must trigger supervisor reconnect, not plain backoff"
295        );
296        let e = TaskWorkerError::from(crate::identity::IdentityError::from(disconnect_storage()));
297        assert!(e.needs_db_reconnect());
298        // Keystore pool loss (rotate/delete paths) must also restart the worker.
299        let e = TaskWorkerError::from(crate::identity::IdentityError::OpenMlsStorageError(
300            xmtp_db::sql_key_store::SqlKeyStoreError::Connection(
301                xmtp_db::ConnectionError::Platform(
302                    xmtp_db::PlatformStorageError::PoolNeedsConnection,
303                ),
304            ),
305        ));
306        assert!(e.needs_db_reconnect());
307        // A non-disconnect storage failure must NOT stop the worker.
308        let e = TaskWorkerError::from(KeyPackageMaintenanceError::Storage(benign_storage()));
309        assert!(
310            !e.needs_db_reconnect(),
311            "benign storage errors must back off, not restart the supervisor"
312        );
313    }
314
315    #[xmtp_common::test(unwrap_try = true)]
316    async fn manual_rotation_nudges_deletion() {
317        tester!(alix, worker_config: no_runner_cfg());
318        let db = alix.context.db();
319        assert!(row_by_hash(&db, kp_deletion_hash()).is_none());
320
321        alix.rotate_and_upload_key_package().await?;
322
323        assert!(
324            row_by_hash(&db, kp_deletion_hash()).is_some(),
325            "manual rotation must self-heal the deletion singleton"
326        );
327        let has_pull_in = db.get_tasks()?.into_iter().any(|t| matches!(
328            TaskProtoDecode::decode(t.data.as_slice()).ok().and_then(|p| p.task),
329            Some(TaskKind::PullInDeadline(p)) if p.target_data_hash == kp_deletion_hash().as_ref()
330        ));
331        assert!(
332            has_pull_in,
333            "manual rotation must enqueue a deletion pull-in"
334        );
335    }
336
337    #[xmtp_common::test(unwrap_try = true)]
338    async fn pending_welcome_preserves_expired_keys_until_completion() {
339        use xmtp_db::diesel::prelude::*;
340        use xmtp_db::incoming_envelope::{
341            IncomingLimits, NetworkEntityKind, NewIncomingEnvelope, PendingBudget, StreamTopic,
342        };
343        use xmtp_db::schema::key_package_history::dsl;
344        use xmtp_proto::types::Cursor;
345
346        tester!(alix, disable_workers);
347        let db = alix.context.db();
348        let history = db
349            .find_key_package_history_entries_before_id(i32::MAX)?
350            .pop()?;
351        let hash =
352            crate::identity::deserialize_key_package_hash_ref(&history.key_package_hash_ref)?;
353        db.raw_query(|conn| {
354            diesel::update(dsl::key_package_history.filter(dsl::id.eq(history.id)))
355                .set(dsl::delete_at_ns.eq(0))
356                .execute(conn)
357        })?;
358        let topic = StreamTopic {
359            entity_id: alix.context.installation_id().to_vec(),
360            kind: NetworkEntityKind::Welcome,
361        };
362        let limit = PendingBudget { rows: 1, bytes: 1 };
363        db.admit_ordered_batch(
364            &topic,
365            Cursor(0),
366            &[NewIncomingEnvelope {
367                sequence_id: Cursor(1),
368                envelope: vec![1],
369            }],
370            IncomingLimits {
371                batch: limit,
372                topic: limit,
373                kind: limit,
374            },
375        )?;
376
377        sweep_expired(&alix.context)?;
378        let retained: Option<openmls::prelude::KeyPackageBundle> =
379            alix.context.mls_storage().key_package(&hash)?;
380        assert!(retained.is_some());
381
382        db.complete_pending_envelope(&topic, Cursor(1))?;
383        sweep_expired(&alix.context)?;
384        let removed: Option<openmls::prelude::KeyPackageBundle> =
385            alix.context.mls_storage().key_package(&hash)?;
386        assert!(removed.is_none());
387        assert!(
388            db.find_key_package_history_entry_by_hash_ref(history.key_package_hash_ref)
389                .is_err()
390        );
391    }
392
393    #[xmtp_common::test(unwrap_try = true)]
394    async fn rotation_task_rotates_and_reschedules() {
395        tester!(alix, worker_config: no_runner_cfg());
396        let db = alix.context.db();
397        let now = xmtp_common::time::now_ns();
398        db.create_or_ignore_task(kp_seed(kp_rotation_proto(), now)?)?;
399        make_rotation_due(&db).await;
400
401        let row = row_by_hash(&db, kp_rotation_hash()).unwrap();
402        TaskWorker::run_and_reschedule_task(row, &alix.context).await?;
403
404        assert!(
405            !db.is_identity_needs_rotation()?,
406            "rotation must have happened"
407        );
408        let after = row_by_hash(&db, kp_rotation_hash()).expect("recurring row survives");
409        let col = db.next_key_package_rotation_ns()?.unwrap();
410        assert_eq!(
411            after.next_attempt_at_ns, col,
412            "reschedule must read the live column"
413        );
414        assert_eq!(after.attempts, 0);
415    }
416
417    #[xmtp_common::test(unwrap_try = true)]
418    async fn rotation_ensures_and_pulls_in_deletion_when_singleton_missing() {
419        tester!(alix, worker_config: no_runner_cfg());
420        let db = alix.context.db();
421        let now = xmtp_common::time::now_ns();
422        db.create_or_ignore_task(kp_seed(kp_rotation_proto(), now)?)?;
423        // Deliberately NO KpDeletion seed: the handler must self-heal it.
424        make_rotation_due(&db).await;
425
426        let row = row_by_hash(&db, kp_rotation_hash()).unwrap();
427        TaskWorker::run_and_reschedule_task(row, &alix.context).await?;
428
429        assert!(
430            row_by_hash(&db, kp_deletion_hash()).is_some(),
431            "rotation must recreate a missing KpDeletion singleton"
432        );
433        let has_pull_in = db.get_tasks()?.iter().any(|t| {
434            matches!(
435                TaskProtoDecode::decode(t.data.as_slice()).ok().and_then(|p| p.task),
436                Some(TaskKind::PullInDeadline(p)) if p.target_data_hash == kp_deletion_hash().as_ref()
437            )
438        });
439        assert!(has_pull_in, "rotation must enqueue a deletion pull-in");
440    }
441
442    #[xmtp_common::test(unwrap_try = true)]
443    async fn deletion_task_sweeps_and_reschedules() {
444        tester!(alix, worker_config: no_runner_cfg());
445        let db = alix.context.db();
446        let now = xmtp_common::time::now_ns();
447        db.create_or_ignore_task(kp_seed(kp_deletion_proto(), now)?)?;
448
449        // A rotation marks the superseded KP delete_at = now + 3s (test cfg).
450        make_rotation_due(&db).await;
451        rotate_if_needed(&alix.context).await?;
452        assert!(db.min_key_package_delete_at_ns()?.is_some());
453        xmtp_common::time::sleep(std::time::Duration::from_secs(4)).await; // pass the grace
454
455        let row = row_by_hash(&db, kp_deletion_hash()).unwrap();
456        TaskWorker::run_and_reschedule_task(row, &alix.context).await?;
457
458        assert!(
459            db.get_expired_key_packages()?.is_empty(),
460            "sweep must delete expired KPs"
461        );
462        let after = row_by_hash(&db, kp_deletion_hash()).expect("recurring row survives");
463        assert!(
464            after.next_attempt_at_ns > xmtp_common::time::now_ns(),
465            "deletion reschedules to next pending deadline or far-future"
466        );
467    }
468
469    #[xmtp_common::test(unwrap_try = true)]
470    async fn kp_tasks_seeded_when_workers_run_absent_when_passive() {
471        tester!(alix); // default: TaskRunner on -> seeds present
472        let db = alix.context.db();
473        assert!(row_by_hash(&db, kp_rotation_hash()).is_some());
474        assert!(row_by_hash(&db, kp_deletion_hash()).is_some());
475
476        tester!(bo, worker_config: no_runner_cfg()); // no TaskRunner -> no seeds
477        let db = bo.context.db();
478        assert!(row_by_hash(&db, kp_rotation_hash()).is_none());
479        assert!(row_by_hash(&db, kp_deletion_hash()).is_none());
480    }
481
482    #[xmtp_common::test(unwrap_try = true)]
483    async fn startup_reconcile_pulls_in_far_scheduled_row() {
484        tester!(alix, worker_config: no_runner_cfg());
485        let db = alix.context.db();
486        let now = xmtp_common::time::now_ns();
487        // Stale persisted row 30d out while the column says due-in-5s
488        // (crash-between-writes scenario).
489        db.create_or_ignore_task(kp_seed(kp_rotation_proto(), now)?)?;
490        let row = row_by_hash(&db, kp_rotation_hash()).unwrap();
491        db.update_task(row.id, 0, now, now + 30 * xmtp_common::NS_IN_DAY)?;
492        db.queue_key_package_rotation()?; // column := now + 5s
493
494        seed_and_reconcile_kp_tasks(&alix.context)?;
495
496        let pull_in = db
497            .get_tasks()?
498            .into_iter()
499            .find(|t| {
500                matches!(
501                    TaskProtoDecode::decode(t.data.as_slice()).ok().and_then(|p| p.task),
502                    Some(TaskKind::PullInDeadline(p)) if p.target_data_hash == kp_rotation_hash().as_ref()
503                )
504            })
505            .expect("reconcile must enqueue a rotation pull-in");
506        TaskWorker::run_and_reschedule_task(pull_in, &alix.context).await?;
507
508        let after = row_by_hash(&db, kp_rotation_hash()).unwrap();
509        let col = db.next_key_package_rotation_ns()?.unwrap();
510        assert_eq!(after.next_attempt_at_ns, col);
511    }
512
513    /// KpRotation firing while NOT due must not rotate or seed deletion — it just
514    /// re-syncs its deadline to the column (spurious-wake safety).
515    #[xmtp_common::test(unwrap_try = true)]
516    async fn rotation_task_not_due_reschedules_without_rotating() {
517        tester!(alix, worker_config: no_runner_cfg());
518        let db = alix.context.db();
519        let now = xmtp_common::time::now_ns();
520        db.create_or_ignore_task(kp_seed(kp_rotation_proto(), now)?)?;
521        // Post-registration column is ~now+30d: not due.
522        let row = row_by_hash(&db, kp_rotation_hash()).unwrap();
523        TaskWorker::run_and_reschedule_task(row, &alix.context).await?;
524
525        assert!(
526            row_by_hash(&db, kp_deletion_hash()).is_none(),
527            "must not seed deletion"
528        );
529        assert!(db.min_key_package_delete_at_ns()?.is_none());
530        let after = row_by_hash(&db, kp_rotation_hash()).unwrap();
531        assert_eq!(
532            after.next_attempt_at_ns,
533            db.next_key_package_rotation_ns()?.unwrap()
534        );
535    }
536
537    /// The welcome nudge must self-heal a missing rotation seed (e.g. startup
538    /// seeding never ran) instead of enqueuing a dropped-on-miss pull-in.
539    #[xmtp_common::test(unwrap_try = true)]
540    async fn welcome_nudge_selfheals_missing_rotation_seed() {
541        tester!(alix, worker_config: no_runner_cfg()); // no TaskRunner -> no seeds
542        let db = alix.context.db();
543        assert!(row_by_hash(&db, kp_rotation_hash()).is_none());
544
545        queue_key_rotation(&alix.context)?;
546
547        assert!(
548            row_by_hash(&db, kp_rotation_hash()).is_some(),
549            "nudge must recreate the missing KpRotation singleton"
550        );
551        let has_pull_in = db.get_tasks()?.iter().any(|t| {
552            matches!(
553                TaskProtoDecode::decode(t.data.as_slice()).ok().and_then(|p| p.task),
554                Some(TaskKind::PullInDeadline(p)) if p.target_data_hash == kp_rotation_hash().as_ref()
555            )
556        });
557        assert!(has_pull_in, "nudge must enqueue a rotation pull-in");
558    }
559
560    #[cfg(not(target_arch = "wasm32"))]
561    #[xmtp_common::test(unwrap_try = true)]
562    async fn queued_rotation_is_visible_before_its_task_fact() {
563        use crate::context::XmtpSharedContext;
564        use crate::subscriptions::internal::InternalEvent;
565        use xmtp_events::EventFilter;
566
567        tester!(alix, worker_config: no_runner_cfg());
568        let subscription = alix.context.events().subscribe(
569            EventFilter::default()
570                .with_internal(|event| matches!(event, InternalEvent::TaskScheduled)),
571            Some(10),
572        );
573        let before = alix.context.events().subscribe(
574            EventFilter::default()
575                .with_internal(|event| matches!(event, InternalEvent::TaskScheduled)),
576            Some(10),
577        );
578        let hook = crate::state_tx::precommit_test_hook::install(move |_| {
579            assert!(before.drain().is_empty());
580        });
581        queue_key_rotation(&alix.context)?;
582        drop(hook);
583        assert!(row_by_hash(&alix.context.db(), kp_rotation_hash()).is_some());
584        assert!(matches!(
585            subscription.next().await.unwrap().internal,
586            Some(InternalEvent::TaskScheduled)
587        ));
588    }
589
590    /// Regression: welcome nudge must pull the parked rotation task in even when
591    /// the seed dispatched BEFORE the column was lowered (the startup race).
592    #[xmtp_common::test(unwrap_try = true)]
593    async fn welcome_nudge_pulls_in_parked_rotation() {
594        tester!(alix, worker_config: no_runner_cfg());
595        let db = alix.context.db();
596        let now = xmtp_common::time::now_ns();
597        db.create_or_ignore_task(kp_seed(kp_rotation_proto(), now)?)?;
598        // Simulate the seed having already dispatched not-due: park it on the column (~+30d).
599        let parked = row_by_hash(&db, kp_rotation_hash()).unwrap();
600        TaskWorker::run_and_reschedule_task(parked, &alix.context).await?;
601        let parked_at = row_by_hash(&db, kp_rotation_hash())
602            .unwrap()
603            .next_attempt_at_ns;
604        assert!(
605            parked_at > now + xmtp_common::NS_IN_DAY,
606            "precondition: parked far out"
607        );
608
609        queue_key_rotation(&alix.context)?; // welcome: column + pull-in, atomically
610
611        let pull_in = db
612            .get_tasks()?
613            .into_iter()
614            .find(|t| {
615                matches!(
616                    TaskProtoDecode::decode(t.data.as_slice()).ok().and_then(|p| p.task),
617                    Some(TaskKind::PullInDeadline(p)) if p.target_data_hash == kp_rotation_hash().as_ref()
618                )
619            })
620            .expect("nudge must enqueue a durable pull-in");
621        TaskWorker::run_and_reschedule_task(pull_in, &alix.context).await?;
622
623        let after = row_by_hash(&db, kp_rotation_hash()).unwrap();
624        let col = db.next_key_package_rotation_ns()?.unwrap();
625        assert_eq!(
626            after.next_attempt_at_ns, col,
627            "rotation row must be pulled in to the lowered column"
628        );
629        // 5s queue debounce + 2s slack for local ops between `now` and the queue call.
630        assert!(after.next_attempt_at_ns <= now + 7 * xmtp_common::NS_IN_SEC);
631    }
632}
xmtp_mls/worker/key_package_maintenance.rs

xmtp_mls/worker/
key_package_maintenance.rs
