xmtp_mls/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    pub(crate) fn increment_metric(&self, metric: Metric) {
126        self.info(metric).increment();
127        tracing::trace!("firing {metric:?}");
128        self.info(metric).fire();
129    }
130
131    pub fn reset_metrics(&self) {
132        *self.metrics.lock() = HashMap::new();
133    }
134
135    /// Register interest in a metric at a certain 'count'.
136    /// Returns a MetricInterest type that can be used to wait for the metric to reach these
137    /// parameters
138    pub fn register_interest(&self, metric_key: Metric, count: usize) -> MetricInterest<Metric> {
139        tracing::info!("registering interest in {metric_key:?}");
140        let info = self.info(metric_key);
141
142        let fut = if self
143            .metrics
144            .lock()
145            .get(&metric_key)
146            .is_some_and(|info| info.count() >= count)
147        {
148            futures::future::ready(()).boxed()
149        } else {
150            info.register_interest(count).boxed()
151        };
152
153        MetricInterest {
154            fut,
155            count,
156            info: self.info(metric_key),
157            metric: metric_key,
158        }
159    }
160
161    pub async fn do_until<F, Fut>(
162        &self,
163        metric: Metric,
164        count: usize,
165        f: F,
166    ) -> Result<(), xmtp_common::time::Expired>
167    where
168        F: Fn() -> Fut,
169        Fut: Future<Output = ()>,
170    {
171        let info = {
172            let mut m = self.metrics.lock();
173            m.entry(metric).or_insert(Info::new()).clone()
174        };
175        xmtp_common::time::timeout(Duration::from_secs(20), async {
176            while info.count() < count {
177                f().await;
178                xmtp_common::task::yield_now().await;
179            }
180        })
181        .await
182    }
183
184    pub fn clear_metric(&self, metric: Metric) {
185        self.metrics
186            .lock()
187            .entry(metric)
188            .or_insert(Info::new())
189            .clear();
190    }
191}
192
193#[cfg(all(test, not(target_arch = "wasm32")))]
194mod tests;
xmtp_mls/worker/metrics.rs

xmtp_mls/worker/
metrics.rs
