xmtp_mls/worker/
hmac_epoch.rs

1//! Report HMAC epoch changes observed while this client is open.
2
3use crate::context::XmtpSharedContext;
4use crate::worker::{BoxedWorker, DynMetrics, Worker, WorkerFactory, WorkerKind, WorkerResult};
5use std::sync::{
6    Arc,
7    atomic::{AtomicI64, Ordering},
8};
9use xmtp_common::time::{self, Duration};
10use xmtp_events::{ClientEvent, EventWriter, HmacKeysUpdated};
11
12pub struct HmacEpochWorker<Context> {
13    context: Context,
14    last_epoch: Arc<AtomicI64>,
15    #[cfg(test)]
16    clock: Option<TestClock>,
17}
18
19#[cfg(test)]
20struct TestClock {
21    now: Arc<AtomicI64>,
22    ticks: futures::channel::mpsc::UnboundedReceiver<i64>,
23    waits: futures::channel::mpsc::UnboundedSender<Duration>,
24}
25
26impl<Context> HmacEpochWorker<Context> {
27    fn now(&self) -> i64 {
28        #[cfg(test)]
29        if let Some(clock) = &self.clock {
30            return clock.now.load(Ordering::Acquire);
31        }
32        time::now_secs()
33    }
34
35    async fn wait(&mut self, duration: Duration) {
36        #[cfg(test)]
37        if let Some(clock) = &mut self.clock {
38            use futures::StreamExt;
39            clock
40                .waits
41                .unbounded_send(duration)
42                .expect("test wait observer must stay open");
43            let now = clock.ticks.next().await.expect("test clock must stay open");
44            clock.now.store(now, Ordering::Release);
45            return;
46        }
47        time::sleep(duration).await;
48    }
49
50    fn wait_until_next_epoch(&self, now: i64) -> Duration {
51        let next =
52            (self.last_epoch.load(Ordering::Acquire) + 1) * xmtp_push_types::HMAC_EPOCH_SECONDS;
53        Duration::from_secs((next - now).max(1) as u64)
54    }
55
56    fn observe_epoch(&mut self, current: i64) -> bool {
57        loop {
58            let previous = self.last_epoch.load(Ordering::Acquire);
59            if current <= previous {
60                return false;
61            }
62            if self
63                .last_epoch
64                .compare_exchange(previous, current, Ordering::AcqRel, Ordering::Acquire)
65                .is_ok()
66            {
67                return true;
68            }
69        }
70    }
71}
72
73impl<Context: XmtpSharedContext> HmacEpochWorker<Context> {
74    // implements: EVENT-026
75    fn report_epoch(&mut self, observed: i64) {
76        if self.observe_epoch(observed) {
77            self.context.task_channels().mark_notification_changed();
78            self.context
79                .events()
80                .emit(Some(ClientEvent::HmacKeysUpdated(HmacKeysUpdated)), None);
81        }
82    }
83}
84
85struct Factory<Context> {
86    context: Context,
87    last_epoch: Arc<AtomicI64>,
88}
89
90impl<Context> Factory<Context> {
91    fn new(context: Context, startup_epoch: i64) -> Self {
92        Self {
93            context,
94            last_epoch: Arc::new(AtomicI64::new(startup_epoch)),
95        }
96    }
97}
98
99impl<Context: XmtpSharedContext + 'static> WorkerFactory for Factory<Context> {
100    fn kind(&self) -> WorkerKind {
101        WorkerKind::HmacEpoch
102    }
103
104    fn create(&self, _metrics: Option<DynMetrics>) -> (BoxedWorker, Option<DynMetrics>) {
105        (
106            Box::new(HmacEpochWorker {
107                context: self.context.clone(),
108                last_epoch: self.last_epoch.clone(),
109                #[cfg(test)]
110                clock: None,
111            }),
112            None,
113        )
114    }
115}
116
117#[xmtp_common::async_trait]
118impl<Context: XmtpSharedContext + 'static> Worker for HmacEpochWorker<Context> {
119    fn kind(&self) -> WorkerKind {
120        WorkerKind::HmacEpoch
121    }
122
123    fn factory<C>(context: C) -> impl WorkerFactory + 'static
124    where
125        Self: Sized,
126        C: XmtpSharedContext + 'static,
127    {
128        Factory::new(context, crate::utils::time::hmac_epoch())
129    }
130
131    async fn run_tasks(&mut self) -> WorkerResult<()> {
132        loop {
133            let now = self.now();
134            self.wait(self.wait_until_next_epoch(now)).await;
135            self.report_epoch(self.now() / xmtp_push_types::HMAC_EPOCH_SECONDS);
136        }
137    }
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143    use crate::tester;
144    use xmtp_events::{EventFilter, EventKind};
145
146    #[xmtp_common::test(unwrap_try = true)]
147    async fn reports_only_new_epochs_after_startup() {
148        let mut worker = HmacEpochWorker {
149            context: (),
150            last_epoch: Arc::new(AtomicI64::new(10)),
151            clock: None,
152        };
153        assert!(!worker.observe_epoch(10));
154        assert!(worker.observe_epoch(11));
155        assert!(!worker.observe_epoch(11));
156        assert!(worker.observe_epoch(12));
157    }
158
159    // verifies: EVENT-001
160    #[xmtp_common::test(unwrap_try = true)]
161    async fn epoch_timer_turn_emits_once_and_invalidates_notifications() {
162        tester!(alix, disable_workers);
163        let current = crate::utils::time::hmac_epoch();
164        let mut worker = HmacEpochWorker {
165            context: alix.context.clone(),
166            last_epoch: Arc::new(AtomicI64::new(current)),
167            clock: None,
168        };
169        let events = alix
170            .context
171            .events()
172            .subscribe_app(EventFilter::new([EventKind::HmacKeysUpdated]))?;
173        let boundary = (current + 1) * xmtp_push_types::HMAC_EPOCH_SECONDS;
174        assert_eq!(
175            worker.wait_until_next_epoch(boundary - 1),
176            Duration::from_secs(1)
177        );
178        let previous_revision = alix.context.task_channels().notification_revision();
179        worker.report_epoch(current);
180        assert!(events.drain().is_empty());
181        worker.report_epoch(current + 1);
182        assert!(matches!(
183            events.drain().as_slice(),
184            [xmtp_events::EventEnvelope {
185                client: Some(ClientEvent::HmacKeysUpdated(_)),
186                ..
187            }]
188        ));
189        assert_eq!(
190            alix.context.task_channels().notification_revision(),
191            previous_revision + 1
192        );
193        worker.report_epoch(current + 1);
194        assert!(events.drain().is_empty());
195    }
196    // verifies: EVENT-026
197    #[xmtp_common::test(unwrap_try = true)]
198    async fn running_epoch_worker_skips_closed_epochs_and_uses_no_network() {
199        tester!(alix, disable_workers);
200        let events = alix
201            .context
202            .events()
203            .subscribe_app(EventFilter::new([EventKind::HmacKeysUpdated]))?;
204        let stats = alix.api_stats();
205        let identity = alix.identity_api_stats();
206        let counts = || {
207            [
208                stats.publish.get_count(),
209                stats.create_upload.get_count(),
210                stats.query.get_count(),
211                stats.query_newest.get_count(),
212                stats.subscribe.get_count(),
213                stats.subscribe_static.get_count(),
214                identity.get_inbox_ids.get_count(),
215                identity.verify_smart_contract_wallet_signatures.get_count(),
216            ]
217        };
218        let before = counts();
219        let revision = alix.context.task_channels().notification_revision();
220        let width = xmtp_push_types::HMAC_EPOCH_SECONDS;
221        let now = Arc::new(AtomicI64::new(10 * width));
222        let (ticks, receiver) = futures::channel::mpsc::unbounded();
223        let (waits, mut scheduled) = futures::channel::mpsc::unbounded();
224        let factory = Factory::new(alix.context.clone(), 10);
225        let mut worker = HmacEpochWorker {
226            context: factory.context.clone(),
227            last_epoch: factory.last_epoch.clone(),
228            clock: Some(TestClock {
229                now: now.clone(),
230                ticks: receiver,
231                waits,
232            }),
233        };
234        let mut running = Box::pin(worker.run_tasks());
235        assert!(futures::poll!(&mut running).is_pending());
236        assert!(events.drain().is_empty());
237        assert_eq!(scheduled.try_recv()?, Duration::from_secs(width as u64));
238        let mut previous = 10;
239        for (epoch, expected) in [(10, 0), (11, 1), (11, 0), (14, 1), (13, 0), (14, 0)] {
240            ticks.unbounded_send(epoch * width)?;
241            assert!(futures::poll!(&mut running).is_pending());
242            let emitted = events.drain();
243            assert_eq!(emitted.len(), expected);
244            previous = previous.max(epoch);
245            assert_eq!(
246                scheduled.try_recv()?,
247                Duration::from_secs(((previous + 1 - epoch) * width) as u64)
248            );
249            assert!(
250                emitted
251                    .iter()
252                    .all(|event| matches!(event.client, Some(ClientEvent::HmacKeysUpdated(_))))
253            );
254        }
255        drop(running);
256        assert_eq!(
257            alix.context.task_channels().notification_revision(),
258            revision + 2
259        );
260        now.store(18 * width, Ordering::Release);
261        let (ticks, receiver) = futures::channel::mpsc::unbounded();
262        let (waits, mut scheduled) = futures::channel::mpsc::unbounded();
263        let factory = Factory::new(alix.context.clone(), 18);
264        let mut reopened = HmacEpochWorker {
265            context: factory.context.clone(),
266            last_epoch: factory.last_epoch.clone(),
267            clock: Some(TestClock {
268                now,
269                ticks: receiver,
270                waits,
271            }),
272        };
273        let mut running = Box::pin(reopened.run_tasks());
274        assert!(futures::poll!(&mut running).is_pending());
275        assert_eq!(scheduled.try_recv()?, Duration::from_secs(width as u64));
276        ticks.unbounded_send(18 * width)?;
277        assert!(futures::poll!(&mut running).is_pending());
278        assert!(events.drain().is_empty());
279        ticks.unbounded_send(19 * width)?;
280        assert!(futures::poll!(&mut running).is_pending());
281        assert_eq!(events.drain().len(), 1);
282        assert_eq!(
283            alix.context.task_channels().notification_revision(),
284            revision + 3
285        );
286        assert_eq!(counts(), before);
287    }
288}
xmtp_mls/worker/hmac_epoch.rs

xmtp_mls/worker/
hmac_epoch.rs
