feat(xmtp_mls): KeyPackageCleaner sleeps to exact rotation deadline (kill 5s poll); span-gated maintain

Author: insipx

crates/xmtp_mls/src/worker/key_package_cleaner.rs

@@ -5,20 +5,31 @@ use crate::worker::BoxedWorker;
 use crate::worker::NeedsDbReconnect;
 use crate::worker::WorkerResult;
 use crate::worker::{Worker, WorkerFactory, WorkerKind};
-use futures::StreamExt;
 use futures::TryFutureExt;
 use openmls_traits::storage::StorageProvider;
-use std::time::Duration;
 use thiserror::Error;
+use tracing::Instrument;
 use xmtp_configuration::CREATE_PQ_KEY_PACKAGE_EXTENSION;
 use xmtp_db::prelude::*;
 use xmtp_db::{
     MlsProviderExt, StorageError,
+    encrypted_store::key_package_history::StoredKeyPackageHistoryEntry,
     sql_key_store::{KEY_PACKAGE_REFERENCES, KEY_PACKAGE_WRAPPER_PRIVATE_KEY},
 };
 
-/// Interval at which the KeyPackagesCleanerWorker runs to delete expired messages.
-pub const INTERVAL_DURATION: Duration = Duration::from_secs(5);
+#[derive(Debug, PartialEq, Eq)]
+enum WakePlan {
+    RunNow,
+    SleepUntil(i64), // absolute deadline ns
+}
+
+fn plan(next_rotation: Option<i64>, now: i64) -> WakePlan {
+    match next_rotation {
+        None => WakePlan::RunNow, // NULL = due now
+        Some(at) if at <= now => WakePlan::RunNow,
+        Some(at) => WakePlan::SleepUntil(at),
+    }
+}
 
 #[derive(Clone)]
 pub struct Factory<Context> {
@@ -116,21 +127,77 @@ where
     Context: XmtpSharedContext + 'static,
 {
     async fn run(&mut self) -> Result<(), KeyPackagesCleanerError> {
-        let (base, jitter) = self
-            .context
-            .worker_interval(WorkerKind::KeyPackageCleaner, INTERVAL_DURATION);
-        let mut intervals = xmtp_common::time::jittered_interval_stream(base, jitter);
-        while (intervals.next().await).is_some() {
-            self.tick().await?;
+        let receiver = self.context.key_package_channels().receiver.clone();
+        let mut receiver = receiver.lock().await;
+        loop {
+            // Drain any pending re-arm signals before computing the plan so we
+            // don't lose a wakeup that arrived while we were working.
+            while receiver.try_recv().is_ok() {}
+
+            match plan(
+                self.context
+                    .db()
+                    .next_key_package_rotation_ns()
+                    .map_err(KeyPackagesCleanerError::Metadata)?,
+                xmtp_common::time::now_ns(),
+            ) {
+                WakePlan::RunNow => {
+                    self.maintain().await?;
+                }
+                WakePlan::SleepUntil(deadline) => {
+                    let dur = std::time::Duration::from_nanos(
+                        (deadline - xmtp_common::time::now_ns()).max(0) as u64,
+                    );
+                    tokio::select! {
+                        // Re-arm signal: recompute the deadline, do NOT run work.
+                        // `None` means every sender was dropped (context torn down) —
+                        // stop rather than busy-spin on a closed channel.
+                        msg = receiver.recv() => { if msg.is_none() { return Ok(()); } }
+                        // Deadline elapsed: time to do maintenance.
+                        () = xmtp_common::time::sleep(dur) => {
+                            self.maintain().await?;
+                        }
+                    }
+                }
+            }
         }
-        Ok(())
     }
 
-    #[tracing::instrument(skip_all, fields(worker = ?self.kind(), operation = "worker_turn"))]
-    async fn tick(&mut self) -> Result<(), KeyPackagesCleanerError> {
-        self.delete_expired_key_packages()?;
-        self.rotate_last_key_package_if_needed().await?;
-        Ok(())
+    /// One maintenance pass. Guards read OUTSIDE the span; if nothing is due,
+    /// returns immediately with no span. Otherwise a single `worker_turn` span
+    /// wraps the work so tracing records the full operation (including errors).
+    async fn maintain(&mut self) -> Result<(), KeyPackagesCleanerError> {
+        let expired = self
+            .context
+            .db()
+            .get_expired_key_packages()
+            .map_err(KeyPackagesCleanerError::Fetch)?;
+        let rotate_due = self
+            .context
+            .db()
+            .is_identity_needs_rotation()
+            .map_err(KeyPackagesCleanerError::Metadata)?;
+
+        if expired.is_empty() && !rotate_due {
+            return Ok(());
+        }
+
+        let span = tracing::info_span!(
+            "worker_turn",
+            worker = ?self.kind(),
+            operation = "worker_turn"
+        );
+        async {
+            if !expired.is_empty() {
+                self.delete_key_packages(expired)?;
+            }
+            if rotate_due {
+                self.rotate_last_key_package_if_needed().await?;
+            }
+            Ok::<(), KeyPackagesCleanerError>(())
+        }
+        .instrument(span)
+        .await
     }
 
     /// Delete a key package from the local database.
@@ -156,65 +223,110 @@ where
         Ok(())
     }
 
-    /// Delete all the expired keys
-    fn delete_expired_key_packages(&mut self) -> Result<(), KeyPackagesCleanerError> {
+    /// Delete an already-fetched list of expired key packages from local state
+    /// and the database history table.
+    fn delete_key_packages(
+        &mut self,
+        expired: Vec<StoredKeyPackageHistoryEntry>,
+    ) -> Result<(), KeyPackagesCleanerError> {
         let conn = self.context.db();
-
-        // Propagate (don't swallow): a swallowed fetch error never triggered the supervisor's
-        // reconnect path, so a pool outage retried silently every 5s.
-        let expired_kps = conn
-            .get_expired_key_packages()
-            .map_err(KeyPackagesCleanerError::Fetch)?;
-        if expired_kps.is_empty() {
-            return Ok(());
-        }
-
-        tracing::info!("Deleting {} expired key packages", expired_kps.len());
-        // Delete from local db
-        for kp in &expired_kps {
+        for kp in &expired {
             self.delete_key_package(
                 kp.key_package_hash_ref.clone(),
                 kp.post_quantum_public_key.clone(),
             )
             .map_err(KeyPackagesCleanerError::DeleteKeyPackage)?;
         }
-
-        // Delete from database using the max expired ID
-        if let Some(max_id) = expired_kps.iter().map(|kp| kp.id).max() {
+        if let Some(max_id) = expired.iter().map(|kp| kp.id).max() {
             conn.delete_key_package_history_up_to_id(max_id)
                 .map_err(KeyPackagesCleanerError::Deletion)?;
             tracing::info!(
-                "Deleted {} expired key packages (up to ID {}) from local DB and state",
-                expired_kps.len(),
+                "Deleted {} expired key packages (up to ID {})",
+                expired.len(),
                 max_id
             );
         }
-
         Ok(())
     }
 
-    /// Check if we need to rotate the keys and upload new keypackage if the las one rotate in has passed
+    /// Upload a fresh key package if the current one has passed its rotation deadline.
     async fn rotate_last_key_package_if_needed(&mut self) -> Result<(), KeyPackagesCleanerError> {
-        let conn = self.context.db();
+        tracing::info!("Rotating key package");
+        self.context
+            .identity()
+            .rotate_and_upload_key_package(
+                self.context.api(),
+                self.context.mls_storage(),
+                CREATE_PQ_KEY_PACKAGE_EXTENSION,
+            )
+            .await
+            .map_err(KeyPackagesCleanerError::Rotation)?;
+        tracing::info!("Key package rotation successful");
+        Ok(())
+    }
+}
 
-        if conn
-            .is_identity_needs_rotation()
-            .map_err(KeyPackagesCleanerError::Metadata)?
-        {
-            tracing::info!("Rotating key package");
-            self.context
-                .identity()
-                .rotate_and_upload_key_package(
-                    self.context.api(),
-                    self.context.mls_storage(),
-                    CREATE_PQ_KEY_PACKAGE_EXTENSION,
-                )
-                .await
-                .map_err(KeyPackagesCleanerError::Rotation)?;
-            tracing::info!("Key package rotation successful");
-            return Ok(());
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn plan_cases() {
+        assert_eq!(super::plan(None, 100), super::WakePlan::RunNow);
+        assert_eq!(super::plan(Some(99), 100), super::WakePlan::RunNow);
+        assert_eq!(super::plan(Some(100), 100), super::WakePlan::RunNow);
+        assert_eq!(
+            super::plan(Some(101), 100),
+            super::WakePlan::SleepUntil(101)
+        );
+    }
+
+    /// Requires a live XMTP node. Verifies that `queue_key_rotation` marks the
+    /// identity for rotation and wakes the KeyPackageCleaner worker, which then
+    /// uploads a new key package within ~12 s.
+    #[xmtp_common::test(unwrap_try = true)]
+    async fn queue_key_rotation_wakes_worker_and_rotates() {
+        use crate::tester;
+
+        tester!(client);
+
+        let installation_id = client.installation_public_key().to_vec();
+
+        // Snapshot the current on-network init key.
+        // `VerifiedKeyPackageV2::hpke_init_key()` returns an owned Vec<u8>.
+        let before = {
+            let mut map = client
+                .get_key_packages_for_installation_ids(vec![installation_id.clone()])
+                .await?;
+            let entry = map
+                .remove(&installation_id)
+                .expect("installation not found in response")?;
+            entry.hpke_init_key()
+        };
+
+        // Queue a rotation and wake the worker.
+        client.queue_key_rotation().await?;
+
+        // Poll up to 12 s (60 × 200 ms) for the on-network key to change.
+        // Uses `xmtp_common::time::sleep` so the test compiles for wasm too
+        // (std::time::Instant panics on wasm).
+        let mut after = None;
+        for _ in 0..60u32 {
+            xmtp_common::time::sleep(std::time::Duration::from_millis(200)).await;
+            let mut map = client
+                .get_key_packages_for_installation_ids(vec![installation_id.clone()])
+                .await?;
+            let entry = map
+                .remove(&installation_id)
+                .expect("installation not found in response")?;
+            let key = entry.hpke_init_key();
+            if key != before {
+                after = Some(key);
+                break;
+            }
         }
 
-        Ok(())
+        assert!(
+            after.is_some(),
+            "key package did not rotate after queue_key_rotation + worker wake"
+        );
     }
 }