xmtp_mls_core/worker/
metrics.rs

1use futures::FutureExt;
2use parking_lot::Mutex;
3use std::{
4    collections::HashMap,
5    fmt::Debug,
6    future::Future,
7    hash::Hash,
8    pin::Pin,
9    sync::{
10        Arc,
11        atomic::{AtomicUsize, Ordering},
12    },
13    time::Duration,
14};
15use tokio::sync::Notify;
16use xmtp_proto::types::InstallationId;
17
18pub struct MetricInterest<Metric> {
19    fut: Pin<Box<dyn Future<Output = ()> + Send>>,
20    count: usize,
21    info: Info,
22    metric: Metric,
23}
24
25impl<Metric> MetricInterest<Metric>
26where
27    Metric: Debug,
28{
29    /// Wait for a metric to be resolved
30    pub async fn wait(self) -> Result<(), xmtp_common::time::Expired> {
31        let Self {
32            fut,
33            count,
34            info,
35            metric,
36        } = self;
37
38        let result = xmtp_common::time::timeout(Duration::from_secs(10), fut).await;
39        tracing::info!("successfully waited for {metric:?}, {result:?}");
40        if info.count() >= count {
41            return Ok(());
42        }
43        tracing::error!(
44            "Timed out waiting for {:?} to be >= {}. Value: {}",
45            metric,
46            count,
47            info.count()
48        );
49        result
50    }
51}
52
53/// Information and interest in a specific metric
54#[derive(Debug, Clone)]
55struct Info {
56    count: Arc<AtomicUsize>,
57    // each Notify is a task waiting on this metric to be incremented
58    notify: Arc<Notify>,
59}
60
61impl Info {
62    fn new() -> Self {
63        Self {
64            count: Arc::new(AtomicUsize::default()),
65            notify: Arc::new(Notify::new()),
66        }
67    }
68
69    fn count(&self) -> usize {
70        self.count.load(Ordering::SeqCst)
71    }
72
73    fn increment(&self) {
74        self.count.fetch_add(1, Ordering::Relaxed);
75    }
76
77    fn clear(&self) {
78        self.count.store(0, Ordering::SeqCst)
79    }
80
81    // Registers interest in the next time this event is fired.
82    // Returns a future that resolves when this event resolves
83    fn register_interest(&self, count: usize) -> impl Future<Output = ()> + 'static {
84        let notify = self.notify.clone();
85        let info_count = self.count.clone();
86        async move {
87            while info_count.load(Ordering::SeqCst) < count {
88                notify.notified().await;
89            }
90        }
91    }
92
93    fn fire(&self) {
94        self.notify.notify_waiters();
95    }
96}
97
98#[derive(Debug)]
99pub struct WorkerMetrics<Metric> {
100    metrics: Mutex<HashMap<Metric, Info>>,
101}
102
103impl<Metric> WorkerMetrics<Metric>
104where
105    Metric: PartialEq + Eq + Hash + Clone + Copy + Debug,
106{
107    pub fn new(_installation_id: InstallationId) -> Self {
108        Self {
109            metrics: Mutex::default(),
110        }
111    }
112
113    fn info(&self, metric: Metric) -> Info {
114        self.metrics
115            .lock()
116            .entry(metric)
117            .or_insert(Info::new())
118            .clone()
119    }
120
121    pub fn get(&self, metric: Metric) -> usize {
122        self.info(metric).count()
123    }
124
125    /// Counts one `metric` and wakes its waiters.
126    pub fn increment_metric(&self, metric: Metric) {
127        self.info(metric).increment();
128        tracing::trace!("firing {metric:?}");
129        self.info(metric).fire();
130    }
131
132    pub fn reset_metrics(&self) {
133        *self.metrics.lock() = HashMap::new();
134    }
135
136    /// Register interest in a metric at a certain 'count'.
137    /// Returns a MetricInterest type that can be used to wait for the metric to reach these
138    /// parameters
139    pub fn register_interest(&self, metric_key: Metric, count: usize) -> MetricInterest<Metric> {
140        tracing::info!("registering interest in {metric_key:?}");
141        let info = self.info(metric_key);
142
143        let fut = if self
144            .metrics
145            .lock()
146            .get(&metric_key)
147            .is_some_and(|info| info.count() >= count)
148        {
149            futures::future::ready(()).boxed()
150        } else {
151            info.register_interest(count).boxed()
152        };
153
154        MetricInterest {
155            fut,
156            count,
157            info: self.info(metric_key),
158            metric: metric_key,
159        }
160    }
161
162    pub async fn do_until<F, Fut>(
163        &self,
164        metric: Metric,
165        count: usize,
166        f: F,
167    ) -> Result<(), xmtp_common::time::Expired>
168    where
169        F: Fn() -> Fut,
170        Fut: Future<Output = ()>,
171    {
172        let info = {
173            let mut m = self.metrics.lock();
174            m.entry(metric).or_insert(Info::new()).clone()
175        };
176        xmtp_common::time::timeout(Duration::from_secs(20), async {
177            while info.count() < count {
178                f().await;
179                xmtp_common::task::yield_now().await;
180            }
181        })
182        .await
183    }
184
185    pub fn clear_metric(&self, metric: Metric) {
186        self.metrics
187            .lock()
188            .entry(metric)
189            .or_insert(Info::new())
190            .clear();
191    }
192}
193
194#[cfg(all(test, not(target_arch = "wasm32")))]
195mod tests;
xmtp_mls_core/worker/metrics.rs

xmtp_mls_core/worker/
metrics.rs
